1package core
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "strings"
9 "sync"
10 "time"
11
12 "github.com/patriceckhart/zot/packages/provider"
13)
14
15// Agent is a stateful conversation bound to a provider client, a model,
16// and a set of tools.
17type Agent struct {
18 Client provider.Client
19 Model string
20 System string
21 Tools Registry
22 MaxSteps int
23 Reasoning string
24
25 // BeforeToolExecute, if set, is called immediately before each
26 // tool runs. Returning (allowed=false, reason) short-circuits
27 // the call with an error result containing reason. Optionally,
28 // returning a non-nil modifiedArgs replaces the JSON args the
29 // tool will see, which lets guards redact / augment / patch the
30 // model's request without rewriting the transcript. Empty or
31 // malformed modifiedArgs is ignored.
32 BeforeToolExecute func(call provider.ToolCallBlock) (allowed bool, reason string, modifiedArgs json.RawMessage)
33
34 // BeforeTurn, if set, is called before each turn's model call.
35 // Returning (allowed=false, reason) aborts the turn; reason is
36 // surfaced as an assistant-like status line. Used for rate-
37 // limiting, business-hour gates, and deny-by-default setups.
38 BeforeTurn func(step int) (allowed bool, reason string)
39
40 // BeforeAssistantMessage, if set, is called after the model's
41 // final assistant message is assembled but before it's appended
42 // to the transcript. Returning (allowed=false) suppresses both
43 // the transcript append and the UI event. A non-empty
44 // replacement rewrites the visible text for the user while
45 // leaving the model's original text in the transcript (so the
46 // model can still see what it said in subsequent turns).
47 BeforeAssistantMessage func(text string) (allowed bool, reason, replacement string)
48
49 // OnEvent, if set, mirrors every AgentEvent the loop emits to
50 // this callback in addition to the per-Prompt sink. Used by the
51 // extension manager to fan events out to subscribed extensions
52 // without each caller having to compose sinks manually.
53 OnEvent func(AgentEvent)
54
55 // OnMessageAppended, if set, fires every time a message is
56 // appended to the in-memory transcript by the agent loop — the
57 // initial user prompt, each finalised assistant message, and
58 // each tool-results message (plus the synthetic OpenAI image
59 // mirror, if any). Hosts wire this to the on-disk session so
60 // that turns are durable as soon as they happen, instead of
61 // only being flushed on a clean exit.
62 OnMessageAppended func(provider.Message)
63
64 // OnUsage, if set, fires after every turn's usage row arrives,
65 // carrying the cumulative usage for the session. Hosts wire
66 // this to the on-disk session so the persisted total stays
67 // current and a crash recovers the right cost figure.
68 OnUsage func(cumulative provider.Usage)
69
70 // OnTranscriptCompacted, if set, fires after Compact replaces the
71 // in-memory transcript with the synthetic summary plus kept tail.
72 // Hosts wire this to append an explicit compaction checkpoint to
73 // the session log; per-message append hooks do not fire for this
74 // wholesale transcript replacement.
75 OnTranscriptCompacted func(messages []provider.Message)
76
77 mu sync.Mutex
78 messages []provider.Message
79 // rev increments whenever the transcript slice is replaced or a
80 // message is appended. The TUI uses it as a cheap redraw cache key
81 // so editor-only typing doesn't copy/rebuild a long transcript on
82 // every keypress.
83 rev uint64
84 cost CostTracker
85
86 // queued holds user messages submitted while the agent is busy.
87 // The loop appends them as normal user messages at safe
88 // boundaries: before the next model call after a tool batch, or
89 // after a text-only assistant turn finishes. It never interrupts
90 // a running tool or cancels an in-flight provider request.
91 queued []string
92}
93
94// NewAgent returns an Agent with sensible defaults.
95func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
96 return &Agent{
97 Client: client,
98 Model: model,
99 System: system,
100 Tools: tools,
101 MaxSteps: 0, // 0 = unlimited
102 }
103}
104
105// QueueMessage queues text to be injected as a user message at the
106// next safe boundary of the active agent loop. It is non-blocking in
107// the sense that it never waits for model/tool work; it only takes
108// the transcript mutex briefly. Empty/whitespace-only messages are
109// ignored.
110func (a *Agent) QueueMessage(text string) bool {
111 text = strings.TrimSpace(text)
112 if text == "" {
113 return false
114 }
115 a.mu.Lock()
116 a.queued = append(a.queued, text)
117 a.mu.Unlock()
118 return true
119}
120
121// PendingQueuedMessages returns a snapshot of user messages waiting
122// to be injected. Used by hosts to render the visible "sliding in"
123// chips without consuming them.
124func (a *Agent) PendingQueuedMessages() []string {
125 a.mu.Lock()
126 defer a.mu.Unlock()
127 out := make([]string, len(a.queued))
128 copy(out, a.queued)
129 return out
130}
131
132// QueuedMessageCount returns the number of messages waiting to be
133// injected at the next safe boundary.
134func (a *Agent) QueuedMessageCount() int {
135 a.mu.Lock()
136 defer a.mu.Unlock()
137 return len(a.queued)
138}
139
140// PopQueuedMessage removes and returns the most recently queued
141// message. Hosts use this for the slide-back keybinding.
142func (a *Agent) PopQueuedMessage() (string, bool) {
143 a.mu.Lock()
144 defer a.mu.Unlock()
145 n := len(a.queued)
146 if n == 0 {
147 return "", false
148 }
149 text := a.queued[n-1]
150 a.queued = a.queued[:n-1]
151 return text, true
152}
153
154// DrainQueuedMessages discards and returns every queued message.
155// Hosts use this on explicit cancel/clear so stale follow-ups do
156// not run after the user aborted the turn.
157func (a *Agent) DrainQueuedMessages() []string {
158 return a.drainQueuedMessages()
159}
160
161func (a *Agent) drainQueuedMessages() []string {
162 a.mu.Lock()
163 defer a.mu.Unlock()
164 out := make([]string, len(a.queued))
165 copy(out, a.queued)
166 a.queued = nil
167 return out
168}
169
170func (a *Agent) appendQueuedAsUser(texts []string, sink func(AgentEvent)) {
171 for _, text := range texts {
172 msg := provider.Message{
173 Role: provider.RoleUser,
174 Content: []provider.Content{provider.TextBlock{Text: text}},
175 Time: time.Now(),
176 }
177 a.mu.Lock()
178 a.messages = append(a.messages, msg)
179 a.rev++
180 a.mu.Unlock()
181 a.fireMessageAppended(msg)
182 if sink != nil {
183 sink(EvUserMessage{Message: msg})
184 }
185 }
186}
187
188// Messages returns a copy of the current transcript.
189func (a *Agent) Messages() []provider.Message {
190 a.mu.Lock()
191 defer a.mu.Unlock()
192 out := make([]provider.Message, len(a.messages))
193 copy(out, a.messages)
194 return out
195}
196
197// Revision returns a monotonically increasing transcript version.
198// It is cheap to query and changes whenever Messages() would return
199// different transcript content because of append/set operations.
200func (a *Agent) Revision() uint64 {
201 a.mu.Lock()
202 defer a.mu.Unlock()
203 return a.rev
204}
205
206// SetTools swaps the tool registry. Used by /reload-ext to hand
207// the agent a fresh registry after extension subprocesses have been
208// respawned (and their freshly-registered tools merged in).
209func (a *Agent) SetTools(reg Registry) {
210 a.mu.Lock()
211 a.Tools = reg
212 a.mu.Unlock()
213}
214
215// SetMessages replaces the transcript (used when resuming a session).
216func (a *Agent) SetMessages(msgs []provider.Message) {
217 a.mu.Lock()
218 defer a.mu.Unlock()
219 a.messages = append(a.messages[:0], msgs...)
220 a.rev++
221}
222
223// Cost returns the cumulative usage.
224func (a *Agent) Cost() provider.Usage {
225 a.mu.Lock()
226 defer a.mu.Unlock()
227 return a.cost.Total
228}
229
230// SeedCost sets the cumulative usage as a baseline before the first
231// turn runs. Used when transferring state from another agent (model
232// or provider switch) so the running cost meter doesn't reset to 0.
233func (a *Agent) SeedCost(u provider.Usage) {
234 a.mu.Lock()
235 defer a.mu.Unlock()
236 a.cost.Total = u
237}
238
239// LastTurnUsage returns the per-turn usage of the most recent
240// completed turn. Drives the "context used" gauge in the status bar
241// without waiting for the next turn to land.
242func (a *Agent) LastTurnUsage() provider.Usage {
243 a.mu.Lock()
244 defer a.mu.Unlock()
245 return a.cost.LastTurn
246}
247
248// SeedLastTurnUsage primes the per-turn snapshot. Used on resume so
249// the gauge reflects the prompt size of the last turn in the session
250// file instead of starting at zero.
251func (a *Agent) SeedLastTurnUsage(u provider.Usage) {
252 a.mu.Lock()
253 defer a.mu.Unlock()
254 a.cost.LastTurn = u
255}
256
257// fireMessageAppended invokes OnMessageAppended without holding the
258// agent mutex, so the host's persistence callback can take its own
259// locks without deadlocking the agent loop. Tolerates a nil hook so
260// non-persisting callers (tests, RPC mode) don't have to set it.
261func (a *Agent) fireMessageAppended(m provider.Message) {
262 if a.OnMessageAppended != nil {
263 a.OnMessageAppended(m)
264 }
265}
266
267// Prompt sends a user message and runs the agent loop until the model
268// stops or an error occurs. Events are delivered via sink in order.
269// sink must not block the caller for long; buffer as needed.
270func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
271 if sink == nil {
272 sink = func(AgentEvent) {}
273 }
274 sink = a.wrapSink(sink)
275 content := []provider.Content{}
276 if text != "" {
277 content = append(content, provider.TextBlock{Text: text})
278 }
279 for _, img := range images {
280 content = append(content, img)
281 }
282 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
283
284 a.mu.Lock()
285 a.messages = append(a.messages, user)
286 a.rev++
287 a.mu.Unlock()
288 a.fireMessageAppended(user)
289 sink(EvUserMessage{Message: user})
290
291 return a.runLoop(ctx, sink)
292}
293
294// Continue runs the agent loop against the existing transcript. Used
295// after appending tool results manually or to retry.
296func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
297 if sink == nil {
298 sink = func(AgentEvent) {}
299 }
300 sink = a.wrapSink(sink)
301 return a.runLoop(ctx, sink)
302}
303
304// wrapSink composes the per-call sink with a.OnEvent (if set) so the
305// extension manager (or any other observer) sees every AgentEvent
306// without having to thread itself through every Prompt callsite.
307func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) {
308 if a.OnEvent == nil {
309 return sink
310 }
311 obs := a.OnEvent
312 return func(ev AgentEvent) {
313 obs(ev)
314 sink(ev)
315 }
316}
317
318func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
319 for step := 1; a.MaxSteps <= 0 || step <= a.MaxSteps; step++ {
320 // Messages queued while the agent was busy are delivered
321 // before the next model call. This is the safe boundary:
322 // any previous tool batch has already completed and its
323 // results have been appended, but no new provider request has
324 // started yet.
325 if pending := a.drainQueuedMessages(); len(pending) > 0 {
326 a.appendQueuedAsUser(pending, sink)
327 }
328
329 sink(EvTurnStart{Step: step})
330 if a.BeforeTurn != nil {
331 if allowed, reason := a.BeforeTurn(step); !allowed {
332 if reason == "" {
333 reason = "turn blocked by extension guard"
334 }
335 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
336 sink(EvDone{})
337 return nil
338 }
339 }
340 stop, assistantMsg, err := a.oneTurn(ctx, sink)
341 sink(EvTurnEnd{Stop: stop, Err: err})
342 if err != nil {
343 return err
344 }
345
346 if stop == provider.StopToolUse {
347 // Execute each tool call, append a single tool-results message, continue.
348 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
349 a.mu.Lock()
350 a.messages = append(a.messages, toolMsg)
351 a.rev++
352 // OpenAI's chat-completions tool message shape is text-centric.
353 // Vision models reliably consume images when they arrive as user
354 // content, so when a tool result contains images we mirror them
355 // into a synthetic user message immediately after the tool result.
356 // This keeps the transcript self-contained for providers that can
357 // see image blocks in tool messages while making OpenAI vision
358 // models actually receive the image bytes.
359 //
360 // The OpenAI Responses route ("openai-codex") has the same
361 // text-centric tool-output shape: a function_call_output only
362 // carries a string, so images in a tool result never reach the
363 // model. Both providers serialize images correctly when they
364 // arrive as user content, so the mirror covers them both.
365 var imageMirror provider.Message
366 if a.Client != nil && (a.Client.Name() == "openai" || a.Client.Name() == "openai-codex") {
367 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
368 a.messages = append(a.messages, mirror)
369 a.rev++
370 imageMirror = mirror
371 }
372 }
373 a.mu.Unlock()
374 a.fireMessageAppended(toolMsg)
375 if len(imageMirror.Content) > 0 {
376 a.fireMessageAppended(imageMirror)
377 }
378 // If context was cancelled during tool execution, bail out.
379 if err := ctx.Err(); err != nil {
380 sink(EvDone{})
381 return err
382 }
383 _ = hadError
384 continue
385 }
386
387 // If the assistant stopped without tool calls but a message was
388 // queued while it was speaking, loop once more so that message
389 // is appended and answered instead of waiting until a later
390 // top-level prompt.
391 if ctx.Err() == nil && a.QueuedMessageCount() > 0 {
392 continue
393 }
394
395 // Terminal stop (end, length, error, aborted).
396 sink(EvDone{})
397 return nil
398 }
399 if a.MaxSteps > 0 {
400 sink(EvDone{})
401 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
402 }
403 return nil
404}
405
406// oneTurn calls the LLM once, forwards events, returns the stop reason
407// and the assembled assistant message (already appended to the transcript).
408func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
409 req := provider.Request{
410 Model: a.Model,
411 System: a.System,
412 Messages: a.Messages(),
413 Tools: a.Tools.Specs(),
414 Reasoning: a.Reasoning,
415 }
416 stream, err := a.Client.Stream(ctx, req)
417 if err != nil {
418 return provider.StopError, provider.Message{}, err
419 }
420
421 sink(EvAssistantStart{})
422
423 var (
424 stop provider.StopReason
425 finalErr error
426 finalMsg provider.Message
427 )
428
429 for ev := range stream {
430 switch e := ev.(type) {
431 case provider.EventStart:
432 // nothing
433 case provider.EventTextDelta:
434 sink(EvTextDelta{Delta: e.Delta})
435 case provider.EventToolStart:
436 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
437 case provider.EventToolArgs:
438 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
439 case provider.EventToolEnd:
440 sink(EvToolUseEnd{ID: e.ID})
441 case provider.EventUsage:
442 cum := a.cost.Add(e.Usage)
443 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
444 if a.OnUsage != nil {
445 a.OnUsage(cum)
446 }
447 case provider.EventDone:
448 stop = e.Stop
449 finalErr = e.Err
450 finalMsg = e.Message
451 }
452 }
453
454 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
455 // produce partial content. When the partial message is text only we
456 // keep whatever was streamed up to the cancel so the user does not
457 // lose visible work (a cut-off summary is still useful). If the
458 // partial message already contained tool-call blocks we drop the
459 // whole thing, because an unmatched tool_use would fail the next
460 // turn with a tool_result mismatch error.
461 keep := len(finalMsg.Content) > 0
462 if stop == provider.StopAborted && keep {
463 hasToolCall := false
464 for _, c := range finalMsg.Content {
465 if _, ok := c.(provider.ToolCallBlock); ok {
466 hasToolCall = true
467 break
468 }
469 }
470 if hasToolCall {
471 keep = false
472 }
473 }
474 if keep {
475 emit := finalMsg
476 suppress := false
477
478 // BeforeAssistantMessage hook: extensions can suppress or
479 // rewrite the visible text. The transcript keeps the
480 // model's original output so the model still sees what it
481 // said on subsequent turns.
482 if a.BeforeAssistantMessage != nil {
483 orig := extractText(finalMsg)
484 if orig != "" {
485 allowed, _, replacement := a.BeforeAssistantMessage(orig)
486 if !allowed {
487 suppress = true
488 } else if replacement != "" && replacement != orig {
489 emit = replaceText(finalMsg, replacement)
490 }
491 }
492 }
493
494 a.mu.Lock()
495 a.messages = append(a.messages, finalMsg)
496 a.rev++
497 a.mu.Unlock()
498 a.fireMessageAppended(finalMsg)
499 if !suppress {
500 sink(EvAssistantMessage{Message: emit})
501 }
502 // Now surface tool calls as EvToolCall events so UIs can render them
503 // in order before the tool results arrive.
504 for _, c := range finalMsg.Content {
505 if tc, ok := c.(provider.ToolCallBlock); ok {
506 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
507 }
508 }
509 }
510
511 return stop, finalMsg, finalErr
512}
513
514// executeTools runs every tool call in the assistant message and returns
515// a single tool-role message carrying all results.
516func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
517 var results []provider.Content
518 hadError := false
519
520 for _, c := range msg.Content {
521 tc, ok := c.(provider.ToolCallBlock)
522 if !ok {
523 continue
524 }
525 res := a.runOneTool(ctx, tc, sink)
526 if res.IsError {
527 hadError = true
528 }
529 results = append(results, provider.ToolResultBlock{
530 CallID: tc.ID,
531 Content: res.Content,
532 IsError: res.IsError,
533 })
534 sink(EvToolResult{ID: tc.ID, Result: res})
535 }
536
537 return provider.Message{
538 Role: provider.RoleTool,
539 Content: results,
540 Time: time.Now(),
541 }, hadError
542}
543
544func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
545 tool, err := a.Tools.Get(tc.Name)
546 if err != nil {
547 return ToolResult{
548 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
549 IsError: true,
550 }
551 }
552
553 args := tc.Arguments
554
555 // Intercept hook: an extension or other guard can refuse the
556 // call before any side effect happens, OR rewrite the args
557 // seen by the tool. The model sees the reason as the tool
558 // error, learns from it, and (typically) proposes a different
559 // action; rewrites are invisible to the model (they apply only
560 // to the execution).
561 if a.BeforeToolExecute != nil {
562 allowed, reason, modified := a.BeforeToolExecute(tc)
563 if !allowed {
564 if reason == "" {
565 reason = "tool call refused by extension guard"
566 }
567 return ToolResult{
568 Content: []provider.Content{provider.TextBlock{Text: reason}},
569 IsError: true,
570 }
571 }
572 if len(modified) > 0 && json.Valid(modified) {
573 args = modified
574 }
575 }
576
577 if len(args) == 0 {
578 args = json.RawMessage("{}")
579 }
580
581 // Recover panics so a buggy tool does not crash the agent.
582 var res ToolResult
583 func() {
584 defer func() {
585 if r := recover(); r != nil {
586 res = ToolResult{
587 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
588 IsError: true,
589 }
590 }
591 }()
592 out, err := tool.Execute(ctx, args, func(text string) {
593 sink(EvToolProgress{ID: tc.ID, Text: text})
594 })
595 if err != nil {
596 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
597 res = ToolResult{
598 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
599 IsError: true,
600 }
601 return
602 }
603 res = ToolResult{
604 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
605 IsError: true,
606 }
607 return
608 }
609 res = out
610 }()
611 return res
612}
613
614// extractText concatenates all TextBlock content in a message. Used
615// by BeforeAssistantMessage so guards see a single string instead of
616// having to walk provider.Content themselves.
617func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
618 var content []provider.Content
619 for _, c := range msg.Content {
620 tr, ok := c.(provider.ToolResultBlock)
621 if !ok {
622 continue
623 }
624 for _, inner := range tr.Content {
625 switch v := inner.(type) {
626 case provider.TextBlock:
627 // Keep short textual context so the model understands why
628 // the images appeared, but don't duplicate giant read
629 // outputs verbatim.
630 if len(v.Text) > 0 && len(v.Text) <= 500 {
631 content = append(content, v)
632 }
633 case provider.ImageBlock:
634 content = append(content, v)
635 }
636 }
637 }
638 if len(content) == 0 {
639 return provider.Message{}
640 }
641 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
642 content = append([]provider.Content{prefix}, content...)
643 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
644}
645
646func extractText(msg provider.Message) string {
647 var out string
648 for _, c := range msg.Content {
649 if tb, ok := c.(provider.TextBlock); ok {
650 if out != "" {
651 out += "\n"
652 }
653 out += tb.Text
654 }
655 }
656 return out
657}
658
659// replaceText returns a copy of msg with every TextBlock replaced by
660// a single TextBlock containing replacement. Non-text content (tool
661// calls, etc.) is preserved in order.
662func replaceText(msg provider.Message, replacement string) provider.Message {
663 out := provider.Message{Role: msg.Role}
664 out.Content = make([]provider.Content, 0, len(msg.Content))
665 replaced := false
666 for _, c := range msg.Content {
667 if _, ok := c.(provider.TextBlock); ok {
668 if !replaced {
669 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
670 replaced = true
671 }
672 continue
673 }
674 out.Content = append(out.Content, c)
675 }
676 if !replaced {
677 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
678 }
679 return out
680}