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