1package core
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "sync"
9 "time"
10
11 "github.com/patriceckhart/zot/internal/provider"
12)
13
14// Agent is a stateful conversation bound to a provider client, a model,
15// and a set of tools.
16type Agent struct {
17 Client provider.Client
18 Model string
19 System string
20 Tools Registry
21 MaxSteps int
22 Reasoning string
23
24 mu sync.Mutex
25 messages []provider.Message
26 cost CostTracker
27}
28
29// NewAgent returns an Agent with sensible defaults.
30func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
31 return &Agent{
32 Client: client,
33 Model: model,
34 System: system,
35 Tools: tools,
36 MaxSteps: 50,
37 }
38}
39
40// Messages returns a copy of the current transcript.
41func (a *Agent) Messages() []provider.Message {
42 a.mu.Lock()
43 defer a.mu.Unlock()
44 out := make([]provider.Message, len(a.messages))
45 copy(out, a.messages)
46 return out
47}
48
49// SetMessages replaces the transcript (used when resuming a session).
50func (a *Agent) SetMessages(msgs []provider.Message) {
51 a.mu.Lock()
52 defer a.mu.Unlock()
53 a.messages = append(a.messages[:0], msgs...)
54}
55
56// Cost returns the cumulative usage.
57func (a *Agent) Cost() provider.Usage {
58 a.mu.Lock()
59 defer a.mu.Unlock()
60 return a.cost.Total
61}
62
63// SeedCost sets the cumulative usage as a baseline before the first
64// turn runs. Used when transferring state from another agent (model
65// or provider switch) so the running cost meter doesn't reset to 0.
66func (a *Agent) SeedCost(u provider.Usage) {
67 a.mu.Lock()
68 defer a.mu.Unlock()
69 a.cost.Total = u
70}
71
72// Prompt sends a user message and runs the agent loop until the model
73// stops or an error occurs. Events are delivered via sink in order.
74// sink must not block the caller for long; buffer as needed.
75func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
76 if sink == nil {
77 sink = func(AgentEvent) {}
78 }
79 content := []provider.Content{}
80 if text != "" {
81 content = append(content, provider.TextBlock{Text: text})
82 }
83 for _, img := range images {
84 content = append(content, img)
85 }
86 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
87
88 a.mu.Lock()
89 a.messages = append(a.messages, user)
90 a.mu.Unlock()
91 sink(EvUserMessage{Message: user})
92
93 return a.runLoop(ctx, sink)
94}
95
96// Continue runs the agent loop against the existing transcript. Used
97// after appending tool results manually or to retry.
98func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
99 if sink == nil {
100 sink = func(AgentEvent) {}
101 }
102 return a.runLoop(ctx, sink)
103}
104
105func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
106 for step := 1; step <= a.MaxSteps; step++ {
107 sink(EvTurnStart{Step: step})
108 stop, assistantMsg, err := a.oneTurn(ctx, sink)
109 sink(EvTurnEnd{Stop: stop, Err: err})
110 if err != nil {
111 return err
112 }
113
114 if stop == provider.StopToolUse {
115 // Execute each tool call, append a single tool-results message, continue.
116 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
117 a.mu.Lock()
118 a.messages = append(a.messages, toolMsg)
119 a.mu.Unlock()
120 // If context was cancelled during tool execution, bail out.
121 if err := ctx.Err(); err != nil {
122 sink(EvDone{})
123 return err
124 }
125 _ = hadError
126 continue
127 }
128
129 // Terminal stop (end, length, error, aborted).
130 sink(EvDone{})
131 return nil
132 }
133 sink(EvDone{})
134 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
135}
136
137// oneTurn calls the LLM once, forwards events, returns the stop reason
138// and the assembled assistant message (already appended to the transcript).
139func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
140 req := provider.Request{
141 Model: a.Model,
142 System: a.System,
143 Messages: a.Messages(),
144 Tools: a.Tools.Specs(),
145 Reasoning: a.Reasoning,
146 }
147 stream, err := a.Client.Stream(ctx, req)
148 if err != nil {
149 return provider.StopError, provider.Message{}, err
150 }
151
152 sink(EvAssistantStart{})
153
154 var (
155 stop provider.StopReason
156 finalErr error
157 finalMsg provider.Message
158 )
159
160 for ev := range stream {
161 switch e := ev.(type) {
162 case provider.EventStart:
163 // nothing
164 case provider.EventTextDelta:
165 sink(EvTextDelta{Delta: e.Delta})
166 case provider.EventToolStart, provider.EventToolArgs, provider.EventToolEnd:
167 // The provider emits these to support partial-JSON UIs.
168 // We surface the complete tool call once per block at EvToolCall
169 // below (after EventDone), so for now we drop them.
170 case provider.EventUsage:
171 cum := a.cost.Add(e.Usage)
172 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
173 case provider.EventDone:
174 stop = e.Stop
175 finalErr = e.Err
176 finalMsg = e.Message
177 }
178 }
179
180 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
181 // produce partial, mid-sentence content that would confuse subsequent
182 // turns if it stayed in the transcript — drop it instead.
183 if len(finalMsg.Content) > 0 && stop != provider.StopAborted {
184 a.mu.Lock()
185 a.messages = append(a.messages, finalMsg)
186 a.mu.Unlock()
187 sink(EvAssistantMessage{Message: finalMsg})
188 // Now surface tool calls as EvToolCall events so UIs can render them
189 // in order before the tool results arrive.
190 for _, c := range finalMsg.Content {
191 if tc, ok := c.(provider.ToolCallBlock); ok {
192 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
193 }
194 }
195 }
196
197 return stop, finalMsg, finalErr
198}
199
200// executeTools runs every tool call in the assistant message and returns
201// a single tool-role message carrying all results.
202func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
203 var results []provider.Content
204 hadError := false
205
206 for _, c := range msg.Content {
207 tc, ok := c.(provider.ToolCallBlock)
208 if !ok {
209 continue
210 }
211 res := a.runOneTool(ctx, tc, sink)
212 if res.IsError {
213 hadError = true
214 }
215 results = append(results, provider.ToolResultBlock{
216 CallID: tc.ID,
217 Content: res.Content,
218 IsError: res.IsError,
219 })
220 sink(EvToolResult{ID: tc.ID, Result: res})
221 }
222
223 return provider.Message{
224 Role: provider.RoleTool,
225 Content: results,
226 Time: time.Now(),
227 }, hadError
228}
229
230func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
231 tool, err := a.Tools.Get(tc.Name)
232 if err != nil {
233 return ToolResult{
234 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
235 IsError: true,
236 }
237 }
238
239 args := tc.Arguments
240 if len(args) == 0 {
241 args = json.RawMessage("{}")
242 }
243
244 // Recover panics so a buggy tool does not crash the agent.
245 var res ToolResult
246 func() {
247 defer func() {
248 if r := recover(); r != nil {
249 res = ToolResult{
250 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
251 IsError: true,
252 }
253 }
254 }()
255 out, err := tool.Execute(ctx, args, func(text string) {
256 sink(EvToolProgress{ID: tc.ID, Text: text})
257 })
258 if err != nil {
259 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
260 res = ToolResult{
261 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
262 IsError: true,
263 }
264 return
265 }
266 res = ToolResult{
267 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
268 IsError: true,
269 }
270 return
271 }
272 res = out
273 }()
274 return res
275}