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