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 Temperature *float32
25
26 // MaxTokens caps the model's output tokens per turn. Zero leaves
27 // the field unset on the provider request, letting each provider
28 // apply its own default (which can be conservative, e.g. Bedrock
29 // defaults to 4096, truncating long writes/edits). Hosts populate
30 // this from the resolved model's MaxOutput so large single-turn
31 // responses aren't silently cut off with stopReason=length.
32 MaxTokens int
33
34 // SessionID is the zot conversation id, forwarded to providers that
35 // support sticky routing. Empty means omitted.
36 SessionID string
37
38 // MaxToolCalls caps provider-executed server-tool calls per turn.
39 // Zero leaves the field unset.
40 MaxToolCalls int
41
42 // BeforeToolExecute, if set, is called immediately before each
43 // tool runs. Returning (allowed=false, reason) short-circuits
44 // the call with an error result containing reason. Optionally,
45 // returning a non-nil modifiedArgs replaces the JSON args the
46 // tool will see, which lets guards redact / augment / patch the
47 // model's request without rewriting the transcript. Empty or
48 // malformed modifiedArgs is ignored.
49 BeforeToolExecute func(call provider.ToolCallBlock) (allowed bool, reason string, modifiedArgs json.RawMessage)
50
51 // BeforeStart replaces the complete system prompt once per runtime session
52 // or explicit prompt/model reset, before any turn starts. It is not called
53 // for ordinary follow-up messages or tool-loop steps.
54 BeforeStart func(context.Context, string) string
55
56 // BeforeTurn, if set, is called before each turn's model call.
57 // Returning (allowed=false, reason) aborts the turn; reason is
58 // surfaced as an assistant-like status line. Used for rate-
59 // limiting, business-hour gates, and deny-by-default setups.
60 BeforeTurn func(step int) (allowed bool, reason string)
61
62 // BeforeAssistantMessage, if set, is called after the model's
63 // final assistant message is assembled but before it's appended
64 // to the transcript. Returning (allowed=false) suppresses both
65 // the transcript append and the UI event. A non-empty
66 // replacement rewrites the visible text for the user while
67 // leaving the model's original text in the transcript (so the
68 // model can still see what it said in subsequent turns).
69 BeforeAssistantMessage func(text string) (allowed bool, reason, replacement string)
70
71 // MaxRetries controls agent-level retries for transient provider
72 // failures that arrive after the HTTP stream opens (for example
73 // Anthropic overloaded_error). Zero disables this retry layer.
74 // RetryBaseDelay is doubled for each attempt; zero uses 2s.
75 MaxRetries int
76 RetryBaseDelay time.Duration
77
78 // OnEvent, if set, mirrors every AgentEvent the loop emits to
79 // this callback in addition to the per-Prompt sink. Used by the
80 // extension manager to fan events out to subscribed extensions
81 // without each caller having to compose sinks manually.
82 OnEvent func(AgentEvent)
83
84 // OnMessageAppended, if set, fires every time a message is
85 // appended to the in-memory transcript by the agent loop — the
86 // initial user prompt, each finalised assistant message, and
87 // each tool-results message (plus the synthetic OpenAI image
88 // mirror, if any). Hosts wire this to the on-disk session so
89 // that turns are durable as soon as they happen, instead of
90 // only being flushed on a clean exit.
91 OnMessageAppended func(provider.Message)
92
93 // OnUsage, if set, fires after every turn's usage row arrives,
94 // carrying the cumulative usage for the session. Hosts wire
95 // this to the on-disk session so the persisted total stays
96 // current and a crash recovers the right cost figure.
97 OnUsage func(cumulative provider.Usage)
98
99 // OnTranscriptCompacted, if set, fires after Compact replaces the
100 // in-memory transcript with the synthetic summary plus kept tail.
101 // Hosts wire this to append an explicit compaction checkpoint to
102 // the session log; per-message append hooks do not fire for this
103 // wholesale transcript replacement.
104 OnTranscriptCompacted func(messages []provider.Message)
105
106 // Preparation caches the effective prompt separately from its unmodified
107 // base so a model or session reset cannot stack extension appendices.
108 startPrepared bool
109 startGeneration uint64
110 startBase, startSystem, startModel, startSession string
111
112 mu sync.Mutex
113 messages []provider.Message
114 // rev increments whenever the transcript slice is replaced or a
115 // message is appended. The TUI uses it as a cheap redraw cache key
116 // so editor-only typing doesn't copy/rebuild a long transcript on
117 // every keypress.
118 rev uint64
119 cost CostTracker
120
121 // queued holds user messages submitted while the agent is busy.
122 // The loop appends them as normal user messages at safe
123 // boundaries: before the next model call after a tool batch, or
124 // after a text-only assistant turn finishes. It never interrupts
125 // a running tool or cancels an in-flight provider request.
126 queued []string
127}
128
129// NewAgent returns an Agent with sensible defaults.
130func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
131 return &Agent{
132 Client: client,
133 Model: model,
134 System: system,
135 Tools: tools,
136 MaxSteps: 0, // 0 = unlimited
137 MaxRetries: 3,
138 RetryBaseDelay: 2 * time.Second,
139 }
140}
141
142// QueueMessage queues text to be injected as a user message at the
143// next safe boundary of the active agent loop. It is non-blocking in
144// the sense that it never waits for model/tool work; it only takes
145// the transcript mutex briefly. Empty/whitespace-only messages are
146// ignored.
147func (a *Agent) QueueMessage(text string) bool {
148 text = strings.TrimSpace(text)
149 if text == "" {
150 return false
151 }
152 a.mu.Lock()
153 a.queued = append(a.queued, text)
154 a.mu.Unlock()
155 return true
156}
157
158// PendingQueuedMessages returns a snapshot of user messages waiting
159// to be injected. Used by hosts to render the visible "sliding in"
160// chips without consuming them.
161func (a *Agent) PendingQueuedMessages() []string {
162 a.mu.Lock()
163 defer a.mu.Unlock()
164 out := make([]string, len(a.queued))
165 copy(out, a.queued)
166 return out
167}
168
169// QueuedMessageCount returns the number of messages waiting to be
170// injected at the next safe boundary.
171func (a *Agent) QueuedMessageCount() int {
172 a.mu.Lock()
173 defer a.mu.Unlock()
174 return len(a.queued)
175}
176
177// PopQueuedMessage removes and returns the most recently queued
178// message. Hosts use this for the slide-back keybinding.
179func (a *Agent) PopQueuedMessage() (string, bool) {
180 a.mu.Lock()
181 defer a.mu.Unlock()
182 n := len(a.queued)
183 if n == 0 {
184 return "", false
185 }
186 text := a.queued[n-1]
187 a.queued = a.queued[:n-1]
188 return text, true
189}
190
191// DrainQueuedMessages discards and returns every queued message.
192// Hosts use this on explicit cancel/clear so stale follow-ups do
193// not run after the user aborted the turn.
194func (a *Agent) DrainQueuedMessages() []string {
195 return a.drainQueuedMessages()
196}
197
198func (a *Agent) drainQueuedMessages() []string {
199 a.mu.Lock()
200 defer a.mu.Unlock()
201 out := make([]string, len(a.queued))
202 copy(out, a.queued)
203 a.queued = nil
204 return out
205}
206
207func (a *Agent) appendQueuedAsUser(texts []string, sink func(AgentEvent)) {
208 for _, text := range texts {
209 msg := provider.Message{
210 Role: provider.RoleUser,
211 Content: []provider.Content{provider.TextBlock{Text: text}},
212 Time: time.Now(),
213 }
214 a.mu.Lock()
215 a.messages = append(a.messages, msg)
216 a.rev++
217 a.mu.Unlock()
218 a.fireMessageAppended(msg)
219 if sink != nil {
220 sink(EvUserMessage{Message: msg})
221 }
222 }
223}
224
225// Messages returns a copy of the current transcript.
226func (a *Agent) Messages() []provider.Message {
227 a.mu.Lock()
228 defer a.mu.Unlock()
229 out := make([]provider.Message, len(a.messages))
230 copy(out, a.messages)
231 return out
232}
233
234// ContextSnapshot returns the provider-neutral inputs that make up the
235// current model context. Hosts use the snapshot for read-only inspection
236// without racing a tool-registry replacement or transcript append.
237func (a *Agent) ContextSnapshot() (system string, tools []provider.Tool, messages []provider.Message) {
238 a.mu.Lock()
239 defer a.mu.Unlock()
240 messages = make([]provider.Message, len(a.messages))
241 copy(messages, a.messages)
242 return a.System, a.Tools.Specs(), messages
243}
244
245// Revision returns a monotonically increasing transcript version.
246// It is cheap to query and changes whenever Messages() would return
247// different transcript content because of append/set operations.
248func (a *Agent) Revision() uint64 {
249 a.mu.Lock()
250 defer a.mu.Unlock()
251 return a.rev
252}
253
254// SetTools swaps the tool registry. Used by /reload-ext to hand
255// the agent a fresh registry after extension subprocesses have been
256// respawned (and their freshly-registered tools merged in).
257func (a *Agent) SetTools(reg Registry) {
258 a.mu.Lock()
259 a.Tools = reg
260 a.mu.Unlock()
261}
262
263// SetMessages replaces the transcript (used when resuming a session).
264func (a *Agent) SetMessages(msgs []provider.Message) {
265 a.mu.Lock()
266 defer a.mu.Unlock()
267 a.messages = append(a.messages[:0], msgs...)
268 if len(msgs) == 0 {
269 a.resetStartLocked()
270 }
271 a.rev++
272}
273
274// AppendUserContext adds a user-role message to the transcript without
275// starting a model turn. Hosts use it for context gathered outside the agent
276// loop, such as the output of an explicitly invoked shell command.
277func (a *Agent) AppendUserContext(text string, meta map[string]string) {
278 msg := provider.Message{
279 Role: provider.RoleUser,
280 Content: []provider.Content{provider.TextBlock{Text: text}},
281 Time: time.Now(),
282 Meta: meta,
283 }
284 a.mu.Lock()
285 a.messages = append(a.messages, msg)
286 a.rev++
287 a.mu.Unlock()
288 a.fireMessageAppended(msg)
289}
290
291// Cost returns the cumulative usage.
292func (a *Agent) Cost() provider.Usage {
293 a.mu.Lock()
294 defer a.mu.Unlock()
295 return a.cost.Total
296}
297
298// SeedCost sets the cumulative usage as a baseline before the first
299// turn runs. Used when transferring state from another agent (model
300// or provider switch) so the running cost meter doesn't reset to 0.
301func (a *Agent) SeedCost(u provider.Usage) {
302 a.mu.Lock()
303 defer a.mu.Unlock()
304 a.cost.Seed(u)
305}
306
307// LastTurnUsage returns the per-turn usage of the most recent
308// completed turn. Drives the "context used" gauge in the status bar
309// without waiting for the next turn to land.
310func (a *Agent) LastTurnUsage() provider.Usage {
311 a.mu.Lock()
312 defer a.mu.Unlock()
313 return a.cost.LastTurn
314}
315
316// SeedLastTurnUsage primes the per-turn snapshot. Used on resume so
317// the gauge reflects the prompt size of the last turn in the session
318// file instead of starting at zero.
319func (a *Agent) SeedLastTurnUsage(u provider.Usage) {
320 a.mu.Lock()
321 defer a.mu.Unlock()
322 a.cost.LastTurn = u
323}
324
325// fireMessageAppended invokes OnMessageAppended without holding the
326// agent mutex, so the host's persistence callback can take its own
327// locks without deadlocking the agent loop. Tolerates a nil hook so
328// non-persisting callers (tests, RPC mode) don't have to set it.
329func (a *Agent) fireMessageAppended(m provider.Message) {
330 if a.OnMessageAppended != nil {
331 a.OnMessageAppended(m)
332 }
333}
334
335// Prompt sends a user message and runs the agent loop until the model
336// stops or an error occurs. Events are delivered via sink in order.
337// sink must not block the caller for long; buffer as needed.
338func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
339 if sink == nil {
340 sink = func(AgentEvent) {}
341 }
342 sink = a.wrapSink(sink)
343 content := []provider.Content{}
344 if text != "" {
345 content = append(content, provider.TextBlock{Text: text})
346 }
347 for _, img := range images {
348 content = append(content, img)
349 }
350 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
351
352 a.mu.Lock()
353 a.messages = append(a.messages, user)
354 a.rev++
355 a.mu.Unlock()
356 a.fireMessageAppended(user)
357 sink(EvUserMessage{Message: user})
358
359 return a.runLoop(ctx, sink)
360}
361
362// Continue runs the agent loop against the existing transcript. Used
363// after appending tool results manually or to retry.
364func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
365 if sink == nil {
366 sink = func(AgentEvent) {}
367 }
368 sink = a.wrapSink(sink)
369 return a.runLoop(ctx, sink)
370}
371
372// wrapSink composes the per-call sink with a.OnEvent (if set) so the
373// extension manager (or any other observer) sees every AgentEvent
374// without having to thread itself through every Prompt callsite.
375func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) {
376 if a.OnEvent == nil {
377 return sink
378 }
379 obs := a.OnEvent
380 return func(ev AgentEvent) {
381 obs(ev)
382 sink(ev)
383 }
384}
385
386func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
387 for step := 1; a.MaxSteps <= 0 || step <= a.MaxSteps; step++ {
388 // Preparation is cached across steps and user messages. Only an
389 // explicit prompt/model/session change invokes BeforeStart again.
390 if err := a.prepareStart(ctx); err != nil {
391 sink(EvDone{})
392 return err
393 }
394 // Messages queued while the agent was busy are delivered
395 // before the next model call. This is the safe boundary:
396 // any previous tool batch has already completed and its
397 // results have been appended, but no new provider request has
398 // started yet.
399 if pending := a.drainQueuedMessages(); len(pending) > 0 {
400 a.appendQueuedAsUser(pending, sink)
401 }
402
403 sink(EvTurnStart{Step: step})
404 if a.BeforeTurn != nil {
405 if allowed, reason := a.BeforeTurn(step); !allowed {
406 if reason == "" {
407 reason = "turn blocked by extension guard"
408 }
409 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
410 sink(EvDone{})
411 return nil
412 }
413 }
414
415 var (
416 stop provider.StopReason
417 assistantMsg provider.Message
418 err error
419 )
420 for attempt := 0; ; attempt++ {
421 stop, assistantMsg, err = a.oneTurn(ctx, sink)
422 sink(EvTurnEnd{Stop: stop, Err: err})
423 if err == nil || !a.canRetryError(err, attempt) {
424 break
425 }
426 a.dropLastAssistantMessage()
427 if sleepErr := sleepRetry(ctx, a.retryDelay(attempt)); sleepErr != nil {
428 return sleepErr
429 }
430 }
431 if err != nil {
432 return err
433 }
434
435 if stop == provider.StopToolUse {
436 // Execute each client tool call, append a single tool-results message, continue.
437 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
438 if len(toolMsg.Content) == 0 {
439 // Provider-executed (server) tools need no client results.
440 continue
441 }
442 a.mu.Lock()
443 a.messages = append(a.messages, toolMsg)
444 a.rev++
445 // OpenAI's chat-completions tool message shape is text-centric.
446 // Vision models reliably consume images when they arrive as user
447 // content, so when a tool result contains images we mirror them
448 // into a synthetic user message immediately after the tool result.
449 // This keeps the transcript self-contained for providers that can
450 // see image blocks in tool messages while making OpenAI vision
451 // models actually receive the image bytes.
452 //
453 // The OpenAI Responses route ("openai-codex") has the same
454 // text-centric tool-output shape: a function_call_output only
455 // carries a string, so images in a tool result never reach the
456 // model. Both providers serialize images correctly when they
457 // arrive as user content, so the mirror covers them both.
458 var imageMirror provider.Message
459 if a.Client != nil && (a.Client.Name() == "openai" || a.Client.Name() == "openai-codex") {
460 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
461 a.messages = append(a.messages, mirror)
462 a.rev++
463 imageMirror = mirror
464 }
465 }
466 a.mu.Unlock()
467 a.fireMessageAppended(toolMsg)
468 if len(imageMirror.Content) > 0 {
469 a.fireMessageAppended(imageMirror)
470 }
471 // If context was cancelled during tool execution, bail out.
472 if err := ctx.Err(); err != nil {
473 sink(EvDone{})
474 return err
475 }
476 _ = hadError
477 continue
478 }
479
480 // If the assistant stopped without tool calls but a message was
481 // queued while it was speaking, loop once more so that message
482 // is appended and answered instead of waiting until a later
483 // top-level prompt.
484 if ctx.Err() == nil && a.QueuedMessageCount() > 0 {
485 continue
486 }
487
488 // Terminal stop (end, length, error, aborted).
489 sink(EvDone{})
490 return nil
491 }
492 if a.MaxSteps > 0 {
493 sink(EvDone{})
494 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
495 }
496 return nil
497}
498
499func (a *Agent) canRetryError(err error, attempt int) bool {
500 if err == nil || a.MaxRetries <= 0 || attempt >= a.MaxRetries {
501 return false
502 }
503 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
504 return false
505 }
506 msg := strings.ToLower(err.Error())
507 if msg == "" || isNonRetryableProviderLimit(msg) {
508 return false
509 }
510 needles := []string{
511 "overloaded", "provider returned error", "rate limit", "ratelimit", "too many requests",
512 "429", "http 429", "500", "http 500", "502", "http 502", "503", "http 503", "504", "http 504",
513 "service unavailable", "server error", "internal error", "network error", "connection error",
514 "connection refused", "connection lost", "fetch failed", "upstream connect", "reset before headers",
515 "socket hang up", "ended without", "stream ended before", "did not get a response", "timed out",
516 "timeout", "terminated", "unexpected eof", "transport failure",
517 // OpenAI's ChatGPT/Codex backend returns this generic message (with a
518 // request ID) for transient server failures and explicitly says
519 // "You can retry your request".
520 "an error occurred while processing your request",
521 // Explicit retry guidance emitted by provider backends (OpenAI
522 // Responses, AWS Bedrock stream exceptions) with varying prefixes.
523 "you can retry your request", "try your request again", "please retry your request",
524 // Capacity messages from the ChatGPT/Codex backend, e.g.
525 // "Our servers are currently overloaded. Please try again later."
526 // The trailing advice also shows up on its own for transient
527 // capacity failures; usage/quota limits are filtered out above by
528 // isNonRetryableProviderLimit before this list is consulted.
529 "servers are currently overloaded", "servers are busy", "try again later",
530 }
531 for _, needle := range needles {
532 if strings.Contains(msg, needle) {
533 return true
534 }
535 }
536 return false
537}
538
539func isNonRetryableProviderLimit(msg string) bool {
540 needles := []string{
541 "usage limit", "monthly usage limit", "freeusagelimit", "gousagelimit",
542 "available balance", "insufficient_quota", "out of budget", "quota exceeded", "billing",
543 }
544 for _, needle := range needles {
545 if strings.Contains(msg, needle) {
546 return true
547 }
548 }
549 return false
550}
551
552func (a *Agent) retryDelay(attempt int) time.Duration {
553 base := a.RetryBaseDelay
554 if base <= 0 {
555 base = 2 * time.Second
556 }
557 return base * time.Duration(1<<attempt)
558}
559
560func sleepRetry(ctx context.Context, d time.Duration) error {
561 if d <= 0 {
562 return ctx.Err()
563 }
564 t := time.NewTimer(d)
565 defer t.Stop()
566 select {
567 case <-ctx.Done():
568 return ctx.Err()
569 case <-t.C:
570 return nil
571 }
572}
573
574func (a *Agent) dropLastAssistantMessage() {
575 a.mu.Lock()
576 defer a.mu.Unlock()
577 if n := len(a.messages); n > 0 && a.messages[n-1].Role == provider.RoleAssistant {
578 a.messages = a.messages[:n-1]
579 a.rev++
580 }
581}
582
583// oneTurn calls the LLM once, forwards events, returns the stop reason
584// and the assembled assistant message (already appended to the transcript).
585func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
586 var req provider.Request
587 for {
588 if err := a.prepareStart(ctx); err != nil {
589 return provider.StopError, provider.Message{}, err
590 }
591 a.mu.Lock()
592 // A reset can also arrive after runLoop's preparation, for example
593 // while BeforeTurn waits. Validate and snapshot under the same lock
594 // so this request never receives an unprepared replacement prompt.
595 if a.BeforeStart != nil && !a.startCurrentLocked() {
596 a.mu.Unlock()
597 continue
598 }
599 req = provider.Request{
600 Model: a.Model,
601 System: a.System,
602 // Repair any dangling tool_use blocks before sending. A turn
603 // aborted mid-flight (cancel, connection drop, ECONNREFUSED to a
604 // dev server, etc.) can leave an assistant tool_use with no
605 // matching tool_result in the live transcript. The load-time
606 // repair in OpenSession only runs on restart, so without this the
607 // next in-process request is rejected by providers like Anthropic
608 // with "tool_use ids were found without tool_result blocks". The
609 // repair is pure and a no-op on already-valid transcripts.
610 Messages: repairToolUseResultPairs(append([]provider.Message(nil), a.messages...)),
611 Tools: a.Tools.Specs(),
612 Reasoning: a.Reasoning,
613 MaxTokens: a.MaxTokens,
614 Temperature: a.Temperature,
615 SessionID: a.SessionID,
616 MaxToolCalls: a.MaxToolCalls,
617 }
618 a.mu.Unlock()
619 break
620 }
621 stream, err := a.Client.Stream(ctx, req)
622 if err != nil {
623 return provider.StopError, provider.Message{}, err
624 }
625
626 sink(EvAssistantStart{})
627
628 var (
629 stop provider.StopReason
630 finalErr error
631 finalMsg provider.Message
632 )
633
634 for ev := range stream {
635 switch e := ev.(type) {
636 case provider.EventStart:
637 // nothing
638 case provider.EventTextDelta:
639 sink(EvTextDelta{Delta: e.Delta})
640 case provider.EventToolStart:
641 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
642 case provider.EventToolArgs:
643 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
644 case provider.EventToolEnd:
645 sink(EvToolUseEnd{ID: e.ID})
646 case provider.EventUsage:
647 cum := a.cost.Add(e.Usage)
648 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
649 if a.OnUsage != nil {
650 a.OnUsage(cum)
651 }
652 case provider.EventDone:
653 stop = e.Stop
654 finalErr = e.Err
655 finalMsg = e.Message
656 }
657 }
658
659 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
660 // produce partial content. When the partial message is text only we
661 // keep whatever was streamed up to the cancel so the user does not
662 // lose visible work (a cut-off summary is still useful). If the
663 // partial message already contained tool-call blocks we drop the
664 // whole thing, because an unmatched tool_use would fail the next
665 // turn with a tool_result mismatch error.
666 keep := len(finalMsg.Content) > 0
667 if stop == provider.StopAborted && keep {
668 hasToolCall := false
669 for _, c := range finalMsg.Content {
670 if _, ok := c.(provider.ToolCallBlock); ok {
671 hasToolCall = true
672 break
673 }
674 }
675 if hasToolCall {
676 keep = false
677 }
678 }
679 if keep {
680 emit := finalMsg
681 suppress := false
682
683 // BeforeAssistantMessage hook: extensions can suppress or
684 // rewrite the visible text. The transcript keeps the
685 // model's original output so the model still sees what it
686 // said on subsequent turns.
687 if a.BeforeAssistantMessage != nil {
688 orig := extractText(finalMsg)
689 if orig != "" {
690 allowed, _, replacement := a.BeforeAssistantMessage(orig)
691 if !allowed {
692 suppress = true
693 } else if replacement != "" && replacement != orig {
694 emit = replaceText(finalMsg, replacement)
695 }
696 }
697 }
698
699 a.mu.Lock()
700 a.messages = append(a.messages, finalMsg)
701 a.rev++
702 a.mu.Unlock()
703 a.fireMessageAppended(finalMsg)
704 if !suppress {
705 sink(EvAssistantMessage{Message: emit})
706 }
707 // Now surface tool calls as EvToolCall events so UIs can render them
708 // in order before the tool results arrive.
709 for _, c := range finalMsg.Content {
710 if tc, ok := c.(provider.ToolCallBlock); ok {
711 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
712 }
713 }
714 }
715
716 return stop, finalMsg, finalErr
717}
718
719// executeTools runs every tool call in the assistant message and returns
720// a single tool-role message carrying all results.
721func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
722 var results []provider.Content
723 var addedTools []string
724 hadError := false
725
726 for _, c := range msg.Content {
727 tc, ok := c.(provider.ToolCallBlock)
728 if !ok || tc.Server {
729 continue
730 }
731 res := a.runOneTool(ctx, tc, sink)
732 if res.IsError {
733 hadError = true
734 }
735 results = append(results, provider.ToolResultBlock{
736 CallID: tc.ID,
737 Content: res.Content,
738 IsError: res.IsError,
739 })
740 for _, name := range res.ActivateTools {
741 if _, err := a.Tools.Get(name); err == nil && !containsString(addedTools, name) {
742 addedTools = append(addedTools, name)
743 }
744 }
745 sink(EvToolResult{ID: tc.ID, Result: res})
746 }
747
748 return provider.Message{
749 Role: provider.RoleTool,
750 Content: results,
751 Time: time.Now(),
752 AddedToolNames: addedTools,
753 }, hadError
754}
755
756func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
757 tool, err := a.Tools.Get(tc.Name)
758 if err != nil {
759 return ToolResult{
760 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
761 IsError: true,
762 }
763 }
764
765 args := tc.Arguments
766
767 // Intercept hook: an extension or other guard can refuse the
768 // call before any side effect happens, OR rewrite the args
769 // seen by the tool. The model sees the reason as the tool
770 // error, learns from it, and (typically) proposes a different
771 // action; rewrites are invisible to the model (they apply only
772 // to the execution).
773 if a.BeforeToolExecute != nil {
774 allowed, reason, modified := a.BeforeToolExecute(tc)
775 if !allowed {
776 if reason == "" {
777 reason = "tool call refused by extension guard"
778 }
779 return ToolResult{
780 Content: []provider.Content{provider.TextBlock{Text: reason}},
781 IsError: true,
782 }
783 }
784 if len(modified) > 0 && json.Valid(modified) {
785 args = modified
786 }
787 }
788
789 if len(args) == 0 {
790 args = json.RawMessage("{}")
791 }
792
793 // Recover panics so a buggy tool does not crash the agent.
794 var res ToolResult
795 func() {
796 defer func() {
797 if r := recover(); r != nil {
798 res = ToolResult{
799 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
800 IsError: true,
801 }
802 }
803 }()
804 out, err := tool.Execute(ctx, args, func(text string) {
805 sink(EvToolProgress{ID: tc.ID, Text: text})
806 })
807 if err != nil {
808 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
809 res = ToolResult{
810 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
811 IsError: true,
812 }
813 return
814 }
815 res = ToolResult{
816 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
817 IsError: true,
818 }
819 return
820 }
821 res = out
822 }()
823 return res
824}
825
826// extractText concatenates all TextBlock content in a message. Used
827// by BeforeAssistantMessage so guards see a single string instead of
828// having to walk provider.Content themselves.
829func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
830 var content []provider.Content
831 for _, c := range msg.Content {
832 tr, ok := c.(provider.ToolResultBlock)
833 if !ok {
834 continue
835 }
836 for _, inner := range tr.Content {
837 switch v := inner.(type) {
838 case provider.TextBlock:
839 // Keep short textual context so the model understands why
840 // the images appeared, but don't duplicate giant read
841 // outputs verbatim.
842 if len(v.Text) > 0 && len(v.Text) <= 500 {
843 content = append(content, v)
844 }
845 case provider.ImageBlock:
846 content = append(content, v)
847 }
848 }
849 }
850 if len(content) == 0 {
851 return provider.Message{}
852 }
853 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
854 content = append([]provider.Content{prefix}, content...)
855 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
856}
857
858func containsString(values []string, want string) bool {
859 for _, value := range values {
860 if value == want {
861 return true
862 }
863 }
864 return false
865}
866
867func extractText(msg provider.Message) string {
868 var out string
869 for _, c := range msg.Content {
870 if tb, ok := c.(provider.TextBlock); ok {
871 if out != "" {
872 out += "\n"
873 }
874 out += tb.Text
875 }
876 }
877 return out
878}
879
880// replaceText returns a copy of msg with every TextBlock replaced by
881// a single TextBlock containing replacement. Non-text content (tool
882// calls, etc.) is preserved in order.
883func replaceText(msg provider.Message, replacement string) provider.Message {
884 out := provider.Message{Role: msg.Role}
885 out.Content = make([]provider.Content, 0, len(msg.Content))
886 replaced := false
887 for _, c := range msg.Content {
888 if _, ok := c.(provider.TextBlock); ok {
889 if !replaced {
890 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
891 replaced = true
892 }
893 continue
894 }
895 out.Content = append(out.Content, c)
896 }
897 if !replaced {
898 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
899 }
900 return out
901}