1// Package extensions implements the host side of zot's subprocess
2// extension protocol. The Manager discovers extensions in well-known
3// directories, spawns each one, completes the hello handshake, and
4// routes slash commands to the right extension.
5//
6// Each extension is its own process, communicating with zot over its
7// stdin/stdout in newline-delimited JSON. Stderr is redirected to a
8// per-extension log file under $ZOT_HOME/logs/. Crashing one
9// extension does not affect the others or the host.
10//
11// See docs/extensions.md for the user-facing reference and
12// internal/extproto for the wire format.
13package extensions
14
15import (
16 "bufio"
17 "context"
18 "encoding/json"
19 "errors"
20 "fmt"
21 "io"
22 "os"
23 "os/exec"
24 "path/filepath"
25 "strings"
26 "sync"
27 "syscall"
28 "time"
29
30 "github.com/patriceckhart/zot/internal/extproto"
31)
32
33// Manifest is the extension.json file shipped alongside an
34// extension's executable. It tells zot how to launch the extension
35// and provides display metadata.
36type Manifest struct {
37 Name string `json:"name"`
38 Version string `json:"version,omitempty"`
39 Exec string `json:"exec"` // executable path, relative to manifest dir
40 Args []string `json:"args,omitempty"` // extra argv passed to exec
41 Language string `json:"language,omitempty"` // informational ("go", "python", "typescript", ...)
42 Enabled *bool `json:"enabled,omitempty"` // nil = enabled
43 Description string `json:"description,omitempty"`
44}
45
46// IsEnabled returns the manifest's effective enabled state. Default
47// is true so adding a new extension folder Just Works without an
48// extra zot ext enable command.
49func (m Manifest) IsEnabled() bool {
50 return m.Enabled == nil || *m.Enabled
51}
52
53// Extension is a running extension subprocess and the metadata zot
54// tracks about it.
55type Extension struct {
56 Manifest Manifest
57 Dir string // absolute path to extension directory
58 LogPath string
59
60 cmd *exec.Cmd
61 stdin io.WriteCloser
62 stdout io.ReadCloser
63 logFile *os.File
64 helloAck bool
65 commands []extproto.RegisterCommandFromExt
66 tools []extproto.RegisterToolFromExt
67
68 // readyCh is closed when the extension sends a ReadyFromExt
69 // frame, or when the host gives up waiting (registrationGrace).
70 readyCh chan struct{}
71 readyOnce sync.Once
72
73 // pending command invocations waiting on a CommandResponseFromExt
74 // keyed by the id we sent in CommandInvokedFromHost.
75 // pendingTool is the same idea for tool calls.
76 mu sync.Mutex
77 pending map[string]chan extproto.CommandResponseFromExt
78 pendingTool map[string]chan extproto.ToolResultFromExt
79}
80
81// HostHooks is the small interface the manager calls back into the
82// running TUI through. Decouples extensions from interactive.go.
83type HostHooks interface {
84 // Notify pushes an ext-originated status message into the chat.
85 // Level is one of "info", "warn", "error", "success".
86 Notify(extName, level, message string)
87
88 // Submit feeds text as if the user had typed and pressed enter,
89 // running it through the agent loop.
90 Submit(text string)
91
92 // Insert places text at the cursor in the editor.
93 Insert(text string)
94
95 // Display appends a one-shot styled note to the chat without
96 // invoking the model and without writing to the transcript.
97 Display(extName, text string)
98}
99
100// Manager owns every extension subprocess for the lifetime of zot.
101type Manager struct {
102 zotHome string
103 cwd string
104 zotVersion string
105 provider string
106 model string
107 hooks HostHooks
108
109 mu sync.RWMutex
110 ext map[string]*Extension // keyed by manifest name
111
112 // commandIndex maps a slash-command name (without the leading /)
113 // to the extension that registered it. First-come-first-served:
114 // later registrations of the same command are dropped with a
115 // warning.
116 commandIndex map[string]*Extension
117
118 // toolIndex maps an extension-defined tool name to its owning
119 // extension. Same first-come-first-served rule as commandIndex.
120 toolIndex map[string]*Extension
121}
122
123// New constructs an empty Manager. Call Discover to populate it from
124// the on-disk extension directories.
125func New(zotHome, cwd, zotVersion, provider, model string, hooks HostHooks) *Manager {
126 return &Manager{
127 zotHome: zotHome,
128 cwd: cwd,
129 zotVersion: zotVersion,
130 provider: provider,
131 model: model,
132 hooks: hooks,
133 ext: map[string]*Extension{},
134 commandIndex: map[string]*Extension{},
135 toolIndex: map[string]*Extension{},
136 }
137}
138
139// Discover scans the global and project extension dirs and starts
140// every extension whose manifest is enabled. Returns a slice of
141// errors encountered (one per extension); a single bad extension
142// doesn't abort the rest.
143func (m *Manager) Discover(ctx context.Context) []error {
144 var errs []error
145 for _, dir := range m.searchDirs() {
146 entries, err := os.ReadDir(dir)
147 if err != nil {
148 continue // missing directory is fine
149 }
150 for _, e := range entries {
151 if !e.IsDir() {
152 continue
153 }
154 extDir := filepath.Join(dir, e.Name())
155 if err := m.loadOne(ctx, extDir); err != nil {
156 errs = append(errs, fmt.Errorf("%s: %w", extDir, err))
157 }
158 }
159 }
160 return errs
161}
162
163// searchDirs returns the directories the discoverer walks, in
164// priority order: project-local first (so a project can override
165// global behavior), then global.
166func (m *Manager) searchDirs() []string {
167 var dirs []string
168 if m.cwd != "" {
169 dirs = append(dirs, filepath.Join(m.cwd, ".zot", "extensions"))
170 }
171 if m.zotHome != "" {
172 dirs = append(dirs, filepath.Join(m.zotHome, "extensions"))
173 }
174 return dirs
175}
176
177// loadOne reads a single extension's manifest and, if enabled,
178// spawns its subprocess + completes the hello handshake.
179func (m *Manager) loadOne(ctx context.Context, dir string) error {
180 manifestPath := filepath.Join(dir, "extension.json")
181 raw, err := os.ReadFile(manifestPath)
182 if err != nil {
183 return fmt.Errorf("read manifest: %w", err)
184 }
185 var mf Manifest
186 if err := json.Unmarshal(raw, &mf); err != nil {
187 return fmt.Errorf("parse manifest: %w", err)
188 }
189 if mf.Name == "" {
190 return errors.New("manifest: name is required")
191 }
192 if mf.Exec == "" {
193 return errors.New("manifest: exec is required")
194 }
195 if !mf.IsEnabled() {
196 // Quietly skip disabled extensions; zot ext list will show them.
197 return nil
198 }
199
200 m.mu.RLock()
201 _, dup := m.ext[mf.Name]
202 m.mu.RUnlock()
203 if dup {
204 // Project-local copy already won; ignore the global one.
205 return nil
206 }
207
208 ext := &Extension{
209 Manifest: mf,
210 Dir: dir,
211 readyCh: make(chan struct{}),
212 pending: map[string]chan extproto.CommandResponseFromExt{},
213 pendingTool: map[string]chan extproto.ToolResultFromExt{},
214 }
215 if err := m.spawn(ctx, ext); err != nil {
216 return err
217 }
218
219 m.mu.Lock()
220 m.ext[mf.Name] = ext
221 // Note: ext.commands and ext.tools may be empty here — they're
222 // populated by the read loop as register_* frames arrive after
223 // hello. Indexing happens in the read loop too. Discover()'s
224 // caller can WaitForReady() before relying on the registries.
225 m.mu.Unlock()
226 return nil
227}
228
229// WaitForReady blocks until every loaded extension has sent its
230// ReadyFromExt sentinel, or the per-extension grace period expires.
231// Call after Discover and before relying on tool registrations.
232func (m *Manager) WaitForReady(grace time.Duration) {
233 m.mu.RLock()
234 exts := make([]*Extension, 0, len(m.ext))
235 for _, e := range m.ext {
236 exts = append(exts, e)
237 }
238 m.mu.RUnlock()
239 for _, ext := range exts {
240 select {
241 case <-ext.readyCh:
242 case <-time.After(grace):
243 fmt.Fprintf(ext.logFile, "[zot] timed out waiting for ready frame; proceeding\n")
244 ext.readyOnce.Do(func() { close(ext.readyCh) })
245 }
246 }
247}
248
249// spawn launches the subprocess, hooks up pipes, logs stderr, and
250// runs the synchronous portion of the hello handshake. Asynchronous
251// frames are processed in a goroutine started here.
252func (m *Manager) spawn(ctx context.Context, ext *Extension) error {
253 logsDir := filepath.Join(m.zotHome, "logs")
254 _ = os.MkdirAll(logsDir, 0o755)
255 logPath := filepath.Join(logsDir, "ext-"+ext.Manifest.Name+".log")
256 logFile, err := os.OpenFile(logPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
257 if err != nil {
258 return fmt.Errorf("open log: %w", err)
259 }
260 ext.LogPath = logPath
261 ext.logFile = logFile
262 fmt.Fprintf(logFile, "\n[zot] starting %s/%s at %s\n", ext.Manifest.Name, ext.Manifest.Version, time.Now().Format(time.RFC3339))
263
264 execPath := ext.Manifest.Exec
265 if !filepath.IsAbs(execPath) {
266 execPath = filepath.Join(ext.Dir, execPath)
267 }
268 cmd := exec.CommandContext(ctx, execPath, ext.Manifest.Args...)
269 cmd.Dir = ext.Dir
270 cmd.Stderr = logFile
271
272 stdin, err := cmd.StdinPipe()
273 if err != nil {
274 return err
275 }
276 stdout, err := cmd.StdoutPipe()
277 if err != nil {
278 return err
279 }
280 if err := cmd.Start(); err != nil {
281 return fmt.Errorf("spawn: %w", err)
282 }
283 ext.cmd = cmd
284 ext.stdin = stdin
285 ext.stdout = stdout
286
287 // Hello handshake. Read the extension's HelloFromExt synchronously
288 // so we can fail fast on a broken extension; everything after is
289 // processed in the read goroutine.
290 scanner := bufio.NewScanner(stdout)
291 scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
292 if !scanner.Scan() {
293 return fmt.Errorf("extension exited before hello: %w", scanner.Err())
294 }
295 var hello extproto.HelloFromExt
296 if err := json.Unmarshal(scanner.Bytes(), &hello); err != nil {
297 return fmt.Errorf("parse hello: %w", err)
298 }
299 if hello.Type != "hello" || hello.Name == "" {
300 return fmt.Errorf("first frame must be hello (got %q)", hello.Type)
301 }
302 // Trust the manifest's name; ignore mismatch from the hello.
303 ext.helloAck = true
304
305 ack, _ := extproto.Encode(extproto.HelloAckFromHost{
306 Type: "hello_ack",
307 ProtocolVersion: extproto.ProtocolVersion,
308 ZotVersion: m.zotVersion,
309 Provider: m.provider,
310 Model: m.model,
311 CWD: m.cwd,
312 })
313 if _, err := stdin.Write(ack); err != nil {
314 return fmt.Errorf("send hello_ack: %w", err)
315 }
316
317 // Spin up the read loop now that the handshake is done.
318 go m.readLoop(ext, scanner)
319 return nil
320}
321
322// readLoop processes every frame the extension sends after hello.
323// Returns when stdout closes.
324func (m *Manager) readLoop(ext *Extension, scanner *bufio.Scanner) {
325 defer func() {
326 // On close, drop every command + tool this extension owned so
327 // future invocations don't dangle. The subprocess is gone; we
328 // won't hear back about its commands or tool calls anymore.
329 m.mu.Lock()
330 for name, owner := range m.commandIndex {
331 if owner == ext {
332 delete(m.commandIndex, name)
333 }
334 }
335 for name, owner := range m.toolIndex {
336 if owner == ext {
337 delete(m.toolIndex, name)
338 }
339 }
340 m.mu.Unlock()
341 ext.readyOnce.Do(func() { close(ext.readyCh) })
342 fmt.Fprintf(ext.logFile, "[zot] extension %s read loop exited at %s\n", ext.Manifest.Name, time.Now().Format(time.RFC3339))
343 }()
344
345 for scanner.Scan() {
346 line := scanner.Bytes()
347 var frame extproto.Frame
348 if err := json.Unmarshal(line, &frame); err != nil {
349 fmt.Fprintf(ext.logFile, "[zot] malformed json from extension: %v\n", err)
350 continue
351 }
352 switch frame.Type {
353 case "register_command":
354 var rc extproto.RegisterCommandFromExt
355 if err := json.Unmarshal(line, &rc); err == nil {
356 ext.commands = append(ext.commands, rc)
357 m.mu.Lock()
358 if _, exists := m.commandIndex[rc.Name]; !exists {
359 m.commandIndex[rc.Name] = ext
360 }
361 m.mu.Unlock()
362 }
363 case "register_tool":
364 var rt extproto.RegisterToolFromExt
365 if err := json.Unmarshal(line, &rt); err != nil {
366 fmt.Fprintf(ext.logFile, "[zot] bad register_tool frame: %v\n", err)
367 continue
368 }
369 // Validate the schema parses as JSON. If not, refuse to
370 // register — a broken schema confuses the model.
371 if len(rt.Schema) > 0 {
372 var tmp any
373 if err := json.Unmarshal(rt.Schema, &tmp); err != nil {
374 fmt.Fprintf(ext.logFile, "[zot] tool %q: schema is not valid json (%v); skipped\n", rt.Name, err)
375 continue
376 }
377 }
378 ext.tools = append(ext.tools, rt)
379 m.mu.Lock()
380 if _, exists := m.toolIndex[rt.Name]; !exists {
381 m.toolIndex[rt.Name] = ext
382 }
383 m.mu.Unlock()
384 case "ready":
385 ext.readyOnce.Do(func() { close(ext.readyCh) })
386 case "tool_result":
387 var tr extproto.ToolResultFromExt
388 if err := json.Unmarshal(line, &tr); err == nil {
389 ext.mu.Lock()
390 ch, ok := ext.pendingTool[tr.ID]
391 if ok {
392 delete(ext.pendingTool, tr.ID)
393 }
394 ext.mu.Unlock()
395 if ok {
396 select {
397 case ch <- tr:
398 default:
399 }
400 }
401 }
402 case "notify":
403 var n extproto.NotifyFromExt
404 if err := json.Unmarshal(line, &n); err == nil {
405 m.hooks.Notify(ext.Manifest.Name, n.Level, n.Message)
406 }
407 case "command_response":
408 var cr extproto.CommandResponseFromExt
409 if err := json.Unmarshal(line, &cr); err == nil {
410 ext.mu.Lock()
411 ch, ok := ext.pending[cr.ID]
412 if ok {
413 delete(ext.pending, cr.ID)
414 }
415 ext.mu.Unlock()
416 if ok {
417 select {
418 case ch <- cr:
419 default:
420 }
421 }
422 }
423 case "shutdown_ack":
424 // Caller of Stop is waiting on the process exit, not this frame.
425 default:
426 fmt.Fprintf(ext.logFile, "[zot] unknown frame type %q\n", frame.Type)
427 }
428 }
429}
430
431// Commands returns a snapshot of every (extension, command) pair
432// currently registered. Used by the slash autocomplete + /help.
433func (m *Manager) Commands() []CommandInfo {
434 m.mu.RLock()
435 defer m.mu.RUnlock()
436 var out []CommandInfo
437 for _, ext := range m.ext {
438 for _, c := range ext.commands {
439 out = append(out, CommandInfo{
440 Extension: ext.Manifest.Name,
441 Name: c.Name,
442 Description: c.Description,
443 })
444 }
445 }
446 return out
447}
448
449// CommandInfo is one extension-registered slash command, surfaced to
450// the rest of zot for display purposes.
451type CommandInfo struct {
452 Extension string
453 Name string
454 Description string
455}
456
457// ToolInfo is one extension-registered tool. Used by the agent's
458// build step to materialise core.Tool wrappers.
459type ToolInfo struct {
460 Extension string
461 Name string
462 Description string
463 Schema json.RawMessage
464}
465
466// Tools returns a snapshot of every (extension, tool) pair currently
467// registered. Used at agent-build time to fold extension tools into
468// the runtime tool registry.
469func (m *Manager) Tools() []ToolInfo {
470 m.mu.RLock()
471 defer m.mu.RUnlock()
472 var out []ToolInfo
473 for _, ext := range m.ext {
474 for _, t := range ext.tools {
475 out = append(out, ToolInfo{
476 Extension: ext.Manifest.Name,
477 Name: t.Name,
478 Description: t.Description,
479 Schema: t.Schema,
480 })
481 }
482 }
483 return out
484}
485
486// HasTool reports whether name is registered by any extension.
487func (m *Manager) HasTool(name string) bool {
488 m.mu.RLock()
489 defer m.mu.RUnlock()
490 _, ok := m.toolIndex[name]
491 return ok
492}
493
494// InvokeTool sends a tool_call to the owning extension and waits for
495// the matching tool_result. Used by the core.Tool wrapper that the
496// agent registers per extension-defined tool.
497func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMessage, timeout time.Duration) (extproto.ToolResultFromExt, error) {
498 m.mu.RLock()
499 ext, ok := m.toolIndex[name]
500 m.mu.RUnlock()
501 if !ok {
502 return extproto.ToolResultFromExt{}, fmt.Errorf("no extension registered for tool %q", name)
503 }
504
505 id := newCorrelationID()
506 ch := make(chan extproto.ToolResultFromExt, 1)
507 ext.mu.Lock()
508 ext.pendingTool[id] = ch
509 ext.mu.Unlock()
510
511 frame, _ := extproto.Encode(extproto.ToolCallFromHost{
512 Type: "tool_call",
513 ID: id,
514 Name: name,
515 Args: args,
516 })
517 if _, err := ext.stdin.Write(frame); err != nil {
518 ext.mu.Lock()
519 delete(ext.pendingTool, id)
520 ext.mu.Unlock()
521 return extproto.ToolResultFromExt{}, fmt.Errorf("write: %w", err)
522 }
523
524 select {
525 case resp := <-ch:
526 return resp, nil
527 case <-time.After(timeout):
528 ext.mu.Lock()
529 delete(ext.pendingTool, id)
530 ext.mu.Unlock()
531 return extproto.ToolResultFromExt{}, fmt.Errorf("timeout waiting for %s/%s", ext.Manifest.Name, name)
532 case <-ctx.Done():
533 ext.mu.Lock()
534 delete(ext.pendingTool, id)
535 ext.mu.Unlock()
536 return extproto.ToolResultFromExt{}, ctx.Err()
537 }
538}
539
540// HasCommand reports whether name is registered by any extension.
541func (m *Manager) HasCommand(name string) bool {
542 m.mu.RLock()
543 defer m.mu.RUnlock()
544 _, ok := m.commandIndex[name]
545 return ok
546}
547
548// Invoke fires the named slash command's handler in the owning
549// extension and waits up to timeout for the response. Returns the
550// extension's CommandResponse so the caller can act on the action
551// (prompt / insert / display).
552func (m *Manager) Invoke(ctx context.Context, name, args string, timeout time.Duration) (extproto.CommandResponseFromExt, error) {
553 m.mu.RLock()
554 ext, ok := m.commandIndex[name]
555 m.mu.RUnlock()
556 if !ok {
557 return extproto.CommandResponseFromExt{}, fmt.Errorf("no extension registered for /%s", name)
558 }
559
560 id := newCorrelationID()
561 ch := make(chan extproto.CommandResponseFromExt, 1)
562 ext.mu.Lock()
563 ext.pending[id] = ch
564 ext.mu.Unlock()
565
566 frame, _ := extproto.Encode(extproto.CommandInvokedFromHost{
567 Type: "command_invoked",
568 ID: id,
569 Name: name,
570 Args: args,
571 })
572 if _, err := ext.stdin.Write(frame); err != nil {
573 ext.mu.Lock()
574 delete(ext.pending, id)
575 ext.mu.Unlock()
576 return extproto.CommandResponseFromExt{}, fmt.Errorf("write: %w", err)
577 }
578
579 select {
580 case resp := <-ch:
581 return resp, nil
582 case <-time.After(timeout):
583 ext.mu.Lock()
584 delete(ext.pending, id)
585 ext.mu.Unlock()
586 return extproto.CommandResponseFromExt{}, fmt.Errorf("timeout waiting for %s/%s", ext.Manifest.Name, name)
587 case <-ctx.Done():
588 ext.mu.Lock()
589 delete(ext.pending, id)
590 ext.mu.Unlock()
591 return extproto.CommandResponseFromExt{}, ctx.Err()
592 }
593}
594
595// Stop cleanly terminates every extension. Sends ShutdownFromHost,
596// waits up to gracePeriod for each subprocess to exit, then SIGTERMs
597// (and SIGKILLs after another second) the holdouts.
598func (m *Manager) Stop(gracePeriod time.Duration) {
599 m.mu.RLock()
600 exts := make([]*Extension, 0, len(m.ext))
601 for _, e := range m.ext {
602 exts = append(exts, e)
603 }
604 m.mu.RUnlock()
605
606 for _, ext := range exts {
607 if frame, err := extproto.Encode(extproto.ShutdownFromHost{Type: "shutdown"}); err == nil {
608 _, _ = ext.stdin.Write(frame)
609 }
610 _ = ext.stdin.Close()
611 }
612
613 deadline := time.Now().Add(gracePeriod)
614 for _, ext := range exts {
615 remaining := time.Until(deadline)
616 if remaining <= 0 {
617 remaining = 100 * time.Millisecond
618 }
619 done := make(chan struct{})
620 go func() { _ = ext.cmd.Wait(); close(done) }()
621 select {
622 case <-done:
623 case <-time.After(remaining):
624 _ = ext.cmd.Process.Signal(syscall.SIGTERM)
625 select {
626 case <-done:
627 case <-time.After(time.Second):
628 _ = ext.cmd.Process.Kill()
629 <-done
630 }
631 }
632 if ext.logFile != nil {
633 _ = ext.logFile.Close()
634 }
635 }
636}
637
638// All returns every extension currently tracked, enabled or not.
639// Used by `zot ext list`.
640func (m *Manager) All() []*Extension {
641 m.mu.RLock()
642 defer m.mu.RUnlock()
643 out := make([]*Extension, 0, len(m.ext))
644 for _, e := range m.ext {
645 out = append(out, e)
646 }
647 return out
648}
649
650// newCorrelationID returns a short non-cryptographic id. We don't
651// need uniqueness across processes, just within the lifetime of one
652// extension's pending map.
653func newCorrelationID() string {
654 return strings.ReplaceAll(time.Now().Format("150405.000000"), ".", "")
655}