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