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. Optionally,
27 // returning a non-nil modifiedArgs replaces the JSON args the
28 // tool will see, which lets guards redact / augment / patch the
29 // model's request without rewriting the transcript. Empty or
30 // malformed modifiedArgs is ignored.
31 BeforeToolExecute func(call provider.ToolCallBlock) (allowed bool, reason string, modifiedArgs json.RawMessage)
32
33 // BeforeTurn, if set, is called before each turn's model call.
34 // Returning (allowed=false, reason) aborts the turn; reason is
35 // surfaced as an assistant-like status line. Used for rate-
36 // limiting, business-hour gates, and deny-by-default setups.
37 BeforeTurn func(step int) (allowed bool, reason string)
38
39 // BeforeAssistantMessage, if set, is called after the model's
40 // final assistant message is assembled but before it's appended
41 // to the transcript. Returning (allowed=false) suppresses both
42 // the transcript append and the UI event. A non-empty
43 // replacement rewrites the visible text for the user while
44 // leaving the model's original text in the transcript (so the
45 // model can still see what it said in subsequent turns).
46 BeforeAssistantMessage func(text string) (allowed bool, reason, replacement string)
47
48 // OnEvent, if set, mirrors every AgentEvent the loop emits to
49 // this callback in addition to the per-Prompt sink. Used by the
50 // extension manager to fan events out to subscribed extensions
51 // without each caller having to compose sinks manually.
52 OnEvent func(AgentEvent)
53
54 mu sync.Mutex
55 messages []provider.Message
56 cost CostTracker
57}
58
59// NewAgent returns an Agent with sensible defaults.
60func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
61 return &Agent{
62 Client: client,
63 Model: model,
64 System: system,
65 Tools: tools,
66 MaxSteps: 50,
67 }
68}
69
70// Messages returns a copy of the current transcript.
71func (a *Agent) Messages() []provider.Message {
72 a.mu.Lock()
73 defer a.mu.Unlock()
74 out := make([]provider.Message, len(a.messages))
75 copy(out, a.messages)
76 return out
77}
78
79// SetTools swaps the tool registry. Used by /reload-ext to hand
80// the agent a fresh registry after extension subprocesses have been
81// respawned (and their freshly-registered tools merged in).
82func (a *Agent) SetTools(reg Registry) {
83 a.mu.Lock()
84 a.Tools = reg
85 a.mu.Unlock()
86}
87
88// SetMessages replaces the transcript (used when resuming a session).
89func (a *Agent) SetMessages(msgs []provider.Message) {
90 a.mu.Lock()
91 defer a.mu.Unlock()
92 a.messages = append(a.messages[:0], msgs...)
93}
94
95// Cost returns the cumulative usage.
96func (a *Agent) Cost() provider.Usage {
97 a.mu.Lock()
98 defer a.mu.Unlock()
99 return a.cost.Total
100}
101
102// SeedCost sets the cumulative usage as a baseline before the first
103// turn runs. Used when transferring state from another agent (model
104// or provider switch) so the running cost meter doesn't reset to 0.
105func (a *Agent) SeedCost(u provider.Usage) {
106 a.mu.Lock()
107 defer a.mu.Unlock()
108 a.cost.Total = u
109}
110
111// Prompt sends a user message and runs the agent loop until the model
112// stops or an error occurs. Events are delivered via sink in order.
113// sink must not block the caller for long; buffer as needed.
114func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
115 if sink == nil {
116 sink = func(AgentEvent) {}
117 }
118 sink = a.wrapSink(sink)
119 content := []provider.Content{}
120 if text != "" {
121 content = append(content, provider.TextBlock{Text: text})
122 }
123 for _, img := range images {
124 content = append(content, img)
125 }
126 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
127
128 a.mu.Lock()
129 a.messages = append(a.messages, user)
130 a.mu.Unlock()
131 sink(EvUserMessage{Message: user})
132
133 return a.runLoop(ctx, sink)
134}
135
136// Continue runs the agent loop against the existing transcript. Used
137// after appending tool results manually or to retry.
138func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
139 if sink == nil {
140 sink = func(AgentEvent) {}
141 }
142 sink = a.wrapSink(sink)
143 return a.runLoop(ctx, sink)
144}
145
146// wrapSink composes the per-call sink with a.OnEvent (if set) so the
147// extension manager (or any other observer) sees every AgentEvent
148// without having to thread itself through every Prompt callsite.
149func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) {
150 if a.OnEvent == nil {
151 return sink
152 }
153 obs := a.OnEvent
154 return func(ev AgentEvent) {
155 obs(ev)
156 sink(ev)
157 }
158}
159
160func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
161 for step := 1; step <= a.MaxSteps; step++ {
162 sink(EvTurnStart{Step: step})
163 if a.BeforeTurn != nil {
164 if allowed, reason := a.BeforeTurn(step); !allowed {
165 if reason == "" {
166 reason = "turn blocked by extension guard"
167 }
168 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
169 sink(EvDone{})
170 return nil
171 }
172 }
173 stop, assistantMsg, err := a.oneTurn(ctx, sink)
174 sink(EvTurnEnd{Stop: stop, Err: err})
175 if err != nil {
176 return err
177 }
178
179 if stop == provider.StopToolUse {
180 // Execute each tool call, append a single tool-results message, continue.
181 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
182 a.mu.Lock()
183 a.messages = append(a.messages, toolMsg)
184 // OpenAI's chat-completions tool message shape is text-centric.
185 // Vision models reliably consume images when they arrive as user
186 // content, so when a tool result contains images we mirror them
187 // into a synthetic user message immediately after the tool result.
188 // This keeps the transcript self-contained for providers that can
189 // see image blocks in tool messages while making OpenAI vision
190 // models actually receive the image bytes.
191 if a.Client != nil && a.Client.Name() == "openai" {
192 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
193 a.messages = append(a.messages, mirror)
194 }
195 }
196 a.mu.Unlock()
197 // If context was cancelled during tool execution, bail out.
198 if err := ctx.Err(); err != nil {
199 sink(EvDone{})
200 return err
201 }
202 _ = hadError
203 continue
204 }
205
206 // Terminal stop (end, length, error, aborted).
207 sink(EvDone{})
208 return nil
209 }
210 sink(EvDone{})
211 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
212}
213
214// oneTurn calls the LLM once, forwards events, returns the stop reason
215// and the assembled assistant message (already appended to the transcript).
216func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
217 req := provider.Request{
218 Model: a.Model,
219 System: a.System,
220 Messages: a.Messages(),
221 Tools: a.Tools.Specs(),
222 Reasoning: a.Reasoning,
223 }
224 stream, err := a.Client.Stream(ctx, req)
225 if err != nil {
226 return provider.StopError, provider.Message{}, err
227 }
228
229 sink(EvAssistantStart{})
230
231 var (
232 stop provider.StopReason
233 finalErr error
234 finalMsg provider.Message
235 )
236
237 for ev := range stream {
238 switch e := ev.(type) {
239 case provider.EventStart:
240 // nothing
241 case provider.EventTextDelta:
242 sink(EvTextDelta{Delta: e.Delta})
243 case provider.EventToolStart:
244 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
245 case provider.EventToolArgs:
246 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
247 case provider.EventToolEnd:
248 sink(EvToolUseEnd{ID: e.ID})
249 case provider.EventUsage:
250 cum := a.cost.Add(e.Usage)
251 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
252 case provider.EventDone:
253 stop = e.Stop
254 finalErr = e.Err
255 finalMsg = e.Message
256 }
257 }
258
259 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
260 // produce partial, mid-sentence content that would confuse subsequent
261 // turns if it stayed in the transcript, drop it instead.
262 if len(finalMsg.Content) > 0 && stop != provider.StopAborted {
263 emit := finalMsg
264 suppress := false
265
266 // BeforeAssistantMessage hook: extensions can suppress or
267 // rewrite the visible text. The transcript keeps the
268 // model's original output so the model still sees what it
269 // said on subsequent turns.
270 if a.BeforeAssistantMessage != nil {
271 orig := extractText(finalMsg)
272 if orig != "" {
273 allowed, _, replacement := a.BeforeAssistantMessage(orig)
274 if !allowed {
275 suppress = true
276 } else if replacement != "" && replacement != orig {
277 emit = replaceText(finalMsg, replacement)
278 }
279 }
280 }
281
282 a.mu.Lock()
283 a.messages = append(a.messages, finalMsg)
284 a.mu.Unlock()
285 if !suppress {
286 sink(EvAssistantMessage{Message: emit})
287 }
288 // Now surface tool calls as EvToolCall events so UIs can render them
289 // in order before the tool results arrive.
290 for _, c := range finalMsg.Content {
291 if tc, ok := c.(provider.ToolCallBlock); ok {
292 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
293 }
294 }
295 }
296
297 return stop, finalMsg, finalErr
298}
299
300// executeTools runs every tool call in the assistant message and returns
301// a single tool-role message carrying all results.
302func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
303 var results []provider.Content
304 hadError := false
305
306 for _, c := range msg.Content {
307 tc, ok := c.(provider.ToolCallBlock)
308 if !ok {
309 continue
310 }
311 res := a.runOneTool(ctx, tc, sink)
312 if res.IsError {
313 hadError = true
314 }
315 results = append(results, provider.ToolResultBlock{
316 CallID: tc.ID,
317 Content: res.Content,
318 IsError: res.IsError,
319 })
320 sink(EvToolResult{ID: tc.ID, Result: res})
321 }
322
323 return provider.Message{
324 Role: provider.RoleTool,
325 Content: results,
326 Time: time.Now(),
327 }, hadError
328}
329
330func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
331 tool, err := a.Tools.Get(tc.Name)
332 if err != nil {
333 return ToolResult{
334 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
335 IsError: true,
336 }
337 }
338
339 args := tc.Arguments
340
341 // Intercept hook: an extension or other guard can refuse the
342 // call before any side effect happens, OR rewrite the args
343 // seen by the tool. The model sees the reason as the tool
344 // error, learns from it, and (typically) proposes a different
345 // action; rewrites are invisible to the model (they apply only
346 // to the execution).
347 if a.BeforeToolExecute != nil {
348 allowed, reason, modified := a.BeforeToolExecute(tc)
349 if !allowed {
350 if reason == "" {
351 reason = "tool call refused by extension guard"
352 }
353 return ToolResult{
354 Content: []provider.Content{provider.TextBlock{Text: reason}},
355 IsError: true,
356 }
357 }
358 if len(modified) > 0 && json.Valid(modified) {
359 args = modified
360 }
361 }
362
363 if len(args) == 0 {
364 args = json.RawMessage("{}")
365 }
366
367 // Recover panics so a buggy tool does not crash the agent.
368 var res ToolResult
369 func() {
370 defer func() {
371 if r := recover(); r != nil {
372 res = ToolResult{
373 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
374 IsError: true,
375 }
376 }
377 }()
378 out, err := tool.Execute(ctx, args, func(text string) {
379 sink(EvToolProgress{ID: tc.ID, Text: text})
380 })
381 if err != nil {
382 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
383 res = ToolResult{
384 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
385 IsError: true,
386 }
387 return
388 }
389 res = ToolResult{
390 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
391 IsError: true,
392 }
393 return
394 }
395 res = out
396 }()
397 return res
398}
399
400// extractText concatenates all TextBlock content in a message. Used
401// by BeforeAssistantMessage so guards see a single string instead of
402// having to walk provider.Content themselves.
403func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
404 var content []provider.Content
405 for _, c := range msg.Content {
406 tr, ok := c.(provider.ToolResultBlock)
407 if !ok {
408 continue
409 }
410 for _, inner := range tr.Content {
411 switch v := inner.(type) {
412 case provider.TextBlock:
413 // Keep short textual context so the model understands why
414 // the images appeared, but don't duplicate giant read
415 // outputs verbatim.
416 if len(v.Text) > 0 && len(v.Text) <= 500 {
417 content = append(content, v)
418 }
419 case provider.ImageBlock:
420 content = append(content, v)
421 }
422 }
423 }
424 if len(content) == 0 {
425 return provider.Message{}
426 }
427 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
428 content = append([]provider.Content{prefix}, content...)
429 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
430}
431
432func extractText(msg provider.Message) string {
433 var out string
434 for _, c := range msg.Content {
435 if tb, ok := c.(provider.TextBlock); ok {
436 if out != "" {
437 out += "\n"
438 }
439 out += tb.Text
440 }
441 }
442 return out
443}
444
445// replaceText returns a copy of msg with every TextBlock replaced by
446// a single TextBlock containing replacement. Non-text content (tool
447// calls, etc.) is preserved in order.
448func replaceText(msg provider.Message, replacement string) provider.Message {
449 out := provider.Message{Role: msg.Role}
450 out.Content = make([]provider.Content, 0, len(msg.Content))
451 replaced := false
452 for _, c := range msg.Content {
453 if _, ok := c.(provider.TextBlock); ok {
454 if !replaced {
455 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
456 replaced = true
457 }
458 continue
459 }
460 out.Content = append(out.Content, c)
461 }
462 if !replaced {
463 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
464 }
465 return out
466}