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