1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "os"
8 "strings"
9 "sync"
10 "time"
11
12 "github.com/patriceckhart/zot/packages/agent/modes"
13 "github.com/patriceckhart/zot/packages/agent/swarm"
14 "github.com/patriceckhart/zot/packages/core"
15 "github.com/patriceckhart/zot/packages/provider"
16)
17
18// runSwarmAgentMode is the daemon-mode entry point used by every
19// swarm-spawned zot subprocess. It's intentionally close in shape to
20// runJSONMode but with two key differences:
21//
22// - Lifetime: the process stays alive across many user turns. The
23// initial positional task (if any) is the first turn; subsequent
24// turns arrive through the inbox unix socket at args.SwarmAgent.
25//
26// - Output: every emitted JSON line is also mirrored verbatim into
27// events.jsonl (see ZOT_SWARM_EVENT_LOG) so a separate zot
28// process can /swarm open this agent and replay its full history
29// even after the parent that spawned us is long gone.
30//
31// The runner in packages/agent/swarm/runner.go is the only caller in
32// production; tests use the stubchild binary under
33// packages/agent/swarm/testdata/cmd/stubchild instead of the real model
34// loop.
35func runSwarmAgentMode(ctx context.Context, args Args, version string) error {
36 if args.SwarmAgent == "" {
37 return fmt.Errorf("--swarm-agent requires a socket path")
38 }
39
40 r, err := Resolve(args, true)
41 if err != nil {
42 return err
43 }
44 extMgr, stopExt := setupNonInteractiveExtensions(ctx, args, &r, version)
45 defer stopExt()
46
47 ag := r.NewAgent()
48 wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
49 sess, _ := openOrCreateSession(args, r, ag, version)
50 defer sess.Close()
51
52 // Open the inbox listener BEFORE emitting agent_ready so the
53 // supervisor can dial through on the very first send. The
54 // swarm supervisor's Inbox.SendInput retries dialing for a
55 // short window, but emitting ready first and then listening
56 // would still race in tight loops.
57 ln, err := swarm.Listen(args.SwarmAgent)
58 if err != nil {
59 return fmt.Errorf("swarm-agent listen: %w", err)
60 }
61 defer ln.Close()
62
63 // Event log is owned by the supervisor's runner via stdout, but
64 // the daemon also writes a redundant copy here when the runner's
65 // pipe is closed (e.g. parent zot exited but the agent is still
66 // running headless). The env var is set by the runner; if it's
67 // empty we silently skip the second mirror.
68 var logMirror *swarm.EventLog
69 if path := os.Getenv("ZOT_SWARM_EVENT_LOG"); path != "" {
70 logMirror, _ = swarm.OpenEventLog(path)
71 }
72 if logMirror != nil {
73 defer logMirror.Close()
74 }
75
76 em := newSwarmEmitter(os.Stdout, logMirror)
77
78 em.emit("agent_ready", map[string]any{
79 "version": version,
80 "cwd": r.CWD,
81 "model": r.Model,
82 })
83
84 // turnCtx is the per-turn context. We keep a per-process
85 // cancel so the "cancel" inbox message can interrupt an
86 // in-flight turn without tearing down the whole daemon.
87 var (
88 mu sync.Mutex
89 turnCtx context.Context = ctx
90 cancelFn context.CancelFunc
91 busyTurn bool
92 turnNo int
93 shutdown = make(chan struct{})
94 )
95
96 runOne := func(prompt string) {
97 mu.Lock()
98 if busyTurn {
99 // Drop concurrent turns rather than queuing. The
100 // supervisor protocol assumes one outstanding turn per
101 // agent; if a user really wants to interrupt and start
102 // another, they should send "cancel" first.
103 mu.Unlock()
104 em.emit("error", map[string]any{"message": "agent busy; send 'cancel' first"})
105 return
106 }
107 busyTurn = true
108 turnNo++
109 step := turnNo
110 turnCtx, cancelFn = context.WithCancel(ctx)
111 c := turnCtx
112 mu.Unlock()
113
114 em.emit("turn_start", map[string]any{"step": step})
115
116 sink := func(ev core.AgentEvent) {
117 em.emit(ev.Type(), modes.EventToJSON(ev))
118 }
119
120 start := len(ag.Messages())
121 err := ag.Prompt(c, prompt, nil, sink)
122 WriteNewTranscript(ag, sess, start)
123
124 em.emit("turn_end", map[string]any{
125 "step": step,
126 "error": errString(err),
127 })
128
129 mu.Lock()
130 busyTurn = false
131 cancelFn = nil
132 mu.Unlock()
133 }
134
135 // Initial task: run before processing the inbox so the agent
136 // "starts working" the moment it boots, matching what users
137 // expect from `/swarm new <task>`.
138 if args.Prompt != "" {
139 go runOne(args.Prompt)
140 }
141
142 // Inbox loop: one supervisor message at a time. We don't spawn
143 // a goroutine per turn because runOne already serialises them
144 // via the busyTurn flag; doing the dispatch on the main
145 // goroutine keeps the daemon's lifecycle easy to follow.
146 for {
147 select {
148 case <-ctx.Done():
149 em.emit("agent_stopped", map[string]any{"reason": "cancelled"})
150 return ctx.Err()
151 case <-shutdown:
152 em.emit("agent_stopped", map[string]any{"reason": "shutdown"})
153 return nil
154 case msg, ok := <-ln.Lines():
155 if !ok {
156 em.emit("agent_stopped", map[string]any{"reason": "inbox-closed"})
157 return nil
158 }
159 switch {
160 case msg == "shutdown":
161 close(shutdown)
162 case msg == "cancel":
163 mu.Lock()
164 if cancelFn != nil {
165 cancelFn()
166 }
167 mu.Unlock()
168 case strings.HasPrefix(msg, "user "):
169 prompt := strings.TrimPrefix(msg, "user ")
170 go runOne(prompt)
171 default:
172 em.emit("error", map[string]any{
173 "message": "unknown supervisor message: " + truncateForLog(msg, 200),
174 })
175 }
176 }
177 }
178}
179
180// swarmEmitter serialises events to stdout and (optionally) to a
181// durable log file. Concurrent goroutines call emit so we have to
182// hold a mutex around the encoder.
183type swarmEmitter struct {
184 mu sync.Mutex
185 w *os.File
186 mirror *swarm.EventLog
187
188 // orphan flips true the first time a stdout write fails (broken
189 // pipe — the supervisor died). Until then the mirror stays
190 // dormant: the supervisor is the canonical writer to events.jsonl
191 // (it parses our stdout and Append()s each event itself). Writing
192 // from both sides used to land every event in the log twice,
193 // which showed up as a fully-duplicated transcript the next time
194 // the agent was reloaded — the exact "why is everything doubled"
195 // bug. Once orphaned the mirror takes over so the events still
196 // land on disk for the next reload.
197 orphan bool
198}
199
200func newSwarmEmitter(w *os.File, mirror *swarm.EventLog) *swarmEmitter {
201 return &swarmEmitter{w: w, mirror: mirror}
202}
203
204func (e *swarmEmitter) emit(typ string, data map[string]any) {
205 if data == nil {
206 data = map[string]any{}
207 }
208 data["type"] = typ
209 data["time"] = time.Now().Format(time.RFC3339Nano)
210
211 e.mu.Lock()
212 defer e.mu.Unlock()
213
214 if !e.orphan {
215 // Encode to a buffer first so we can detect the stdout write
216 // error — json.Encoder swallows partial writes silently, but
217 // a direct Write on *os.File returns broken-pipe immediately.
218 line, err := json.Marshal(data)
219 if err == nil {
220 line = append(line, '\n')
221 if _, werr := e.w.Write(line); werr != nil {
222 // Supervisor's stdout pipe is gone (parent zot exited
223 // but we kept running). Switch to mirror-only mode so
224 // subsequent events still get persisted; also retro-
225 // actively log this very event to the mirror so it
226 // doesn't get lost in the handoff.
227 e.orphan = true
228 }
229 }
230 }
231
232 if e.orphan && e.mirror != nil {
233 // Drop "type" + "time" out of data into Event fields so the
234 // mirror file matches what the supervisor would have written.
235 flat := map[string]any{}
236 for k, v := range data {
237 if k == "type" || k == "time" {
238 continue
239 }
240 flat[k] = v
241 }
242 _ = e.mirror.Append(swarm.NewEvent(typ, flat))
243 }
244}
245
246func errString(err error) string {
247 if err == nil {
248 return ""
249 }
250 return err.Error()
251}
252
253func truncateForLog(s string, n int) string {
254 if len(s) <= n {
255 return s
256 }
257 if n <= 3 {
258 return strings.Repeat(".", n)
259 }
260 return s[:n-3] + "..."
261}
262
263// _ keeps the provider import used; provider types may surface
264// through ag.OnEvent / modes.EventToJSON in future iterations.
265var _ provider.Content = provider.TextBlock{}