1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "time"
8
9 "github.com/patriceckhart/zot/packages/agent/extproto"
10)
11
12// EmitEvent fires a one-way lifecycle event to every extension that
13// subscribed to it via SubscribeFromExt.events. Notification enqueue is
14// ordered and non-blocking; a slow consumer is disconnected on overflow.
15//
16// Event names are documented on extproto.EventFromHost. Unknown event
17// names are still routed (subscribers can use any string they want).
18func (m *Manager) EmitEvent(ev extproto.EventFromHost) {
19 m.lifecycleMu.Lock()
20 defer m.lifecycleMu.Unlock()
21 if m.stopping {
22 return
23 }
24 m.emitEventLocked(ev)
25}
26
27func (m *Manager) emitEventLocked(ev extproto.EventFromHost) {
28 ev.Type = "event"
29 m.sequence++
30 ev.Sequence = m.sequence
31 if ev.SessionID == "" {
32 ev.SessionID = m.sessionID
33 }
34 if ev.CWD == "" {
35 ev.CWD = m.sessionCWD
36 }
37 if ev.CWD == "" {
38 ev.CWD = m.cwd
39 }
40 // A malformed model call must still have an observable terminal result.
41 if len(ev.ToolArgs) > 0 && !json.Valid(ev.ToolArgs) {
42 ev.ToolArgsRaw = string(ev.ToolArgs)
43 ev.ToolArgs = nil
44 }
45 frame, err := extproto.Encode(ev)
46 if err != nil {
47 return
48 }
49 m.mu.RLock()
50 subs := make([]*Extension, 0, len(m.ext))
51 for _, ext := range m.ext {
52 ext.mu.Lock()
53 _, subscribed := ext.eventSubs[ev.Event]
54 ext.mu.Unlock()
55 if subscribed {
56 subs = append(subs, ext)
57 }
58 }
59 m.mu.RUnlock()
60
61 for _, ext := range subs {
62 if ext.stdin == nil {
63 continue
64 }
65 var err error
66 if pipe, ok := ext.stdin.(*orderedPipe); ok {
67 err = pipe.enqueue(frame, nil)
68 } else {
69 // Tests may provide a synchronous in-memory writer.
70 _, err = ext.stdin.Write(frame)
71 }
72 if err != nil && ext.logFile != nil {
73 fmt.Fprintf(ext.logFile, "[zot] lifecycle delivery failed: %v\n", err)
74 }
75 }
76}
77
78// InterceptResult aggregates the outcome of walking every subscribed
79// interceptor for one event. The zero value means "allow, no
80// modification". Callers check Block first; if allowed, use the
81// optional rewrite fields (ModifiedArgs for tool_call, ReplaceText
82// for assistant_message) to carry the rewrite into the action.
83type InterceptResult struct {
84 SystemPrompt json.RawMessage
85 failure string
86 Block bool
87 Reason string
88 ModifiedArgs json.RawMessage
89 ReplaceText string
90}
91
92const interceptTimeout = 5 * time.Second
93
94// InterceptToolCall is the typed entry point for tool_call
95// interception. Subscribers are invoked serially; the first to return
96// Block=true wins. Rewrites (ModifiedArgs) from earlier subscribers
97// flow into later ones, so a chain of guards can successively redact
98// / patch the args.
99func (m *Manager) InterceptToolCall(ctx context.Context, toolID, toolName string, args json.RawMessage) InterceptResult {
100 subs := m.interceptSubsFor("tool_call")
101 if len(subs) == 0 {
102 return InterceptResult{}
103 }
104 current := args
105 for _, ext := range subs {
106 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
107 Event: "tool_call",
108 ToolID: toolID,
109 ToolName: toolName,
110 ToolArgs: current,
111 })
112 if r.Block {
113 r.ModifiedArgs = current
114 return r
115 }
116 if len(r.ModifiedArgs) > 0 && json.Valid(r.ModifiedArgs) {
117 current = r.ModifiedArgs
118 }
119 }
120 out := InterceptResult{}
121 if !jsonEqual(current, args) {
122 out.ModifiedArgs = current
123 }
124 return out
125}
126
127// InterceptTurnStart asks every subscriber whether the upcoming turn
128// may run. Block=true aborts the turn with Reason shown to the user.
129// Rewrites are not supported for this event.
130func (m *Manager) InterceptTurnStart(ctx context.Context, step int) InterceptResult {
131 subs := m.interceptSubsFor("turn_start")
132 if len(subs) == 0 {
133 return InterceptResult{}
134 }
135 for _, ext := range subs {
136 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
137 Event: "turn_start",
138 Step: step,
139 })
140 if r.Block {
141 return r
142 }
143 }
144 return InterceptResult{}
145}
146
147// InterceptAssistantMessage asks every subscriber to approve, block,
148// or rewrite the assistant's final visible text. Block=true hides the
149// message from the user entirely; a non-empty ReplaceText rewrites
150// what the user sees while keeping the model's original text in the
151// transcript. Successive rewrites chain: each subscriber sees the
152// previous subscriber's output.
153func (m *Manager) InterceptAssistantMessage(ctx context.Context, text string) InterceptResult {
154 subs := m.interceptSubsFor("assistant_message")
155 if len(subs) == 0 {
156 return InterceptResult{}
157 }
158 current := text
159 for _, ext := range subs {
160 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
161 Event: "assistant_message",
162 Text: current,
163 })
164 if r.Block {
165 return r
166 }
167 if r.ReplaceText != "" {
168 current = r.ReplaceText
169 }
170 }
171 out := InterceptResult{}
172 if current != text {
173 out.ReplaceText = current
174 }
175 return out
176}
177
178// interceptSubsFor returns the snapshot of extensions that subscribed
179// to intercepting the named event.
180func (m *Manager) interceptSubsFor(event string) []*Extension {
181 m.mu.RLock()
182 defer m.mu.RUnlock()
183 subs := make([]*Extension, 0, len(m.ext))
184 for _, ext := range m.ext {
185 ext.mu.Lock()
186 _, subscribed := ext.interceptSubs[event]
187 ext.mu.Unlock()
188 if subscribed {
189 subs = append(subs, ext)
190 }
191 }
192 return subs
193}
194
195// askIntercept sends one EventInterceptFromHost to ext and waits for
196// the reply, a timeout, or context cancellation. Returns a typed
197// result. before_agent_start callers bound the complete operation with a
198// context deadline; other events retain the legacy response-only timeout.
199func (m *Manager) askIntercept(ctx context.Context, ext *Extension, payload extproto.EventInterceptFromHost) InterceptResult {
200 id := newCorrelationID()
201 ch := make(chan extproto.EventInterceptResponseFromExt, 1)
202 ext.mu.Lock()
203 ext.pendingIntercept[id] = ch
204 ext.mu.Unlock()
205
206 payload.Type = "event_intercept"
207 payload.ID = id
208 frame, err := extproto.Encode(payload)
209 if err != nil {
210 ext.mu.Lock()
211 delete(ext.pendingIntercept, id)
212 ext.mu.Unlock()
213 return InterceptResult{}
214 }
215 var writeErr error
216 if payload.Event == "before_agent_start" {
217 // Bound pipe writes too. Closing the pipe releases a blocked writer.
218 done := make(chan error, 1)
219 go func() { _, err := ext.stdin.Write(frame); done <- err }()
220 select {
221 case writeErr = <-done:
222 case <-ctx.Done():
223 _ = ext.stdin.Close()
224 <-done
225 writeErr = ctx.Err()
226 }
227 } else {
228 _, writeErr = ext.stdin.Write(frame)
229 }
230 if err := writeErr; err != nil {
231 ext.mu.Lock()
232 delete(ext.pendingIntercept, id)
233 ext.mu.Unlock()
234 fmt.Fprintf(ext.logFile, "[zot] intercept write failed: %v\n", err)
235 return InterceptResult{failure: "request write failed"}
236 }
237
238 select {
239 case resp := <-ch:
240 return InterceptResult{
241 SystemPrompt: resp.SystemPrompt,
242 Block: resp.Block,
243 Reason: resp.Reason,
244 ModifiedArgs: resp.ModifiedArgs,
245 ReplaceText: resp.ReplaceText,
246 }
247 case <-time.After(interceptTimeout):
248 ext.mu.Lock()
249 delete(ext.pendingIntercept, id)
250 ext.mu.Unlock()
251 fmt.Fprintf(ext.logFile, "[zot] intercept %s timed out; allowing\n", payload.Event)
252 return InterceptResult{failure: "response timed out"}
253 case <-ctx.Done():
254 ext.mu.Lock()
255 delete(ext.pendingIntercept, id)
256 ext.mu.Unlock()
257 return InterceptResult{failure: "response canceled or timed out"}
258 }
259}
260
261// jsonEqual reports whether a and b encode the same JSON value. Used
262// to detect whether the interceptor chain actually mutated the args.
263func jsonEqual(a, b json.RawMessage) bool {
264 if len(a) == len(b) {
265 same := true
266 for i := range a {
267 if a[i] != b[i] {
268 same = false
269 break
270 }
271 }
272 if same {
273 return true
274 }
275 }
276 // Fallback to structural compare so whitespace differences don't
277 // register as a mutation. (We don't actually rely on this to be
278 // cheap, callers only use it to pick a code path.)
279 var av, bv any
280 if json.Unmarshal(a, &av) != nil || json.Unmarshal(b, &bv) != nil {
281 return false
282 }
283 ae, _ := json.Marshal(av)
284 be, _ := json.Marshal(bv)
285 return string(ae) == string(be)
286}