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