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 // OnMessageAppended, if set, fires every time a message is
55 // appended to the in-memory transcript by the agent loop — the
56 // initial user prompt, each finalised assistant message, and
57 // each tool-results message (plus the synthetic OpenAI image
58 // mirror, if any). Hosts wire this to the on-disk session so
59 // that turns are durable as soon as they happen, instead of
60 // only being flushed on a clean exit.
61 OnMessageAppended func(provider.Message)
62
63 // OnUsage, if set, fires after every turn's usage row arrives,
64 // carrying the cumulative usage for the session. Hosts wire
65 // this to the on-disk session so the persisted total stays
66 // current and a crash recovers the right cost figure.
67 OnUsage func(cumulative provider.Usage)
68
69 mu sync.Mutex
70 messages []provider.Message
71 // rev increments whenever the transcript slice is replaced or a
72 // message is appended. The TUI uses it as a cheap redraw cache key
73 // so editor-only typing doesn't copy/rebuild a long transcript on
74 // every keypress.
75 rev uint64
76 cost CostTracker
77}
78
79// NewAgent returns an Agent with sensible defaults.
80func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
81 return &Agent{
82 Client: client,
83 Model: model,
84 System: system,
85 Tools: tools,
86 MaxSteps: 0, // 0 = unlimited
87 }
88}
89
90// Messages returns a copy of the current transcript.
91func (a *Agent) Messages() []provider.Message {
92 a.mu.Lock()
93 defer a.mu.Unlock()
94 out := make([]provider.Message, len(a.messages))
95 copy(out, a.messages)
96 return out
97}
98
99// Revision returns a monotonically increasing transcript version.
100// It is cheap to query and changes whenever Messages() would return
101// different transcript content because of append/set operations.
102func (a *Agent) Revision() uint64 {
103 a.mu.Lock()
104 defer a.mu.Unlock()
105 return a.rev
106}
107
108// SetTools swaps the tool registry. Used by /reload-ext to hand
109// the agent a fresh registry after extension subprocesses have been
110// respawned (and their freshly-registered tools merged in).
111func (a *Agent) SetTools(reg Registry) {
112 a.mu.Lock()
113 a.Tools = reg
114 a.mu.Unlock()
115}
116
117// SetMessages replaces the transcript (used when resuming a session).
118func (a *Agent) SetMessages(msgs []provider.Message) {
119 a.mu.Lock()
120 defer a.mu.Unlock()
121 a.messages = append(a.messages[:0], msgs...)
122 a.rev++
123}
124
125// Cost returns the cumulative usage.
126func (a *Agent) Cost() provider.Usage {
127 a.mu.Lock()
128 defer a.mu.Unlock()
129 return a.cost.Total
130}
131
132// SeedCost sets the cumulative usage as a baseline before the first
133// turn runs. Used when transferring state from another agent (model
134// or provider switch) so the running cost meter doesn't reset to 0.
135func (a *Agent) SeedCost(u provider.Usage) {
136 a.mu.Lock()
137 defer a.mu.Unlock()
138 a.cost.Total = u
139}
140
141// fireMessageAppended invokes OnMessageAppended without holding the
142// agent mutex, so the host's persistence callback can take its own
143// locks without deadlocking the agent loop. Tolerates a nil hook so
144// non-persisting callers (tests, RPC mode) don't have to set it.
145func (a *Agent) fireMessageAppended(m provider.Message) {
146 if a.OnMessageAppended != nil {
147 a.OnMessageAppended(m)
148 }
149}
150
151// Prompt sends a user message and runs the agent loop until the model
152// stops or an error occurs. Events are delivered via sink in order.
153// sink must not block the caller for long; buffer as needed.
154func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
155 if sink == nil {
156 sink = func(AgentEvent) {}
157 }
158 sink = a.wrapSink(sink)
159 content := []provider.Content{}
160 if text != "" {
161 content = append(content, provider.TextBlock{Text: text})
162 }
163 for _, img := range images {
164 content = append(content, img)
165 }
166 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
167
168 a.mu.Lock()
169 a.messages = append(a.messages, user)
170 a.rev++
171 a.mu.Unlock()
172 a.fireMessageAppended(user)
173 sink(EvUserMessage{Message: user})
174
175 return a.runLoop(ctx, sink)
176}
177
178// Continue runs the agent loop against the existing transcript. Used
179// after appending tool results manually or to retry.
180func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
181 if sink == nil {
182 sink = func(AgentEvent) {}
183 }
184 sink = a.wrapSink(sink)
185 return a.runLoop(ctx, sink)
186}
187
188// wrapSink composes the per-call sink with a.OnEvent (if set) so the
189// extension manager (or any other observer) sees every AgentEvent
190// without having to thread itself through every Prompt callsite.
191func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) {
192 if a.OnEvent == nil {
193 return sink
194 }
195 obs := a.OnEvent
196 return func(ev AgentEvent) {
197 obs(ev)
198 sink(ev)
199 }
200}
201
202func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
203 for step := 1; a.MaxSteps <= 0 || step <= a.MaxSteps; step++ {
204 sink(EvTurnStart{Step: step})
205 if a.BeforeTurn != nil {
206 if allowed, reason := a.BeforeTurn(step); !allowed {
207 if reason == "" {
208 reason = "turn blocked by extension guard"
209 }
210 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
211 sink(EvDone{})
212 return nil
213 }
214 }
215 stop, assistantMsg, err := a.oneTurn(ctx, sink)
216 sink(EvTurnEnd{Stop: stop, Err: err})
217 if err != nil {
218 return err
219 }
220
221 if stop == provider.StopToolUse {
222 // Execute each tool call, append a single tool-results message, continue.
223 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
224 a.mu.Lock()
225 a.messages = append(a.messages, toolMsg)
226 a.rev++
227 // OpenAI's chat-completions tool message shape is text-centric.
228 // Vision models reliably consume images when they arrive as user
229 // content, so when a tool result contains images we mirror them
230 // into a synthetic user message immediately after the tool result.
231 // This keeps the transcript self-contained for providers that can
232 // see image blocks in tool messages while making OpenAI vision
233 // models actually receive the image bytes.
234 var imageMirror provider.Message
235 if a.Client != nil && a.Client.Name() == "openai" {
236 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
237 a.messages = append(a.messages, mirror)
238 a.rev++
239 imageMirror = mirror
240 }
241 }
242 a.mu.Unlock()
243 a.fireMessageAppended(toolMsg)
244 if len(imageMirror.Content) > 0 {
245 a.fireMessageAppended(imageMirror)
246 }
247 // If context was cancelled during tool execution, bail out.
248 if err := ctx.Err(); err != nil {
249 sink(EvDone{})
250 return err
251 }
252 _ = hadError
253 continue
254 }
255
256 // Terminal stop (end, length, error, aborted).
257 sink(EvDone{})
258 return nil
259 }
260 if a.MaxSteps > 0 {
261 sink(EvDone{})
262 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
263 }
264 return nil
265}
266
267// oneTurn calls the LLM once, forwards events, returns the stop reason
268// and the assembled assistant message (already appended to the transcript).
269func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
270 req := provider.Request{
271 Model: a.Model,
272 System: a.System,
273 Messages: a.Messages(),
274 Tools: a.Tools.Specs(),
275 Reasoning: a.Reasoning,
276 }
277 stream, err := a.Client.Stream(ctx, req)
278 if err != nil {
279 return provider.StopError, provider.Message{}, err
280 }
281
282 sink(EvAssistantStart{})
283
284 var (
285 stop provider.StopReason
286 finalErr error
287 finalMsg provider.Message
288 )
289
290 for ev := range stream {
291 switch e := ev.(type) {
292 case provider.EventStart:
293 // nothing
294 case provider.EventTextDelta:
295 sink(EvTextDelta{Delta: e.Delta})
296 case provider.EventToolStart:
297 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
298 case provider.EventToolArgs:
299 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
300 case provider.EventToolEnd:
301 sink(EvToolUseEnd{ID: e.ID})
302 case provider.EventUsage:
303 cum := a.cost.Add(e.Usage)
304 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
305 if a.OnUsage != nil {
306 a.OnUsage(cum)
307 }
308 case provider.EventDone:
309 stop = e.Stop
310 finalErr = e.Err
311 finalMsg = e.Message
312 }
313 }
314
315 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
316 // produce partial content. When the partial message is text only we
317 // keep whatever was streamed up to the cancel so the user does not
318 // lose visible work (a cut-off summary is still useful). If the
319 // partial message already contained tool-call blocks we drop the
320 // whole thing, because an unmatched tool_use would fail the next
321 // turn with a tool_result mismatch error.
322 keep := len(finalMsg.Content) > 0
323 if stop == provider.StopAborted && keep {
324 hasToolCall := false
325 for _, c := range finalMsg.Content {
326 if _, ok := c.(provider.ToolCallBlock); ok {
327 hasToolCall = true
328 break
329 }
330 }
331 if hasToolCall {
332 keep = false
333 }
334 }
335 if keep {
336 emit := finalMsg
337 suppress := false
338
339 // BeforeAssistantMessage hook: extensions can suppress or
340 // rewrite the visible text. The transcript keeps the
341 // model's original output so the model still sees what it
342 // said on subsequent turns.
343 if a.BeforeAssistantMessage != nil {
344 orig := extractText(finalMsg)
345 if orig != "" {
346 allowed, _, replacement := a.BeforeAssistantMessage(orig)
347 if !allowed {
348 suppress = true
349 } else if replacement != "" && replacement != orig {
350 emit = replaceText(finalMsg, replacement)
351 }
352 }
353 }
354
355 a.mu.Lock()
356 a.messages = append(a.messages, finalMsg)
357 a.rev++
358 a.mu.Unlock()
359 a.fireMessageAppended(finalMsg)
360 if !suppress {
361 sink(EvAssistantMessage{Message: emit})
362 }
363 // Now surface tool calls as EvToolCall events so UIs can render them
364 // in order before the tool results arrive.
365 for _, c := range finalMsg.Content {
366 if tc, ok := c.(provider.ToolCallBlock); ok {
367 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
368 }
369 }
370 }
371
372 return stop, finalMsg, finalErr
373}
374
375// executeTools runs every tool call in the assistant message and returns
376// a single tool-role message carrying all results.
377func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
378 var results []provider.Content
379 hadError := false
380
381 for _, c := range msg.Content {
382 tc, ok := c.(provider.ToolCallBlock)
383 if !ok {
384 continue
385 }
386 res := a.runOneTool(ctx, tc, sink)
387 if res.IsError {
388 hadError = true
389 }
390 results = append(results, provider.ToolResultBlock{
391 CallID: tc.ID,
392 Content: res.Content,
393 IsError: res.IsError,
394 })
395 sink(EvToolResult{ID: tc.ID, Result: res})
396 }
397
398 return provider.Message{
399 Role: provider.RoleTool,
400 Content: results,
401 Time: time.Now(),
402 }, hadError
403}
404
405func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
406 tool, err := a.Tools.Get(tc.Name)
407 if err != nil {
408 return ToolResult{
409 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
410 IsError: true,
411 }
412 }
413
414 args := tc.Arguments
415
416 // Intercept hook: an extension or other guard can refuse the
417 // call before any side effect happens, OR rewrite the args
418 // seen by the tool. The model sees the reason as the tool
419 // error, learns from it, and (typically) proposes a different
420 // action; rewrites are invisible to the model (they apply only
421 // to the execution).
422 if a.BeforeToolExecute != nil {
423 allowed, reason, modified := a.BeforeToolExecute(tc)
424 if !allowed {
425 if reason == "" {
426 reason = "tool call refused by extension guard"
427 }
428 return ToolResult{
429 Content: []provider.Content{provider.TextBlock{Text: reason}},
430 IsError: true,
431 }
432 }
433 if len(modified) > 0 && json.Valid(modified) {
434 args = modified
435 }
436 }
437
438 if len(args) == 0 {
439 args = json.RawMessage("{}")
440 }
441
442 // Recover panics so a buggy tool does not crash the agent.
443 var res ToolResult
444 func() {
445 defer func() {
446 if r := recover(); r != nil {
447 res = ToolResult{
448 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
449 IsError: true,
450 }
451 }
452 }()
453 out, err := tool.Execute(ctx, args, func(text string) {
454 sink(EvToolProgress{ID: tc.ID, Text: text})
455 })
456 if err != nil {
457 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
458 res = ToolResult{
459 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
460 IsError: true,
461 }
462 return
463 }
464 res = ToolResult{
465 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
466 IsError: true,
467 }
468 return
469 }
470 res = out
471 }()
472 return res
473}
474
475// extractText concatenates all TextBlock content in a message. Used
476// by BeforeAssistantMessage so guards see a single string instead of
477// having to walk provider.Content themselves.
478func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
479 var content []provider.Content
480 for _, c := range msg.Content {
481 tr, ok := c.(provider.ToolResultBlock)
482 if !ok {
483 continue
484 }
485 for _, inner := range tr.Content {
486 switch v := inner.(type) {
487 case provider.TextBlock:
488 // Keep short textual context so the model understands why
489 // the images appeared, but don't duplicate giant read
490 // outputs verbatim.
491 if len(v.Text) > 0 && len(v.Text) <= 500 {
492 content = append(content, v)
493 }
494 case provider.ImageBlock:
495 content = append(content, v)
496 }
497 }
498 }
499 if len(content) == 0 {
500 return provider.Message{}
501 }
502 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
503 content = append([]provider.Content{prefix}, content...)
504 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
505}
506
507func extractText(msg provider.Message) string {
508 var out string
509 for _, c := range msg.Content {
510 if tb, ok := c.(provider.TextBlock); ok {
511 if out != "" {
512 out += "\n"
513 }
514 out += tb.Text
515 }
516 }
517 return out
518}
519
520// replaceText returns a copy of msg with every TextBlock replaced by
521// a single TextBlock containing replacement. Non-text content (tool
522// calls, etc.) is preserved in order.
523func replaceText(msg provider.Message, replacement string) provider.Message {
524 out := provider.Message{Role: msg.Role}
525 out.Content = make([]provider.Content, 0, len(msg.Content))
526 replaced := false
527 for _, c := range msg.Content {
528 if _, ok := c.(provider.TextBlock); ok {
529 if !replaced {
530 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
531 replaced = true
532 }
533 continue
534 }
535 out.Content = append(out.Content, c)
536 }
537 if !replaced {
538 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
539 }
540 return out
541}