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/internal/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// fireMessageAppended invokes OnMessageAppended without holding the
240// agent mutex, so the host's persistence callback can take its own
241// locks without deadlocking the agent loop. Tolerates a nil hook so
242// non-persisting callers (tests, RPC mode) don't have to set it.
243func (a *Agent) fireMessageAppended(m provider.Message) {
244 if a.OnMessageAppended != nil {
245 a.OnMessageAppended(m)
246 }
247}
248
249// Prompt sends a user message and runs the agent loop until the model
250// stops or an error occurs. Events are delivered via sink in order.
251// sink must not block the caller for long; buffer as needed.
252func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
253 if sink == nil {
254 sink = func(AgentEvent) {}
255 }
256 sink = a.wrapSink(sink)
257 content := []provider.Content{}
258 if text != "" {
259 content = append(content, provider.TextBlock{Text: text})
260 }
261 for _, img := range images {
262 content = append(content, img)
263 }
264 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
265
266 a.mu.Lock()
267 a.messages = append(a.messages, user)
268 a.rev++
269 a.mu.Unlock()
270 a.fireMessageAppended(user)
271 sink(EvUserMessage{Message: user})
272
273 return a.runLoop(ctx, sink)
274}
275
276// Continue runs the agent loop against the existing transcript. Used
277// after appending tool results manually or to retry.
278func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
279 if sink == nil {
280 sink = func(AgentEvent) {}
281 }
282 sink = a.wrapSink(sink)
283 return a.runLoop(ctx, sink)
284}
285
286// wrapSink composes the per-call sink with a.OnEvent (if set) so the
287// extension manager (or any other observer) sees every AgentEvent
288// without having to thread itself through every Prompt callsite.
289func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) {
290 if a.OnEvent == nil {
291 return sink
292 }
293 obs := a.OnEvent
294 return func(ev AgentEvent) {
295 obs(ev)
296 sink(ev)
297 }
298}
299
300func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
301 for step := 1; a.MaxSteps <= 0 || step <= a.MaxSteps; step++ {
302 // Messages queued while the agent was busy are delivered
303 // before the next model call. This is the safe boundary:
304 // any previous tool batch has already completed and its
305 // results have been appended, but no new provider request has
306 // started yet.
307 if pending := a.drainQueuedMessages(); len(pending) > 0 {
308 a.appendQueuedAsUser(pending, sink)
309 }
310
311 sink(EvTurnStart{Step: step})
312 if a.BeforeTurn != nil {
313 if allowed, reason := a.BeforeTurn(step); !allowed {
314 if reason == "" {
315 reason = "turn blocked by extension guard"
316 }
317 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
318 sink(EvDone{})
319 return nil
320 }
321 }
322 stop, assistantMsg, err := a.oneTurn(ctx, sink)
323 sink(EvTurnEnd{Stop: stop, Err: err})
324 if err != nil {
325 return err
326 }
327
328 if stop == provider.StopToolUse {
329 // Execute each tool call, append a single tool-results message, continue.
330 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
331 a.mu.Lock()
332 a.messages = append(a.messages, toolMsg)
333 a.rev++
334 // OpenAI's chat-completions tool message shape is text-centric.
335 // Vision models reliably consume images when they arrive as user
336 // content, so when a tool result contains images we mirror them
337 // into a synthetic user message immediately after the tool result.
338 // This keeps the transcript self-contained for providers that can
339 // see image blocks in tool messages while making OpenAI vision
340 // models actually receive the image bytes.
341 var imageMirror provider.Message
342 if a.Client != nil && a.Client.Name() == "openai" {
343 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
344 a.messages = append(a.messages, mirror)
345 a.rev++
346 imageMirror = mirror
347 }
348 }
349 a.mu.Unlock()
350 a.fireMessageAppended(toolMsg)
351 if len(imageMirror.Content) > 0 {
352 a.fireMessageAppended(imageMirror)
353 }
354 // If context was cancelled during tool execution, bail out.
355 if err := ctx.Err(); err != nil {
356 sink(EvDone{})
357 return err
358 }
359 _ = hadError
360 continue
361 }
362
363 // If the assistant stopped without tool calls but a message was
364 // queued while it was speaking, loop once more so that message
365 // is appended and answered instead of waiting until a later
366 // top-level prompt.
367 if ctx.Err() == nil && a.QueuedMessageCount() > 0 {
368 continue
369 }
370
371 // Terminal stop (end, length, error, aborted).
372 sink(EvDone{})
373 return nil
374 }
375 if a.MaxSteps > 0 {
376 sink(EvDone{})
377 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
378 }
379 return nil
380}
381
382// oneTurn calls the LLM once, forwards events, returns the stop reason
383// and the assembled assistant message (already appended to the transcript).
384func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
385 req := provider.Request{
386 Model: a.Model,
387 System: a.System,
388 Messages: a.Messages(),
389 Tools: a.Tools.Specs(),
390 Reasoning: a.Reasoning,
391 }
392 stream, err := a.Client.Stream(ctx, req)
393 if err != nil {
394 return provider.StopError, provider.Message{}, err
395 }
396
397 sink(EvAssistantStart{})
398
399 var (
400 stop provider.StopReason
401 finalErr error
402 finalMsg provider.Message
403 )
404
405 for ev := range stream {
406 switch e := ev.(type) {
407 case provider.EventStart:
408 // nothing
409 case provider.EventTextDelta:
410 sink(EvTextDelta{Delta: e.Delta})
411 case provider.EventToolStart:
412 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
413 case provider.EventToolArgs:
414 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
415 case provider.EventToolEnd:
416 sink(EvToolUseEnd{ID: e.ID})
417 case provider.EventUsage:
418 cum := a.cost.Add(e.Usage)
419 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
420 if a.OnUsage != nil {
421 a.OnUsage(cum)
422 }
423 case provider.EventDone:
424 stop = e.Stop
425 finalErr = e.Err
426 finalMsg = e.Message
427 }
428 }
429
430 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
431 // produce partial content. When the partial message is text only we
432 // keep whatever was streamed up to the cancel so the user does not
433 // lose visible work (a cut-off summary is still useful). If the
434 // partial message already contained tool-call blocks we drop the
435 // whole thing, because an unmatched tool_use would fail the next
436 // turn with a tool_result mismatch error.
437 keep := len(finalMsg.Content) > 0
438 if stop == provider.StopAborted && keep {
439 hasToolCall := false
440 for _, c := range finalMsg.Content {
441 if _, ok := c.(provider.ToolCallBlock); ok {
442 hasToolCall = true
443 break
444 }
445 }
446 if hasToolCall {
447 keep = false
448 }
449 }
450 if keep {
451 emit := finalMsg
452 suppress := false
453
454 // BeforeAssistantMessage hook: extensions can suppress or
455 // rewrite the visible text. The transcript keeps the
456 // model's original output so the model still sees what it
457 // said on subsequent turns.
458 if a.BeforeAssistantMessage != nil {
459 orig := extractText(finalMsg)
460 if orig != "" {
461 allowed, _, replacement := a.BeforeAssistantMessage(orig)
462 if !allowed {
463 suppress = true
464 } else if replacement != "" && replacement != orig {
465 emit = replaceText(finalMsg, replacement)
466 }
467 }
468 }
469
470 a.mu.Lock()
471 a.messages = append(a.messages, finalMsg)
472 a.rev++
473 a.mu.Unlock()
474 a.fireMessageAppended(finalMsg)
475 if !suppress {
476 sink(EvAssistantMessage{Message: emit})
477 }
478 // Now surface tool calls as EvToolCall events so UIs can render them
479 // in order before the tool results arrive.
480 for _, c := range finalMsg.Content {
481 if tc, ok := c.(provider.ToolCallBlock); ok {
482 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
483 }
484 }
485 }
486
487 return stop, finalMsg, finalErr
488}
489
490// executeTools runs every tool call in the assistant message and returns
491// a single tool-role message carrying all results.
492func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
493 var results []provider.Content
494 hadError := false
495
496 for _, c := range msg.Content {
497 tc, ok := c.(provider.ToolCallBlock)
498 if !ok {
499 continue
500 }
501 res := a.runOneTool(ctx, tc, sink)
502 if res.IsError {
503 hadError = true
504 }
505 results = append(results, provider.ToolResultBlock{
506 CallID: tc.ID,
507 Content: res.Content,
508 IsError: res.IsError,
509 })
510 sink(EvToolResult{ID: tc.ID, Result: res})
511 }
512
513 return provider.Message{
514 Role: provider.RoleTool,
515 Content: results,
516 Time: time.Now(),
517 }, hadError
518}
519
520func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
521 tool, err := a.Tools.Get(tc.Name)
522 if err != nil {
523 return ToolResult{
524 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
525 IsError: true,
526 }
527 }
528
529 args := tc.Arguments
530
531 // Intercept hook: an extension or other guard can refuse the
532 // call before any side effect happens, OR rewrite the args
533 // seen by the tool. The model sees the reason as the tool
534 // error, learns from it, and (typically) proposes a different
535 // action; rewrites are invisible to the model (they apply only
536 // to the execution).
537 if a.BeforeToolExecute != nil {
538 allowed, reason, modified := a.BeforeToolExecute(tc)
539 if !allowed {
540 if reason == "" {
541 reason = "tool call refused by extension guard"
542 }
543 return ToolResult{
544 Content: []provider.Content{provider.TextBlock{Text: reason}},
545 IsError: true,
546 }
547 }
548 if len(modified) > 0 && json.Valid(modified) {
549 args = modified
550 }
551 }
552
553 if len(args) == 0 {
554 args = json.RawMessage("{}")
555 }
556
557 // Recover panics so a buggy tool does not crash the agent.
558 var res ToolResult
559 func() {
560 defer func() {
561 if r := recover(); r != nil {
562 res = ToolResult{
563 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
564 IsError: true,
565 }
566 }
567 }()
568 out, err := tool.Execute(ctx, args, func(text string) {
569 sink(EvToolProgress{ID: tc.ID, Text: text})
570 })
571 if err != nil {
572 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
573 res = ToolResult{
574 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
575 IsError: true,
576 }
577 return
578 }
579 res = ToolResult{
580 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
581 IsError: true,
582 }
583 return
584 }
585 res = out
586 }()
587 return res
588}
589
590// extractText concatenates all TextBlock content in a message. Used
591// by BeforeAssistantMessage so guards see a single string instead of
592// having to walk provider.Content themselves.
593func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
594 var content []provider.Content
595 for _, c := range msg.Content {
596 tr, ok := c.(provider.ToolResultBlock)
597 if !ok {
598 continue
599 }
600 for _, inner := range tr.Content {
601 switch v := inner.(type) {
602 case provider.TextBlock:
603 // Keep short textual context so the model understands why
604 // the images appeared, but don't duplicate giant read
605 // outputs verbatim.
606 if len(v.Text) > 0 && len(v.Text) <= 500 {
607 content = append(content, v)
608 }
609 case provider.ImageBlock:
610 content = append(content, v)
611 }
612 }
613 }
614 if len(content) == 0 {
615 return provider.Message{}
616 }
617 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
618 content = append([]provider.Content{prefix}, content...)
619 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
620}
621
622func extractText(msg provider.Message) string {
623 var out string
624 for _, c := range msg.Content {
625 if tb, ok := c.(provider.TextBlock); ok {
626 if out != "" {
627 out += "\n"
628 }
629 out += tb.Text
630 }
631 }
632 return out
633}
634
635// replaceText returns a copy of msg with every TextBlock replaced by
636// a single TextBlock containing replacement. Non-text content (tool
637// calls, etc.) is preserved in order.
638func replaceText(msg provider.Message, replacement string) provider.Message {
639 out := provider.Message{Role: msg.Role}
640 out.Content = make([]provider.Content, 0, len(msg.Content))
641 replaced := false
642 for _, c := range msg.Content {
643 if _, ok := c.(provider.TextBlock); ok {
644 if !replaced {
645 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
646 replaced = true
647 }
648 continue
649 }
650 out.Content = append(out.Content, c)
651 }
652 if !replaced {
653 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
654 }
655 return out
656}