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