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.EvTextDelta:
86 m["delta"] = e.Delta
87 case core.EvToolUseStart:
88 m["id"] = e.ID
89 m["name"] = e.Name
90 case core.EvToolUseArgs:
91 m["id"] = e.ID
92 m["delta"] = e.Delta
93 case core.EvToolUseEnd:
94 m["id"] = e.ID
95 case core.EvToolCall:
96 m["id"] = e.ID
97 m["name"] = e.Name
98 var args any
99 _ = json.Unmarshal(e.Args, &args)
100 m["args"] = args
101 case core.EvToolExecutionStarted:
102 m["id"] = e.ID
103 m["name"] = e.Name
104 case core.EvToolProgress:
105 m["id"] = e.ID
106 m["text"] = e.Text
107 case core.EvToolResult:
108 m["id"] = e.ID
109 m["is_error"] = e.Result.IsError
110 m["content"] = ContentToJSON(e.Result.Content)
111 case core.EvUsage:
112 m["input"] = e.Usage.InputTokens
113 m["output"] = e.Usage.OutputTokens
114 m["reasoning"] = nullableReasoningTokens(e.Usage)
115 m["cache_read"] = e.Usage.CacheReadTokens
116 m["cache_write"] = e.Usage.CacheWriteTokens
117 m["cost_usd"] = e.Usage.CostUSD
118 m["cumulative"] = map[string]any{
119 "input": e.Cumulative.InputTokens,
120 "output": e.Cumulative.OutputTokens,
121 "reasoning": nullableReasoningTokens(e.Cumulative),
122 "cache_read": e.Cumulative.CacheReadTokens,
123 "cache_write": e.Cumulative.CacheWriteTokens,
124 "cost_usd": e.Cumulative.CostUSD,
125 }
126 case core.EvTurnEnd:
127 m["stop"] = string(e.Stop)
128 if e.Err != nil {
129 m["error"] = e.Err.Error()
130 }
131 }
132 return m
133}
134
135// ContentToJSON serialises a transcript content slice into the same
136// shape used in EventToJSON. Exported alongside EventToJSON for the
137// RPC mode.
138func ContentToJSON(blocks []provider.Content) []map[string]any {
139 out := make([]map[string]any, 0, len(blocks))
140 for _, b := range blocks {
141 switch v := b.(type) {
142 case provider.TextBlock:
143 out = append(out, map[string]any{"type": "text", "text": v.Text})
144 case provider.ImageBlock:
145 out = append(out, map[string]any{"type": "image", "mime_type": v.MimeType, "bytes": len(v.Data)})
146 case provider.ToolCallBlock:
147 var args any
148 _ = json.Unmarshal(v.Arguments, &args)
149 out = append(out, map[string]any{"type": "tool_call", "id": v.ID, "name": v.Name, "args": args})
150 case provider.ToolResultBlock:
151 out = append(out, map[string]any{
152 "type": "tool_result",
153 "call_id": v.CallID,
154 "is_error": v.IsError,
155 "content": ContentToJSON(v.Content),
156 })
157 }
158 }
159 return out
160}
161
162func nullableReasoningTokens(usage provider.Usage) any {
163 if !usage.ReasoningTokensKnown {
164 return nil
165 }
166 return usage.ReasoningTokens
167}