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