1// Package swarm implements zot's multi-agent supervisor.
2//
3// A Swarm manages a set of headless zot subprocesses ("agents")
4// working in their own git worktrees. The interactive TUI exposes
5// the supervisor through the /swarm slash command and a dashboard
6// dialog; non-TUI code can drive it directly through this package.
7//
8// Each Agent has:
9// - a unique id (short slug + nanoseconds)
10// - a branch name on a fresh git worktree under <root>/swarm/<id>
11// - a Runner (the thing that actually executes the task)
12// - a Status string + Activity string that the dashboard reads
13//
14// The Runner abstraction means tests can swap a fake in instead of
15// really spawning a subprocess; the production Runner shells out to
16// `zot --print` so we reuse zot's own model resolution and tooling
17// without re-implementing the agent loop.
18package swarm
19
20import (
21 "context"
22 "errors"
23 "fmt"
24 "os"
25 "path/filepath"
26 "sort"
27 "strings"
28 "sync"
29 "time"
30)
31
32// Status is the high-level lifecycle state of an Agent.
33type Status string
34
35const (
36 StatusPending Status = "pending" // created, not started yet
37 StatusRunning Status = "running" // Runner.Run is in flight
38 StatusDone Status = "done" // Runner.Run returned nil
39 StatusFailed Status = "failed" // Runner.Run returned an error
40 StatusKilled Status = "killed" // Stop() called before completion
41 StatusDetached Status = "detached" // reloaded from disk; no live runner
42)
43
44// Config configures a Swarm.
45type Config struct {
46 // Root is the directory under which worktrees + state files live.
47 // Typically <ZotHome>/swarm, but tests pass a tempdir.
48 Root string
49
50 // RepoRoot is the path to the user's git repo (CWD of the parent
51 // zot). New worktrees branch off this repo.
52 RepoRoot string
53
54 // Worktree creates the per-agent working directory. If nil, the
55 // default git-worktree implementation is used.
56 Worktree WorktreeManager
57
58 // NewRunner produces the Runner for an Agent. If nil, the default
59 // `zot --print` exec runner is used. Tests inject a fake here.
60 NewRunner func(a *Agent) Runner
61
62 // Now is a clock seam for tests; defaults to time.Now.
63 Now func() time.Time
64}
65
66// WorktreeManager creates and removes per-agent working directories.
67type WorktreeManager interface {
68 Create(id, branch, base string) (dir string, err error)
69 Remove(id, dir string) error
70}
71
72// Runner executes one agent task. Run blocks until the task finishes,
73// is cancelled via ctx, or hits an unrecoverable error.
74//
75// Run should report progress by writing short human-readable strings
76// to the activity channel and final transcript text to transcript.
77// Both channels are non-blocking sinks owned by the Swarm; if the
78// dashboard isn't reading, sends are dropped.
79type Runner interface {
80 Run(ctx context.Context, sink Sink) error
81}
82
83// Sink is how a Runner reports activity and transcript back to the
84// supervisor. All methods are safe to call from any goroutine and
85// never block.
86type Sink interface {
87 // Activity sets the one-line "what is this agent doing right now"
88 // string shown in the dashboard.
89 Activity(msg string)
90 // Transcript appends a chunk of agent output (typically a final
91 // assistant message) to the agent's running transcript.
92 Transcript(chunk string)
93}
94
95// Swarm supervises a set of Agents.
96type Swarm struct {
97 cfg Config
98
99 mu sync.Mutex
100 agents map[string]*Agent
101 order []string // creation order for stable listing
102
103 // activeSession is the host session id the dashboard is
104 // currently scoped to. When non-empty, SnapshotAll filters out
105 // agents whose SessionID doesn't match (legacy / unscoped
106 // agents with SessionID == "" are always shown). Spawn stamps
107 // new agents with this value so they appear only in the session
108 // that created them. When empty (the default), the historical
109 // "show everything" behaviour is preserved — important for
110 // tests and any scripted use of the Swarm that doesn't bother
111 // with sessions.
112 activeSession string
113}
114
115// New constructs a Swarm from cfg. Missing config fields are filled
116// with defaults.
117func New(cfg Config) *Swarm {
118 if cfg.Now == nil {
119 cfg.Now = time.Now
120 }
121 if cfg.Worktree == nil {
122 cfg.Worktree = &gitWorktree{root: filepath.Join(cfg.Root, "worktrees"), repo: cfg.RepoRoot}
123 }
124 if cfg.NewRunner == nil {
125 cfg.NewRunner = func(a *Agent) Runner { return &execRunner{agent: a} }
126 }
127 return &Swarm{
128 cfg: cfg,
129 agents: map[string]*Agent{},
130 }
131}
132
133// SetActiveSession scopes the dashboard view (and Spawn stamping)
134// to a particular host zot session id. Pass empty to clear the
135// scope and revert to "show every agent" (the original behaviour).
136//
137// Existing in-memory agents keep their SessionID; only the filter
138// applied at snapshot time changes. So swapping the active session
139// with /sessions instantly re-narrows the dashboard without
140// touching any agent state.
141func (f *Swarm) SetActiveSession(id string) {
142 f.mu.Lock()
143 defer f.mu.Unlock()
144 f.activeSession = id
145}
146
147// ActiveSession returns the current scope, mostly for tests and
148// diagnostics. Empty means "no scope; show everything".
149func (f *Swarm) ActiveSession() string {
150 f.mu.Lock()
151 defer f.mu.Unlock()
152 return f.activeSession
153}
154
155// agentStateDir is the per-agent state directory laid out as:
156//
157// <root>/agents/<id>/
158// events.jsonl durable event log (runner-owned)
159// in.sock unix socket inbox (child-owned)
160// session.json persistent agent session (child-owned)
161// meta.json static metadata (id, task, branch, dir)
162func (f *Swarm) agentStateDir(id string) string {
163 return filepath.Join(f.cfg.Root, "agents", id)
164}
165
166// SpawnRequest configures a Spawn. Only Task is required; the rest
167// are optional. Model + Provider, when set, get baked into the
168// child argv as --model / --provider so the agent runs against the
169// chosen model regardless of the parent's current selection.
170type SpawnRequest struct {
171 Task string
172 Model string // optional override; child resolves default if empty
173 Provider string // optional override; usually paired with Model
174}
175
176// Spawn creates a new Agent for the given task, sets up its worktree,
177// allocates the on-disk state directory (events log, inbox socket
178// path, session file path), and starts the Runner on a background
179// goroutine. The returned Agent is already in StatusRunning (or
180// StatusFailed if worktree setup failed before the goroutine
181// started). This is the historical signature; callers that want to
182// override the child's model use SpawnReq instead.
183func (f *Swarm) Spawn(ctx context.Context, task string) (*Agent, error) {
184 return f.SpawnReq(ctx, SpawnRequest{Task: task})
185}
186
187// SpawnReq is the full-fat variant of Spawn that accepts a
188// SpawnRequest. Existing callers can keep using Spawn; new code that
189// wants to pin the child's model uses this.
190func (f *Swarm) SpawnReq(ctx context.Context, req SpawnRequest) (*Agent, error) {
191 task := strings.TrimSpace(req.Task)
192 if task == "" {
193 return nil, errors.New("swarm: empty task")
194 }
195 id := newAgentID(task, f.cfg.Now())
196 branch := "swarm/" + id
197 dir, err := f.cfg.Worktree.Create(id, branch, f.cfg.RepoRoot)
198 if err != nil {
199 return nil, fmt.Errorf("worktree create: %w", err)
200 }
201
202 stateDir := f.agentStateDir(id)
203 if err := os.MkdirAll(stateDir, 0o755); err != nil {
204 return nil, fmt.Errorf("swarm state dir: %w", err)
205 }
206 logPath := filepath.Join(stateDir, "events.jsonl")
207 sessionPath := filepath.Join(stateDir, "session.json")
208 // Unix sockets have a hard 104-byte path limit on darwin and 108
209 // on linux. Long ZOT_HOME paths plus an agent slug blow that cap
210 // quickly. Pick the shortest path that still keeps sockets
211 // per-swarm-root so two zot instances on the same machine don't
212 // collide. inboxSocketPath falls back from $TMPDIR to /tmp if
213 // neither is short enough.
214 inboxPath, err := inboxSocketPath(f.cfg.Root, id)
215 if err != nil {
216 return nil, fmt.Errorf("swarm inbox path: %w", err)
217 }
218
219 // Snapshot activeSession under the lock; we use it twice (struct
220 // init below + persistence). Reading it without the lock would
221 // race a concurrent SetActiveSession call.
222 f.mu.Lock()
223 sessionID := f.activeSession
224 f.mu.Unlock()
225
226 a := &Agent{
227 ID: id,
228 Task: task,
229 Branch: branch,
230 Dir: dir,
231 Started: f.cfg.Now(),
232 Model: strings.TrimSpace(req.Model),
233 Provider: strings.TrimSpace(req.Provider),
234 SessionID: sessionID,
235 InboxPath: inboxPath,
236 EventLogPath: logPath,
237 SessionPath: sessionPath,
238 inbox: NewInbox(inboxPath),
239 status: StatusPending,
240 activity: "queued",
241 done: make(chan struct{}),
242 }
243 a.ctx, a.cancel = context.WithCancel(ctx)
244 a.runner = f.cfg.NewRunner(a)
245
246 f.mu.Lock()
247 f.agents[id] = a
248 f.order = append(f.order, id)
249 f.mu.Unlock()
250
251 // Persist the agent's identity so a later `zot` invocation can
252 // reload it from disk via Swarm.Reload. Best-effort: if the disk
253 // is read-only we still let the runner start, the user just won't
254 // see this agent on the next launch.
255 _ = writeAgentMeta(stateDir, a)
256
257 go f.run(a)
258 return a, nil
259}
260
261// SendInput delivers msg to the agent's inbox. msg should already
262// include the leading kind ("user " + text, "cancel", "shutdown").
263// Returns an error wrapping ErrNotReady if the child hasn't opened
264// its listener yet; the caller can retry or surface the wait.
265func (f *Swarm) SendInput(id, msg string) error {
266 a := f.Get(id)
267 if a == nil {
268 return fmt.Errorf("no such agent %q", id)
269 }
270 if a.inbox == nil {
271 return fmt.Errorf("agent %s has no inbox", a.ID)
272 }
273 return a.inbox.SendInput(msg)
274}
275
276// SendUserTurn is sugar for the common "send the next user turn"
277// case. It quotes nothing and forwards verbatim; callers are
278// expected to have already trimmed and expanded the text.
279func (f *Swarm) SendUserTurn(id, text string) error {
280 return f.SendInput(id, "user "+text)
281}
282
283func (f *Swarm) run(a *Agent) {
284 a.setStatus(StatusRunning)
285 a.setActivity("starting")
286 err := a.runner.Run(a.ctx, agentSink{a: a})
287 a.mu.Lock()
288 a.finished = f.cfg.Now()
289 switch {
290 case a.status == StatusKilled:
291 // Already finalised by Stop.
292 case errors.Is(err, context.Canceled):
293 a.status = StatusKilled
294 a.activity = "cancelled"
295 case err != nil:
296 a.status = StatusFailed
297 a.activity = "error: " + truncate(err.Error(), 120)
298 a.lastErr = err
299 default:
300 a.status = StatusDone
301 a.activity = "done"
302 }
303 a.mu.Unlock()
304 close(a.done)
305}
306
307// List returns a snapshot of every agent in creation order. The
308// returned slice is a copy; callers may iterate without holding the
309// swarm lock. Agent fields are read under their own mutex during
310// formatting in Snapshot.
311func (f *Swarm) List() []*Agent {
312 f.mu.Lock()
313 defer f.mu.Unlock()
314 out := make([]*Agent, 0, len(f.order))
315 for _, id := range f.order {
316 out = append(out, f.agents[id])
317 }
318 return out
319}
320
321// Get returns the agent with the given (possibly truncated) id, or
322// nil. Matching is prefix-based so the user can type the first few
323// characters of a long id.
324func (f *Swarm) Get(id string) *Agent {
325 f.mu.Lock()
326 defer f.mu.Unlock()
327 if a, ok := f.agents[id]; ok {
328 return a
329 }
330 // Prefix match.
331 var hits []*Agent
332 for k, a := range f.agents {
333 if strings.HasPrefix(k, id) {
334 hits = append(hits, a)
335 }
336 }
337 if len(hits) == 1 {
338 return hits[0]
339 }
340 return nil
341}
342
343// Stop cancels the agent's context. The Runner should observe the
344// cancellation and return promptly; the goroutine then finalises the
345// agent in StatusKilled. Also closes the supervisor-side inbox handle
346// so any pending SendInput retries fail fast instead of dialing a
347// socket about to be unlinked.
348//
349// Stop is a no-op for any agent that's not in a live runnable state
350// — Done / Failed / Killed (already finalised) and Detached (no
351// in-process runner; reloaded from disk). Calling Stop on a detached
352// agent must not crash: buildDetachedAgent doesn't allocate a
353// context/cancel pair because there's nothing to cancel.
354func (f *Swarm) Stop(id string) error {
355 a := f.Get(id)
356 if a == nil {
357 return fmt.Errorf("no such agent %q", id)
358 }
359 a.mu.Lock()
360 switch a.status {
361 case StatusDone, StatusFailed, StatusKilled, StatusDetached:
362 a.mu.Unlock()
363 return nil
364 }
365 a.status = StatusKilled
366 a.activity = "stopped"
367 a.mu.Unlock()
368 // Belt-and-braces guard against the cancel-less case (detached
369 // agents skip this branch above, but a future code path that
370 // builds an Agent without a runner shouldn't crash the
371 // supervisor).
372 if a.cancel != nil {
373 a.cancel()
374 }
375 if a.inbox != nil {
376 _ = a.inbox.Close()
377 }
378 return nil
379}
380
381// StopAll cancels every running agent. Used on shutdown.
382func (f *Swarm) StopAll() {
383 for _, a := range f.List() {
384 _ = f.Stop(a.ID)
385 }
386}
387
388// Remove tears down the worktree for a terminated agent. It is an
389// error to remove an agent that's still running; call Stop first and
390// wait for the status to settle. Detached agents (reloaded from
391// disk) remove cleanly because they have no live runner racing for
392// the same files.
393func (f *Swarm) Remove(id string) error {
394 a := f.Get(id)
395 if a == nil {
396 return fmt.Errorf("no such agent %q", id)
397 }
398 a.mu.Lock()
399 st := a.status
400 dir := a.Dir
401 a.mu.Unlock()
402 if st == StatusRunning || st == StatusPending {
403 return fmt.Errorf("agent %s still %s", a.ID, st)
404 }
405 if err := f.cfg.Worktree.Remove(a.ID, dir); err != nil {
406 return err
407 }
408 // Best-effort cleanup of the per-agent state directory (meta.json,
409 // events.jsonl, session.json, in.sock if it's local). Worktree
410 // removal already succeeded; failing here would leave the user
411 // with no recourse, so swallow the error.
412 _ = os.RemoveAll(f.agentStateDir(a.ID))
413 f.mu.Lock()
414 delete(f.agents, a.ID)
415 for i, k := range f.order {
416 if k == a.ID {
417 f.order = append(f.order[:i], f.order[i+1:]...)
418 break
419 }
420 }
421 f.mu.Unlock()
422 return nil
423}
424
425// Snapshot returns a read-only view of one agent. Safe for the TUI
426// goroutine to call repeatedly; never blocks on the Runner.
427type AgentSnapshot struct {
428 ID string
429 Task string
430 Branch string
431 Dir string
432 Status Status
433 Activity string
434 Started time.Time
435 Finished time.Time
436 Err string
437 Tail string // last few transcript lines, joined with "\n"
438 Lines []string // full transcript (already capped by Agent.appendTranscript)
439
440 // Model and Provider expose the per-agent overrides set at
441 // Spawn time (empty when the agent inherits the child's default
442 // resolution). The dashboard surfaces these so the user can
443 // confirm which model an agent is running against.
444 Model string
445 Provider string
446
447 // Paths to the agent's durable state. Surface them in the
448 // snapshot so the dashboard / /swarm open can read events.jsonl
449 // or resume the session without going back through the Agent.
450 InboxPath string
451 EventLogPath string
452 SessionPath string
453}
454
455// Snapshot copies the live agent state into a value the caller can
456// inspect at leisure.
457func (a *Agent) Snapshot() AgentSnapshot {
458 a.mu.Lock()
459 defer a.mu.Unlock()
460 tail := strings.Join(lastN(a.transcript, 6), "\n")
461 lines := make([]string, len(a.transcript))
462 copy(lines, a.transcript)
463 errStr := ""
464 if a.lastErr != nil {
465 errStr = a.lastErr.Error()
466 }
467 return AgentSnapshot{
468 ID: a.ID, Task: a.Task, Branch: a.Branch, Dir: a.Dir,
469 Status: a.status, Activity: a.activity,
470 Started: a.Started, Finished: a.finished,
471 Err: errStr, Tail: tail, Lines: lines,
472 Model: a.Model,
473 Provider: a.Provider,
474 InboxPath: a.InboxPath,
475 EventLogPath: a.EventLogPath,
476 SessionPath: a.SessionPath,
477 }
478}
479
480// SnapshotAll returns snapshots of every agent in creation order,
481// scoped to the active session when one is set.
482//
483// Scoping rules:
484// - activeSession == "": no filter, every agent is returned
485// (historical behaviour; used by tests and scripted callers).
486// - activeSession != "": include only agents whose SessionID
487// matches activeSession OR is empty. The empty-id pass-through
488// keeps pre-upgrade agents (their meta.json was written before
489// session_id existed) visible from any session so the user
490// doesn't lose access after the schema bump.
491func (f *Swarm) SnapshotAll() []AgentSnapshot {
492 agents := f.List()
493 f.mu.Lock()
494 active := f.activeSession
495 f.mu.Unlock()
496
497 out := make([]AgentSnapshot, 0, len(agents))
498 for _, a := range agents {
499 if active != "" && a.SessionID != "" && a.SessionID != active {
500 continue
501 }
502 out = append(out, a.Snapshot())
503 }
504 // Sort by start time for a stable, deterministic listing.
505 sort.SliceStable(out, func(i, j int) bool { return out[i].Started.Before(out[j].Started) })
506 return out
507}
508
509// agentSink is the Sink the Swarm hands to each Runner.
510type agentSink struct{ a *Agent }
511
512func (s agentSink) Activity(msg string) { s.a.setActivity(msg) }
513func (s agentSink) Transcript(chunk string) { s.a.appendTranscript(chunk) }
514
515func truncate(s string, n int) string {
516 if len(s) <= n {
517 return s
518 }
519 return s[:n-1] + "…"
520}
521
522func lastN(lines []string, n int) []string {
523 if len(lines) <= n {
524 return lines
525 }
526 return lines[len(lines)-n:]
527}