1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "strings"
8 "time"
9 "unicode/utf8"
10
11 "github.com/bnema/zut/packages/agent/extproto"
12)
13
14// EmitEvent fires a one-way lifecycle event to every extension that
15// subscribed to it via SubscribeFromExt.events. Non-blocking: each
16// extension's pipe write happens on a per-call goroutine so a slow
17// extension can't stall the agent loop.
18//
19// Event names are documented on extproto.EventFromHost. Unknown event
20// names are still routed (subscribers can use any string they want).
21func (m *Manager) EmitEvent(ev extproto.EventFromHost) {
22 ev.Type = "event"
23 frame, err := extproto.Encode(ev)
24 if err != nil {
25 return
26 }
27 m.mu.RLock()
28 subs := make([]*Extension, 0, len(m.ext))
29 for _, ext := range m.ext {
30 ext.mu.Lock()
31 _, subscribed := ext.eventSubs[ev.Event]
32 ext.mu.Unlock()
33 if subscribed {
34 subs = append(subs, ext)
35 }
36 }
37 m.mu.RUnlock()
38
39 for _, ext := range subs {
40 go func(ext *Extension) {
41 defer func() { _ = recover() }()
42 _, _ = ext.writeFrame(frame)
43 }(ext)
44 }
45}
46
47// UpdateSessionState updates the cached state used when an extension is
48// reloaded during the active session. The session file remains the durable
49// source of truth; this only keeps the in-memory lifecycle replay current.
50func (m *Manager) UpdateSessionState(extension string, state json.RawMessage) {
51 if m == nil || strings.TrimSpace(extension) == "" || len(state) > maxSessionStateBytes {
52 return
53 }
54 m.mu.Lock()
55 defer m.mu.Unlock()
56 if m.activeSession == nil {
57 return
58 }
59 if m.activeSession.states == nil {
60 m.activeSession.states = map[string]json.RawMessage{}
61 }
62 if len(state) == 0 || strings.TrimSpace(string(state)) == "null" {
63 delete(m.activeSession.states, extension)
64 return
65 }
66 m.activeSession.states[extension] = append(json.RawMessage(nil), state...)
67}
68
69// EmitSessionEvent broadcasts a session lifecycle event with the active
70// session identity. Each extension receives only its own opaque state
71// snapshot, preventing one extension from reading another extension's data.
72// The latest event is retained and replayed to extensions loaded later.
73func (m *Manager) EmitSessionEvent(event string, session *extproto.SessionContext, states map[string]json.RawMessage) {
74 if event == "" {
75 return
76 }
77 m.sessionEmitMu.Lock()
78 defer m.sessionEmitMu.Unlock()
79 base := extproto.EventFromHost{Type: "event", Event: event, Session: cloneSessionContext(session)}
80 storedStates := cloneExtensionStates(states)
81 m.mu.Lock()
82 m.activeSession = &sessionEventState{event: base, states: storedStates}
83 subs := make([]*Extension, 0, len(m.ext))
84 for _, ext := range m.ext {
85 ext.mu.Lock()
86 _, subscribed := ext.eventSubs[event]
87 ext.mu.Unlock()
88 if subscribed {
89 subs = append(subs, ext)
90 }
91 }
92 m.mu.Unlock()
93 for _, ext := range subs {
94 m.sendSessionEvent(ext, base, storedStates[ext.Manifest.Name])
95 }
96}
97
98func (m *Manager) sendSessionEvent(ext *Extension, base extproto.EventFromHost, state json.RawMessage) {
99 if ext == nil {
100 return
101 }
102 frame := base
103 frame.State = append(json.RawMessage(nil), state...)
104 encoded, err := extproto.Encode(frame)
105 if err != nil {
106 return
107 }
108 if !ext.enqueueLifecycleFrame(encoded) {
109 ext.recordDiagnostic("session lifecycle frame dropped because the extension queue is full or stopped")
110 }
111}
112
113func (m *Manager) replaySessionEvent(ext *Extension) {
114 m.sessionEmitMu.Lock()
115 defer m.sessionEmitMu.Unlock()
116 m.mu.RLock()
117 if m.activeSession == nil {
118 m.mu.RUnlock()
119 return
120 }
121 base := m.activeSession.event
122 state := append(json.RawMessage(nil), m.activeSession.states[ext.Manifest.Name]...)
123 m.mu.RUnlock()
124 m.sendSessionEvent(ext, base, state)
125}
126
127func cloneSessionContext(in *extproto.SessionContext) *extproto.SessionContext {
128 if in == nil {
129 return nil
130 }
131 out := *in
132 return &out
133}
134
135func cloneExtensionStates(in map[string]json.RawMessage) map[string]json.RawMessage {
136 if len(in) == 0 {
137 return nil
138 }
139 out := make(map[string]json.RawMessage, len(in))
140 for name, state := range in {
141 out[name] = append(json.RawMessage(nil), state...)
142 }
143 return out
144}
145
146// InterceptResult aggregates the outcome of walking every subscribed
147// interceptor for one event. The zero value means "allow, no
148// modification". Callers check Block first; if allowed, use the
149// optional rewrite fields (ModifiedArgs for tool_call, ReplaceText
150// for assistant_message) to carry the rewrite into the action.
151type InterceptResult struct {
152 Block bool
153 Reason string
154 ModifiedArgs json.RawMessage
155 ReplaceText string
156 Context string
157}
158
159const (
160 interceptTimeout = 5 * time.Second
161 maxInterceptContext = 16 * 1024
162 maxSessionStateBytes = 256 * 1024
163)
164
165const interceptContextTruncatedMarker = "\n[extension context truncated]"
166
167func boundInterceptContext(text string) string {
168 text = strings.TrimSpace(text)
169 if len(text) <= maxInterceptContext {
170 return text
171 }
172 limit := maxInterceptContext - len(interceptContextTruncatedMarker)
173 cut := limit
174 for cut > 0 && !utf8.RuneStart(text[cut]) {
175 cut--
176 }
177 return strings.TrimSpace(text[:cut]) + interceptContextTruncatedMarker
178}
179
180// InterceptToolCall is the typed entry point for tool_call
181// interception. Subscribers are invoked serially; the first to return
182// Block=true wins. Rewrites (ModifiedArgs) from earlier subscribers
183// flow into later ones, so a chain of guards can successively redact
184// / patch the args.
185func (m *Manager) InterceptToolCall(ctx context.Context, toolID, toolName string, args json.RawMessage) InterceptResult {
186 subs := m.interceptSubsFor("tool_call")
187 if len(subs) == 0 {
188 return InterceptResult{}
189 }
190 current := args
191 for _, ext := range subs {
192 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
193 Event: "tool_call",
194 ToolID: toolID,
195 ToolName: toolName,
196 ToolArgs: current,
197 })
198 if r.Block {
199 return r
200 }
201 if len(r.ModifiedArgs) > 0 && json.Valid(r.ModifiedArgs) {
202 current = r.ModifiedArgs
203 }
204 }
205 out := InterceptResult{}
206 if !jsonEqual(current, args) {
207 out.ModifiedArgs = current
208 }
209 return out
210}
211
212// InterceptTurnStart asks every subscriber whether the upcoming turn
213// may run. Block=true aborts the turn with Reason shown to the user.
214// Rewrites are not supported for this event.
215func (m *Manager) InterceptTurnStart(ctx context.Context, step int) InterceptResult {
216 subs := m.interceptSubsFor("turn_start")
217 if len(subs) == 0 {
218 return InterceptResult{}
219 }
220 var contexts []string
221 for _, ext := range subs {
222 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
223 Event: "turn_start",
224 Step: step,
225 })
226 if r.Block {
227 return r
228 }
229 if text := strings.TrimSpace(r.Context); text != "" {
230 contexts = append(contexts, "["+ext.Manifest.Name+"]\n"+text)
231 }
232 }
233 return InterceptResult{Context: boundInterceptContext(strings.Join(contexts, "\n\n"))}
234}
235
236// InterceptAssistantMessage asks every subscriber to approve, block,
237// or rewrite the assistant's final visible text. Block=true hides the
238// message from the user entirely; a non-empty ReplaceText rewrites
239// what the user sees while keeping the model's original text in the
240// transcript. Successive rewrites chain: each subscriber sees the
241// previous subscriber's output.
242func (m *Manager) InterceptAssistantMessage(ctx context.Context, text string) InterceptResult {
243 subs := m.interceptSubsFor("assistant_message")
244 if len(subs) == 0 {
245 return InterceptResult{}
246 }
247 current := text
248 for _, ext := range subs {
249 r := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{
250 Event: "assistant_message",
251 Text: current,
252 })
253 if r.Block {
254 return r
255 }
256 if r.ReplaceText != "" {
257 current = r.ReplaceText
258 }
259 }
260 out := InterceptResult{}
261 if current != text {
262 out.ReplaceText = current
263 }
264 return out
265}
266
267// interceptSubsFor returns the snapshot of extensions that subscribed
268// to intercepting the named event.
269func (m *Manager) interceptSubsFor(event string) []*Extension {
270 m.mu.RLock()
271 defer m.mu.RUnlock()
272 subs := make([]*Extension, 0, len(m.ext))
273 for _, ext := range m.ext {
274 ext.mu.Lock()
275 _, subscribed := ext.interceptSubs[event]
276 ext.mu.Unlock()
277 if subscribed {
278 subs = append(subs, ext)
279 }
280 }
281 return subs
282}
283
284// askIntercept sends one EventInterceptFromHost to ext and waits for
285// the reply, a timeout, or context cancellation. Returns a typed
286// result. Never blocks for longer than interceptTimeout.
287func (m *Manager) askIntercept(ctx context.Context, ext *Extension, payload extproto.EventInterceptFromHost) InterceptResult {
288 id := newCorrelationID()
289 ch := make(chan extproto.EventInterceptResponseFromExt, 1)
290 ext.mu.Lock()
291 ext.pendingIntercept[id] = ch
292 ext.mu.Unlock()
293
294 payload.Type = "event_intercept"
295 payload.ID = id
296 frame, err := extproto.Encode(payload)
297 if err != nil {
298 ext.mu.Lock()
299 delete(ext.pendingIntercept, id)
300 ext.mu.Unlock()
301 return InterceptResult{}
302 }
303 if _, err := ext.writeFrame(frame); err != nil {
304 ext.mu.Lock()
305 delete(ext.pendingIntercept, id)
306 ext.mu.Unlock()
307 fmt.Fprintf(ext.logFile, "[zut] intercept write failed: %v\n", err)
308 return InterceptResult{}
309 }
310
311 select {
312 case resp := <-ch:
313 return InterceptResult{
314 Block: resp.Block,
315 Reason: resp.Reason,
316 ModifiedArgs: resp.ModifiedArgs,
317 ReplaceText: resp.ReplaceText,
318 Context: resp.Context,
319 }
320 case <-time.After(interceptTimeout):
321 ext.mu.Lock()
322 delete(ext.pendingIntercept, id)
323 ext.mu.Unlock()
324 fmt.Fprintf(ext.logFile, "[zut] intercept %s timed out; allowing\n", payload.Event)
325 return InterceptResult{}
326 case <-ctx.Done():
327 ext.mu.Lock()
328 delete(ext.pendingIntercept, id)
329 ext.mu.Unlock()
330 return InterceptResult{}
331 }
332}
333
334// jsonEqual reports whether a and b encode the same JSON value. Used
335// to detect whether the interceptor chain actually mutated the args.
336func jsonEqual(a, b json.RawMessage) bool {
337 if len(a) == len(b) {
338 same := true
339 for i := range a {
340 if a[i] != b[i] {
341 same = false
342 break
343 }
344 }
345 if same {
346 return true
347 }
348 }
349 // Fallback to structural compare so whitespace differences don't
350 // register as a mutation. (We don't actually rely on this to be
351 // cheap, callers only use it to pick a code path.)
352 var av, bv any
353 if json.Unmarshal(a, &av) != nil || json.Unmarshal(b, &bv) != nil {
354 return false
355 }
356 ae, _ := json.Marshal(av)
357 be, _ := json.Marshal(bv)
358 return string(ae) == string(be)
359}