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 var imageMirror provider.Message
360 if a.Client != nil && a.Client.Name() == "openai" {
361 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
362 a.messages = append(a.messages, mirror)
363 a.rev++
364 imageMirror = mirror
365 }
366 }
367 a.mu.Unlock()
368 a.fireMessageAppended(toolMsg)
369 if len(imageMirror.Content) > 0 {
370 a.fireMessageAppended(imageMirror)
371 }
372 // If context was cancelled during tool execution, bail out.
373 if err := ctx.Err(); err != nil {
374 sink(EvDone{})
375 return err
376 }
377 _ = hadError
378 continue
379 }
380
381 // If the assistant stopped without tool calls but a message was
382 // queued while it was speaking, loop once more so that message
383 // is appended and answered instead of waiting until a later
384 // top-level prompt.
385 if ctx.Err() == nil && a.QueuedMessageCount() > 0 {
386 continue
387 }
388
389 // Terminal stop (end, length, error, aborted).
390 sink(EvDone{})
391 return nil
392 }
393 if a.MaxSteps > 0 {
394 sink(EvDone{})
395 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
396 }
397 return nil
398}
399
400// oneTurn calls the LLM once, forwards events, returns the stop reason
401// and the assembled assistant message (already appended to the transcript).
402func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
403 req := provider.Request{
404 Model: a.Model,
405 System: a.System,
406 Messages: a.Messages(),
407 Tools: a.Tools.Specs(),
408 Reasoning: a.Reasoning,
409 }
410 stream, err := a.Client.Stream(ctx, req)
411 if err != nil {
412 return provider.StopError, provider.Message{}, err
413 }
414
415 sink(EvAssistantStart{})
416
417 var (
418 stop provider.StopReason
419 finalErr error
420 finalMsg provider.Message
421 )
422
423 for ev := range stream {
424 switch e := ev.(type) {
425 case provider.EventStart:
426 // nothing
427 case provider.EventTextDelta:
428 sink(EvTextDelta{Delta: e.Delta})
429 case provider.EventToolStart:
430 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
431 case provider.EventToolArgs:
432 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
433 case provider.EventToolEnd:
434 sink(EvToolUseEnd{ID: e.ID})
435 case provider.EventUsage:
436 cum := a.cost.Add(e.Usage)
437 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
438 if a.OnUsage != nil {
439 a.OnUsage(cum)
440 }
441 case provider.EventDone:
442 stop = e.Stop
443 finalErr = e.Err
444 finalMsg = e.Message
445 }
446 }
447
448 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
449 // produce partial content. When the partial message is text only we
450 // keep whatever was streamed up to the cancel so the user does not
451 // lose visible work (a cut-off summary is still useful). If the
452 // partial message already contained tool-call blocks we drop the
453 // whole thing, because an unmatched tool_use would fail the next
454 // turn with a tool_result mismatch error.
455 keep := len(finalMsg.Content) > 0
456 if stop == provider.StopAborted && keep {
457 hasToolCall := false
458 for _, c := range finalMsg.Content {
459 if _, ok := c.(provider.ToolCallBlock); ok {
460 hasToolCall = true
461 break
462 }
463 }
464 if hasToolCall {
465 keep = false
466 }
467 }
468 if keep {
469 emit := finalMsg
470 suppress := false
471
472 // BeforeAssistantMessage hook: extensions can suppress or
473 // rewrite the visible text. The transcript keeps the
474 // model's original output so the model still sees what it
475 // said on subsequent turns.
476 if a.BeforeAssistantMessage != nil {
477 orig := extractText(finalMsg)
478 if orig != "" {
479 allowed, _, replacement := a.BeforeAssistantMessage(orig)
480 if !allowed {
481 suppress = true
482 } else if replacement != "" && replacement != orig {
483 emit = replaceText(finalMsg, replacement)
484 }
485 }
486 }
487
488 a.mu.Lock()
489 a.messages = append(a.messages, finalMsg)
490 a.rev++
491 a.mu.Unlock()
492 a.fireMessageAppended(finalMsg)
493 if !suppress {
494 sink(EvAssistantMessage{Message: emit})
495 }
496 // Now surface tool calls as EvToolCall events so UIs can render them
497 // in order before the tool results arrive.
498 for _, c := range finalMsg.Content {
499 if tc, ok := c.(provider.ToolCallBlock); ok {
500 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
501 }
502 }
503 }
504
505 return stop, finalMsg, finalErr
506}
507
508// executeTools runs every tool call in the assistant message and returns
509// a single tool-role message carrying all results.
510func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
511 var results []provider.Content
512 hadError := false
513
514 for _, c := range msg.Content {
515 tc, ok := c.(provider.ToolCallBlock)
516 if !ok {
517 continue
518 }
519 res := a.runOneTool(ctx, tc, sink)
520 if res.IsError {
521 hadError = true
522 }
523 results = append(results, provider.ToolResultBlock{
524 CallID: tc.ID,
525 Content: res.Content,
526 IsError: res.IsError,
527 })
528 sink(EvToolResult{ID: tc.ID, Result: res})
529 }
530
531 return provider.Message{
532 Role: provider.RoleTool,
533 Content: results,
534 Time: time.Now(),
535 }, hadError
536}
537
538func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
539 tool, err := a.Tools.Get(tc.Name)
540 if err != nil {
541 return ToolResult{
542 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
543 IsError: true,
544 }
545 }
546
547 args := tc.Arguments
548
549 // Intercept hook: an extension or other guard can refuse the
550 // call before any side effect happens, OR rewrite the args
551 // seen by the tool. The model sees the reason as the tool
552 // error, learns from it, and (typically) proposes a different
553 // action; rewrites are invisible to the model (they apply only
554 // to the execution).
555 if a.BeforeToolExecute != nil {
556 allowed, reason, modified := a.BeforeToolExecute(tc)
557 if !allowed {
558 if reason == "" {
559 reason = "tool call refused by extension guard"
560 }
561 return ToolResult{
562 Content: []provider.Content{provider.TextBlock{Text: reason}},
563 IsError: true,
564 }
565 }
566 if len(modified) > 0 && json.Valid(modified) {
567 args = modified
568 }
569 }
570
571 if len(args) == 0 {
572 args = json.RawMessage("{}")
573 }
574
575 // Recover panics so a buggy tool does not crash the agent.
576 var res ToolResult
577 func() {
578 defer func() {
579 if r := recover(); r != nil {
580 res = ToolResult{
581 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
582 IsError: true,
583 }
584 }
585 }()
586 out, err := tool.Execute(ctx, args, func(text string) {
587 sink(EvToolProgress{ID: tc.ID, Text: text})
588 })
589 if err != nil {
590 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
591 res = ToolResult{
592 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
593 IsError: true,
594 }
595 return
596 }
597 res = ToolResult{
598 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
599 IsError: true,
600 }
601 return
602 }
603 res = out
604 }()
605 return res
606}
607
608// extractText concatenates all TextBlock content in a message. Used
609// by BeforeAssistantMessage so guards see a single string instead of
610// having to walk provider.Content themselves.
611func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
612 var content []provider.Content
613 for _, c := range msg.Content {
614 tr, ok := c.(provider.ToolResultBlock)
615 if !ok {
616 continue
617 }
618 for _, inner := range tr.Content {
619 switch v := inner.(type) {
620 case provider.TextBlock:
621 // Keep short textual context so the model understands why
622 // the images appeared, but don't duplicate giant read
623 // outputs verbatim.
624 if len(v.Text) > 0 && len(v.Text) <= 500 {
625 content = append(content, v)
626 }
627 case provider.ImageBlock:
628 content = append(content, v)
629 }
630 }
631 }
632 if len(content) == 0 {
633 return provider.Message{}
634 }
635 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
636 content = append([]provider.Content{prefix}, content...)
637 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
638}
639
640func extractText(msg provider.Message) string {
641 var out string
642 for _, c := range msg.Content {
643 if tb, ok := c.(provider.TextBlock); ok {
644 if out != "" {
645 out += "\n"
646 }
647 out += tb.Text
648 }
649 }
650 return out
651}
652
653// replaceText returns a copy of msg with every TextBlock replaced by
654// a single TextBlock containing replacement. Non-text content (tool
655// calls, etc.) is preserved in order.
656func replaceText(msg provider.Message, replacement string) provider.Message {
657 out := provider.Message{Role: msg.Role}
658 out.Content = make([]provider.Content, 0, len(msg.Content))
659 replaced := false
660 for _, c := range msg.Content {
661 if _, ok := c.(provider.TextBlock); ok {
662 if !replaced {
663 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
664 replaced = true
665 }
666 continue
667 }
668 out.Content = append(out.Content, c)
669 }
670 if !replaced {
671 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
672 }
673 return out
674}