1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "time"
8
9 "github.com/patriceckhart/zot/internal/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 Block bool
56 Reason string
57 ModifiedArgs json.RawMessage
58 ReplaceText string
59}
60
61const interceptTimeout = 5 * time.Second
62
63// InterceptToolCall is the typed entry point for tool_call
64// interception. Subscribers are invoked serially; the first to return
65// Block=true wins. Rewrites (ModifiedArgs) from earlier subscribers
66// flow into later ones, so a chain of guards can successively redact
67// / patch the args.
68func (m *Manager) InterceptToolCall(ctx context.Context, toolID, toolName string, args json.RawMessage) InterceptResult {
69 subs := m.interceptSubsFor("tool_call")
70 if len(subs) == 0 {
71 return InterceptResult{}
72 }
73 current := args
74 for _, ext := range subs {
75 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
76 Event: "tool_call",
77 ToolID: toolID,
78 ToolName: toolName,
79 ToolArgs: current,
80 })
81 if r.Block {
82 return r
83 }
84 if len(r.ModifiedArgs) > 0 && json.Valid(r.ModifiedArgs) {
85 current = r.ModifiedArgs
86 }
87 }
88 out := InterceptResult{}
89 if !jsonEqual(current, args) {
90 out.ModifiedArgs = current
91 }
92 return out
93}
94
95// InterceptTurnStart asks every subscriber whether the upcoming turn
96// may run. Block=true aborts the turn with Reason shown to the user.
97// Rewrites are not supported for this event.
98func (m *Manager) InterceptTurnStart(ctx context.Context, step int) InterceptResult {
99 subs := m.interceptSubsFor("turn_start")
100 if len(subs) == 0 {
101 return InterceptResult{}
102 }
103 for _, ext := range subs {
104 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
105 Event: "turn_start",
106 Step: step,
107 })
108 if r.Block {
109 return r
110 }
111 }
112 return InterceptResult{}
113}
114
115// InterceptAssistantMessage asks every subscriber to approve, block,
116// or rewrite the assistant's final visible text. Block=true hides the
117// message from the user entirely; a non-empty ReplaceText rewrites
118// what the user sees while keeping the model's original text in the
119// transcript. Successive rewrites chain: each subscriber sees the
120// previous subscriber's output.
121func (m *Manager) InterceptAssistantMessage(ctx context.Context, text string) InterceptResult {
122 subs := m.interceptSubsFor("assistant_message")
123 if len(subs) == 0 {
124 return InterceptResult{}
125 }
126 current := text
127 for _, ext := range subs {
128 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
129 Event: "assistant_message",
130 Text: current,
131 })
132 if r.Block {
133 return r
134 }
135 if r.ReplaceText != "" {
136 current = r.ReplaceText
137 }
138 }
139 out := InterceptResult{}
140 if current != text {
141 out.ReplaceText = current
142 }
143 return out
144}
145
146// interceptSubsFor returns the snapshot of extensions that subscribed
147// to intercepting the named event.
148func (m *Manager) interceptSubsFor(event string) []*Extension {
149 m.mu.RLock()
150 defer m.mu.RUnlock()
151 subs := make([]*Extension, 0, len(m.ext))
152 for _, ext := range m.ext {
153 ext.mu.Lock()
154 _, subscribed := ext.interceptSubs[event]
155 ext.mu.Unlock()
156 if subscribed {
157 subs = append(subs, ext)
158 }
159 }
160 return subs
161}
162
163// askIntercept sends one EventInterceptFromHost to ext and waits for
164// the reply, a timeout, or context cancellation. Returns a typed
165// result. Never blocks for longer than interceptTimeout.
166func (m *Manager) askIntercept(ctx context.Context, ext *Extension, payload extproto.EventInterceptFromHost) InterceptResult {
167 id := newCorrelationID()
168 ch := make(chan extproto.EventInterceptResponseFromExt, 1)
169 ext.mu.Lock()
170 ext.pendingIntercept[id] = ch
171 ext.mu.Unlock()
172
173 payload.Type = "event_intercept"
174 payload.ID = id
175 frame, err := extproto.Encode(payload)
176 if err != nil {
177 ext.mu.Lock()
178 delete(ext.pendingIntercept, id)
179 ext.mu.Unlock()
180 return InterceptResult{}
181 }
182 if _, err := ext.stdin.Write(frame); err != nil {
183 ext.mu.Lock()
184 delete(ext.pendingIntercept, id)
185 ext.mu.Unlock()
186 fmt.Fprintf(ext.logFile, "[zot] intercept write failed: %v\n", err)
187 return InterceptResult{}
188 }
189
190 select {
191 case resp := <-ch:
192 return InterceptResult{
193 Block: resp.Block,
194 Reason: resp.Reason,
195 ModifiedArgs: resp.ModifiedArgs,
196 ReplaceText: resp.ReplaceText,
197 }
198 case <-time.After(interceptTimeout):
199 ext.mu.Lock()
200 delete(ext.pendingIntercept, id)
201 ext.mu.Unlock()
202 fmt.Fprintf(ext.logFile, "[zot] intercept %s timed out; allowing\n", payload.Event)
203 return InterceptResult{}
204 case <-ctx.Done():
205 ext.mu.Lock()
206 delete(ext.pendingIntercept, id)
207 ext.mu.Unlock()
208 return InterceptResult{}
209 }
210}
211
212// jsonEqual reports whether a and b encode the same JSON value. Used
213// to detect whether the interceptor chain actually mutated the args.
214func jsonEqual(a, b json.RawMessage) bool {
215 if len(a) == len(b) {
216 same := true
217 for i := range a {
218 if a[i] != b[i] {
219 same = false
220 break
221 }
222 }
223 if same {
224 return true
225 }
226 }
227 // Fallback to structural compare so whitespace differences don't
228 // register as a mutation. (We don't actually rely on this to be
229 // cheap, callers only use it to pick a code path.)
230 var av, bv any
231 if json.Unmarshal(a, &av) != nil || json.Unmarshal(b, &bv) != nil {
232 return false
233 }
234 ae, _ := json.Marshal(av)
235 be, _ := json.Marshal(bv)
236 return string(ae) == string(be)
237}