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", "tool_prompt", "insert", "display", "open_panel", "noop"
259 Prompt string
260 ToolName string
261 ToolArgs json.RawMessage
262 Insert string
263 Display string
264 OpenPanel *Panel
265 Error string
266}
267
268// Prompt returns a Response that submits text as a fresh user message
269// to the agent (running it through the model loop as if the user had
270// typed and pressed enter).
271func Prompt(text string) Response { return Response{Action: "prompt", Prompt: text} }
272
273// ToolPrompt asks the host to run a tool and show its result to the model
274// after text has been submitted as a user prompt. Unlike CallTool, execution
275// takes place after the command handler returns and is not returned to it.
276func ToolPrompt(name string, args any, text string) Response {
277 if name == "" || text == "" {
278 return Errorf("tool_prompt requires a name and prompt")
279 }
280 data, err := json.Marshal(args)
281 if err != nil || len(data) == 0 || data[0] != '{' {
282 return Errorf("tool_prompt args must be a JSON object")
283 }
284 return Response{Action: "tool_prompt", Prompt: text, ToolName: name, ToolArgs: data}
285}
286
287// Insert returns a Response that drops text into the editor at the
288// cursor without submitting.
289func Insert(text string) Response { return Response{Action: "insert", Insert: text} }
290
291// Display returns a Response that adds a one-shot styled note to the
292// chat without invoking the model. Useful for showing a result without
293// burning tokens.
294func Display(text string) Response { return Response{Action: "display", Display: text} }
295
296// OpenPanel returns a Response that opens an interactive extension-owned
297// panel inside zot.
298func OpenPanel(id, title string, lines []string, footer string) Response {
299 return Response{Action: "open_panel", OpenPanel: &Panel{ID: id, Title: title, Lines: lines, Footer: footer}}
300}
301
302// Noop returns a Response that signals "I handled it, no UI change".
303// Use after pushing your own state or notifications.
304func Noop() Response { return Response{Action: "noop"} }
305
306// Errorf returns a Response that surfaces the error in the chat as a
307// red status line.
308func Errorf(format string, args ...any) Response {
309 return Response{Action: "noop", Error: fmt.Sprintf(format, args...)}
310}
311
312// Extension is one zot extension. Construct with New, register
313// commands, then call Run.
314type Extension struct {
315 name string
316 version string
317
318 in io.Reader
319 out io.Writer
320 stderr io.Writer
321 writeMu sync.Mutex
322
323 mu sync.Mutex
324 commands map[string]CommandHandler
325 descriptions []descTuple // ordered so register frames arrive in registration order
326 tools map[string]InteractiveToolHandler
327 toolCancels map[string]context.CancelFunc
328 pendingCalls map[string]chan extproto.ToolResultFromHost
329 callSeq atomic.Uint64
330 callToolReady atomic.Bool
331 done chan struct{}
332 toolDefs []toolDef // ordered so register frames arrive in registration order
333 eventHandlers map[string]EventHandler
334 eventNames []string // declared subscription order
335
336 interceptTool InterceptHandler
337 interceptToolRich ToolCallHandler
338 interceptOn bool
339 interceptStart BeforeAgentStartHandler
340 interceptTurn TurnStartHandler
341 interceptAssistant AssistantMessageHandler
342 panelKeys map[string]func(key, text string)
343 panelCloses map[string]func()
344 onHello func(HostInfo)
345
346 // Caps reported in the hello frame.
347 caps []string
348
349 // hostInfo is filled in once HelloAck arrives.
350 host HostInfo
351}
352
353type descTuple struct {
354 name, desc string
355}
356
357type toolDef struct {
358 name string
359 description string
360 schema json.RawMessage
361 deferred bool
362 interactive bool
363}
364
365// HostInfo is what the host (zot) tells us in HelloAck. Useful for
366// extensions that want to behave differently per provider.
367type HostInfo struct {
368 ProtocolVersion int
369 ZotVersion string
370 Provider string
371 Model string
372 CWD string
373 ExtensionDir string
374 DataDir string
375}
376
377// New constructs an Extension with the given identifier. name should
378// match the name field in extension.json.
379func New(name, version string) *Extension {
380 return &Extension{
381 name: name,
382 version: version,
383 in: os.Stdin,
384 out: os.Stdout,
385 stderr: os.Stderr,
386 commands: map[string]CommandHandler{},
387 tools: map[string]InteractiveToolHandler{},
388 toolCancels: map[string]context.CancelFunc{},
389 pendingCalls: map[string]chan extproto.ToolResultFromHost{},
390 done: make(chan struct{}),
391 eventHandlers: map[string]EventHandler{},
392 panelKeys: map[string]func(key, text string){},
393 panelCloses: map[string]func(){},
394 caps: []string{"commands", "tools", "events", "panels", "tool_cancel", "call_tool"},
395 }
396}
397
398// Host returns the HostInfo received during the hello handshake.
399// Returns the zero value if Run hasn't started yet.
400func (e *Extension) Host() HostInfo { return e.host }
401
402// OnHello registers a callback that runs after zot acknowledges the
403// hello frame and before commands/tools are announced. Use it when
404// registrations depend on host metadata such as CWD or ZotHome-derived
405// paths. The callback may call Command, Tool, On, and interceptor
406// registration methods.
407func (e *Extension) OnHello(fn func(HostInfo)) {
408 e.mu.Lock()
409 e.onHello = fn
410 e.mu.Unlock()
411}
412
413// Logf writes a line to the extension's stderr, which zot captures to
414// $ZOT_HOME/logs/ext-<name>.log. Use this for debug output: anything
415// you print to stdout would corrupt the JSON wire protocol.
416func (e *Extension) Logf(format string, args ...any) {
417 fmt.Fprintf(e.stderr, "["+e.name+"] "+format+"\n", args...)
418}
419
420// OnPanelKey registers callbacks for panel key + close events.
421func (e *Extension) OnPanelKey(panelID string, onKey func(key, text string), onClose func()) {
422 e.mu.Lock()
423 defer e.mu.Unlock()
424 if onKey != nil {
425 e.panelKeys[panelID] = onKey
426 }
427 if onClose != nil {
428 e.panelCloses[panelID] = onClose
429 }
430}
431
432// OpenPanel opens an interactive panel spontaneously from extension code
433// without requiring a slash command. Safe to call from a tool handler goroutine
434// or any background context. The panel receives panel_key events via OnPanelKey
435// and can be dismissed with ClosePanel, exactly as with command-response panels.
436func (e *Extension) OpenPanel(id, title string, lines []string, footer string) {
437 _ = e.send(extproto.OpenPanelFromExt{
438 Type: "open_panel",
439 Panel: extproto.PanelSpec{ID: id, Title: title, Lines: lines, Footer: footer},
440 })
441}
442
443// RenderPanel pushes a fresh frame for panelID.
444func (e *Extension) RenderPanel(panelID, title string, lines []string, footer string) {
445 _ = e.send(extproto.PanelRenderFromExt{Type: "panel_render", PanelID: panelID, Title: title, Lines: lines, Footer: footer})
446}
447
448// ClosePanel tells zot to close panelID.
449func (e *Extension) ClosePanel(panelID string) {
450 _ = e.send(extproto.PanelCloseFromExt{Type: "panel_close", PanelID: panelID})
451}
452
453// Command registers a slash-command handler. Call this BEFORE Run().
454// Once Run is going, when the user runs /name in zot, fn is invoked
455// with the remaining args.
456//
457// Naming conflicts with built-in commands (e.g. /help) are silently
458// shadowed by the built-in; check the user's "ext logs" if a command
459// you registered isn't taking effect.
460func (e *Extension) Command(name, description string, fn CommandHandler) {
461 e.mu.Lock()
462 e.commands[name] = fn
463 e.descriptions = append(e.descriptions, descTuple{name: name, desc: description})
464 e.mu.Unlock()
465}
466
467// Tool registers an LLM-callable tool. schema is a JSON Schema
468// object describing the tool's args (the same shape Anthropic /
469// OpenAI accept). Call this BEFORE Run(); zot folds extension tools
470// into the agent's registry once the extension's ready frame fires.
471//
472// Naming conflicts with built-in tools (read, write, edit, bash,
473// glob, skill) are silently shadowed by the built-in.
474func (e *Extension) Tool(name, description string, schema json.RawMessage, fn ToolHandler) {
475 e.registerTool(name, description, schema, false, fn)
476}
477
478// DeferredTool registers a tool whose definition stays hidden until another
479// tool result names it in ActivateTools.
480func (e *Extension) DeferredTool(name, description string, schema json.RawMessage, fn ToolHandler) {
481 e.registerTool(name, description, schema, true, fn)
482}
483
484// InteractiveTool registers a tool that intentionally waits for user input.
485// The interactive host imposes no reply deadline. Headless hosts reject it.
486// The handler must honor ctx cancellation and close any panel it owns.
487func (e *Extension) InteractiveTool(name, description string, schema json.RawMessage, fn InteractiveToolHandler) {
488 e.mu.Lock()
489 defer e.mu.Unlock()
490 e.tools[name] = fn
491 e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema, interactive: true})
492}
493
494func (e *Extension) registerTool(name, description string, schema json.RawMessage, deferred bool, fn ToolHandler) {
495 e.mu.Lock()
496 e.tools[name] = func(_ context.Context, args json.RawMessage) ToolResult { return fn(args) }
497 e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema, deferred: deferred})
498 e.mu.Unlock()
499}
500
501// On subscribes to a lifecycle event. fn is called for each
502// notification; the same name can only have one handler (later
503// registrations replace earlier ones). Recognised names:
504// session_start, session_end, user_prompt_submit, turn_start, turn_end,
505// tool_call, tool_result, tool_confirmation_requested, permission_decision,
506// assistant_message, pre_compact, post_compact, subagent_start, subagent_stop.
507func (e *Extension) On(name string, fn EventHandler) {
508 e.mu.Lock()
509 if _, exists := e.eventHandlers[name]; !exists {
510 e.eventNames = append(e.eventNames, name)
511 }
512 e.eventHandlers[name] = fn
513 e.mu.Unlock()
514}
515
516// InterceptToolCall registers a guard that runs immediately before
517// each tool call. Returning (false, reason) refuses the call; the
518// model sees reason as the tool error text. Multiple extensions can
519// install interceptors; if any one refuses, the call is blocked.
520//
521// Use this to build permission gates: refuse `bash` calls containing
522// `rm -rf`, ask the user for confirmation on dangerous patterns,
523// audit-log every call, etc.
524//
525// For the richer form (arg rewrites, structured decisions) use
526// InterceptToolCallX.
527func (e *Extension) InterceptToolCall(fn InterceptHandler) {
528 e.mu.Lock()
529 e.interceptTool = fn
530 e.interceptOn = true
531 e.mu.Unlock()
532}
533
534// InterceptToolCallX is the richer variant. Return a ToolCallDecision
535// to block, allow, or rewrite args mid-flight. If both
536// InterceptToolCall and InterceptToolCallX are set, the X form wins.
537func (e *Extension) InterceptToolCallX(fn ToolCallHandler) {
538 e.mu.Lock()
539 e.interceptToolRich = fn
540 e.interceptOn = true
541 e.mu.Unlock()
542}
543
544// InterceptBeforeAgentStart registers a session-scoped prompt replacement hook.
545func (e *Extension) InterceptBeforeAgentStart(fn BeforeAgentStartHandler) {
546 e.mu.Lock()
547 e.interceptStart = fn
548 e.mu.Unlock()
549}
550
551// InterceptTurnStart registers a guard that runs before every turn's
552// model call. Return Block=true with Reason to abort the turn.
553// Useful for deny-by-default gates and usage quotas.
554func (e *Extension) InterceptTurnStart(fn TurnStartHandler) {
555 e.mu.Lock()
556 e.interceptTurn = fn
557 e.mu.Unlock()
558}
559
560// InterceptAssistantMessage registers a guard that runs after the
561// model's final text is assembled but before it's shown. Return
562// Block=true to suppress entirely, or ReplaceText to rewrite what
563// the user sees. The model's original output stays in the transcript.
564func (e *Extension) InterceptAssistantMessage(fn AssistantMessageHandler) {
565 e.mu.Lock()
566 e.interceptAssistant = fn
567 e.mu.Unlock()
568}
569
570// Notify pushes an info-level status note into zot's chat without
571// requiring a slash command from the user.
572func (e *Extension) Notify(level, message string) {
573 _ = e.send(extproto.NotifyFromExt{
574 Type: "notify",
575 Level: level,
576 Message: message,
577 })
578}
579
580// Run starts the protocol loop. Blocks until stdin closes (zot has
581// shut us down). Returns the first fatal error, or nil on clean exit.
582func (e *Extension) Run() error {
583 defer close(e.done)
584 defer func() {
585 e.mu.Lock()
586 defer e.mu.Unlock()
587 for id, cancel := range e.toolCancels {
588 cancel()
589 delete(e.toolCancels, id)
590 }
591 }()
592 scanner := bufio.NewScanner(e.in)
593 scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
594
595 // Send hello first. zot answers with hello_ack immediately; read it
596 // synchronously so extensions can use HostInfo while registering
597 // commands/tools before the ready frame.
598 if err := e.send(extproto.HelloFromExt{
599 Type: "hello",
600 Name: e.name,
601 Version: e.version,
602 Capabilities: e.caps,
603 }); err != nil {
604 return err
605 }
606 if !scanner.Scan() {
607 return scanner.Err()
608 }
609 var ack extproto.HelloAckFromHost
610 if err := json.Unmarshal(scanner.Bytes(), &ack); err != nil {
611 return fmt.Errorf("parse hello_ack: %w", err)
612 }
613 if ack.Type != "hello_ack" {
614 return fmt.Errorf("first host frame must be hello_ack (got %q)", ack.Type)
615 }
616 e.callToolReady.Store(slices.Contains(ack.Capabilities, "call_tool"))
617 e.host = HostInfo{
618 ProtocolVersion: ack.ProtocolVersion,
619 ZotVersion: ack.ZotVersion,
620 Provider: ack.Provider,
621 Model: ack.Model,
622 CWD: ack.CWD,
623 ExtensionDir: ack.ExtensionDir,
624 DataDir: ack.DataDir,
625 }
626 e.mu.Lock()
627 onHello := e.onHello
628 e.mu.Unlock()
629 if onHello != nil {
630 onHello(e.host)
631 }
632
633 // Re-announce all commands (covers Command calls made before Run and
634 // from OnHello, and is also fine for those made after via Command()'s
635 // direct send).
636 e.mu.Lock()
637 descs := append([]descTuple(nil), e.descriptions...)
638 toolDefs := append([]toolDef(nil), e.toolDefs...)
639 eventNames := append([]string(nil), e.eventNames...)
640 interceptStart := e.interceptStart != nil
641 interceptTool := e.interceptOn
642 interceptTurn := e.interceptTurn != nil
643 interceptAsst := e.interceptAssistant != nil
644 e.mu.Unlock()
645 for _, d := range descs {
646 _ = e.send(extproto.RegisterCommandFromExt{
647 Type: "register_command",
648 Name: d.name,
649 Description: d.desc,
650 })
651 }
652 for _, td := range toolDefs {
653 _ = e.send(extproto.RegisterToolFromExt{
654 Type: "register_tool",
655 Name: td.name,
656 Description: td.description,
657 Schema: td.schema,
658 Deferred: td.deferred,
659 Interactive: td.interactive,
660 })
661 }
662 var intercepts []string
663 if interceptStart {
664 intercepts = append(intercepts, "before_agent_start")
665 }
666 if interceptTool {
667 intercepts = append(intercepts, "tool_call")
668 }
669 if interceptTurn {
670 intercepts = append(intercepts, "turn_start")
671 }
672 if interceptAsst {
673 intercepts = append(intercepts, "assistant_message")
674 }
675 if len(eventNames) > 0 || len(intercepts) > 0 {
676 _ = e.send(extproto.SubscribeFromExt{
677 Type: "subscribe",
678 Events: eventNames,
679 Intercept: intercepts,
680 })
681 }
682 // Sentinel: tells the host that all initial registrations have
683 // been flushed and the agent registry can be built. Never block
684 // on this; if the host doesn't act on it, registrations still
685 // land in time for the typical use case.
686 _ = e.send(extproto.ReadyFromExt{Type: "ready"})
687
688 for scanner.Scan() {
689 line := scanner.Bytes()
690 var frame extproto.Frame
691 if err := json.Unmarshal(line, &frame); err != nil {
692 e.Logf("malformed frame from host: %v", err)
693 continue
694 }
695 switch frame.Type {
696 case "hello_ack":
697 // Already handled before registrations; ignore duplicates.
698 case "command_invoked":
699 var ci extproto.CommandInvokedFromHost
700 if err := json.Unmarshal(line, &ci); err != nil {
701 continue
702 }
703 e.mu.Lock()
704 fn := e.commands[ci.Name]
705 e.mu.Unlock()
706 if fn == nil {
707 e.respond(ci.ID, Errorf("no handler for /%s", ci.Name))
708 continue
709 }
710 // Run the handler on its own goroutine so a slow handler
711 // doesn't block subsequent commands. Order isn't promised.
712 go func(id string, fn CommandHandler, args string) {
713 defer func() {
714 if r := recover(); r != nil {
715 e.respond(id, Errorf("panic: %v", r))
716 }
717 }()
718 resp := fn(args)
719 e.respond(id, resp)
720 }(ci.ID, fn, ci.Args)
721 case "tool_call":
722 var tc extproto.ToolCallFromHost
723 if err := json.Unmarshal(line, &tc); err != nil {
724 continue
725 }
726 e.mu.Lock()
727 fn := e.tools[tc.Name]
728 e.mu.Unlock()
729 if fn == nil {
730 e.respondTool(tc.ID, TextErrorResult(fmt.Sprintf("no handler for tool %q", tc.Name)))
731 continue
732 }
733 ctx, cancel := context.WithCancel(context.WithValue(context.Background(), parentCallKey{}, tc.ID))
734 e.mu.Lock()
735 e.toolCancels[tc.ID] = cancel
736 e.mu.Unlock()
737 go func(id string, fn InteractiveToolHandler, args json.RawMessage) {
738 defer func() {
739 cancel()
740 e.mu.Lock()
741 delete(e.toolCancels, id)
742 e.mu.Unlock()
743 }()
744 defer func() {
745 if r := recover(); r != nil && ctx.Err() == nil {
746 e.respondTool(id, TextErrorResult(fmt.Sprintf("panic: %v", r)))
747 }
748 }()
749 res := fn(ctx, args)
750 if ctx.Err() == nil {
751 e.respondTool(id, res)
752 }
753 }(tc.ID, fn, tc.Args)
754 case "tool_result":
755 var tr extproto.ToolResultFromHost
756 if json.Unmarshal(line, &tr) == nil {
757 e.mu.Lock()
758 ch := e.pendingCalls[tr.ID]
759 e.mu.Unlock()
760 if ch != nil {
761 select {
762 case ch <- tr:
763 default:
764 }
765 }
766 }
767 case "tool_cancel":
768 var tc extproto.ToolCancelFromHost
769 if json.Unmarshal(line, &tc) != nil {
770 continue
771 }
772 e.mu.Lock()
773 if cancel := e.toolCancels[tc.ID]; cancel != nil {
774 cancel()
775 }
776 e.mu.Unlock()
777 case "event":
778 var ef extproto.EventFromHost
779 if err := json.Unmarshal(line, &ef); err != nil {
780 continue
781 }
782 e.mu.Lock()
783 handler := e.eventHandlers[ef.Event]
784 e.mu.Unlock()
785 if handler != nil {
786 func() {
787 defer func() {
788 if r := recover(); r != nil {
789 e.Logf("event %s handler panicked: %v", ef.Event, r)
790 }
791 }()
792 handler(Event{
793 Name: ef.Event, Step: ef.Step, Stop: ef.Stop,
794 SessionID: ef.SessionID, AgentRunID: ef.AgentRunID, Queued: ef.Queued, ImageCount: ef.ImageCount, CWD: ef.CWD, Sequence: ef.Sequence,
795 Status: ef.Status, Reason: ef.Reason, Source: ef.Source, Decision: ef.Decision, Stage: ef.Stage,
796 Result: ef.Result, Executed: ef.Executed, CompactionID: ef.CompactionID,
797 MessageCount: ef.MessageCount, TokenEstimate: ef.TokenEstimate,
798 AgentID: ef.AgentID, AgentName: ef.Name,
799 Error: ef.Error, ToolID: ef.ToolID, ToolName: ef.ToolName,
800 ToolArgs: ef.ToolArgs, ToolArgsRaw: ef.ToolArgsRaw, ToolPreview: ef.ToolPreview, Text: ef.Text,
801 })
802 }()
803 }
804 case "event_intercept":
805 var ei extproto.EventInterceptFromHost
806 if err := json.Unmarshal(line, &ei); err != nil {
807 continue
808 }
809 go e.dispatchIntercept(ei)
810 case "panel_key":
811 var pk extproto.PanelKeyFromHost
812 if err := json.Unmarshal(line, &pk); err != nil {
813 continue
814 }
815 e.mu.Lock()
816 h := e.panelKeys[pk.PanelID]
817 e.mu.Unlock()
818 if h != nil {
819 go h(pk.Key, pk.Text)
820 }
821 case "panel_close":
822 var pc extproto.PanelCloseFromHost
823 if err := json.Unmarshal(line, &pc); err != nil {
824 continue
825 }
826 e.mu.Lock()
827 h := e.panelCloses[pc.PanelID]
828 e.mu.Unlock()
829 if h != nil {
830 go h()
831 }
832 case "shutdown":
833 _ = e.send(extproto.ShutdownAckFromExt{Type: "shutdown_ack"})
834 return nil
835 default:
836 e.Logf("unknown frame type %q", frame.Type)
837 }
838 }
839 return scanner.Err()
840}
841
842// respond serialises a CommandResponseFromExt for the given id.
843func (e *Extension) respond(id string, r Response) {
844 if r.Action == "" {
845 r.Action = "noop"
846 }
847 var panel *extproto.PanelSpec
848 if r.OpenPanel != nil {
849 panel = &extproto.PanelSpec{ID: r.OpenPanel.ID, Title: r.OpenPanel.Title, Lines: r.OpenPanel.Lines, Footer: r.OpenPanel.Footer}
850 }
851 _ = e.send(extproto.CommandResponseFromExt{
852 Type: "command_response",
853 ID: id,
854 Action: r.Action,
855 Prompt: r.Prompt,
856 ToolName: r.ToolName,
857 ToolArgs: r.ToolArgs,
858 Insert: r.Insert,
859 Display: r.Display,
860 OpenPanel: panel,
861 Error: r.Error,
862 })
863}
864
865// respondTool serialises a ToolResultFromExt for the given id.
866func (e *Extension) respondTool(id string, r ToolResult) {
867 blocks := make([]extproto.ContentBlock, 0, len(r.Content))
868 for _, c := range r.Content {
869 blocks = append(blocks, extproto.ContentBlock{
870 Type: c.Type,
871 Text: c.Text,
872 MimeType: c.MimeType,
873 Data: c.Data,
874 })
875 }
876 _ = e.send(extproto.ToolResultFromExt{
877 Type: "tool_result",
878 ID: id,
879 Content: blocks,
880 IsError: r.IsError,
881 ActivateTools: r.ActivateTools,
882 })
883}
884
885// send marshals v + LF and writes it under a mutex (so concurrent
886// goroutines don't interleave bytes on stdout).
887func (e *Extension) send(v any) error {
888 b, err := extproto.Encode(v)
889 if err != nil {
890 return err
891 }
892 e.writeMu.Lock()
893 defer e.writeMu.Unlock()
894 _, err = e.out.Write(b)
895 return err
896}
897
898// dispatchIntercept runs the per-event handler (tool_call / turn_start
899// / assistant_message) on its own goroutine, catches panics, and
900// always emits exactly one event_intercept_response. Called from the
901// Run loop.
902func (e *Extension) dispatchIntercept(ei extproto.EventInterceptFromHost) {
903 defer func() {
904 if r := recover(); r != nil {
905 _ = e.send(extproto.EventInterceptResponseFromExt{
906 Type: "event_intercept_response",
907 ID: ei.ID,
908 Block: true,
909 Reason: fmt.Sprintf("intercept panic: %v", r),
910 })
911 }
912 }()
913
914 resp := extproto.EventInterceptResponseFromExt{
915 Type: "event_intercept_response",
916 ID: ei.ID,
917 }
918
919 switch ei.Event {
920 case "before_agent_start":
921 e.mu.Lock()
922 fn := e.interceptStart
923 e.mu.Unlock()
924 if fn != nil && ei.SystemPrompt != nil {
925 if replacement := fn(BeforeAgentStartEvent{
926 SystemPrompt: *ei.SystemPrompt, SessionID: ei.SessionID, AgentRunID: ei.AgentRunID,
927 CWD: ei.CWD, Provider: ei.Provider, Model: ei.Model,
928 }); replacement != nil {
929 resp.SystemPrompt, _ = json.Marshal(*replacement)
930 }
931 }
932 case "tool_call":
933 e.mu.Lock()
934 rich := e.interceptToolRich
935 plain := e.interceptTool
936 e.mu.Unlock()
937 if rich != nil {
938 d := rich(ei.ToolName, ei.ToolArgs)
939 resp.Block = d.Block
940 resp.Reason = d.Reason
941 if !d.Block && len(d.ModifiedArgs) > 0 && json.Valid(d.ModifiedArgs) {
942 resp.ModifiedArgs = d.ModifiedArgs
943 }
944 } else if plain != nil {
945 allow, reason := plain(ei.ToolName, ei.ToolArgs)
946 resp.Block = !allow
947 resp.Reason = reason
948 }
949 case "turn_start":
950 e.mu.Lock()
951 fn := e.interceptTurn
952 e.mu.Unlock()
953 if fn != nil {
954 d := fn(ei.Step)
955 resp.Block = d.Block
956 resp.Reason = d.Reason
957 }
958 case "assistant_message":
959 e.mu.Lock()
960 fn := e.interceptAssistant
961 e.mu.Unlock()
962 if fn != nil {
963 d := fn(ei.Text)
964 resp.Block = d.Block
965 resp.Reason = d.Reason
966 if !d.Block && d.ReplaceText != "" && d.ReplaceText != ei.Text {
967 resp.ReplaceText = d.ReplaceText
968 }
969 }
970 }
971
972 _ = e.send(resp)
973}