1package modes
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "io"
8 "sync"
9
10 "github.com/bnema/zut/packages/core"
11 "github.com/bnema/zut/packages/provider"
12)
13
14// RunJSON runs the agent to completion, writing one JSON object per
15// AgentEvent as newline-delimited JSON.
16func RunJSON(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, out io.Writer) error {
17 _, err := RunJSONWithContextRecovery(ctx, ag, prompt, images, out, nil)
18 return err
19}
20
21// RunJSONWithContextRecovery is RunJSON with an optional checkpoint writer
22// for a one-shot context-overflow compaction. It suppresses the recoverable
23// first turn_end so stdout never contains a terminal error frame before the
24// continued turn succeeds. OutputStart identifies the suffix to persist after
25// a successful checkpoint.
26func RunJSONWithContextRecovery(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, out io.Writer, persistCompaction func([]provider.Message) error) (ContextRecoveryResult, error) {
27 enc := json.NewEncoder(out)
28 var writeMu sync.Mutex
29 var writeErr error
30 write := func(v any) {
31 writeMu.Lock()
32 defer writeMu.Unlock()
33 if writeErr == nil {
34 writeErr = enc.Encode(v)
35 }
36 }
37
38 var runErr error
39 sink := func(ev core.AgentEvent) {
40 write(EventToJSON(ev))
41 }
42
43 recovery, err := PromptWithContextRecovery(ctx, ag, prompt, images, sink, ContextRecoveryOptions{
44 PersistCompaction: persistCompaction,
45 SuppressInitialOverflowEvent: true,
46 })
47 if err != nil {
48 runErr = err
49 }
50
51 if runErr != nil {
52 write(map[string]any{"type": "error", "message": runErr.Error()})
53 }
54 writeMu.Lock()
55 outputErr := writeErr
56 writeMu.Unlock()
57 return recovery, errors.Join(runErr, outputErr)
58}
59
60// EventToJSON converts an AgentEvent to a JSON-friendly map. The on-wire
61// schema is deliberately simple and flat. Exported so the RPC mode can
62// reuse the same serialisation as `zut --json`.
63func EventToJSON(ev core.AgentEvent) map[string]any {
64 m := map[string]any{"type": ev.Type()}
65 switch e := ev.(type) {
66 case core.EvTurnStart:
67 m["step"] = e.Step
68 case core.EvRequestStarted:
69 m["provider"] = e.Provider
70 m["model"] = e.Model
71 m["scope"] = string(e.Scope)
72 m["attempt"] = e.Attempt
73 m["max_attempts"] = e.MaxAttempts
74 case core.EvRetryScheduled:
75 m["scope"] = string(e.Scope)
76 m["attempt"] = e.Attempt
77 m["max_attempts"] = e.MaxAttempts
78 m["delay_ms"] = e.Delay.Milliseconds()
79 case core.EvUserMessage:
80 m["content"] = ContentToJSON(e.Message.Content)
81 m["time"] = e.Message.Time
82 case core.EvAssistantMessage:
83 m["content"] = ContentToJSON(e.Message.Content)
84 m["time"] = e.Message.Time
85 case core.EvCacheDiagnostics:
86 m["eligible"] = e.Eligible
87 m["mode"] = e.Mode
88 m["transport"] = e.Transport
89 m["continuation"] = e.Continuation
90 case core.EvTextDelta:
91 m["delta"] = e.Delta
92 case core.EvToolUseStart:
93 m["id"] = e.ID
94 m["name"] = e.Name
95 case core.EvToolUseArgs:
96 m["id"] = e.ID
97 m["delta"] = e.Delta
98 case core.EvToolUseEnd:
99 m["id"] = e.ID
100 case core.EvToolCall:
101 m["id"] = e.ID
102 m["name"] = e.Name
103 var args any
104 _ = json.Unmarshal(e.Args, &args)
105 m["args"] = args
106 case core.EvToolExecutionStarted:
107 m["id"] = e.ID
108 m["name"] = e.Name
109 case core.EvToolProgress:
110 m["id"] = e.ID
111 m["text"] = e.Text
112 case core.EvPlanUpdate:
113 if e.Update.Explanation != nil {
114 m["explanation"] = *e.Update.Explanation
115 } else {
116 m["explanation"] = nil
117 }
118 m["plan"] = e.Update.Plan
119 case core.EvToolResult:
120 m["id"] = e.ID
121 m["is_error"] = e.Result.IsError
122 m["content"] = ContentToJSON(e.Result.Content)
123 case core.EvUsage:
124 m["input"] = e.Usage.InputTokens
125 m["output"] = e.Usage.OutputTokens
126 m["reasoning"] = nullableReasoningTokens(e.Usage)
127 m["cache_read"] = e.Usage.CacheReadTokens
128 m["cache_write"] = e.Usage.CacheWriteTokens
129 m["cache_measured_prompt"] = e.Usage.CacheMeasuredPromptTokens
130 m["cache_measured_read"] = e.Usage.CacheMeasuredReadTokens
131 m["cost_usd"] = e.Usage.CostUSD
132 m["cumulative"] = map[string]any{
133 "input": e.Cumulative.InputTokens,
134 "output": e.Cumulative.OutputTokens,
135 "reasoning": nullableReasoningTokens(e.Cumulative),
136 "cache_read": e.Cumulative.CacheReadTokens,
137 "cache_write": e.Cumulative.CacheWriteTokens,
138 "cache_measured_prompt": e.Cumulative.CacheMeasuredPromptTokens,
139 "cache_measured_read": e.Cumulative.CacheMeasuredReadTokens,
140 "cost_usd": e.Cumulative.CostUSD,
141 }
142 case core.EvTurnRecovery:
143 m["reason"] = e.Reason
144 m["attempt"] = e.Attempt
145 case core.EvTurnEnd:
146 m["stop"] = string(e.Stop)
147 if e.Err != nil {
148 m["error"] = e.Err.Error()
149 }
150 }
151 return m
152}
153
154// ContentToJSON serialises a transcript content slice into the same
155// shape used in EventToJSON. Exported alongside EventToJSON for the
156// RPC mode.
157func ContentToJSON(blocks []provider.Content) []map[string]any {
158 out := make([]map[string]any, 0, len(blocks))
159 for _, b := range blocks {
160 switch v := b.(type) {
161 case provider.TextBlock:
162 out = append(out, map[string]any{"type": "text", "text": v.Text})
163 case provider.ImageBlock:
164 out = append(out, map[string]any{"type": "image", "mime_type": v.MimeType, "bytes": len(v.Data)})
165 case provider.ToolCallBlock:
166 var args any
167 _ = json.Unmarshal(v.Arguments, &args)
168 out = append(out, map[string]any{"type": "tool_call", "id": v.ID, "name": v.Name, "args": args})
169 case provider.ToolResultBlock:
170 out = append(out, map[string]any{
171 "type": "tool_result",
172 "call_id": v.CallID,
173 "is_error": v.IsError,
174 "content": ContentToJSON(v.Content),
175 })
176 }
177 }
178 return out
179}
180
181func nullableReasoningTokens(usage provider.Usage) any {
182 if !usage.ReasoningTokensKnown {
183 return nil
184 }
185 return usage.ReasoningTokens
186}