1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "os"
9 "os/signal"
10 "path/filepath"
11 "strings"
12 "syscall"
13
14 "github.com/bnema/zut/packages/agent/modes"
15 "github.com/bnema/zut/packages/agent/subagents"
16 "github.com/bnema/zut/packages/core"
17 "github.com/bnema/zut/packages/provider"
18)
19
20type orchestratedModeHooks struct {
21 promptLabel string
22
23 // startupSink is used for entry.pre agent turns. The completion updates
24 // themselves go through runPrimary and therefore use the same renderer as
25 // the initial parent turn.
26 startupSink func() (func(core.AgentEvent), func())
27 runPrimary func(context.Context, *core.Agent, string, func([]provider.Message) error) (string, modes.ContextRecoveryResult, error)
28 emitError func(error)
29 finish func(string) error
30}
31
32// runOrchestratedMode owns the mode-neutral headless lifecycle. Modes provide
33// only their existing renderer and wire format; the tracker, worker runtime,
34// startup hooks, transcript persistence, and completion-wave loop are shared.
35func runOrchestratedMode(parentCtx context.Context, args Args, version string, hooks orchestratedModeHooks) (runErr error) {
36 ctx, stopSignal := signal.NotifyContext(parentCtx, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
37 defer stopSignal()
38
39 if args.NoYolo {
40 fmt.Fprintf(os.Stderr, "warning: --no-yolo has no effect in %s mode (no interactive prompt available); tools will run without confirmation\n", hooks.promptLabel)
41 }
42 r, err := Resolve(args, true)
43 if err != nil {
44 return err
45 }
46 initialCfg, err := LoadConfig()
47 if err != nil {
48 return fmt.Errorf("load config for orchestrated %s: %w", hooks.promptLabel, err)
49 }
50 prompt := args.Prompt
51 if prompt == "" {
52 piped, readErr := readAllStdin()
53 if readErr != nil {
54 return fmt.Errorf("read %s prompt from stdin: %w", hooks.promptLabel, readErr)
55 }
56 prompt = strings.TrimSpace(piped)
57 }
58 if prompt == "" {
59 return fmt.Errorf("%s mode requires a prompt (arg or stdin)", hooks.promptLabel)
60 }
61 extMgr, stopExt := setupNonInteractiveExtensions(ctx, args, &r, version)
62 defer stopExt()
63
64 tracker := subagents.NewCompletionTracker()
65 runtime := newOrchestratedRuntime(ctx, args, r, initialCfg, tracker)
66 defer func() {
67 if closeErr := closeSubagentRuntimeFresh(runtime); closeErr != nil {
68 runErr = errors.Join(runErr, closeErr)
69 }
70 }()
71 if resident := runtime.ResidentManager(); resident != nil {
72 for _, reconcileErr := range resident.Reconcile() {
73 fmt.Fprintf(os.Stderr, "zut: reconcile resident subagent: %v\n", reconcileErr)
74 }
75 }
76
77 prepareRegistry := func(reg core.Registry) core.Registry {
78 return runtime.PrepareRegistry(reg)
79 }
80 runtime.SetProvider(r.Provider)
81 runtime.SetProviderSettings(r.BaseURL, r.InsecureTLS)
82 runtime.SetFastMode(r.FastMode)
83 runtime.PrepareResolvedRegistry(r.ToolRegistry, r.WebSearchPolicy)
84 ag := r.NewAgent()
85 initialAg := ag
86 defer func() {
87 closeAgentLSP(ag)
88 if ag != initialAg {
89 closeAgentLSP(initialAg)
90 }
91 }()
92 wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
93
94 sess, err := openOrCreateSession(ctx, args, r, ag, version)
95 if err != nil {
96 return err
97 }
98 if sess != nil {
99 defer func() { runErr = joinSessionCloseError(runErr, sess) }()
100 var providerName, model string
101 sess, ag, providerName, model, err = applyInitialSessionResumeWithRuntime(ctx, args, r, extMgr, sess, ag, runtime)
102 if err != nil {
103 return err
104 }
105 r.Provider, r.Model = providerName, model
106 ag.CommitToolResult = func(_ string, result core.ToolResult) error { return persistToolResultState(extMgr, sess, result) }
107 runtime.SetActiveSession(sess.ID)
108 }
109 announceSession(extMgr, sess)
110 wireRetryLifecyclePersistence(ag, sess)
111 runtime.WireRequiredWorkerGate(ag)
112
113 start := len(ag.Messages())
114 preSink, finishPre := hooks.startupSink()
115 preErr := runZutfileStartupPre(ctx, args.StartupPre, r.CWD, r.Sandbox, ag, preSink, os.Stderr)
116 finishPre()
117 if preErr != nil {
118 if hooks.emitError != nil {
119 hooks.emitError(preErr)
120 }
121 return preErr
122 }
123 if strings.TrimSpace(args.StartupPre) != "" {
124 refreshedPolicy, refreshErr := reloadResourcesAfterStartupPreWithRegistry(ctx, args, extMgr, r.Sandbox, ag, prepareRegistry)
125 if refreshErr != nil {
126 if hooks.emitError != nil {
127 hooks.emitError(refreshErr)
128 }
129 return refreshErr
130 }
131 // entry.pre may have changed the effective permission set. Keep the
132 // resident child factory in step with the refreshed parent registry.
133 runtime.SetWebSearchPolicy(refreshedPolicy)
134 }
135
136 var usagePersistErr error
137 if sess != nil {
138 ag.OnUsage = func(cumulative provider.Usage) {
139 if usagePersistErr == nil {
140 usagePersistErr = sess.AppendUsage(ag.LastTurnUsage(), cumulative)
141 }
142 }
143 }
144 var persistCompaction func([]provider.Message) error
145 if sess != nil {
146 persistCompaction = func(messages []provider.Message) error {
147 return sess.AppendCompactionWithUsage(messages, ag.Cost())
148 }
149 }
150
151 transcriptStart := start
152 runParent := func(turnCtx context.Context, text string) (string, error) {
153 output, recovery, turnErr := hooks.runPrimary(turnCtx, ag, text, persistCompaction)
154 if recovery.Compacted {
155 transcriptStart = recovery.OutputStart
156 }
157 if usagePersistErr != nil {
158 // Usage persistence happens after the mode renderer has completed.
159 // JSON therefore needs an explicit terminal error for this failure;
160 // unlike provider errors, no renderer event has reported it yet.
161 if turnErr == nil && hooks.emitError != nil {
162 hooks.emitError(usagePersistErr)
163 }
164 return output, errors.Join(turnErr, usagePersistErr)
165 }
166 return output, turnErr
167 }
168
169 finalText, err := runParent(ctx, prompt)
170 if err == nil {
171 finalText, err = runHeadlessContinuation(ctx, tracker, finalText, runParent)
172 if err != nil && hooks.emitError != nil {
173 // A completion wait has no renderer call in which to report its
174 // failure, so JSON must add its single terminal error object here.
175 // Errors from the follow-up parent turn were already rendered by
176 // that mode and must not be emitted a second time.
177 var waitErr *headlessCompletionWaitError
178 if errors.As(err, &waitErr) {
179 hooks.emitError(waitErr)
180 }
181 }
182 }
183 if persistErr := WriteNewTranscript(ag, sess, transcriptStart); persistErr != nil {
184 if err == nil && hooks.emitError != nil {
185 hooks.emitError(persistErr)
186 }
187 err = errors.Join(err, persistErr)
188 }
189 if err != nil {
190 return err
191 }
192 return hooks.finish(finalText)
193}
194
195func newOrchestratedRuntime(_ context.Context, args Args, r Resolved, cfg Config, tracker *subagents.CompletionTracker) *subagentRuntime {
196 return newSubagentRuntime(subagentRuntimeConfig{
197 Args: args,
198 Root: filepath.Join(ZutHome(), "subagents"),
199 RepoRoot: r.CWD,
200 Provider: r.Provider,
201 Model: r.Model,
202 Reasoning: r.Reasoning,
203 BaseURL: r.BaseURL,
204 InsecureTLS: r.InsecureTLS,
205 FastMode: r.FastMode,
206 Policy: subagentPolicyFromConfig(cfg.Subagents),
207 WebSearchPolicy: webSearchPolicyForRegistry(r.WebSearchPolicy, r.ToolRegistry),
208 ResidentCompletion: func(completion subagents.ResidentCompletion) {
209 status := "completed"
210 errText := ""
211 if completion.Err != nil {
212 status, errText = "failed", completion.Err.Error()
213 if errors.Is(completion.Err, context.Canceled) {
214 status = "interrupted"
215 }
216 }
217 tracker.Report(subagents.Completion{AgentID: completion.ChildID, TurnID: completion.TurnID, Status: status, Task: completion.Task, Error: errText, Summary: completion.Summary})
218 },
219 OnResidentSpawned: func(spec subagents.ResidentChildSpec, turnID, _ string) { tracker.TrackResident(spec.ID, turnID) },
220 })
221}
222
223func runOrchestratedStreamMode(ctx context.Context, args Args, version string) error {
224 return runOrchestratedMode(ctx, args, version, orchestratedModeHooks{
225 promptLabel: "stream",
226 startupSink: func() (func(core.AgentEvent), func()) {
227 return newStreamTextSink(os.Stdout)
228 },
229 runPrimary: func(turnCtx context.Context, ag *core.Agent, prompt string, persist func([]provider.Message) error) (string, modes.ContextRecoveryResult, error) {
230 recovery, err := modes.RunStreamWithContextRecovery(turnCtx, ag, prompt, nil, os.Stdout, os.Stderr, persist)
231 return "", recovery, err
232 },
233 finish: func(string) error { return nil },
234 })
235}
236
237func runOrchestratedJSONMode(ctx context.Context, args Args, version string) error {
238 enc := json.NewEncoder(os.Stdout)
239 var outputErr error
240 encode := func(value any) {
241 if outputErr == nil {
242 outputErr = enc.Encode(value)
243 }
244 }
245 runErr := runOrchestratedMode(ctx, args, version, orchestratedModeHooks{
246 promptLabel: "json",
247 startupSink: func() (func(core.AgentEvent), func()) {
248 return func(ev core.AgentEvent) {
249 encode(modes.EventToJSON(ev))
250 }, func() {}
251 },
252 runPrimary: func(turnCtx context.Context, ag *core.Agent, prompt string, persist func([]provider.Message) error) (string, modes.ContextRecoveryResult, error) {
253 recovery, err := modes.RunJSONWithContextRecovery(turnCtx, ag, prompt, nil, os.Stdout, persist)
254 return "", recovery, err
255 },
256 emitError: func(err error) {
257 if err != nil {
258 encode(map[string]any{"type": "error", "message": err.Error()})
259 }
260 },
261 finish: func(string) error { return nil },
262 })
263 return errors.Join(runErr, outputErr)
264}