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