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