1// Package zotext 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/pkg/zotext"
13//
14// func main() {
15// ext := zotext.New("hello", "1.0.0")
16// ext.Command("hello", "say hi", func(args string) zotext.Response {
17// return zotext.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// The same wire format also has reference clients in TypeScript and
26// Python under examples/extensions/. Use whichever language fits.
27package zotext
28
29import (
30 "bufio"
31 "encoding/base64"
32 "encoding/json"
33 "fmt"
34 "io"
35 "os"
36 "sync"
37
38 "github.com/patriceckhart/zot/internal/extproto"
39)
40
41func base64Encode(b []byte) string { return base64.StdEncoding.EncodeToString(b) }
42
43// CommandHandler is invoked when the user runs the extension's
44// registered slash command. args is everything the user typed after
45// the command name (already trimmed). Return a Response describing
46// what zot should do next.
47type CommandHandler func(args string) Response
48
49// ToolHandler is invoked when the LLM calls a tool the extension
50// registered. args is the raw JSON object the model produced; the
51// handler is responsible for parsing/validating it. Return a
52// ToolResult describing what zot should send back to the model.
53type ToolHandler func(args json.RawMessage) ToolResult
54
55// EventHandler is called for each lifecycle event the extension
56// subscribed to via Subscribe. The handler is invoked synchronously
57// on the read goroutine; keep it quick or hand off to your own
58// worker.
59type EventHandler func(ev Event)
60
61// Event is a lifecycle notification from zot. The fields populated
62// depend on Name (the host's event_name string):
63//
64// session_start : (no extra fields)
65// turn_start : Step
66// turn_end : Stop, optional Error
67// tool_call : ToolID, ToolName, ToolArgs
68// assistant_message: Text
69type Event struct {
70 Name string
71
72 Step int
73 Stop string
74 Error string
75
76 ToolID string
77 ToolName string
78 ToolArgs json.RawMessage
79
80 Text string
81}
82
83// InterceptHandler decides whether a tool call may proceed. Return
84// (allow=true) to permit, (allow=false, reason) to refuse. The
85// reason is shown to the model as the tool's error text.
86type InterceptHandler func(toolName string, args json.RawMessage) (allow bool, reason string)
87
88// ToolResult is the extension's reply to a tool invocation. Build
89// one with TextResult, ImageResult, or directly when you need to
90// combine multiple blocks.
91type ToolResult struct {
92 Content []ToolContent
93 IsError bool
94}
95
96// ToolContent is one block of tool output. Either Text is set, or
97// MimeType+Data (base64-encoded). Use the Text/Image helpers below.
98type ToolContent struct {
99 Type string // "text" | "image"
100 Text string
101 MimeType string
102 Data string // base64
103}
104
105// Text returns a text content block.
106func Text(s string) ToolContent { return ToolContent{Type: "text", Text: s} }
107
108// Image returns an image content block. data must already be
109// base64-encoded; use ImageBytes to encode raw bytes.
110func Image(mimeType, base64Data string) ToolContent {
111 return ToolContent{Type: "image", MimeType: mimeType, Data: base64Data}
112}
113
114// ImageBytes returns an image content block from raw bytes,
115// encoding them to base64 for the wire.
116func ImageBytes(mimeType string, data []byte) ToolContent {
117 return ToolContent{Type: "image", MimeType: mimeType, Data: base64Encode(data)}
118}
119
120// TextResult returns a tool result with one text block, success.
121func TextResult(s string) ToolResult { return ToolResult{Content: []ToolContent{Text(s)}} }
122
123// TextErrorResult returns a tool result with one text block, marked
124// as an error to the model.
125func TextErrorResult(s string) ToolResult {
126 return ToolResult{Content: []ToolContent{Text(s)}, IsError: true}
127}
128
129// Response tells zot how to react to a command invocation. Construct
130// one with Prompt(), Insert(), Display(), or Noop().
131type Response struct {
132 Action string // "prompt", "insert", "display", "noop"
133 Prompt string
134 Insert string
135 Display string
136 Error string
137}
138
139// Prompt returns a Response that submits text as a fresh user message
140// to the agent (running it through the model loop as if the user had
141// typed and pressed enter).
142func Prompt(text string) Response { return Response{Action: "prompt", Prompt: text} }
143
144// Insert returns a Response that drops text into the editor at the
145// cursor without submitting.
146func Insert(text string) Response { return Response{Action: "insert", Insert: text} }
147
148// Display returns a Response that adds a one-shot styled note to the
149// chat without invoking the model. Useful for showing a result without
150// burning tokens.
151func Display(text string) Response { return Response{Action: "display", Display: text} }
152
153// Noop returns a Response that signals "I handled it, no UI change".
154// Use after pushing your own state or notifications.
155func Noop() Response { return Response{Action: "noop"} }
156
157// Errorf returns a Response that surfaces the error in the chat as a
158// red status line.
159func Errorf(format string, args ...any) Response {
160 return Response{Action: "noop", Error: fmt.Sprintf(format, args...)}
161}
162
163// Extension is one zot extension. Construct with New, register
164// commands, then call Run.
165type Extension struct {
166 name string
167 version string
168
169 in io.Reader
170 out io.Writer
171 stderr io.Writer
172 writeMu sync.Mutex
173
174 mu sync.Mutex
175 commands map[string]CommandHandler
176 descriptions []descTuple // ordered so register frames arrive in registration order
177 tools map[string]ToolHandler
178 toolDefs []toolDef // ordered so register frames arrive in registration order
179 eventHandlers map[string]EventHandler
180 eventNames []string // declared subscription order
181 interceptTool InterceptHandler
182 interceptOn bool
183
184 // Caps reported in the hello frame.
185 caps []string
186
187 // hostInfo is filled in once HelloAck arrives.
188 host HostInfo
189}
190
191type descTuple struct {
192 name, desc string
193}
194
195type toolDef struct {
196 name string
197 description string
198 schema json.RawMessage
199}
200
201// HostInfo is what the host (zot) tells us in HelloAck. Useful for
202// extensions that want to behave differently per provider.
203type HostInfo struct {
204 ProtocolVersion int
205 ZotVersion string
206 Provider string
207 Model string
208 CWD string
209}
210
211// New constructs an Extension with the given identifier. name should
212// match the name field in extension.json.
213func New(name, version string) *Extension {
214 return &Extension{
215 name: name,
216 version: version,
217 in: os.Stdin,
218 out: os.Stdout,
219 stderr: os.Stderr,
220 commands: map[string]CommandHandler{},
221 tools: map[string]ToolHandler{},
222 eventHandlers: map[string]EventHandler{},
223 caps: []string{"commands", "tools", "events"},
224 }
225}
226
227// Host returns the HostInfo received during the hello handshake.
228// Returns the zero value if Run hasn't started yet.
229func (e *Extension) Host() HostInfo { return e.host }
230
231// Logf writes a line to the extension's stderr, which zot captures to
232// $ZOT_HOME/logs/ext-<name>.log. Use this for debug output: anything
233// you print to stdout would corrupt the JSON wire protocol.
234func (e *Extension) Logf(format string, args ...any) {
235 fmt.Fprintf(e.stderr, "["+e.name+"] "+format+"\n", args...)
236}
237
238// Command registers a slash-command handler. Call this BEFORE Run().
239// Once Run is going, when the user runs /name in zot, fn is invoked
240// with the remaining args.
241//
242// Naming conflicts with built-in commands (e.g. /help) are silently
243// shadowed by the built-in; check the user's "ext logs" if a command
244// you registered isn't taking effect.
245func (e *Extension) Command(name, description string, fn CommandHandler) {
246 e.mu.Lock()
247 e.commands[name] = fn
248 e.descriptions = append(e.descriptions, descTuple{name: name, desc: description})
249 e.mu.Unlock()
250}
251
252// Tool registers an LLM-callable tool. schema is a JSON Schema
253// object describing the tool's args (the same shape Anthropic /
254// OpenAI accept). Call this BEFORE Run(); zot folds extension tools
255// into the agent's registry once the extension's ready frame fires.
256//
257// Naming conflicts with built-in tools (read, write, edit, bash,
258// skill) are silently shadowed by the built-in.
259func (e *Extension) Tool(name, description string, schema json.RawMessage, fn ToolHandler) {
260 e.mu.Lock()
261 e.tools[name] = fn
262 e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema})
263 e.mu.Unlock()
264}
265
266// On subscribes to a lifecycle event. fn is called for each
267// notification; the same name can only have one handler (later
268// registrations replace earlier ones). Recognised names:
269// session_start, turn_start, turn_end, tool_call,
270// assistant_message.
271func (e *Extension) On(name string, fn EventHandler) {
272 e.mu.Lock()
273 if _, exists := e.eventHandlers[name]; !exists {
274 e.eventNames = append(e.eventNames, name)
275 }
276 e.eventHandlers[name] = fn
277 e.mu.Unlock()
278}
279
280// InterceptToolCall registers a guard that runs immediately before
281// each tool call. Returning (false, reason) refuses the call; the
282// model sees reason as the tool error text. Multiple extensions can
283// install interceptors; if any one refuses, the call is blocked.
284//
285// Use this to build permission gates: refuse `bash` calls containing
286// `rm -rf`, ask the user for confirmation on dangerous patterns,
287// audit-log every call, etc.
288func (e *Extension) InterceptToolCall(fn InterceptHandler) {
289 e.mu.Lock()
290 e.interceptTool = fn
291 e.interceptOn = true
292 e.mu.Unlock()
293}
294
295// Notify pushes an info-level status note into zot's chat without
296// requiring a slash command from the user.
297func (e *Extension) Notify(level, message string) {
298 _ = e.send(extproto.NotifyFromExt{
299 Type: "notify",
300 Level: level,
301 Message: message,
302 })
303}
304
305// Run starts the protocol loop. Blocks until stdin closes (zot has
306// shut us down). Returns the first fatal error, or nil on clean exit.
307func (e *Extension) Run() error {
308 // Send hello, then re-announce all commands (covers Command calls
309 // made before Run, and is also fine for those made after via
310 // Command()'s direct send).
311 if err := e.send(extproto.HelloFromExt{
312 Type: "hello",
313 Name: e.name,
314 Version: e.version,
315 Capabilities: e.caps,
316 }); err != nil {
317 return err
318 }
319 e.mu.Lock()
320 descs := append([]descTuple(nil), e.descriptions...)
321 toolDefs := append([]toolDef(nil), e.toolDefs...)
322 eventNames := append([]string(nil), e.eventNames...)
323 interceptOn := e.interceptOn
324 e.mu.Unlock()
325 for _, d := range descs {
326 _ = e.send(extproto.RegisterCommandFromExt{
327 Type: "register_command",
328 Name: d.name,
329 Description: d.desc,
330 })
331 }
332 for _, td := range toolDefs {
333 _ = e.send(extproto.RegisterToolFromExt{
334 Type: "register_tool",
335 Name: td.name,
336 Description: td.description,
337 Schema: td.schema,
338 })
339 }
340 if len(eventNames) > 0 || interceptOn {
341 sub := extproto.SubscribeFromExt{Type: "subscribe", Events: eventNames}
342 if interceptOn {
343 sub.Intercept = []string{"tool_call"}
344 }
345 _ = e.send(sub)
346 }
347 // Sentinel: tells the host that all initial registrations have
348 // been flushed and the agent registry can be built. Never block
349 // on this; if the host doesn't act on it, registrations still
350 // land in time for the typical use case.
351 _ = e.send(extproto.ReadyFromExt{Type: "ready"})
352
353 scanner := bufio.NewScanner(e.in)
354 scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
355 for scanner.Scan() {
356 line := scanner.Bytes()
357 var frame extproto.Frame
358 if err := json.Unmarshal(line, &frame); err != nil {
359 e.Logf("malformed frame from host: %v", err)
360 continue
361 }
362 switch frame.Type {
363 case "hello_ack":
364 var ack extproto.HelloAckFromHost
365 if err := json.Unmarshal(line, &ack); err == nil {
366 e.host = HostInfo{
367 ProtocolVersion: ack.ProtocolVersion,
368 ZotVersion: ack.ZotVersion,
369 Provider: ack.Provider,
370 Model: ack.Model,
371 CWD: ack.CWD,
372 }
373 }
374 case "command_invoked":
375 var ci extproto.CommandInvokedFromHost
376 if err := json.Unmarshal(line, &ci); err != nil {
377 continue
378 }
379 e.mu.Lock()
380 fn := e.commands[ci.Name]
381 e.mu.Unlock()
382 if fn == nil {
383 e.respond(ci.ID, Errorf("no handler for /%s", ci.Name))
384 continue
385 }
386 // Run the handler on its own goroutine so a slow handler
387 // doesn't block subsequent commands. Order isn't promised.
388 go func(id string, fn CommandHandler, args string) {
389 defer func() {
390 if r := recover(); r != nil {
391 e.respond(id, Errorf("panic: %v", r))
392 }
393 }()
394 resp := fn(args)
395 e.respond(id, resp)
396 }(ci.ID, fn, ci.Args)
397 case "tool_call":
398 var tc extproto.ToolCallFromHost
399 if err := json.Unmarshal(line, &tc); err != nil {
400 continue
401 }
402 e.mu.Lock()
403 fn := e.tools[tc.Name]
404 e.mu.Unlock()
405 if fn == nil {
406 e.respondTool(tc.ID, TextErrorResult(fmt.Sprintf("no handler for tool %q", tc.Name)))
407 continue
408 }
409 go func(id string, fn ToolHandler, args json.RawMessage) {
410 defer func() {
411 if r := recover(); r != nil {
412 e.respondTool(id, TextErrorResult(fmt.Sprintf("panic: %v", r)))
413 }
414 }()
415 res := fn(args)
416 e.respondTool(id, res)
417 }(tc.ID, fn, tc.Args)
418 case "event":
419 var ef extproto.EventFromHost
420 if err := json.Unmarshal(line, &ef); err != nil {
421 continue
422 }
423 e.mu.Lock()
424 handler := e.eventHandlers[ef.Event]
425 e.mu.Unlock()
426 if handler != nil {
427 func() {
428 defer func() {
429 if r := recover(); r != nil {
430 e.Logf("event %s handler panicked: %v", ef.Event, r)
431 }
432 }()
433 handler(Event{
434 Name: ef.Event, Step: ef.Step, Stop: ef.Stop,
435 Error: ef.Error, ToolID: ef.ToolID, ToolName: ef.ToolName,
436 ToolArgs: ef.ToolArgs, Text: ef.Text,
437 })
438 }()
439 }
440 case "event_intercept":
441 var ei extproto.EventInterceptFromHost
442 if err := json.Unmarshal(line, &ei); err != nil {
443 continue
444 }
445 e.mu.Lock()
446 fn := e.interceptTool
447 e.mu.Unlock()
448 if fn == nil || ei.Event != "tool_call" {
449 _ = e.send(extproto.EventInterceptResponseFromExt{
450 Type: "event_intercept_response", ID: ei.ID,
451 })
452 continue
453 }
454 go func(id, name string, args json.RawMessage) {
455 defer func() {
456 if r := recover(); r != nil {
457 _ = e.send(extproto.EventInterceptResponseFromExt{
458 Type: "event_intercept_response", ID: id,
459 Block: true, Reason: fmt.Sprintf("intercept panic: %v", r),
460 })
461 }
462 }()
463 allow, reason := fn(name, args)
464 _ = e.send(extproto.EventInterceptResponseFromExt{
465 Type: "event_intercept_response", ID: id,
466 Block: !allow, Reason: reason,
467 })
468 }(ei.ID, ei.ToolName, ei.ToolArgs)
469 case "shutdown":
470 _ = e.send(extproto.ShutdownAckFromExt{Type: "shutdown_ack"})
471 return nil
472 default:
473 e.Logf("unknown frame type %q", frame.Type)
474 }
475 }
476 return scanner.Err()
477}
478
479// respond serialises a CommandResponseFromExt for the given id.
480func (e *Extension) respond(id string, r Response) {
481 if r.Action == "" {
482 r.Action = "noop"
483 }
484 _ = e.send(extproto.CommandResponseFromExt{
485 Type: "command_response",
486 ID: id,
487 Action: r.Action,
488 Prompt: r.Prompt,
489 Insert: r.Insert,
490 Display: r.Display,
491 Error: r.Error,
492 })
493}
494
495// respondTool serialises a ToolResultFromExt for the given id.
496func (e *Extension) respondTool(id string, r ToolResult) {
497 blocks := make([]extproto.ContentBlock, 0, len(r.Content))
498 for _, c := range r.Content {
499 blocks = append(blocks, extproto.ContentBlock{
500 Type: c.Type,
501 Text: c.Text,
502 MimeType: c.MimeType,
503 Data: c.Data,
504 })
505 }
506 _ = e.send(extproto.ToolResultFromExt{
507 Type: "tool_result",
508 ID: id,
509 Content: blocks,
510 IsError: r.IsError,
511 })
512}
513
514// send marshals v + LF and writes it under a mutex (so concurrent
515// goroutines don't interleave bytes on stdout).
516func (e *Extension) send(v any) error {
517 b, err := extproto.Encode(v)
518 if err != nil {
519 return err
520 }
521 e.writeMu.Lock()
522 defer e.writeMu.Unlock()
523 _, err = e.out.Write(b)
524 return err
525}