1// Package ext is the Go SDK for writing zot extensions.
2//
3// An extension is a subprocess that talks to zot over its stdin/stdout
4// in newline-delimited JSON. This package wraps the wire format so
5// extension authors can write straightforward Go without reimplementing
6// the protocol.
7//
8// Minimal example (registers /hello and replies with a static prompt):
9//
10// package main
11//
12// import "github.com/patriceckhart/zot/packages/agent/ext"
13//
14// func main() {
15// ext := ext.New("hello", "1.0.0")
16// ext.Command("hello", "say hi", func(args string) ext.Response {
17// return ext.Prompt("Greet me in three different languages.")
18// })
19// ext.Run()
20// }
21//
22// Build it, drop the binary + an extension.json next to it under
23// `$ZOT_HOME/extensions/hello/`, and zot picks it up on next launch.
24//
25// Use OnHello when registration or configuration depends on host metadata
26// such as HostInfo.CWD, Provider, Model, ZotVersion, ExtensionDir, or DataDir.
27// Run sends hello, waits for hello_ack, runs OnHello, announces registrations,
28// then sends ready.
29//
30// The same wire format also has reference clients in TypeScript and
31// Python under examples/extensions/. Use whichever language fits.
32package ext
33
34import (
35 "bufio"
36 "encoding/base64"
37 "encoding/json"
38 "fmt"
39 "io"
40 "os"
41 "sync"
42
43 "github.com/patriceckhart/zot/packages/agent/extproto"
44)
45
46func base64Encode(b []byte) string { return base64.StdEncoding.EncodeToString(b) }
47
48// CommandHandler is invoked when the user runs the extension's
49// registered slash command. args is everything the user typed after
50// the command name (already trimmed). Return a Response describing
51// what zot should do next.
52type CommandHandler func(args string) Response
53
54// ToolHandler is invoked when the LLM calls a tool the extension
55// registered. args is the raw JSON object the model produced; the
56// handler is responsible for parsing/validating it. Return a
57// ToolResult describing what zot should send back to the model.
58type ToolHandler func(args json.RawMessage) ToolResult
59
60// EventHandler is called for each lifecycle event the extension
61// subscribed to via Subscribe. The handler is invoked synchronously
62// on the read goroutine; keep it quick or hand off to your own
63// worker.
64type EventHandler func(ev Event)
65
66// EventResult contains bounded text/base64-image tool output for observers.
67// Truncated reports omitted or shortened content, not execution failure.
68type EventResult = extproto.EventResult
69
70// Event is an observational lifecycle notification. SessionID, CWD and Sequence
71// identify its context. Name selects the remaining payload:
72//
73// session_start: no additional fields
74// session_end: Reason
75// user_prompt_submit: Text, Queued, ImageCount
76// turn_start: Step
77// turn_end: Stop, Error (model-response boundary, before client tools)
78// tool_call: ToolID, ToolName, ToolArgs (original arguments)
79// tool_result: ToolID, ToolName, ToolArgs (effective), Status, Executed, Result
80// tool_confirmation_requested: ToolID, ToolName, ToolPreview
81// permission_decision: ToolID, ToolName, Decision, Source, Stage, Reason
82// assistant_message: Text
83// pre_compact, post_compact: CompactionID, MessageCount, TokenEstimate, Status, Error
84// subagent_start, subagent_stop: AgentID, AgentRunID, AgentName, Status, Error
85//
86// Result and count pointers distinguish absent fields from false/zero values.
87// Notifications are best-effort, not a durable audit log. See docs/extensions.md
88// for payload limits, privacy considerations, and delivery semantics.
89type Event struct {
90 Name string
91 AgentRunID string
92 SessionID string
93 Queued bool
94 ImageCount int
95 CWD string
96 Sequence uint64
97 Status string
98 Reason string
99 Source string
100 Decision string
101 Stage string
102 Result *EventResult
103 Executed *bool
104 CompactionID string
105 MessageCount *int
106 TokenEstimate *int
107 AgentID string
108 AgentName string
109
110 Step int
111 Stop string
112 Error string
113
114 ToolID string
115 ToolName string
116 ToolArgs json.RawMessage
117 ToolPreview string
118 ToolArgsRaw string
119
120 Text string
121}
122
123// InterceptHandler decides whether a tool call may proceed. Return
124// (allow=true) to permit, (allow=false, reason) to refuse. The
125// reason is shown to the model as the tool's error text.
126//
127// This is the original 2-value handler. For richer control (arg
128// rewrites, turn_start / assistant_message interception, etc.)
129// register via InterceptToolCallX, InterceptTurnStart, or
130// InterceptAssistantMessage.
131type InterceptHandler func(toolName string, args json.RawMessage) (allow bool, reason string)
132
133// ToolCallDecision is the richer reply an InterceptToolCallHandler
134// can return. Zero value means "allow, pass through".
135type ToolCallDecision struct {
136 // Block refuses the call. The model sees Reason as the tool's
137 // error text and typically proposes a different action.
138 Block bool
139 // Reason is the refusal text shown to the model on Block, or a
140 // logged note when passing through (optional).
141 Reason string
142 // ModifiedArgs, when non-nil and Block is false, replaces the
143 // JSON args the tool actually runs with. Lets guards redact /
144 // augment / patch the model's request without rewriting the
145 // transcript. Must encode to a JSON object.
146 ModifiedArgs json.RawMessage
147}
148
149// ToolCallHandler is the richer form of an interceptor. Use this
150// when you want to rewrite args instead of only blocking.
151type ToolCallHandler func(toolName string, args json.RawMessage) ToolCallDecision
152
153// BeforeAgentStartEvent contains the complete prompt and effective runtime context.
154type BeforeAgentStartEvent struct {
155 SystemPrompt string
156 SessionID string
157 AgentRunID string
158 CWD string
159 Provider string
160 Model string
161}
162
163// BeforeAgentStartHandler returns an exact replacement, or nil to keep the prompt.
164// A pointer to an empty string intentionally removes the system prompt.
165type BeforeAgentStartHandler func(BeforeAgentStartEvent) *string
166
167// TurnStartDecision controls whether the next model call runs. Zero
168// value means "allow".
169type TurnStartDecision struct {
170 Block bool
171 Reason string
172}
173
174// TurnStartHandler is called before every turn's model call. Return
175// Block=true with Reason to abort the turn (shown to the user as a
176// status line). Useful for rate-limiting and business-hour gates.
177type TurnStartHandler func(step int) TurnStartDecision
178
179// AssistantMessageDecision controls the final assistant-text
180// rendering. Zero value means "allow, show as-is".
181type AssistantMessageDecision struct {
182 // Block suppresses the message from the user entirely. The
183 // transcript still records the model's original output so the
184 // model sees what it said on the next turn.
185 Block bool
186 // Reason is logged when blocking (optional).
187 Reason string
188 // ReplaceText, when non-empty and Block is false, is what the
189 // user sees. The model's original text still lives in the
190 // transcript.
191 ReplaceText string
192}
193
194// AssistantMessageHandler is called after the model's final text is
195// assembled but before it's shown. Use it to scrub secrets, expand
196// templates, enforce tone, or suppress responses entirely.
197type AssistantMessageHandler func(text string) AssistantMessageDecision
198
199// ToolResult is the extension's reply to a tool invocation. Build
200// one with TextResult, ImageResult, or directly when you need to
201// combine multiple blocks.
202type ToolResult struct {
203 Content []ToolContent
204 IsError bool
205 ActivateTools []string
206}
207
208// ToolContent is one block of tool output. Either Text is set, or
209// MimeType+Data (base64-encoded). Use the Text/Image helpers below.
210type ToolContent struct {
211 Type string // "text" | "image"
212 Text string
213 MimeType string
214 Data string // base64
215}
216
217// Text returns a text content block.
218func Text(s string) ToolContent { return ToolContent{Type: "text", Text: s} }
219
220// Image returns an image content block. data must already be
221// base64-encoded; use ImageBytes to encode raw bytes.
222func Image(mimeType, base64Data string) ToolContent {
223 return ToolContent{Type: "image", MimeType: mimeType, Data: base64Data}
224}
225
226// ImageBytes returns an image content block from raw bytes,
227// encoding them to base64 for the wire.
228func ImageBytes(mimeType string, data []byte) ToolContent {
229 return ToolContent{Type: "image", MimeType: mimeType, Data: base64Encode(data)}
230}
231
232// TextResult returns a tool result with one text block, success.
233func TextResult(s string) ToolResult { return ToolResult{Content: []ToolContent{Text(s)}} }
234
235// TextErrorResult returns a tool result with one text block, marked
236// as an error to the model.
237func TextErrorResult(s string) ToolResult {
238 return ToolResult{Content: []ToolContent{Text(s)}, IsError: true}
239}
240
241// Response tells zot how to react to a command invocation. Construct
242// one with Prompt(), Insert(), Display(), or Noop().
243type Panel struct {
244 ID string
245 Title string
246 Lines []string
247 Footer string
248}
249
250type Response struct {
251 Action string // "prompt", "insert", "display", "open_panel", "noop"
252 Prompt string
253 Insert string
254 Display string
255 OpenPanel *Panel
256 Error string
257}
258
259// Prompt returns a Response that submits text as a fresh user message
260// to the agent (running it through the model loop as if the user had
261// typed and pressed enter).
262func Prompt(text string) Response { return Response{Action: "prompt", Prompt: text} }
263
264// Insert returns a Response that drops text into the editor at the
265// cursor without submitting.
266func Insert(text string) Response { return Response{Action: "insert", Insert: text} }
267
268// Display returns a Response that adds a one-shot styled note to the
269// chat without invoking the model. Useful for showing a result without
270// burning tokens.
271func Display(text string) Response { return Response{Action: "display", Display: text} }
272
273// OpenPanel returns a Response that opens an interactive extension-owned
274// panel inside zot.
275func OpenPanel(id, title string, lines []string, footer string) Response {
276 return Response{Action: "open_panel", OpenPanel: &Panel{ID: id, Title: title, Lines: lines, Footer: footer}}
277}
278
279// Noop returns a Response that signals "I handled it, no UI change".
280// Use after pushing your own state or notifications.
281func Noop() Response { return Response{Action: "noop"} }
282
283// Errorf returns a Response that surfaces the error in the chat as a
284// red status line.
285func Errorf(format string, args ...any) Response {
286 return Response{Action: "noop", Error: fmt.Sprintf(format, args...)}
287}
288
289// Extension is one zot extension. Construct with New, register
290// commands, then call Run.
291type Extension struct {
292 name string
293 version string
294
295 in io.Reader
296 out io.Writer
297 stderr io.Writer
298 writeMu sync.Mutex
299
300 mu sync.Mutex
301 commands map[string]CommandHandler
302 descriptions []descTuple // ordered so register frames arrive in registration order
303 tools map[string]ToolHandler
304 toolDefs []toolDef // ordered so register frames arrive in registration order
305 eventHandlers map[string]EventHandler
306 eventNames []string // declared subscription order
307
308 interceptTool InterceptHandler
309 interceptToolRich ToolCallHandler
310 interceptOn bool
311 interceptStart BeforeAgentStartHandler
312 interceptTurn TurnStartHandler
313 interceptAssistant AssistantMessageHandler
314 panelKeys map[string]func(key, text string)
315 panelCloses map[string]func()
316 onHello func(HostInfo)
317
318 // Caps reported in the hello frame.
319 caps []string
320
321 // hostInfo is filled in once HelloAck arrives.
322 host HostInfo
323}
324
325type descTuple struct {
326 name, desc string
327}
328
329type toolDef struct {
330 name string
331 description string
332 schema json.RawMessage
333 deferred bool
334}
335
336// HostInfo is what the host (zot) tells us in HelloAck. Useful for
337// extensions that want to behave differently per provider.
338type HostInfo struct {
339 ProtocolVersion int
340 ZotVersion string
341 Provider string
342 Model string
343 CWD string
344 ExtensionDir string
345 DataDir string
346}
347
348// New constructs an Extension with the given identifier. name should
349// match the name field in extension.json.
350func New(name, version string) *Extension {
351 return &Extension{
352 name: name,
353 version: version,
354 in: os.Stdin,
355 out: os.Stdout,
356 stderr: os.Stderr,
357 commands: map[string]CommandHandler{},
358 tools: map[string]ToolHandler{},
359 eventHandlers: map[string]EventHandler{},
360 panelKeys: map[string]func(key, text string){},
361 panelCloses: map[string]func(){},
362 caps: []string{"commands", "tools", "events", "panels"},
363 }
364}
365
366// Host returns the HostInfo received during the hello handshake.
367// Returns the zero value if Run hasn't started yet.
368func (e *Extension) Host() HostInfo { return e.host }
369
370// OnHello registers a callback that runs after zot acknowledges the
371// hello frame and before commands/tools are announced. Use it when
372// registrations depend on host metadata such as CWD or ZotHome-derived
373// paths. The callback may call Command, Tool, On, and interceptor
374// registration methods.
375func (e *Extension) OnHello(fn func(HostInfo)) {
376 e.mu.Lock()
377 e.onHello = fn
378 e.mu.Unlock()
379}
380
381// Logf writes a line to the extension's stderr, which zot captures to
382// $ZOT_HOME/logs/ext-<name>.log. Use this for debug output: anything
383// you print to stdout would corrupt the JSON wire protocol.
384func (e *Extension) Logf(format string, args ...any) {
385 fmt.Fprintf(e.stderr, "["+e.name+"] "+format+"\n", args...)
386}
387
388// OnPanelKey registers callbacks for panel key + close events.
389func (e *Extension) OnPanelKey(panelID string, onKey func(key, text string), onClose func()) {
390 e.mu.Lock()
391 defer e.mu.Unlock()
392 if onKey != nil {
393 e.panelKeys[panelID] = onKey
394 }
395 if onClose != nil {
396 e.panelCloses[panelID] = onClose
397 }
398}
399
400// OpenPanel opens an interactive panel spontaneously from extension code
401// without requiring a slash command. Safe to call from a tool handler goroutine
402// or any background context. The panel receives panel_key events via OnPanelKey
403// and can be dismissed with ClosePanel, exactly as with command-response panels.
404func (e *Extension) OpenPanel(id, title string, lines []string, footer string) {
405 _ = e.send(extproto.OpenPanelFromExt{
406 Type: "open_panel",
407 Panel: extproto.PanelSpec{ID: id, Title: title, Lines: lines, Footer: footer},
408 })
409}
410
411// RenderPanel pushes a fresh frame for panelID.
412func (e *Extension) RenderPanel(panelID, title string, lines []string, footer string) {
413 _ = e.send(extproto.PanelRenderFromExt{Type: "panel_render", PanelID: panelID, Title: title, Lines: lines, Footer: footer})
414}
415
416// ClosePanel tells zot to close panelID.
417func (e *Extension) ClosePanel(panelID string) {
418 _ = e.send(extproto.PanelCloseFromExt{Type: "panel_close", PanelID: panelID})
419}
420
421// Command registers a slash-command handler. Call this BEFORE Run().
422// Once Run is going, when the user runs /name in zot, fn is invoked
423// with the remaining args.
424//
425// Naming conflicts with built-in commands (e.g. /help) are silently
426// shadowed by the built-in; check the user's "ext logs" if a command
427// you registered isn't taking effect.
428func (e *Extension) Command(name, description string, fn CommandHandler) {
429 e.mu.Lock()
430 e.commands[name] = fn
431 e.descriptions = append(e.descriptions, descTuple{name: name, desc: description})
432 e.mu.Unlock()
433}
434
435// Tool registers an LLM-callable tool. schema is a JSON Schema
436// object describing the tool's args (the same shape Anthropic /
437// OpenAI accept). Call this BEFORE Run(); zot folds extension tools
438// into the agent's registry once the extension's ready frame fires.
439//
440// Naming conflicts with built-in tools (read, write, edit, bash,
441// glob, skill) are silently shadowed by the built-in.
442func (e *Extension) Tool(name, description string, schema json.RawMessage, fn ToolHandler) {
443 e.registerTool(name, description, schema, false, fn)
444}
445
446// DeferredTool registers a tool whose definition stays hidden until another
447// tool result names it in ActivateTools.
448func (e *Extension) DeferredTool(name, description string, schema json.RawMessage, fn ToolHandler) {
449 e.registerTool(name, description, schema, true, fn)
450}
451
452func (e *Extension) registerTool(name, description string, schema json.RawMessage, deferred bool, fn ToolHandler) {
453 e.mu.Lock()
454 e.tools[name] = fn
455 e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema, deferred: deferred})
456 e.mu.Unlock()
457}
458
459// On subscribes to a lifecycle event. fn is called for each
460// notification; the same name can only have one handler (later
461// registrations replace earlier ones). Recognised names:
462// session_start, session_end, user_prompt_submit, turn_start, turn_end,
463// tool_call, tool_result, tool_confirmation_requested, permission_decision,
464// assistant_message, pre_compact, post_compact, subagent_start, subagent_stop.
465func (e *Extension) On(name string, fn EventHandler) {
466 e.mu.Lock()
467 if _, exists := e.eventHandlers[name]; !exists {
468 e.eventNames = append(e.eventNames, name)
469 }
470 e.eventHandlers[name] = fn
471 e.mu.Unlock()
472}
473
474// InterceptToolCall registers a guard that runs immediately before
475// each tool call. Returning (false, reason) refuses the call; the
476// model sees reason as the tool error text. Multiple extensions can
477// install interceptors; if any one refuses, the call is blocked.
478//
479// Use this to build permission gates: refuse `bash` calls containing
480// `rm -rf`, ask the user for confirmation on dangerous patterns,
481// audit-log every call, etc.
482//
483// For the richer form (arg rewrites, structured decisions) use
484// InterceptToolCallX.
485func (e *Extension) InterceptToolCall(fn InterceptHandler) {
486 e.mu.Lock()
487 e.interceptTool = fn
488 e.interceptOn = true
489 e.mu.Unlock()
490}
491
492// InterceptToolCallX is the richer variant. Return a ToolCallDecision
493// to block, allow, or rewrite args mid-flight. If both
494// InterceptToolCall and InterceptToolCallX are set, the X form wins.
495func (e *Extension) InterceptToolCallX(fn ToolCallHandler) {
496 e.mu.Lock()
497 e.interceptToolRich = fn
498 e.interceptOn = true
499 e.mu.Unlock()
500}
501
502// InterceptBeforeAgentStart registers a session-scoped prompt replacement hook.
503func (e *Extension) InterceptBeforeAgentStart(fn BeforeAgentStartHandler) {
504 e.mu.Lock()
505 e.interceptStart = fn
506 e.mu.Unlock()
507}
508
509// InterceptTurnStart registers a guard that runs before every turn's
510// model call. Return Block=true with Reason to abort the turn.
511// Useful for deny-by-default gates and usage quotas.
512func (e *Extension) InterceptTurnStart(fn TurnStartHandler) {
513 e.mu.Lock()
514 e.interceptTurn = fn
515 e.mu.Unlock()
516}
517
518// InterceptAssistantMessage registers a guard that runs after the
519// model's final text is assembled but before it's shown. Return
520// Block=true to suppress entirely, or ReplaceText to rewrite what
521// the user sees. The model's original output stays in the transcript.
522func (e *Extension) InterceptAssistantMessage(fn AssistantMessageHandler) {
523 e.mu.Lock()
524 e.interceptAssistant = fn
525 e.mu.Unlock()
526}
527
528// Notify pushes an info-level status note into zot's chat without
529// requiring a slash command from the user.
530func (e *Extension) Notify(level, message string) {
531 _ = e.send(extproto.NotifyFromExt{
532 Type: "notify",
533 Level: level,
534 Message: message,
535 })
536}
537
538// Run starts the protocol loop. Blocks until stdin closes (zot has
539// shut us down). Returns the first fatal error, or nil on clean exit.
540func (e *Extension) Run() error {
541 scanner := bufio.NewScanner(e.in)
542 scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
543
544 // Send hello first. zot answers with hello_ack immediately; read it
545 // synchronously so extensions can use HostInfo while registering
546 // commands/tools before the ready frame.
547 if err := e.send(extproto.HelloFromExt{
548 Type: "hello",
549 Name: e.name,
550 Version: e.version,
551 Capabilities: e.caps,
552 }); err != nil {
553 return err
554 }
555 if !scanner.Scan() {
556 return scanner.Err()
557 }
558 var ack extproto.HelloAckFromHost
559 if err := json.Unmarshal(scanner.Bytes(), &ack); err != nil {
560 return fmt.Errorf("parse hello_ack: %w", err)
561 }
562 if ack.Type != "hello_ack" {
563 return fmt.Errorf("first host frame must be hello_ack (got %q)", ack.Type)
564 }
565 e.host = HostInfo{
566 ProtocolVersion: ack.ProtocolVersion,
567 ZotVersion: ack.ZotVersion,
568 Provider: ack.Provider,
569 Model: ack.Model,
570 CWD: ack.CWD,
571 ExtensionDir: ack.ExtensionDir,
572 DataDir: ack.DataDir,
573 }
574 e.mu.Lock()
575 onHello := e.onHello
576 e.mu.Unlock()
577 if onHello != nil {
578 onHello(e.host)
579 }
580
581 // Re-announce all commands (covers Command calls made before Run and
582 // from OnHello, and is also fine for those made after via Command()'s
583 // direct send).
584 e.mu.Lock()
585 descs := append([]descTuple(nil), e.descriptions...)
586 toolDefs := append([]toolDef(nil), e.toolDefs...)
587 eventNames := append([]string(nil), e.eventNames...)
588 interceptStart := e.interceptStart != nil
589 interceptTool := e.interceptOn
590 interceptTurn := e.interceptTurn != nil
591 interceptAsst := e.interceptAssistant != nil
592 e.mu.Unlock()
593 for _, d := range descs {
594 _ = e.send(extproto.RegisterCommandFromExt{
595 Type: "register_command",
596 Name: d.name,
597 Description: d.desc,
598 })
599 }
600 for _, td := range toolDefs {
601 _ = e.send(extproto.RegisterToolFromExt{
602 Type: "register_tool",
603 Name: td.name,
604 Description: td.description,
605 Schema: td.schema,
606 Deferred: td.deferred,
607 })
608 }
609 var intercepts []string
610 if interceptStart {
611 intercepts = append(intercepts, "before_agent_start")
612 }
613 if interceptTool {
614 intercepts = append(intercepts, "tool_call")
615 }
616 if interceptTurn {
617 intercepts = append(intercepts, "turn_start")
618 }
619 if interceptAsst {
620 intercepts = append(intercepts, "assistant_message")
621 }
622 if len(eventNames) > 0 || len(intercepts) > 0 {
623 _ = e.send(extproto.SubscribeFromExt{
624 Type: "subscribe",
625 Events: eventNames,
626 Intercept: intercepts,
627 })
628 }
629 // Sentinel: tells the host that all initial registrations have
630 // been flushed and the agent registry can be built. Never block
631 // on this; if the host doesn't act on it, registrations still
632 // land in time for the typical use case.
633 _ = e.send(extproto.ReadyFromExt{Type: "ready"})
634
635 for scanner.Scan() {
636 line := scanner.Bytes()
637 var frame extproto.Frame
638 if err := json.Unmarshal(line, &frame); err != nil {
639 e.Logf("malformed frame from host: %v", err)
640 continue
641 }
642 switch frame.Type {
643 case "hello_ack":
644 // Already handled before registrations; ignore duplicates.
645 case "command_invoked":
646 var ci extproto.CommandInvokedFromHost
647 if err := json.Unmarshal(line, &ci); err != nil {
648 continue
649 }
650 e.mu.Lock()
651 fn := e.commands[ci.Name]
652 e.mu.Unlock()
653 if fn == nil {
654 e.respond(ci.ID, Errorf("no handler for /%s", ci.Name))
655 continue
656 }
657 // Run the handler on its own goroutine so a slow handler
658 // doesn't block subsequent commands. Order isn't promised.
659 go func(id string, fn CommandHandler, args string) {
660 defer func() {
661 if r := recover(); r != nil {
662 e.respond(id, Errorf("panic: %v", r))
663 }
664 }()
665 resp := fn(args)
666 e.respond(id, resp)
667 }(ci.ID, fn, ci.Args)
668 case "tool_call":
669 var tc extproto.ToolCallFromHost
670 if err := json.Unmarshal(line, &tc); err != nil {
671 continue
672 }
673 e.mu.Lock()
674 fn := e.tools[tc.Name]
675 e.mu.Unlock()
676 if fn == nil {
677 e.respondTool(tc.ID, TextErrorResult(fmt.Sprintf("no handler for tool %q", tc.Name)))
678 continue
679 }
680 go func(id string, fn ToolHandler, args json.RawMessage) {
681 defer func() {
682 if r := recover(); r != nil {
683 e.respondTool(id, TextErrorResult(fmt.Sprintf("panic: %v", r)))
684 }
685 }()
686 res := fn(args)
687 e.respondTool(id, res)
688 }(tc.ID, fn, tc.Args)
689 case "event":
690 var ef extproto.EventFromHost
691 if err := json.Unmarshal(line, &ef); err != nil {
692 continue
693 }
694 e.mu.Lock()
695 handler := e.eventHandlers[ef.Event]
696 e.mu.Unlock()
697 if handler != nil {
698 func() {
699 defer func() {
700 if r := recover(); r != nil {
701 e.Logf("event %s handler panicked: %v", ef.Event, r)
702 }
703 }()
704 handler(Event{
705 Name: ef.Event, Step: ef.Step, Stop: ef.Stop,
706 SessionID: ef.SessionID, AgentRunID: ef.AgentRunID, Queued: ef.Queued, ImageCount: ef.ImageCount, CWD: ef.CWD, Sequence: ef.Sequence,
707 Status: ef.Status, Reason: ef.Reason, Source: ef.Source, Decision: ef.Decision, Stage: ef.Stage,
708 Result: ef.Result, Executed: ef.Executed, CompactionID: ef.CompactionID,
709 MessageCount: ef.MessageCount, TokenEstimate: ef.TokenEstimate,
710 AgentID: ef.AgentID, AgentName: ef.Name,
711 Error: ef.Error, ToolID: ef.ToolID, ToolName: ef.ToolName,
712 ToolArgs: ef.ToolArgs, ToolArgsRaw: ef.ToolArgsRaw, ToolPreview: ef.ToolPreview, Text: ef.Text,
713 })
714 }()
715 }
716 case "event_intercept":
717 var ei extproto.EventInterceptFromHost
718 if err := json.Unmarshal(line, &ei); err != nil {
719 continue
720 }
721 go e.dispatchIntercept(ei)
722 case "panel_key":
723 var pk extproto.PanelKeyFromHost
724 if err := json.Unmarshal(line, &pk); err != nil {
725 continue
726 }
727 e.mu.Lock()
728 h := e.panelKeys[pk.PanelID]
729 e.mu.Unlock()
730 if h != nil {
731 go h(pk.Key, pk.Text)
732 }
733 case "panel_close":
734 var pc extproto.PanelCloseFromHost
735 if err := json.Unmarshal(line, &pc); err != nil {
736 continue
737 }
738 e.mu.Lock()
739 h := e.panelCloses[pc.PanelID]
740 e.mu.Unlock()
741 if h != nil {
742 go h()
743 }
744 case "shutdown":
745 _ = e.send(extproto.ShutdownAckFromExt{Type: "shutdown_ack"})
746 return nil
747 default:
748 e.Logf("unknown frame type %q", frame.Type)
749 }
750 }
751 return scanner.Err()
752}
753
754// respond serialises a CommandResponseFromExt for the given id.
755func (e *Extension) respond(id string, r Response) {
756 if r.Action == "" {
757 r.Action = "noop"
758 }
759 var panel *extproto.PanelSpec
760 if r.OpenPanel != nil {
761 panel = &extproto.PanelSpec{ID: r.OpenPanel.ID, Title: r.OpenPanel.Title, Lines: r.OpenPanel.Lines, Footer: r.OpenPanel.Footer}
762 }
763 _ = e.send(extproto.CommandResponseFromExt{
764 Type: "command_response",
765 ID: id,
766 Action: r.Action,
767 Prompt: r.Prompt,
768 Insert: r.Insert,
769 Display: r.Display,
770 OpenPanel: panel,
771 Error: r.Error,
772 })
773}
774
775// respondTool serialises a ToolResultFromExt for the given id.
776func (e *Extension) respondTool(id string, r ToolResult) {
777 blocks := make([]extproto.ContentBlock, 0, len(r.Content))
778 for _, c := range r.Content {
779 blocks = append(blocks, extproto.ContentBlock{
780 Type: c.Type,
781 Text: c.Text,
782 MimeType: c.MimeType,
783 Data: c.Data,
784 })
785 }
786 _ = e.send(extproto.ToolResultFromExt{
787 Type: "tool_result",
788 ID: id,
789 Content: blocks,
790 IsError: r.IsError,
791 ActivateTools: r.ActivateTools,
792 })
793}
794
795// send marshals v + LF and writes it under a mutex (so concurrent
796// goroutines don't interleave bytes on stdout).
797func (e *Extension) send(v any) error {
798 b, err := extproto.Encode(v)
799 if err != nil {
800 return err
801 }
802 e.writeMu.Lock()
803 defer e.writeMu.Unlock()
804 _, err = e.out.Write(b)
805 return err
806}
807
808// dispatchIntercept runs the per-event handler (tool_call / turn_start
809// / assistant_message) on its own goroutine, catches panics, and
810// always emits exactly one event_intercept_response. Called from the
811// Run loop.
812func (e *Extension) dispatchIntercept(ei extproto.EventInterceptFromHost) {
813 defer func() {
814 if r := recover(); r != nil {
815 _ = e.send(extproto.EventInterceptResponseFromExt{
816 Type: "event_intercept_response",
817 ID: ei.ID,
818 Block: true,
819 Reason: fmt.Sprintf("intercept panic: %v", r),
820 })
821 }
822 }()
823
824 resp := extproto.EventInterceptResponseFromExt{
825 Type: "event_intercept_response",
826 ID: ei.ID,
827 }
828
829 switch ei.Event {
830 case "before_agent_start":
831 e.mu.Lock()
832 fn := e.interceptStart
833 e.mu.Unlock()
834 if fn != nil && ei.SystemPrompt != nil {
835 if replacement := fn(BeforeAgentStartEvent{
836 SystemPrompt: *ei.SystemPrompt, SessionID: ei.SessionID, AgentRunID: ei.AgentRunID,
837 CWD: ei.CWD, Provider: ei.Provider, Model: ei.Model,
838 }); replacement != nil {
839 resp.SystemPrompt, _ = json.Marshal(*replacement)
840 }
841 }
842 case "tool_call":
843 e.mu.Lock()
844 rich := e.interceptToolRich
845 plain := e.interceptTool
846 e.mu.Unlock()
847 if rich != nil {
848 d := rich(ei.ToolName, ei.ToolArgs)
849 resp.Block = d.Block
850 resp.Reason = d.Reason
851 if !d.Block && len(d.ModifiedArgs) > 0 && json.Valid(d.ModifiedArgs) {
852 resp.ModifiedArgs = d.ModifiedArgs
853 }
854 } else if plain != nil {
855 allow, reason := plain(ei.ToolName, ei.ToolArgs)
856 resp.Block = !allow
857 resp.Reason = reason
858 }
859 case "turn_start":
860 e.mu.Lock()
861 fn := e.interceptTurn
862 e.mu.Unlock()
863 if fn != nil {
864 d := fn(ei.Step)
865 resp.Block = d.Block
866 resp.Reason = d.Reason
867 }
868 case "assistant_message":
869 e.mu.Lock()
870 fn := e.interceptAssistant
871 e.mu.Unlock()
872 if fn != nil {
873 d := fn(ei.Text)
874 resp.Block = d.Block
875 resp.Reason = d.Reason
876 if !d.Block && d.ReplaceText != "" && d.ReplaceText != ei.Text {
877 resp.ReplaceText = d.ReplaceText
878 }
879 }
880 }
881
882 _ = e.send(resp)
883}