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