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 runtime.Supervisor() != nil {
72 if _, reloadErrs := runtime.Supervisor().Reload(); len(reloadErrs) != 0 {
73 return fmt.Errorf("reload subagents: %w", errors.Join(reloadErrs...))
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 runtime.WireRequiredWorkerGate(ag)
111
112 start := len(ag.Messages())
113 preSink, finishPre := hooks.startupSink()
114 preErr := runZutfileStartupPre(ctx, args.StartupPre, r.CWD, r.Sandbox, ag, preSink, os.Stderr)
115 finishPre()
116 if preErr != nil {
117 if hooks.emitError != nil {
118 hooks.emitError(preErr)
119 }
120 return preErr
121 }
122 if strings.TrimSpace(args.StartupPre) != "" {
123 refreshedPolicy, refreshErr := reloadResourcesAfterStartupPreWithRegistry(ctx, args, extMgr, r.Sandbox, ag, prepareRegistry)
124 if refreshErr != nil {
125 if hooks.emitError != nil {
126 hooks.emitError(refreshErr)
127 }
128 return refreshErr
129 }
130 // entry.pre may have changed the effective permission set. Keep the
131 // worker supervisor in step with the refreshed parent registry.
132 runtime.SetWebSearchPolicy(refreshedPolicy)
133 }
134
135 var usagePersistErr error
136 if sess != nil {
137 ag.OnUsage = func(cumulative provider.Usage) {
138 if usagePersistErr == nil {
139 usagePersistErr = sess.AppendUsage(ag.LastTurnUsage(), cumulative)
140 }
141 }
142 }
143 var persistCompaction func([]provider.Message) error
144 if sess != nil {
145 persistCompaction = func(messages []provider.Message) error {
146 return sess.AppendCompactionWithUsage(messages, ag.Cost())
147 }
148 }
149
150 transcriptStart := start
151 runParent := func(turnCtx context.Context, text string) (string, error) {
152 output, recovery, turnErr := hooks.runPrimary(turnCtx, ag, text, persistCompaction)
153 if recovery.Compacted {
154 transcriptStart = recovery.OutputStart
155 }
156 if usagePersistErr != nil {
157 // Usage persistence happens after the mode renderer has completed.
158 // JSON therefore needs an explicit terminal error for this failure;
159 // unlike provider errors, no renderer event has reported it yet.
160 if turnErr == nil && hooks.emitError != nil {
161 hooks.emitError(usagePersistErr)
162 }
163 return output, errors.Join(turnErr, usagePersistErr)
164 }
165 return output, turnErr
166 }
167
168 finalText, err := runParent(ctx, prompt)
169 if err == nil {
170 finalText, err = runHeadlessContinuation(ctx, tracker, finalText, runParent)
171 if err != nil && hooks.emitError != nil {
172 // A completion wait has no renderer call in which to report its
173 // failure, so JSON must add its single terminal error object here.
174 // Errors from the follow-up parent turn were already rendered by
175 // that mode and must not be emitted a second time.
176 var waitErr *headlessCompletionWaitError
177 if errors.As(err, &waitErr) {
178 hooks.emitError(waitErr)
179 }
180 }
181 }
182 if persistErr := WriteNewTranscript(ag, sess, transcriptStart); persistErr != nil {
183 if err == nil && hooks.emitError != nil {
184 hooks.emitError(persistErr)
185 }
186 err = errors.Join(err, persistErr)
187 }
188 if err != nil {
189 return err
190 }
191 return hooks.finish(finalText)
192}
193
194func newOrchestratedRuntime(ctx context.Context, args Args, r Resolved, cfg Config, tracker *subagents.CompletionTracker) *subagentRuntime {
195 onSpawned := func(a *subagents.Agent, task string, required bool) {
196 if !required {
197 tracker.TrackTurn(a, task, false)
198 }
199 }
200 onResumed := func(a *subagents.Agent, prompt string, required bool) {
201 // BeforeResumed normally owns future-turn registration. Keep this
202 // fallback for direct runtime callers that do not install the pre-hook;
203 // SubagentResumeTool suppresses it when BeforeResumed accepted delivery.
204 if !required {
205 tracker.TrackTurn(a, prompt, true)
206 }
207 }
208 onBeforeResumed := func(a *subagents.Agent, prompt string, required bool) func() {
209 if required {
210 return nil
211 }
212 return tracker.TrackFutureTurn(a, prompt, true)
213 }
214 onStopRequested := func(a *subagents.Agent) {
215 tracker.TrackExit(a, "stopped")
216 }
217 return newSubagentRuntime(subagentRuntimeConfig{
218 Context: ctx,
219 Args: args,
220 Root: filepath.Join(ZutHome(), "subagents"),
221 RepoRoot: r.CWD,
222 Provider: r.Provider,
223 Model: r.Model,
224 Reasoning: r.Reasoning,
225 BaseURL: r.BaseURL,
226 InsecureTLS: r.InsecureTLS,
227 FastMode: r.FastMode,
228 APIKey: args.APIKey,
229 Policy: subagentPolicyFromConfig(cfg.Subagents),
230 WebSearchPolicy: webSearchPolicyForRegistry(r.WebSearchPolicy, r.ToolRegistry),
231 OnSpawned: onSpawned,
232 OnResumed: onResumed,
233 BeforeResumed: onBeforeResumed,
234 OnStopRequested: onStopRequested,
235 })
236}
237
238func runOrchestratedStreamMode(ctx context.Context, args Args, version string) error {
239 return runOrchestratedMode(ctx, args, version, orchestratedModeHooks{
240 promptLabel: "stream",
241 startupSink: func() (func(core.AgentEvent), func()) {
242 return newStreamTextSink(os.Stdout)
243 },
244 runPrimary: func(turnCtx context.Context, ag *core.Agent, prompt string, persist func([]provider.Message) error) (string, modes.ContextRecoveryResult, error) {
245 recovery, err := modes.RunStreamWithContextRecovery(turnCtx, ag, prompt, nil, os.Stdout, os.Stderr, persist)
246 return "", recovery, err
247 },
248 finish: func(string) error { return nil },
249 })
250}
251
252func runOrchestratedJSONMode(ctx context.Context, args Args, version string) error {
253 enc := json.NewEncoder(os.Stdout)
254 var outputErr error
255 encode := func(value any) {
256 if outputErr == nil {
257 outputErr = enc.Encode(value)
258 }
259 }
260 runErr := runOrchestratedMode(ctx, args, version, orchestratedModeHooks{
261 promptLabel: "json",
262 startupSink: func() (func(core.AgentEvent), func()) {
263 return func(ev core.AgentEvent) {
264 encode(modes.EventToJSON(ev))
265 }, func() {}
266 },
267 runPrimary: func(turnCtx context.Context, ag *core.Agent, prompt string, persist func([]provider.Message) error) (string, modes.ContextRecoveryResult, error) {
268 recovery, err := modes.RunJSONWithContextRecovery(turnCtx, ag, prompt, nil, os.Stdout, persist)
269 return "", recovery, err
270 },
271 emitError: func(err error) {
272 if err != nil {
273 encode(map[string]any{"type": "error", "message": err.Error()})
274 }
275 },
276 finish: func(string) error { return nil },
277 })
278 return errors.Join(runErr, outputErr)
279}