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// Prompt sends a user message and runs the agent loop until the model
64// stops or an error occurs. Events are delivered via sink in order.
65// sink must not block the caller for long; buffer as needed.
66func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
67 if sink == nil {
68 sink = func(AgentEvent) {}
69 }
70 content := []provider.Content{}
71 if text != "" {
72 content = append(content, provider.TextBlock{Text: text})
73 }
74 for _, img := range images {
75 content = append(content, img)
76 }
77 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
78
79 a.mu.Lock()
80 a.messages = append(a.messages, user)
81 a.mu.Unlock()
82 sink(EvUserMessage{Message: user})
83
84 return a.runLoop(ctx, sink)
85}
86
87// Continue runs the agent loop against the existing transcript. Used
88// after appending tool results manually or to retry.
89func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
90 if sink == nil {
91 sink = func(AgentEvent) {}
92 }
93 return a.runLoop(ctx, sink)
94}
95
96func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
97 for step := 1; step <= a.MaxSteps; step++ {
98 sink(EvTurnStart{Step: step})
99 stop, assistantMsg, err := a.oneTurn(ctx, sink)
100 sink(EvTurnEnd{Stop: stop, Err: err})
101 if err != nil {
102 return err
103 }
104
105 if stop == provider.StopToolUse {
106 // Execute each tool call, append a single tool-results message, continue.
107 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
108 a.mu.Lock()
109 a.messages = append(a.messages, toolMsg)
110 a.mu.Unlock()
111 // If context was cancelled during tool execution, bail out.
112 if err := ctx.Err(); err != nil {
113 sink(EvDone{})
114 return err
115 }
116 _ = hadError
117 continue
118 }
119
120 // Terminal stop (end, length, error, aborted).
121 sink(EvDone{})
122 return nil
123 }
124 sink(EvDone{})
125 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
126}
127
128// oneTurn calls the LLM once, forwards events, returns the stop reason
129// and the assembled assistant message (already appended to the transcript).
130func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
131 req := provider.Request{
132 Model: a.Model,
133 System: a.System,
134 Messages: a.Messages(),
135 Tools: a.Tools.Specs(),
136 Reasoning: a.Reasoning,
137 }
138 stream, err := a.Client.Stream(ctx, req)
139 if err != nil {
140 return provider.StopError, provider.Message{}, err
141 }
142
143 sink(EvAssistantStart{})
144
145 var (
146 stop provider.StopReason
147 finalErr error
148 finalMsg provider.Message
149 )
150
151 for ev := range stream {
152 switch e := ev.(type) {
153 case provider.EventStart:
154 // nothing
155 case provider.EventTextDelta:
156 sink(EvTextDelta{Delta: e.Delta})
157 case provider.EventToolStart, provider.EventToolArgs, provider.EventToolEnd:
158 // The provider emits these to support partial-JSON UIs.
159 // We surface the complete tool call once per block at EvToolCall
160 // below (after EventDone), so for now we drop them.
161 case provider.EventUsage:
162 cum := a.cost.Add(e.Usage)
163 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
164 case provider.EventDone:
165 stop = e.Stop
166 finalErr = e.Err
167 finalMsg = e.Message
168 }
169 }
170
171 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
172 // produce partial, mid-sentence content that would confuse subsequent
173 // turns if it stayed in the transcript — drop it instead.
174 if len(finalMsg.Content) > 0 && stop != provider.StopAborted {
175 a.mu.Lock()
176 a.messages = append(a.messages, finalMsg)
177 a.mu.Unlock()
178 sink(EvAssistantMessage{Message: finalMsg})
179 // Now surface tool calls as EvToolCall events so UIs can render them
180 // in order before the tool results arrive.
181 for _, c := range finalMsg.Content {
182 if tc, ok := c.(provider.ToolCallBlock); ok {
183 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
184 }
185 }
186 }
187
188 return stop, finalMsg, finalErr
189}
190
191// executeTools runs every tool call in the assistant message and returns
192// a single tool-role message carrying all results.
193func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
194 var results []provider.Content
195 hadError := false
196
197 for _, c := range msg.Content {
198 tc, ok := c.(provider.ToolCallBlock)
199 if !ok {
200 continue
201 }
202 res := a.runOneTool(ctx, tc, sink)
203 if res.IsError {
204 hadError = true
205 }
206 results = append(results, provider.ToolResultBlock{
207 CallID: tc.ID,
208 Content: res.Content,
209 IsError: res.IsError,
210 })
211 sink(EvToolResult{ID: tc.ID, Result: res})
212 }
213
214 return provider.Message{
215 Role: provider.RoleTool,
216 Content: results,
217 Time: time.Now(),
218 }, hadError
219}
220
221func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
222 tool, err := a.Tools.Get(tc.Name)
223 if err != nil {
224 return ToolResult{
225 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
226 IsError: true,
227 }
228 }
229
230 args := tc.Arguments
231 if len(args) == 0 {
232 args = json.RawMessage("{}")
233 }
234
235 // Recover panics so a buggy tool does not crash the agent.
236 var res ToolResult
237 func() {
238 defer func() {
239 if r := recover(); r != nil {
240 res = ToolResult{
241 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
242 IsError: true,
243 }
244 }
245 }()
246 out, err := tool.Execute(ctx, args, func(text string) {
247 sink(EvToolProgress{ID: tc.ID, Text: text})
248 })
249 if err != nil {
250 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
251 res = ToolResult{
252 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
253 IsError: true,
254 }
255 return
256 }
257 res = ToolResult{
258 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
259 IsError: true,
260 }
261 return
262 }
263 res = out
264 }()
265 return res
266}