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