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