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 OriginExtension: m.ToolCallOrigin(toolID),
108 Event: "tool_call",
109 ToolID: toolID,
110 ToolName: toolName,
111 ToolArgs: current,
112 })
113 if r.Block {
114 r.ModifiedArgs = current
115 return r
116 }
117 if len(r.ModifiedArgs) > 0 && json.Valid(r.ModifiedArgs) {
118 current = r.ModifiedArgs
119 }
120 }
121 out := InterceptResult{}
122 if !jsonEqual(current, args) {
123 out.ModifiedArgs = current
124 }
125 return out
126}
127
128// InterceptTurnStart asks every subscriber whether the upcoming turn
129// may run. Block=true aborts the turn with Reason shown to the user.
130// Rewrites are not supported for this event.
131func (m *Manager) InterceptTurnStart(ctx context.Context, step int) InterceptResult {
132 subs := m.interceptSubsFor("turn_start")
133 if len(subs) == 0 {
134 return InterceptResult{}
135 }
136 for _, ext := range subs {
137 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
138 Event: "turn_start",
139 Step: step,
140 })
141 if r.Block {
142 return r
143 }
144 }
145 return InterceptResult{}
146}
147
148// InterceptAssistantMessage asks every subscriber to approve, block,
149// or rewrite the assistant's final visible text. Block=true hides the
150// message from the user entirely; a non-empty ReplaceText rewrites
151// what the user sees while keeping the model's original text in the
152// transcript. Successive rewrites chain: each subscriber sees the
153// previous subscriber's output.
154func (m *Manager) InterceptAssistantMessage(ctx context.Context, text string) InterceptResult {
155 subs := m.interceptSubsFor("assistant_message")
156 if len(subs) == 0 {
157 return InterceptResult{}
158 }
159 current := text
160 for _, ext := range subs {
161 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
162 Event: "assistant_message",
163 Text: current,
164 })
165 if r.Block {
166 return r
167 }
168 if r.ReplaceText != "" {
169 current = r.ReplaceText
170 }
171 }
172 out := InterceptResult{}
173 if current != text {
174 out.ReplaceText = current
175 }
176 return out
177}
178
179// interceptSubsFor returns the snapshot of extensions that subscribed
180// to intercepting the named event.
181func (m *Manager) interceptSubsFor(event string) []*Extension {
182 m.mu.RLock()
183 defer m.mu.RUnlock()
184 subs := make([]*Extension, 0, len(m.ext))
185 for _, ext := range m.ext {
186 ext.mu.Lock()
187 _, subscribed := ext.interceptSubs[event]
188 ext.mu.Unlock()
189 if subscribed {
190 subs = append(subs, ext)
191 }
192 }
193 return subs
194}
195
196// askIntercept sends one EventInterceptFromHost to ext and waits for
197// the reply, a timeout, or context cancellation. Returns a typed
198// result. before_agent_start callers bound the complete operation with a
199// context deadline; other events retain the legacy response-only timeout.
200func (m *Manager) askIntercept(ctx context.Context, ext *Extension, payload extproto.EventInterceptFromHost) InterceptResult {
201 id := newCorrelationID()
202 ch := make(chan extproto.EventInterceptResponseFromExt, 1)
203 ext.mu.Lock()
204 ext.pendingIntercept[id] = ch
205 ext.mu.Unlock()
206
207 payload.Type = "event_intercept"
208 payload.ID = id
209 frame, err := extproto.Encode(payload)
210 if err != nil {
211 ext.mu.Lock()
212 delete(ext.pendingIntercept, id)
213 ext.mu.Unlock()
214 return InterceptResult{}
215 }
216 var writeErr error
217 if payload.Event == "before_agent_start" {
218 // Bound pipe writes too. Closing the pipe releases a blocked writer.
219 done := make(chan error, 1)
220 go func() { _, err := ext.stdin.Write(frame); done <- err }()
221 select {
222 case writeErr = <-done:
223 case <-ctx.Done():
224 _ = ext.stdin.Close()
225 <-done
226 writeErr = ctx.Err()
227 }
228 } else {
229 _, writeErr = ext.stdin.Write(frame)
230 }
231 if err := writeErr; err != nil {
232 ext.mu.Lock()
233 delete(ext.pendingIntercept, id)
234 ext.mu.Unlock()
235 fmt.Fprintf(ext.logFile, "[zot] intercept write failed: %v\n", err)
236 return InterceptResult{failure: "request write failed"}
237 }
238
239 select {
240 case resp := <-ch:
241 return InterceptResult{
242 SystemPrompt: resp.SystemPrompt,
243 Block: resp.Block,
244 Reason: resp.Reason,
245 ModifiedArgs: resp.ModifiedArgs,
246 ReplaceText: resp.ReplaceText,
247 }
248 case <-time.After(interceptTimeout):
249 ext.mu.Lock()
250 delete(ext.pendingIntercept, id)
251 ext.mu.Unlock()
252 fmt.Fprintf(ext.logFile, "[zot] intercept %s timed out; allowing\n", payload.Event)
253 return InterceptResult{failure: "response timed out"}
254 case <-ctx.Done():
255 ext.mu.Lock()
256 delete(ext.pendingIntercept, id)
257 ext.mu.Unlock()
258 return InterceptResult{failure: "response canceled or timed out"}
259 }
260}
261
262// jsonEqual reports whether a and b encode the same JSON value. Used
263// to detect whether the interceptor chain actually mutated the args.
264func jsonEqual(a, b json.RawMessage) bool {
265 if len(a) == len(b) {
266 same := true
267 for i := range a {
268 if a[i] != b[i] {
269 same = false
270 break
271 }
272 }
273 if same {
274 return true
275 }
276 }
277 // Fallback to structural compare so whitespace differences don't
278 // register as a mutation. (We don't actually rely on this to be
279 // cheap, callers only use it to pick a code path.)
280 var av, bv any
281 if json.Unmarshal(a, &av) != nil || json.Unmarshal(b, &bv) != nil {
282 return false
283 }
284 ae, _ := json.Marshal(av)
285 be, _ := json.Marshal(bv)
286 return string(ae) == string(be)
287}