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// ToolResult is the extension's reply to a tool invocation. Build
56// one with TextResult, ImageResult, or directly when you need to
57// combine multiple blocks.
58type ToolResult struct {
59 Content []ToolContent
60 IsError bool
61}
62
63// ToolContent is one block of tool output. Either Text is set, or
64// MimeType+Data (base64-encoded). Use the Text/Image helpers below.
65type ToolContent struct {
66 Type string // "text" | "image"
67 Text string
68 MimeType string
69 Data string // base64
70}
71
72// Text returns a text content block.
73func Text(s string) ToolContent { return ToolContent{Type: "text", Text: s} }
74
75// Image returns an image content block. data must already be
76// base64-encoded; use ImageBytes to encode raw bytes.
77func Image(mimeType, base64Data string) ToolContent {
78 return ToolContent{Type: "image", MimeType: mimeType, Data: base64Data}
79}
80
81// ImageBytes returns an image content block from raw bytes,
82// encoding them to base64 for the wire.
83func ImageBytes(mimeType string, data []byte) ToolContent {
84 return ToolContent{Type: "image", MimeType: mimeType, Data: base64Encode(data)}
85}
86
87// TextResult returns a tool result with one text block, success.
88func TextResult(s string) ToolResult { return ToolResult{Content: []ToolContent{Text(s)}} }
89
90// TextErrorResult returns a tool result with one text block, marked
91// as an error to the model.
92func TextErrorResult(s string) ToolResult {
93 return ToolResult{Content: []ToolContent{Text(s)}, IsError: true}
94}
95
96// Response tells zot how to react to a command invocation. Construct
97// one with Prompt(), Insert(), Display(), or Noop().
98type Response struct {
99 Action string // "prompt", "insert", "display", "noop"
100 Prompt string
101 Insert string
102 Display string
103 Error string
104}
105
106// Prompt returns a Response that submits text as a fresh user message
107// to the agent (running it through the model loop as if the user had
108// typed and pressed enter).
109func Prompt(text string) Response { return Response{Action: "prompt", Prompt: text} }
110
111// Insert returns a Response that drops text into the editor at the
112// cursor without submitting.
113func Insert(text string) Response { return Response{Action: "insert", Insert: text} }
114
115// Display returns a Response that adds a one-shot styled note to the
116// chat without invoking the model. Useful for showing a result without
117// burning tokens.
118func Display(text string) Response { return Response{Action: "display", Display: text} }
119
120// Noop returns a Response that signals "I handled it, no UI change".
121// Use after pushing your own state or notifications.
122func Noop() Response { return Response{Action: "noop"} }
123
124// Errorf returns a Response that surfaces the error in the chat as a
125// red status line.
126func Errorf(format string, args ...any) Response {
127 return Response{Action: "noop", Error: fmt.Sprintf(format, args...)}
128}
129
130// Extension is one zot extension. Construct with New, register
131// commands, then call Run.
132type Extension struct {
133 name string
134 version string
135
136 in io.Reader
137 out io.Writer
138 stderr io.Writer
139 writeMu sync.Mutex
140
141 mu sync.Mutex
142 commands map[string]CommandHandler
143 descriptions []descTuple // ordered so register frames arrive in registration order
144 tools map[string]ToolHandler
145 toolDefs []toolDef // ordered so register frames arrive in registration order
146
147 // Caps reported in the hello frame. The SDK enables "commands"
148 // and "tools" automatically; future capabilities (events) will
149 // live here when those phases land.
150 caps []string
151
152 // hostInfo is filled in once HelloAck arrives.
153 host HostInfo
154}
155
156type descTuple struct {
157 name, desc string
158}
159
160type toolDef struct {
161 name string
162 description string
163 schema json.RawMessage
164}
165
166// HostInfo is what the host (zot) tells us in HelloAck. Useful for
167// extensions that want to behave differently per provider.
168type HostInfo struct {
169 ProtocolVersion int
170 ZotVersion string
171 Provider string
172 Model string
173 CWD string
174}
175
176// New constructs an Extension with the given identifier. name should
177// match the name field in extension.json.
178func New(name, version string) *Extension {
179 return &Extension{
180 name: name,
181 version: version,
182 in: os.Stdin,
183 out: os.Stdout,
184 stderr: os.Stderr,
185 commands: map[string]CommandHandler{},
186 tools: map[string]ToolHandler{},
187 caps: []string{"commands", "tools"},
188 }
189}
190
191// Host returns the HostInfo received during the hello handshake.
192// Returns the zero value if Run hasn't started yet.
193func (e *Extension) Host() HostInfo { return e.host }
194
195// Logf writes a line to the extension's stderr, which zot captures to
196// $ZOT_HOME/logs/ext-<name>.log. Use this for debug output: anything
197// you print to stdout would corrupt the JSON wire protocol.
198func (e *Extension) Logf(format string, args ...any) {
199 fmt.Fprintf(e.stderr, "["+e.name+"] "+format+"\n", args...)
200}
201
202// Command registers a slash-command handler. Call this BEFORE Run().
203// Once Run is going, when the user runs /name in zot, fn is invoked
204// with the remaining args.
205//
206// Naming conflicts with built-in commands (e.g. /help) are silently
207// shadowed by the built-in; check the user's "ext logs" if a command
208// you registered isn't taking effect.
209func (e *Extension) Command(name, description string, fn CommandHandler) {
210 e.mu.Lock()
211 e.commands[name] = fn
212 e.descriptions = append(e.descriptions, descTuple{name: name, desc: description})
213 e.mu.Unlock()
214}
215
216// Tool registers an LLM-callable tool. schema is a JSON Schema
217// object describing the tool's args (the same shape Anthropic /
218// OpenAI accept). Call this BEFORE Run(); zot folds extension tools
219// into the agent's registry once the extension's ready frame fires.
220//
221// Naming conflicts with built-in tools (read, write, edit, bash,
222// skill) are silently shadowed by the built-in.
223func (e *Extension) Tool(name, description string, schema json.RawMessage, fn ToolHandler) {
224 e.mu.Lock()
225 e.tools[name] = fn
226 e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema})
227 e.mu.Unlock()
228}
229
230// Notify pushes an info-level status note into zot's chat without
231// requiring a slash command from the user.
232func (e *Extension) Notify(level, message string) {
233 _ = e.send(extproto.NotifyFromExt{
234 Type: "notify",
235 Level: level,
236 Message: message,
237 })
238}
239
240// Run starts the protocol loop. Blocks until stdin closes (zot has
241// shut us down). Returns the first fatal error, or nil on clean exit.
242func (e *Extension) Run() error {
243 // Send hello, then re-announce all commands (covers Command calls
244 // made before Run, and is also fine for those made after via
245 // Command()'s direct send).
246 if err := e.send(extproto.HelloFromExt{
247 Type: "hello",
248 Name: e.name,
249 Version: e.version,
250 Capabilities: e.caps,
251 }); err != nil {
252 return err
253 }
254 e.mu.Lock()
255 descs := append([]descTuple(nil), e.descriptions...)
256 toolDefs := append([]toolDef(nil), e.toolDefs...)
257 e.mu.Unlock()
258 for _, d := range descs {
259 _ = e.send(extproto.RegisterCommandFromExt{
260 Type: "register_command",
261 Name: d.name,
262 Description: d.desc,
263 })
264 }
265 for _, td := range toolDefs {
266 _ = e.send(extproto.RegisterToolFromExt{
267 Type: "register_tool",
268 Name: td.name,
269 Description: td.description,
270 Schema: td.schema,
271 })
272 }
273 // Sentinel: tells the host that all initial registrations have
274 // been flushed and the agent registry can be built. Never block
275 // on this; if the host doesn't act on it, registrations still
276 // land in time for the typical use case.
277 _ = e.send(extproto.ReadyFromExt{Type: "ready"})
278
279 scanner := bufio.NewScanner(e.in)
280 scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
281 for scanner.Scan() {
282 line := scanner.Bytes()
283 var frame extproto.Frame
284 if err := json.Unmarshal(line, &frame); err != nil {
285 e.Logf("malformed frame from host: %v", err)
286 continue
287 }
288 switch frame.Type {
289 case "hello_ack":
290 var ack extproto.HelloAckFromHost
291 if err := json.Unmarshal(line, &ack); err == nil {
292 e.host = HostInfo{
293 ProtocolVersion: ack.ProtocolVersion,
294 ZotVersion: ack.ZotVersion,
295 Provider: ack.Provider,
296 Model: ack.Model,
297 CWD: ack.CWD,
298 }
299 }
300 case "command_invoked":
301 var ci extproto.CommandInvokedFromHost
302 if err := json.Unmarshal(line, &ci); err != nil {
303 continue
304 }
305 e.mu.Lock()
306 fn := e.commands[ci.Name]
307 e.mu.Unlock()
308 if fn == nil {
309 e.respond(ci.ID, Errorf("no handler for /%s", ci.Name))
310 continue
311 }
312 // Run the handler on its own goroutine so a slow handler
313 // doesn't block subsequent commands. Order isn't promised.
314 go func(id string, fn CommandHandler, args string) {
315 defer func() {
316 if r := recover(); r != nil {
317 e.respond(id, Errorf("panic: %v", r))
318 }
319 }()
320 resp := fn(args)
321 e.respond(id, resp)
322 }(ci.ID, fn, ci.Args)
323 case "tool_call":
324 var tc extproto.ToolCallFromHost
325 if err := json.Unmarshal(line, &tc); err != nil {
326 continue
327 }
328 e.mu.Lock()
329 fn := e.tools[tc.Name]
330 e.mu.Unlock()
331 if fn == nil {
332 e.respondTool(tc.ID, TextErrorResult(fmt.Sprintf("no handler for tool %q", tc.Name)))
333 continue
334 }
335 go func(id string, fn ToolHandler, args json.RawMessage) {
336 defer func() {
337 if r := recover(); r != nil {
338 e.respondTool(id, TextErrorResult(fmt.Sprintf("panic: %v", r)))
339 }
340 }()
341 res := fn(args)
342 e.respondTool(id, res)
343 }(tc.ID, fn, tc.Args)
344 case "shutdown":
345 _ = e.send(extproto.ShutdownAckFromExt{Type: "shutdown_ack"})
346 return nil
347 default:
348 e.Logf("unknown frame type %q", frame.Type)
349 }
350 }
351 return scanner.Err()
352}
353
354// respond serialises a CommandResponseFromExt for the given id.
355func (e *Extension) respond(id string, r Response) {
356 if r.Action == "" {
357 r.Action = "noop"
358 }
359 _ = e.send(extproto.CommandResponseFromExt{
360 Type: "command_response",
361 ID: id,
362 Action: r.Action,
363 Prompt: r.Prompt,
364 Insert: r.Insert,
365 Display: r.Display,
366 Error: r.Error,
367 })
368}
369
370// respondTool serialises a ToolResultFromExt for the given id.
371func (e *Extension) respondTool(id string, r ToolResult) {
372 blocks := make([]extproto.ContentBlock, 0, len(r.Content))
373 for _, c := range r.Content {
374 blocks = append(blocks, extproto.ContentBlock{
375 Type: c.Type,
376 Text: c.Text,
377 MimeType: c.MimeType,
378 Data: c.Data,
379 })
380 }
381 _ = e.send(extproto.ToolResultFromExt{
382 Type: "tool_result",
383 ID: id,
384 Content: blocks,
385 IsError: r.IsError,
386 })
387}
388
389// send marshals v + LF and writes it under a mutex (so concurrent
390// goroutines don't interleave bytes on stdout).
391func (e *Extension) send(v any) error {
392 b, err := extproto.Encode(v)
393 if err != nil {
394 return err
395 }
396 e.writeMu.Lock()
397 defer e.writeMu.Unlock()
398 _, err = e.out.Write(b)
399 return err
400}