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.EvToolResult:
113 m["id"] = e.ID
114 m["is_error"] = e.Result.IsError
115 m["content"] = ContentToJSON(e.Result.Content)
116 case core.EvUsage:
117 m["input"] = e.Usage.InputTokens
118 m["output"] = e.Usage.OutputTokens
119 m["reasoning"] = nullableReasoningTokens(e.Usage)
120 m["cache_read"] = e.Usage.CacheReadTokens
121 m["cache_write"] = e.Usage.CacheWriteTokens
122 m["cache_measured_prompt"] = e.Usage.CacheMeasuredPromptTokens
123 m["cache_measured_read"] = e.Usage.CacheMeasuredReadTokens
124 m["cost_usd"] = e.Usage.CostUSD
125 m["cumulative"] = map[string]any{
126 "input": e.Cumulative.InputTokens,
127 "output": e.Cumulative.OutputTokens,
128 "reasoning": nullableReasoningTokens(e.Cumulative),
129 "cache_read": e.Cumulative.CacheReadTokens,
130 "cache_write": e.Cumulative.CacheWriteTokens,
131 "cache_measured_prompt": e.Cumulative.CacheMeasuredPromptTokens,
132 "cache_measured_read": e.Cumulative.CacheMeasuredReadTokens,
133 "cost_usd": e.Cumulative.CostUSD,
134 }
135 case core.EvTurnEnd:
136 m["stop"] = string(e.Stop)
137 if e.Err != nil {
138 m["error"] = e.Err.Error()
139 }
140 }
141 return m
142}
143
144// ContentToJSON serialises a transcript content slice into the same
145// shape used in EventToJSON. Exported alongside EventToJSON for the
146// RPC mode.
147func ContentToJSON(blocks []provider.Content) []map[string]any {
148 out := make([]map[string]any, 0, len(blocks))
149 for _, b := range blocks {
150 switch v := b.(type) {
151 case provider.TextBlock:
152 out = append(out, map[string]any{"type": "text", "text": v.Text})
153 case provider.ImageBlock:
154 out = append(out, map[string]any{"type": "image", "mime_type": v.MimeType, "bytes": len(v.Data)})
155 case provider.ToolCallBlock:
156 var args any
157 _ = json.Unmarshal(v.Arguments, &args)
158 out = append(out, map[string]any{"type": "tool_call", "id": v.ID, "name": v.Name, "args": args})
159 case provider.ToolResultBlock:
160 out = append(out, map[string]any{
161 "type": "tool_result",
162 "call_id": v.CallID,
163 "is_error": v.IsError,
164 "content": ContentToJSON(v.Content),
165 })
166 }
167 }
168 return out
169}
170
171func nullableReasoningTokens(usage provider.Usage) any {
172 if !usage.ReasoningTokensKnown {
173 return nil
174 }
175 return usage.ReasoningTokens
176}