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 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 // worker supervisor 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(ctx context.Context, args Args, r Resolved, cfg Config, tracker *subagents.CompletionTracker) *subagentRuntime {
196 onSpawned := func(a *subagents.Agent, task string, _ bool) {
197 tracker.TrackTurn(a, task, false)
198 }
199 onResumed := func(a *subagents.Agent, prompt string, _ bool) {
200 // BeforeResumed normally owns future-turn registration. Keep this
201 // fallback for direct runtime callers that do not install the pre-hook;
202 // SubagentResumeTool suppresses it when BeforeResumed accepted delivery.
203 tracker.TrackTurn(a, prompt, true)
204 }
205 onBeforeResumed := func(a *subagents.Agent, prompt string, _ bool) func() {
206 return tracker.TrackFutureTurn(a, prompt, true)
207 }
208 onStopRequested := func(a *subagents.Agent) {
209 tracker.TrackExit(a, "stopped")
210 }
211 return newSubagentRuntime(subagentRuntimeConfig{
212 Context: ctx,
213 Args: args,
214 Root: filepath.Join(ZutHome(), "subagents"),
215 RepoRoot: r.CWD,
216 Provider: r.Provider,
217 Model: r.Model,
218 Reasoning: r.Reasoning,
219 BaseURL: r.BaseURL,
220 InsecureTLS: r.InsecureTLS,
221 FastMode: r.FastMode,
222 APIKey: args.APIKey,
223 Policy: subagentPolicyFromConfig(cfg.Subagents),
224 WebSearchPolicy: webSearchPolicyForRegistry(r.WebSearchPolicy, r.ToolRegistry),
225 OnSpawned: onSpawned,
226 OnResumed: onResumed,
227 BeforeResumed: onBeforeResumed,
228 OnStopRequested: onStopRequested,
229 })
230}
231
232func runOrchestratedStreamMode(ctx context.Context, args Args, version string) error {
233 return runOrchestratedMode(ctx, args, version, orchestratedModeHooks{
234 promptLabel: "stream",
235 startupSink: func() (func(core.AgentEvent), func()) {
236 return newStreamTextSink(os.Stdout)
237 },
238 runPrimary: func(turnCtx context.Context, ag *core.Agent, prompt string, persist func([]provider.Message) error) (string, modes.ContextRecoveryResult, error) {
239 recovery, err := modes.RunStreamWithContextRecovery(turnCtx, ag, prompt, nil, os.Stdout, os.Stderr, persist)
240 return "", recovery, err
241 },
242 finish: func(string) error { return nil },
243 })
244}
245
246func runOrchestratedJSONMode(ctx context.Context, args Args, version string) error {
247 enc := json.NewEncoder(os.Stdout)
248 var outputErr error
249 encode := func(value any) {
250 if outputErr == nil {
251 outputErr = enc.Encode(value)
252 }
253 }
254 runErr := runOrchestratedMode(ctx, args, version, orchestratedModeHooks{
255 promptLabel: "json",
256 startupSink: func() (func(core.AgentEvent), func()) {
257 return func(ev core.AgentEvent) {
258 encode(modes.EventToJSON(ev))
259 }, func() {}
260 },
261 runPrimary: func(turnCtx context.Context, ag *core.Agent, prompt string, persist func([]provider.Message) error) (string, modes.ContextRecoveryResult, error) {
262 recovery, err := modes.RunJSONWithContextRecovery(turnCtx, ag, prompt, nil, os.Stdout, persist)
263 return "", recovery, err
264 },
265 emitError: func(err error) {
266 if err != nil {
267 encode(map[string]any{"type": "error", "message": err.Error()})
268 }
269 },
270 finish: func(string) error { return nil },
271 })
272 return errors.Join(runErr, outputErr)
273}