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