1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "io"
9 "os"
10 "strconv"
11 "strings"
12 "sync"
13 "time"
14 "unicode/utf8"
15
16 "github.com/bnema/zut/packages/agent/modes"
17 "github.com/bnema/zut/packages/agent/subagents"
18 "github.com/bnema/zut/packages/core"
19 "github.com/bnema/zut/packages/provider"
20)
21
22// runSubagentWorkerMode is the daemon-mode entry point used by every
23// subagent-spawned zut subprocess. It's intentionally close in shape to
24// runJSONMode but with two key differences:
25//
26// - Lifetime: the process stays alive across many user turns. The
27// initial positional task (if any) is the first turn; subsequent
28// turns arrive through the inbox unix socket at args.SubagentWorker.
29//
30// - Output: every emitted JSON line is also mirrored verbatim into
31// events.jsonl (see ZUT_SUBAGENT_EVENT_LOG) so a separate zut
32// process can /subagents open this agent and replay its full history
33// even after the parent that spawned us is long gone.
34//
35// The runner in packages/agent/subagents/runner.go is the only caller in
36// production; tests use the stubchild binary under
37// packages/agent/subagents/testdata/cmd/stubchild instead of the real model
38// loop.
39func runSubagentWorkerMode(ctx context.Context, args Args, version string) error {
40 if args.SubagentWorker == "" {
41 return fmt.Errorf("--subagent-worker requires a socket path")
42 }
43 if os.Getenv("ZUT_SUBAGENT_CREDENTIAL_STDIN") == "1" {
44 var inherited subagents.Credential
45 dec := json.NewDecoder(io.LimitReader(os.Stdin, 1<<20))
46 if err := dec.Decode(&inherited); err != nil {
47 return fmt.Errorf("read inherited subagent credential: %w", err)
48 }
49 if inherited.Value == "" || inherited.Method == "" {
50 return fmt.Errorf("read inherited subagent credential: missing value or method")
51 }
52 args.inheritedCredential = inherited.Value
53 args.inheritedAuthMethod = inherited.Method
54 args.inheritedAccountID = inherited.AccountID
55 }
56
57 r, err := Resolve(args, true)
58 if err != nil {
59 return err
60 }
61 extMgr, stopExt := setupNonInteractiveExtensions(ctx, args, &r, version)
62 defer stopExt()
63
64 ag := r.NewAgent()
65 initialAg := ag
66 defer func() {
67 closeAgentLSP(ag)
68 if ag != initialAg {
69 closeAgentLSP(initialAg)
70 }
71 }()
72 wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
73 sess, err := openOrCreateSession(args, r, ag, version)
74 if err != nil {
75 return err
76 }
77 if sess != nil {
78 var providerName, model string
79 sess, ag, providerName, model, err = applyInitialSessionResume(ctx, args, r, extMgr, sess, ag)
80 if err != nil {
81 return err
82 }
83 r.Provider, r.Model = providerName, model
84 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
85 defer sess.Close()
86 }
87 announceSession(extMgr, sess)
88
89 // Open the inbox listener BEFORE emitting agent_ready so the
90 // supervisor can dial through on the very first send. The
91 // subagent supervisor's Inbox retries dialing for a
92 // short window, but emitting ready first and then listening
93 // would still race in tight loops.
94 ln, err := subagents.Listen(args.SubagentWorker)
95 if err != nil {
96 return fmt.Errorf("subagent-worker listen: %w", err)
97 }
98 defer ln.Close()
99
100 // Event log is owned by the supervisor's runner via stdout, but
101 // the daemon also writes a redundant copy here when the runner's
102 // pipe is closed (e.g. parent zut exited but the agent is still
103 // running headless). The env var is set by the runner; if it's
104 // empty we silently skip the second mirror.
105 var logMirror *subagents.EventLog
106 if path := os.Getenv("ZUT_SUBAGENT_EVENT_LOG"); path != "" {
107 logMirror, _ = subagents.OpenEventLog(path)
108 }
109 if logMirror != nil {
110 defer logMirror.Close()
111 }
112
113 em := newSubagentEmitter(os.Stdout, logMirror)
114 em.setProtocolIdentity(os.Getenv("ZUT_SUBAGENT_AGENT_ID"))
115 em.emit("agent.ready", map[string]any{
116 "version": version,
117 "cwd": r.CWD,
118 "model": r.Model,
119 })
120
121 // Heartbeats make a live-but-detached child distinguishable from a
122 // failed turn after supervisor reconnect. The parent still treats
123 // unknown event types as forward-compatible.
124 heartbeatDone := make(chan struct{})
125 defer close(heartbeatDone)
126 go func() {
127 interval := 10 * time.Second
128 if raw := os.Getenv("ZUT_SUBAGENT_HEARTBEAT_INTERVAL"); raw != "" {
129 if parsed, parseErr := time.ParseDuration(raw); parseErr == nil && parsed > 0 {
130 interval = parsed
131 }
132 }
133 ticker := time.NewTicker(interval)
134 defer ticker.Stop()
135 for {
136 select {
137 case <-heartbeatDone:
138 return
139 case <-ticker.C:
140 em.emit("agent.heartbeat", map[string]any{"activity": "alive"})
141 }
142 }
143 }()
144
145 // Keep a per-turn cancel so the "cancel" inbox message can interrupt
146 // an in-flight turn without tearing down the whole daemon.
147 var (
148 mu sync.Mutex
149 cancelFn context.CancelFunc
150 busyTurn bool
151 turnNo = initialTurnNumber(os.Getenv("ZUT_SUBAGENT_EVENT_LOG"))
152 turnWG sync.WaitGroup
153 closing bool
154 shutdown = make(chan struct{})
155 shutdownOnce sync.Once
156 )
157 maxOutputBytes, maxOutputLines := workerOutputLimits()
158
159 runOne := func(prompt string) {
160 mu.Lock()
161 // launchTurn admits the goroutine to turnWG before it starts so
162 // shutdown can join it. A shutdown can win the scheduling race
163 // after that admission but before this goroutine acquires mu; in
164 // that case do not start a new Prompt after waitForActiveTurn has
165 // declared the worker closing.
166 if closing {
167 mu.Unlock()
168 return
169 }
170 if args.SubagentMaxTurns > 0 && turnNo >= args.SubagentMaxTurns {
171 // Consume a unique rejected attempt so repeated prompts after
172 // the limit cannot reuse the same turn id in the event log.
173 turnNo++
174 step := turnNo
175 mu.Unlock()
176 turnID := fmt.Sprintf("turn-%d", step)
177 errPayload := map[string]any{"code": "max_turns", "message": "maximum subagent turns reached"}
178 em.emit("turn.result", map[string]any{"status": "failed", "turn_id": turnID, "error": errPayload})
179 em.emit("turn.failed", map[string]any{"turn_id": turnID, "error": errPayload})
180 em.emit("turn_end", map[string]any{"step": step, "turn_id": turnID, "error": errPayload["message"]})
181 em.emit("agent.idle", map[string]any{"turn_id": turnID})
182 return
183 }
184 if busyTurn {
185 // Drop concurrent turns rather than queuing. The
186 // supervisor protocol assumes one outstanding turn per
187 // agent; if a user really wants to interrupt and start
188 // another, they should send "cancel" first.
189 mu.Unlock()
190 em.emit("error", map[string]any{"message": "agent busy; send 'cancel' first"})
191 return
192 }
193 busyTurn = true
194 turnNo++
195 step := turnNo
196 c, cancel := context.WithCancel(ctx)
197 cancelFn = cancel
198 mu.Unlock()
199
200 turnID := fmt.Sprintf("turn-%d", step)
201 em.emit("turn.started", map[string]any{"step": step, "turn_id": turnID})
202
203 sink := func(ev core.AgentEvent) {
204 em.emit(ev.Type(), modes.EventToJSON(ev))
205 }
206
207 start := len(ag.Messages())
208 err := ag.Prompt(c, prompt, nil, sink)
209 WriteNewTranscript(ag, sess, start)
210
211 output := finalAssistantText(ag.Messages()[start:])
212 resultStatus := "succeeded"
213 if err != nil {
214 if errors.Is(err, context.Canceled) {
215 resultStatus = "canceled"
216 } else {
217 resultStatus = "failed"
218 }
219 }
220 resultPayload := map[string]any{
221 "status": resultStatus,
222 "turn_id": turnID,
223 "summary": boundedResultSummary(output),
224 "output": boundedResultOutput(output, maxOutputBytes, maxOutputLines),
225 "error": resultErrorPayload(err),
226 }
227 em.emit("turn.result", resultPayload)
228 if err != nil {
229 em.emit("turn.failed", map[string]any{"turn_id": turnID, "error": resultErrorPayload(err)})
230 }
231 em.emit("turn_end", map[string]any{
232 "step": step,
233 "turn_id": turnID,
234 "error": errString(err),
235 })
236 em.emit("agent.idle", map[string]any{"turn_id": turnID})
237
238 cancel()
239 mu.Lock()
240 busyTurn = false
241 cancelFn = nil
242 mu.Unlock()
243 }
244
245 launchTurn := func(prompt string) {
246 mu.Lock()
247 if closing {
248 mu.Unlock()
249 return
250 }
251 // Register the goroutine before launching it. Shutdown can then
252 // close the gate and wait for every turn that was admitted, even
253 // if runOne has not acquired mu and set busyTurn yet.
254 turnWG.Add(1)
255 mu.Unlock()
256 go func() {
257 defer turnWG.Done()
258 runOne(prompt)
259 }()
260 }
261
262 waitForActiveTurn := func(cancel bool) {
263 mu.Lock()
264 closing = true
265 activeCancel := cancelFn
266 mu.Unlock()
267 if cancel && activeCancel != nil {
268 activeCancel()
269 }
270 turnWG.Wait()
271 }
272
273 requestShutdown := func() {
274 mu.Lock()
275 closing = true
276 mu.Unlock()
277 shutdownOnce.Do(func() { close(shutdown) })
278 }
279
280 // Initial task: run before processing the inbox so the agent
281 // "starts working" the moment it boots, matching what users
282 // expect from `/subagents new <task>`.
283 if args.Prompt != "" {
284 launchTurn(args.Prompt)
285 }
286
287 // Inbox loop: one supervisor message at a time. We don't spawn
288 // a goroutine per turn because runOne already serialises them
289 // via the busyTurn flag; doing the dispatch on the main
290 // goroutine keeps the daemon's lifecycle easy to follow.
291 for {
292 select {
293 case <-ctx.Done():
294 waitForActiveTurn(true)
295 em.emit("agent.exited", map[string]any{"reason": "cancelled"})
296 return ctx.Err()
297 case <-shutdown:
298 // Shutdown is an explicit supervisor stop. Cancel an active
299 // turn before joining it so a detached worker cannot remain
300 // alive indefinitely while the host waits for its socket to
301 // disappear.
302 waitForActiveTurn(true)
303 em.emit("agent.exited", map[string]any{"reason": "shutdown"})
304 return nil
305 case msg, ok := <-ln.Lines():
306 if !ok {
307 waitForActiveTurn(true)
308 em.emit("agent.exited", map[string]any{"reason": "inbox-closed"})
309 return nil
310 }
311 command, parseErr := subagents.ParseCommand(msg)
312 if parseErr != nil {
313 em.emit("error", map[string]any{"message": "invalid supervisor command"})
314 continue
315 }
316 switch command.Type {
317 case subagents.CommandAgentShutdown:
318 requestShutdown()
319 case subagents.CommandTurnCancel:
320 mu.Lock()
321 if cancelFn != nil {
322 cancelFn()
323 }
324 mu.Unlock()
325 case subagents.CommandAgentPing:
326 em.emit("agent.heartbeat", map[string]any{"activity": "alive"})
327 case subagents.CommandTurnStart:
328 var payload subagents.TurnStartPayload
329 if err := command.DecodePayload(&payload); err != nil {
330 em.emit("error", map[string]any{"message": "invalid turn.start payload"})
331 continue
332 }
333 launchTurn(payload.Prompt)
334 default:
335 // Unknown commands are rejected rather than interpreted as
336 // child-originated control requests.
337 em.emit("error", map[string]any{"message": "unknown supervisor command"})
338 }
339 }
340 }
341}
342
343// subagentEmitter serialises events to stdout and (optionally) to a
344// durable log file. Concurrent goroutines call emit so we have to
345// hold a mutex around the encoder.
346type subagentEmitter struct {
347 mu sync.Mutex
348 w *os.File
349 mirror *subagents.EventLog
350 agentID string
351
352 // orphan flips true the first time a stdout write fails (broken
353 // pipe — the supervisor died). Until then the mirror stays
354 // dormant: the supervisor is the canonical writer to events.jsonl
355 // (it parses our stdout and Append()s each event itself). Writing
356 // from both sides used to land every event in the log twice,
357 // which showed up as a fully-duplicated transcript the next time
358 // the agent was reloaded — the exact "why is everything doubled"
359 // bug. Once orphaned the mirror takes over so the events still
360 // land on disk for the next reload.
361 orphan bool
362}
363
364func newSubagentEmitter(w *os.File, mirror *subagents.EventLog) *subagentEmitter {
365 return &subagentEmitter{w: w, mirror: mirror}
366}
367
368func (e *subagentEmitter) setProtocolIdentity(agentID string) {
369 e.mu.Lock()
370 e.agentID = agentID
371 e.mu.Unlock()
372}
373
374func (e *subagentEmitter) emit(typ string, data map[string]any) {
375 if data == nil {
376 data = map[string]any{}
377 }
378 payload := make(map[string]any, len(data))
379 for key, value := range data {
380 payload[key] = value
381 }
382 turnID, _ := payload["turn_id"].(string)
383
384 e.mu.Lock()
385 defer e.mu.Unlock()
386
387 wireType := canonicalWorkerEvent(typ)
388 envelope := subagents.NewEventEnvelope(wireType, e.agentID, turnID, payload)
389 line, err := subagents.MarshalJSONL(envelope)
390
391 if !e.orphan && err == nil {
392 if _, werr := e.w.Write(line); werr != nil {
393 // Supervisor's stdout pipe is gone (parent zut exited but we
394 // kept running). Switch to mirror-only mode and preserve this
395 // event in the durable log.
396 e.orphan = true
397 }
398 }
399
400 if e.orphan && e.mirror != nil {
401 _ = e.mirror.Append(subagents.NewEvent(wireType, payload))
402 }
403}
404
405func canonicalWorkerEvent(typ string) string {
406 switch typ {
407 case "agent_ready":
408 return subagents.EventAgentReady
409 case "agent_stopped":
410 return subagents.EventAgentExited
411 case "turn_start":
412 return subagents.EventTurnStarted
413 case "turn_progress":
414 return subagents.EventTurnProgress
415 case "tool_call":
416 return subagents.EventToolStarted
417 case "tool_result":
418 return subagents.EventToolFinished
419 case "text_delta":
420 return subagents.EventMessageDelta
421 default:
422 return typ
423 }
424}
425
426func finalAssistantText(messages []provider.Message) string {
427 var out strings.Builder
428 for _, message := range messages {
429 if message.Role != provider.RoleAssistant {
430 continue
431 }
432 for _, block := range message.Content {
433 if text, ok := block.(provider.TextBlock); ok {
434 if out.Len() > 0 {
435 out.WriteByte('\n')
436 }
437 out.WriteString(text.Text)
438 }
439 }
440 }
441 return out.String()
442}
443
444func firstLine(text string) string {
445 return strings.Split(strings.TrimSpace(text), "\n")[0]
446}
447
448func boundedResultSummary(text string) string {
449 return boundedResultOutput(firstLine(text), 4*1024, 1)
450}
451
452func workerOutputLimits() (maxBytes, maxLines int) {
453 maxBytes, maxLines = 500_000, 5_000
454 if value, err := strconv.Atoi(strings.TrimSpace(os.Getenv("ZUT_SUBAGENT_MAX_OUTPUT_BYTES"))); err == nil && value > 0 {
455 maxBytes = value
456 }
457 if value, err := strconv.Atoi(strings.TrimSpace(os.Getenv("ZUT_SUBAGENT_MAX_OUTPUT_LINES"))); err == nil && value > 0 {
458 maxLines = value
459 }
460 return maxBytes, maxLines
461}
462
463func boundedResultOutput(text string, maxBytes, maxLines int) string {
464 const marker = "...[output truncated]"
465 lines := strings.Split(text, "\n")
466 if maxLines > 0 && len(lines) > maxLines {
467 if maxLines == 1 {
468 text = marker
469 } else {
470 text = strings.Join(append(lines[:maxLines-1], marker), "\n")
471 }
472 }
473 if maxBytes > 0 && len([]byte(text)) > maxBytes {
474 if maxBytes <= len(marker) {
475 return truncateOutputUTF8(text, maxBytes)
476 }
477 return truncateOutputUTF8(text, maxBytes-len(marker)) + marker
478 }
479 return text
480}
481
482func truncateOutputUTF8(value string, maxBytes int) string {
483 if maxBytes <= 0 {
484 return ""
485 }
486 data := []byte(value)
487 if len(data) <= maxBytes {
488 return value
489 }
490 data = data[:maxBytes]
491 for len(data) > 0 && !utf8.Valid(data) {
492 data = data[:len(data)-1]
493 }
494 return string(data)
495}
496
497func resultErrorPayload(err error) map[string]any {
498 if err == nil {
499 return nil
500 }
501 return map[string]any{"code": "turn_failed", "message": truncateForLog(err.Error(), 500)}
502}
503
504func errString(err error) string {
505 if err == nil {
506 return ""
507 }
508 return err.Error()
509}
510
511func initialTurnNumber(path string) int {
512 if strings.TrimSpace(path) == "" {
513 return 0
514 }
515 events, err := subagents.ReadEventLog(path)
516 if err != nil {
517 return 0
518 }
519 maxStep := 0
520 for _, ev := range events {
521 if step, ok := ev.Data["step"].(float64); ok && int(step) > maxStep {
522 maxStep = int(step)
523 }
524 if strings.HasPrefix(ev.TurnID, "turn-") {
525 if n, parseErr := strconv.Atoi(strings.TrimPrefix(ev.TurnID, "turn-")); parseErr == nil && n > maxStep {
526 maxStep = n
527 }
528 }
529 }
530 return maxStep
531}
532
533func truncateForLog(s string, n int) string {
534 if len(s) <= n {
535 return s
536 }
537 if n <= 3 {
538 return strings.Repeat(".", n)
539 }
540 return s[:n-3] + "..."
541}