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