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 // lastFrameTime is updated by the read loop on every frame it
83 // processes. Used by the auto-ready idle watchdog so legacy
84 // extensions (no `ready` frame) don't pin the WaitForReady wait
85 // to its full grace.
86 lastFrameTime time.Time
87
88 // eventSubs and interceptSubs are the sets of event names this
89 // extension subscribed to via SubscribeFromExt. Used by
90 // EmitEvent / InterceptToolCall to filter recipients.
91 eventSubs map[string]struct{}
92 interceptSubs map[string]struct{}
93}
94
95// HostHooks is the small interface the manager calls back into the
96// running TUI through. Decouples extensions from interactive.go.
97type HostHooks interface {
98 // Notify pushes an ext-originated status message into the chat.
99 // Level is one of "info", "warn", "error", "success".
100 Notify(extName, level, message string)
101
102 // Submit feeds text as if the user had typed and pressed enter,
103 // running it through the agent loop.
104 Submit(text string)
105
106 // Insert places text at the cursor in the editor.
107 Insert(text string)
108
109 // Display appends a one-shot styled note to the chat without
110 // invoking the model and without writing to the transcript.
111 Display(extName, text string)
112}
113
114// Manager owns every extension subprocess for the lifetime of zot.
115type Manager struct {
116 zotHome string
117 cwd string
118 zotVersion string
119 provider string
120 model string
121 hooks HostHooks
122
123 mu sync.RWMutex
124 ext map[string]*Extension // keyed by manifest name
125
126 // commandIndex maps a slash-command name (without the leading /)
127 // to the extension that registered it. First-come-first-served:
128 // later registrations of the same command are dropped with a
129 // warning.
130 commandIndex map[string]*Extension
131
132 // toolIndex maps an extension-defined tool name to its owning
133 // extension. Same first-come-first-served rule as commandIndex.
134 toolIndex map[string]*Extension
135
136 // explicitPaths remembers ad-hoc paths passed via --ext so
137 // Reload can respawn them alongside the discovered set.
138 explicitPaths []string
139
140 // onReload, if set, is invoked after a successful Reload. Used
141 // by the host so it can rebuild the agent's tool registry with
142 // the freshly-registered extension tools.
143 onReload func()
144}
145
146// New constructs an empty Manager. Call Discover to populate it from
147// the on-disk extension directories.
148func New(zotHome, cwd, zotVersion, provider, model string, hooks HostHooks) *Manager {
149 return &Manager{
150 zotHome: zotHome,
151 cwd: cwd,
152 zotVersion: zotVersion,
153 provider: provider,
154 model: model,
155 hooks: hooks,
156 ext: map[string]*Extension{},
157 commandIndex: map[string]*Extension{},
158 toolIndex: map[string]*Extension{},
159 }
160}
161
162// Discover scans the global and project extension dirs and starts
163// every extension whose manifest is enabled. Spawns happen in
164// parallel so a slow runtime (e.g. `npx tsx` cold-start, ~1.5s)
165// doesn't block other extensions from starting. Returns a slice of
166// errors encountered (one per extension); a single bad extension
167// doesn't abort the rest.
168func (m *Manager) Discover(ctx context.Context) []error {
169 type loadJob struct {
170 dir string
171 }
172 var jobs []loadJob
173 seenDirs := map[string]bool{} // dedup by basename so project wins
174 for _, dir := range m.searchDirs() {
175 entries, err := os.ReadDir(dir)
176 if err != nil {
177 continue // missing directory is fine
178 }
179 for _, e := range entries {
180 if !e.IsDir() {
181 continue
182 }
183 if seenDirs[e.Name()] {
184 continue // higher-priority location already queued
185 }
186 seenDirs[e.Name()] = true
187 jobs = append(jobs, loadJob{dir: filepath.Join(dir, e.Name())})
188 }
189 }
190
191 var wg sync.WaitGroup
192 errCh := make(chan error, len(jobs))
193 for _, j := range jobs {
194 wg.Add(1)
195 go func(extDir string) {
196 defer wg.Done()
197 if err := m.loadOne(ctx, extDir); err != nil {
198 errCh <- fmt.Errorf("%s: %w", extDir, err)
199 }
200 }(j.dir)
201 }
202 wg.Wait()
203 close(errCh)
204
205 var errs []error
206 for e := range errCh {
207 errs = append(errs, e)
208 }
209 return errs
210}
211
212// searchDirs returns the directories the discoverer walks, in
213// priority order: project-local first (so a project can override
214// global behavior), then global.
215func (m *Manager) searchDirs() []string {
216 var dirs []string
217 if m.cwd != "" {
218 dirs = append(dirs, filepath.Join(m.cwd, ".zot", "extensions"))
219 }
220 if m.zotHome != "" {
221 dirs = append(dirs, filepath.Join(m.zotHome, "extensions"))
222 }
223 return dirs
224}
225
226// loadOne reads a single extension's manifest and, if enabled,
227// spawns its subprocess + completes the hello handshake.
228func (m *Manager) loadOne(ctx context.Context, dir string) error {
229 manifestPath := filepath.Join(dir, "extension.json")
230 raw, err := os.ReadFile(manifestPath)
231 if err != nil {
232 return fmt.Errorf("read manifest: %w", err)
233 }
234 var mf Manifest
235 if err := json.Unmarshal(raw, &mf); err != nil {
236 return fmt.Errorf("parse manifest: %w", err)
237 }
238 if mf.Name == "" {
239 return errors.New("manifest: name is required")
240 }
241 if mf.Exec == "" {
242 return errors.New("manifest: exec is required")
243 }
244 if !mf.IsEnabled() {
245 // Quietly skip disabled extensions; zot ext list will show them.
246 return nil
247 }
248
249 m.mu.RLock()
250 _, dup := m.ext[mf.Name]
251 m.mu.RUnlock()
252 if dup {
253 // Project-local copy already won; ignore the global one.
254 return nil
255 }
256
257 ext := &Extension{
258 Manifest: mf,
259 Dir: dir,
260 readyCh: make(chan struct{}),
261 pending: map[string]chan extproto.CommandResponseFromExt{},
262 pendingTool: map[string]chan extproto.ToolResultFromExt{},
263 pendingIntercept: map[string]chan extproto.EventInterceptResponseFromExt{},
264 eventSubs: map[string]struct{}{},
265 interceptSubs: map[string]struct{}{},
266 }
267 if err := m.spawn(ctx, ext); err != nil {
268 return err
269 }
270
271 m.mu.Lock()
272 m.ext[mf.Name] = ext
273 // Note: ext.commands and ext.tools may be empty here — they're
274 // populated by the read loop as register_* frames arrive after
275 // hello. Indexing happens in the read loop too. Discover()'s
276 // caller can WaitForReady() before relying on the registries.
277 m.mu.Unlock()
278 return nil
279}
280
281// LoadExplicit loads each path as an ad-hoc extension. Used for
282// `zot --ext <path>` so extension authors can iterate on a working
283// copy without having to `zot ext install` after every change.
284//
285// Loaded BEFORE Discover so explicit paths win on name conflicts
286// against installed extensions. Spawns happen in parallel like the
287// regular discovery path; errors are returned per path.
288func (m *Manager) LoadExplicit(ctx context.Context, paths []string) []error {
289 if len(paths) == 0 {
290 return nil
291 }
292
293 var wg sync.WaitGroup
294 errCh := make(chan error, len(paths))
295 absPaths := make([]string, 0, len(paths))
296 for _, p := range paths {
297 abs, err := filepath.Abs(p)
298 if err != nil {
299 errCh <- fmt.Errorf("%s: %w", p, err)
300 continue
301 }
302 absPaths = append(absPaths, abs)
303 wg.Add(1)
304 go func(extDir string) {
305 defer wg.Done()
306 if err := m.loadOne(ctx, extDir); err != nil {
307 errCh <- fmt.Errorf("%s: %w", extDir, err)
308 }
309 }(abs)
310 }
311 wg.Wait()
312 close(errCh)
313
314 m.mu.Lock()
315 m.explicitPaths = append(m.explicitPaths, absPaths...)
316 m.mu.Unlock()
317
318 var errs []error
319 for e := range errCh {
320 errs = append(errs, e)
321 }
322 return errs
323}
324
325// SetOnReload registers a callback fired after a successful Reload.
326// Hosts use it to rebuild the agent's tool registry with freshly-
327// registered extension tools.
328func (m *Manager) SetOnReload(fn func()) {
329 m.mu.Lock()
330 m.onReload = fn
331 m.mu.Unlock()
332}
333
334// ReloadStats summarises the outcome of Reload.
335type ReloadStats struct {
336 Stopped int // how many old processes were torn down
337 Loaded int // how many new processes reached spawn
338 Ready int // how many of those signalled ready in time
339 Errors []error // non-fatal per-extension errors
340}
341
342// Reload tears down every running extension, re-reads the manifests
343// from disk, respawns everyone (including the --ext paths remembered
344// from LoadExplicit), waits up to grace for ready signals, and
345// invokes the SetOnReload callback so the host can rebuild its tool
346// registry. The manager's internal maps are cleared before the
347// new load to ensure a clean slate.
348//
349// Safe to call concurrently with normal host operations: the lock is
350// released between stop and respawn so pending InvokeTool / Invoke
351// calls on the old processes get a clean error as their stdin
352// closes.
353func (m *Manager) Reload(ctx context.Context, grace time.Duration) ReloadStats {
354 stats := ReloadStats{}
355
356 // Snapshot and remember the explicit paths before we wipe state.
357 m.mu.Lock()
358 old := m.ext
359 explicit := append([]string(nil), m.explicitPaths...)
360 stats.Stopped = len(old)
361 m.ext = map[string]*Extension{}
362 m.commandIndex = map[string]*Extension{}
363 m.toolIndex = map[string]*Extension{}
364 m.explicitPaths = nil
365 callback := m.onReload
366 m.mu.Unlock()
367
368 // Graceful stop of the old set (reuses Stop's shutdown logic,
369 // but Stop re-reads m.ext which is now empty, so we replicate
370 // the small shutdown loop here on the snapshot).
371 for _, ext := range old {
372 if frame, err := extproto.Encode(extproto.ShutdownFromHost{Type: "shutdown"}); err == nil {
373 _, _ = ext.stdin.Write(frame)
374 }
375 _ = ext.stdin.Close()
376 }
377 deadline := time.Now().Add(grace)
378 for _, ext := range old {
379 remaining := time.Until(deadline)
380 if remaining <= 0 {
381 remaining = 100 * time.Millisecond
382 }
383 done := make(chan struct{})
384 go func(e *Extension) { _ = e.cmd.Wait(); close(done) }(ext)
385 select {
386 case <-done:
387 case <-time.After(remaining):
388 _ = ext.cmd.Process.Signal(syscall.SIGTERM)
389 select {
390 case <-done:
391 case <-time.After(time.Second):
392 _ = ext.cmd.Process.Kill()
393 <-done
394 }
395 }
396 if ext.logFile != nil {
397 _ = ext.logFile.Close()
398 }
399 }
400
401 // Fresh load. Explicit paths first (they still win on conflict).
402 if errs := m.LoadExplicit(ctx, explicit); len(errs) > 0 {
403 stats.Errors = append(stats.Errors, errs...)
404 }
405 if errs := m.Discover(ctx); len(errs) > 0 {
406 stats.Errors = append(stats.Errors, errs...)
407 }
408
409 m.mu.RLock()
410 stats.Loaded = len(m.ext)
411 m.mu.RUnlock()
412
413 // Wait for ready frames. Use the same 3s grace zot uses at
414 // startup so the reload feels no slower than a cold boot.
415 readyDeadline := time.Now().Add(grace)
416 if time.Until(readyDeadline) < 3*time.Second {
417 readyDeadline = time.Now().Add(3 * time.Second)
418 }
419 m.WaitForReady(time.Until(readyDeadline))
420
421 m.mu.RLock()
422 for _, ext := range m.ext {
423 select {
424 case <-ext.readyCh:
425 stats.Ready++
426 default:
427 }
428 }
429 m.mu.RUnlock()
430
431 if callback != nil {
432 callback()
433 }
434
435 return stats
436}
437
438// WaitForReady blocks until every loaded extension has signalled
439// ReadyFromExt, or the grace period expires for the slowest one.
440//
441// Waits run in parallel: total time is max(per-extension wait), not
442// sum. Without this, a single slow extension (e.g. `npx tsx` cold)
443// would gate every other extension's wait too and zot startup would
444// scale linearly with the number of slow runtimes installed.
445//
446// Call after Discover and before relying on tool registrations.
447func (m *Manager) WaitForReady(grace time.Duration) {
448 m.mu.RLock()
449 exts := make([]*Extension, 0, len(m.ext))
450 for _, e := range m.ext {
451 exts = append(exts, e)
452 }
453 m.mu.RUnlock()
454
455 deadline := time.After(grace)
456 var wg sync.WaitGroup
457 for _, ext := range exts {
458 wg.Add(1)
459 go func(ext *Extension) {
460 defer wg.Done()
461 select {
462 case <-ext.readyCh:
463 case <-deadline:
464 fmt.Fprintf(ext.logFile, "[zot] timed out waiting for ready frame; proceeding\n")
465 ext.readyOnce.Do(func() { close(ext.readyCh) })
466 }
467 }(ext)
468 }
469 wg.Wait()
470}
471
472// spawn launches the subprocess, hooks up pipes, logs stderr, and
473// runs the synchronous portion of the hello handshake. Asynchronous
474// frames are processed in a goroutine started here.
475func (m *Manager) spawn(ctx context.Context, ext *Extension) error {
476 logsDir := filepath.Join(m.zotHome, "logs")
477 _ = os.MkdirAll(logsDir, 0o755)
478 logPath := filepath.Join(logsDir, "ext-"+ext.Manifest.Name+".log")
479 logFile, err := os.OpenFile(logPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
480 if err != nil {
481 return fmt.Errorf("open log: %w", err)
482 }
483 ext.LogPath = logPath
484 ext.logFile = logFile
485 fmt.Fprintf(logFile, "\n[zot] starting %s/%s at %s\n", ext.Manifest.Name, ext.Manifest.Version, time.Now().Format(time.RFC3339))
486
487 // Exec resolution rules:
488 // - absolute path: used as-is.
489 // - starts with "." (./ or ../): resolved relative to ext.Dir.
490 // - bare name (no path separator): looked up via $PATH so
491 // "node", "npx", "python3",
492 // "tsx" etc. work without
493 // forcing absolute paths.
494 // - other relative form (foo/bar): resolved relative to ext.Dir.
495 execPath := ext.Manifest.Exec
496 switch {
497 case filepath.IsAbs(execPath):
498 // keep
499 case strings.HasPrefix(execPath, "."+string(filepath.Separator)) ||
500 strings.HasPrefix(execPath, ".."+string(filepath.Separator)) ||
501 execPath == "." || execPath == "..":
502 execPath = filepath.Join(ext.Dir, execPath)
503 case strings.ContainsRune(execPath, filepath.Separator):
504 execPath = filepath.Join(ext.Dir, execPath)
505 default:
506 // bare name: leave as-is for exec.LookPath via exec.Command.
507 }
508 cmd := exec.CommandContext(ctx, execPath, ext.Manifest.Args...)
509 cmd.Dir = ext.Dir
510 cmd.Stderr = logFile
511
512 stdin, err := cmd.StdinPipe()
513 if err != nil {
514 return err
515 }
516 stdout, err := cmd.StdoutPipe()
517 if err != nil {
518 return err
519 }
520 if err := cmd.Start(); err != nil {
521 return fmt.Errorf("spawn: %w", err)
522 }
523 ext.cmd = cmd
524 ext.stdin = stdin
525 ext.stdout = stdout
526
527 // Hello handshake. Read the extension's HelloFromExt synchronously
528 // so we can fail fast on a broken extension; everything after is
529 // processed in the read goroutine.
530 scanner := bufio.NewScanner(stdout)
531 scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
532 if !scanner.Scan() {
533 return fmt.Errorf("extension exited before hello: %w", scanner.Err())
534 }
535 var hello extproto.HelloFromExt
536 if err := json.Unmarshal(scanner.Bytes(), &hello); err != nil {
537 return fmt.Errorf("parse hello: %w", err)
538 }
539 if hello.Type != "hello" || hello.Name == "" {
540 return fmt.Errorf("first frame must be hello (got %q)", hello.Type)
541 }
542 // Trust the manifest's name; ignore mismatch from the hello.
543 ext.helloAck = true
544
545 ack, _ := extproto.Encode(extproto.HelloAckFromHost{
546 Type: "hello_ack",
547 ProtocolVersion: extproto.ProtocolVersion,
548 ZotVersion: m.zotVersion,
549 Provider: m.provider,
550 Model: m.model,
551 CWD: m.cwd,
552 })
553 if _, err := stdin.Write(ack); err != nil {
554 return fmt.Errorf("send hello_ack: %w", err)
555 }
556
557 // Spin up the read loop now that the handshake is done.
558 go m.readLoop(ext, scanner)
559
560 // Compatibility shim: extensions built against the phase-1 SDK
561 // don't send a ready frame. Watch the read loop's frame arrival
562 // rate; if nothing's arrived for readyIdleWindow we treat the
563 // extension as ready so WaitForReady doesn't burn the full grace
564 // on every startup. Newer extensions still trigger the explicit
565 // path on their own ready frame.
566 go m.assumeReadyAfterIdle(ext)
567
568 return nil
569}
570
571// readyIdleWindow is how long the manager waits for a frame after
572// hello before assuming an extension that doesn't send `ready` is
573// nevertheless ready. 250ms is enough for any well-behaved native
574// binary to flush its register frames; slow runtimes (npx tsx) flush
575// even faster once they've started, so this rarely affects them.
576const readyIdleWindow = 250 * time.Millisecond
577
578func (m *Manager) assumeReadyAfterIdle(ext *Extension) {
579 ext.mu.Lock()
580 last := ext.lastFrameTime
581 ext.mu.Unlock()
582 for {
583 select {
584 case <-ext.readyCh:
585 return
586 case <-time.After(readyIdleWindow):
587 }
588 ext.mu.Lock()
589 current := ext.lastFrameTime
590 ext.mu.Unlock()
591 if current.Equal(last) {
592 // No new frame in the idle window. Treat as ready.
593 ext.readyOnce.Do(func() {
594 fmt.Fprintf(ext.logFile, "[zot] no ready frame; auto-readying after idle (legacy SDK?)\n")
595 close(ext.readyCh)
596 })
597 return
598 }
599 last = current
600 }
601}
602
603// readLoop processes every frame the extension sends after hello.
604// Returns when stdout closes.
605func (m *Manager) readLoop(ext *Extension, scanner *bufio.Scanner) {
606 defer func() {
607 // On close, drop every command + tool this extension owned so
608 // future invocations don't dangle. The subprocess is gone; we
609 // won't hear back about its commands or tool calls anymore.
610 m.mu.Lock()
611 for name, owner := range m.commandIndex {
612 if owner == ext {
613 delete(m.commandIndex, name)
614 }
615 }
616 for name, owner := range m.toolIndex {
617 if owner == ext {
618 delete(m.toolIndex, name)
619 }
620 }
621 m.mu.Unlock()
622 ext.readyOnce.Do(func() { close(ext.readyCh) })
623 fmt.Fprintf(ext.logFile, "[zot] extension %s read loop exited at %s\n", ext.Manifest.Name, time.Now().Format(time.RFC3339))
624 }()
625
626 for scanner.Scan() {
627 line := scanner.Bytes()
628 ext.mu.Lock()
629 ext.lastFrameTime = time.Now()
630 ext.mu.Unlock()
631 var frame extproto.Frame
632 if err := json.Unmarshal(line, &frame); err != nil {
633 fmt.Fprintf(ext.logFile, "[zot] malformed json from extension: %v\n", err)
634 continue
635 }
636 switch frame.Type {
637 case "register_command":
638 var rc extproto.RegisterCommandFromExt
639 if err := json.Unmarshal(line, &rc); err == nil {
640 // Both ext.commands and m.commandIndex are read by
641 // the public Commands() / HasCommand() helpers under
642 // m.mu, so the writes have to take the same lock to
643 // keep the race detector happy.
644 m.mu.Lock()
645 ext.commands = append(ext.commands, rc)
646 if _, exists := m.commandIndex[rc.Name]; !exists {
647 m.commandIndex[rc.Name] = ext
648 }
649 m.mu.Unlock()
650 }
651 case "register_tool":
652 var rt extproto.RegisterToolFromExt
653 if err := json.Unmarshal(line, &rt); err != nil {
654 fmt.Fprintf(ext.logFile, "[zot] bad register_tool frame: %v\n", err)
655 continue
656 }
657 // Validate the schema parses as JSON. If not, refuse to
658 // register — a broken schema confuses the model.
659 if len(rt.Schema) > 0 {
660 var tmp any
661 if err := json.Unmarshal(rt.Schema, &tmp); err != nil {
662 fmt.Fprintf(ext.logFile, "[zot] tool %q: schema is not valid json (%v); skipped\n", rt.Name, err)
663 continue
664 }
665 }
666 m.mu.Lock()
667 ext.tools = append(ext.tools, rt)
668 if _, exists := m.toolIndex[rt.Name]; !exists {
669 m.toolIndex[rt.Name] = ext
670 }
671 m.mu.Unlock()
672 case "ready":
673 ext.readyOnce.Do(func() { close(ext.readyCh) })
674 case "subscribe":
675 var sub extproto.SubscribeFromExt
676 if err := json.Unmarshal(line, &sub); err == nil {
677 ext.mu.Lock()
678 for _, ev := range sub.Events {
679 ext.eventSubs[ev] = struct{}{}
680 }
681 for _, ev := range sub.Intercept {
682 switch ev {
683 case "tool_call", "turn_start", "assistant_message":
684 ext.interceptSubs[ev] = struct{}{}
685 }
686 }
687 ext.mu.Unlock()
688 }
689 case "event_intercept_response":
690 var er extproto.EventInterceptResponseFromExt
691 if err := json.Unmarshal(line, &er); err == nil {
692 ext.mu.Lock()
693 ch, ok := ext.pendingIntercept[er.ID]
694 if ok {
695 delete(ext.pendingIntercept, er.ID)
696 }
697 ext.mu.Unlock()
698 if ok {
699 select {
700 case ch <- er:
701 default:
702 }
703 }
704 }
705 case "tool_result":
706 var tr extproto.ToolResultFromExt
707 if err := json.Unmarshal(line, &tr); err == nil {
708 ext.mu.Lock()
709 ch, ok := ext.pendingTool[tr.ID]
710 if ok {
711 delete(ext.pendingTool, tr.ID)
712 }
713 ext.mu.Unlock()
714 if ok {
715 select {
716 case ch <- tr:
717 default:
718 }
719 }
720 }
721 case "notify":
722 var n extproto.NotifyFromExt
723 if err := json.Unmarshal(line, &n); err == nil {
724 m.hooks.Notify(ext.Manifest.Name, n.Level, n.Message)
725 }
726 case "command_response":
727 var cr extproto.CommandResponseFromExt
728 if err := json.Unmarshal(line, &cr); err == nil {
729 ext.mu.Lock()
730 ch, ok := ext.pending[cr.ID]
731 if ok {
732 delete(ext.pending, cr.ID)
733 }
734 ext.mu.Unlock()
735 if ok {
736 select {
737 case ch <- cr:
738 default:
739 }
740 }
741 }
742 case "shutdown_ack":
743 // Caller of Stop is waiting on the process exit, not this frame.
744 default:
745 fmt.Fprintf(ext.logFile, "[zot] unknown frame type %q\n", frame.Type)
746 }
747 }
748}
749
750// Commands returns a snapshot of every (extension, command) pair
751// currently registered. Used by the slash autocomplete + /help.
752func (m *Manager) Commands() []CommandInfo {
753 m.mu.RLock()
754 defer m.mu.RUnlock()
755 var out []CommandInfo
756 for _, ext := range m.ext {
757 for _, c := range ext.commands {
758 out = append(out, CommandInfo{
759 Extension: ext.Manifest.Name,
760 Name: c.Name,
761 Description: c.Description,
762 })
763 }
764 }
765 return out
766}
767
768// CommandInfo is one extension-registered slash command, surfaced to
769// the rest of zot for display purposes.
770type CommandInfo struct {
771 Extension string
772 Name string
773 Description string
774}
775
776// ToolInfo is one extension-registered tool. Used by the agent's
777// build step to materialise core.Tool wrappers.
778type ToolInfo struct {
779 Extension string
780 Name string
781 Description string
782 Schema json.RawMessage
783}
784
785// Tools returns a snapshot of every (extension, tool) pair currently
786// registered. Used at agent-build time to fold extension tools into
787// the runtime tool registry.
788func (m *Manager) Tools() []ToolInfo {
789 m.mu.RLock()
790 defer m.mu.RUnlock()
791 var out []ToolInfo
792 for _, ext := range m.ext {
793 for _, t := range ext.tools {
794 out = append(out, ToolInfo{
795 Extension: ext.Manifest.Name,
796 Name: t.Name,
797 Description: t.Description,
798 Schema: t.Schema,
799 })
800 }
801 }
802 return out
803}
804
805// HasTool reports whether name is registered by any extension.
806func (m *Manager) HasTool(name string) bool {
807 m.mu.RLock()
808 defer m.mu.RUnlock()
809 _, ok := m.toolIndex[name]
810 return ok
811}
812
813// InvokeTool sends a tool_call to the owning extension and waits for
814// the matching tool_result. Used by the core.Tool wrapper that the
815// agent registers per extension-defined tool.
816func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMessage, timeout time.Duration) (extproto.ToolResultFromExt, error) {
817 m.mu.RLock()
818 ext, ok := m.toolIndex[name]
819 m.mu.RUnlock()
820 if !ok {
821 return extproto.ToolResultFromExt{}, fmt.Errorf("no extension registered for tool %q", name)
822 }
823
824 id := newCorrelationID()
825 ch := make(chan extproto.ToolResultFromExt, 1)
826 ext.mu.Lock()
827 ext.pendingTool[id] = ch
828 ext.mu.Unlock()
829
830 frame, _ := extproto.Encode(extproto.ToolCallFromHost{
831 Type: "tool_call",
832 ID: id,
833 Name: name,
834 Args: args,
835 })
836 if _, err := ext.stdin.Write(frame); err != nil {
837 ext.mu.Lock()
838 delete(ext.pendingTool, id)
839 ext.mu.Unlock()
840 return extproto.ToolResultFromExt{}, fmt.Errorf("write: %w", err)
841 }
842
843 select {
844 case resp := <-ch:
845 return resp, nil
846 case <-time.After(timeout):
847 ext.mu.Lock()
848 delete(ext.pendingTool, id)
849 ext.mu.Unlock()
850 return extproto.ToolResultFromExt{}, fmt.Errorf("timeout waiting for %s/%s", ext.Manifest.Name, name)
851 case <-ctx.Done():
852 ext.mu.Lock()
853 delete(ext.pendingTool, id)
854 ext.mu.Unlock()
855 return extproto.ToolResultFromExt{}, ctx.Err()
856 }
857}
858
859// HasCommand reports whether name is registered by any extension.
860func (m *Manager) HasCommand(name string) bool {
861 m.mu.RLock()
862 defer m.mu.RUnlock()
863 _, ok := m.commandIndex[name]
864 return ok
865}
866
867// Invoke fires the named slash command's handler in the owning
868// extension and waits up to timeout for the response. Returns the
869// extension's CommandResponse so the caller can act on the action
870// (prompt / insert / display).
871func (m *Manager) Invoke(ctx context.Context, name, args string, timeout time.Duration) (extproto.CommandResponseFromExt, error) {
872 m.mu.RLock()
873 ext, ok := m.commandIndex[name]
874 m.mu.RUnlock()
875 if !ok {
876 return extproto.CommandResponseFromExt{}, fmt.Errorf("no extension registered for /%s", name)
877 }
878
879 id := newCorrelationID()
880 ch := make(chan extproto.CommandResponseFromExt, 1)
881 ext.mu.Lock()
882 ext.pending[id] = ch
883 ext.mu.Unlock()
884
885 frame, _ := extproto.Encode(extproto.CommandInvokedFromHost{
886 Type: "command_invoked",
887 ID: id,
888 Name: name,
889 Args: args,
890 })
891 if _, err := ext.stdin.Write(frame); err != nil {
892 ext.mu.Lock()
893 delete(ext.pending, id)
894 ext.mu.Unlock()
895 return extproto.CommandResponseFromExt{}, fmt.Errorf("write: %w", err)
896 }
897
898 select {
899 case resp := <-ch:
900 return resp, nil
901 case <-time.After(timeout):
902 ext.mu.Lock()
903 delete(ext.pending, id)
904 ext.mu.Unlock()
905 return extproto.CommandResponseFromExt{}, fmt.Errorf("timeout waiting for %s/%s", ext.Manifest.Name, name)
906 case <-ctx.Done():
907 ext.mu.Lock()
908 delete(ext.pending, id)
909 ext.mu.Unlock()
910 return extproto.CommandResponseFromExt{}, ctx.Err()
911 }
912}
913
914// Stop cleanly terminates every extension. Sends ShutdownFromHost,
915// waits up to gracePeriod for each subprocess to exit, then SIGTERMs
916// (and SIGKILLs after another second) the holdouts.
917func (m *Manager) Stop(gracePeriod time.Duration) {
918 m.mu.RLock()
919 exts := make([]*Extension, 0, len(m.ext))
920 for _, e := range m.ext {
921 exts = append(exts, e)
922 }
923 m.mu.RUnlock()
924
925 for _, ext := range exts {
926 if frame, err := extproto.Encode(extproto.ShutdownFromHost{Type: "shutdown"}); err == nil {
927 _, _ = ext.stdin.Write(frame)
928 }
929 _ = ext.stdin.Close()
930 }
931
932 deadline := time.Now().Add(gracePeriod)
933 for _, ext := range exts {
934 remaining := time.Until(deadline)
935 if remaining <= 0 {
936 remaining = 100 * time.Millisecond
937 }
938 done := make(chan struct{})
939 go func() { _ = ext.cmd.Wait(); close(done) }()
940 select {
941 case <-done:
942 case <-time.After(remaining):
943 _ = ext.cmd.Process.Signal(syscall.SIGTERM)
944 select {
945 case <-done:
946 case <-time.After(time.Second):
947 _ = ext.cmd.Process.Kill()
948 <-done
949 }
950 }
951 if ext.logFile != nil {
952 _ = ext.logFile.Close()
953 }
954 }
955}
956
957// All returns every extension currently tracked, enabled or not.
958// Used by `zot ext list`.
959func (m *Manager) All() []*Extension {
960 m.mu.RLock()
961 defer m.mu.RUnlock()
962 out := make([]*Extension, 0, len(m.ext))
963 for _, e := range m.ext {
964 out = append(out, e)
965 }
966 return out
967}
968
969// newCorrelationID returns a short non-cryptographic id. We don't
970// need uniqueness across processes, just within the lifetime of one
971// extension's pending map.
972func newCorrelationID() string {
973 return strings.ReplaceAll(time.Now().Format("150405.000000"), ".", "")
974}