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 }
475 for _, needle := range needles {
476 if strings.Contains(msg, needle) {
477 return true
478 }
479 }
480 return false
481}
482
483func isNonRetryableProviderLimit(msg string) bool {
484 needles := []string{
485 "usage limit", "monthly usage limit", "freeusagelimit", "gousagelimit",
486 "available balance", "insufficient_quota", "out of budget", "quota exceeded", "billing",
487 }
488 for _, needle := range needles {
489 if strings.Contains(msg, needle) {
490 return true
491 }
492 }
493 return false
494}
495
496func (a *Agent) retryDelay(attempt int) time.Duration {
497 base := a.RetryBaseDelay
498 if base <= 0 {
499 base = 2 * time.Second
500 }
501 return base * time.Duration(1<<attempt)
502}
503
504func sleepRetry(ctx context.Context, d time.Duration) error {
505 if d <= 0 {
506 return ctx.Err()
507 }
508 t := time.NewTimer(d)
509 defer t.Stop()
510 select {
511 case <-ctx.Done():
512 return ctx.Err()
513 case <-t.C:
514 return nil
515 }
516}
517
518func (a *Agent) dropLastAssistantMessage() {
519 a.mu.Lock()
520 defer a.mu.Unlock()
521 if n := len(a.messages); n > 0 && a.messages[n-1].Role == provider.RoleAssistant {
522 a.messages = a.messages[:n-1]
523 a.rev++
524 }
525}
526
527// oneTurn calls the LLM once, forwards events, returns the stop reason
528// and the assembled assistant message (already appended to the transcript).
529func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
530 req := provider.Request{
531 Model: a.Model,
532 System: a.System,
533 // Repair any dangling tool_use blocks before sending. A turn
534 // aborted mid-flight (cancel, connection drop, ECONNREFUSED to a
535 // dev server, etc.) can leave an assistant tool_use with no
536 // matching tool_result in the live transcript. The load-time
537 // repair in OpenSession only runs on restart, so without this the
538 // next in-process request is rejected by providers like Anthropic
539 // with "tool_use ids were found without tool_result blocks". The
540 // repair is pure and a no-op on already-valid transcripts.
541 Messages: repairToolUseResultPairs(a.Messages()),
542 Tools: a.Tools.Specs(),
543 Reasoning: a.Reasoning,
544 MaxTokens: a.MaxTokens,
545 Temperature: a.Temperature,
546 }
547 stream, err := a.Client.Stream(ctx, req)
548 if err != nil {
549 return provider.StopError, provider.Message{}, err
550 }
551
552 sink(EvAssistantStart{})
553
554 var (
555 stop provider.StopReason
556 finalErr error
557 finalMsg provider.Message
558 )
559
560 for ev := range stream {
561 switch e := ev.(type) {
562 case provider.EventStart:
563 // nothing
564 case provider.EventTextDelta:
565 sink(EvTextDelta{Delta: e.Delta})
566 case provider.EventToolStart:
567 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
568 case provider.EventToolArgs:
569 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
570 case provider.EventToolEnd:
571 sink(EvToolUseEnd{ID: e.ID})
572 case provider.EventUsage:
573 cum := a.cost.Add(e.Usage)
574 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
575 if a.OnUsage != nil {
576 a.OnUsage(cum)
577 }
578 case provider.EventDone:
579 stop = e.Stop
580 finalErr = e.Err
581 finalMsg = e.Message
582 }
583 }
584
585 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
586 // produce partial content. When the partial message is text only we
587 // keep whatever was streamed up to the cancel so the user does not
588 // lose visible work (a cut-off summary is still useful). If the
589 // partial message already contained tool-call blocks we drop the
590 // whole thing, because an unmatched tool_use would fail the next
591 // turn with a tool_result mismatch error.
592 keep := len(finalMsg.Content) > 0
593 if stop == provider.StopAborted && keep {
594 hasToolCall := false
595 for _, c := range finalMsg.Content {
596 if _, ok := c.(provider.ToolCallBlock); ok {
597 hasToolCall = true
598 break
599 }
600 }
601 if hasToolCall {
602 keep = false
603 }
604 }
605 if keep {
606 emit := finalMsg
607 suppress := false
608
609 // BeforeAssistantMessage hook: extensions can suppress or
610 // rewrite the visible text. The transcript keeps the
611 // model's original output so the model still sees what it
612 // said on subsequent turns.
613 if a.BeforeAssistantMessage != nil {
614 orig := extractText(finalMsg)
615 if orig != "" {
616 allowed, _, replacement := a.BeforeAssistantMessage(orig)
617 if !allowed {
618 suppress = true
619 } else if replacement != "" && replacement != orig {
620 emit = replaceText(finalMsg, replacement)
621 }
622 }
623 }
624
625 a.mu.Lock()
626 a.messages = append(a.messages, finalMsg)
627 a.rev++
628 a.mu.Unlock()
629 a.fireMessageAppended(finalMsg)
630 if !suppress {
631 sink(EvAssistantMessage{Message: emit})
632 }
633 // Now surface tool calls as EvToolCall events so UIs can render them
634 // in order before the tool results arrive.
635 for _, c := range finalMsg.Content {
636 if tc, ok := c.(provider.ToolCallBlock); ok {
637 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
638 }
639 }
640 }
641
642 return stop, finalMsg, finalErr
643}
644
645// executeTools runs every tool call in the assistant message and returns
646// a single tool-role message carrying all results.
647func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
648 var results []provider.Content
649 var addedTools []string
650 hadError := false
651
652 for _, c := range msg.Content {
653 tc, ok := c.(provider.ToolCallBlock)
654 if !ok {
655 continue
656 }
657 res := a.runOneTool(ctx, tc, sink)
658 if res.IsError {
659 hadError = true
660 }
661 results = append(results, provider.ToolResultBlock{
662 CallID: tc.ID,
663 Content: res.Content,
664 IsError: res.IsError,
665 })
666 for _, name := range res.ActivateTools {
667 if _, err := a.Tools.Get(name); err == nil && !containsString(addedTools, name) {
668 addedTools = append(addedTools, name)
669 }
670 }
671 sink(EvToolResult{ID: tc.ID, Result: res})
672 }
673
674 return provider.Message{
675 Role: provider.RoleTool,
676 Content: results,
677 Time: time.Now(),
678 AddedToolNames: addedTools,
679 }, hadError
680}
681
682func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
683 tool, err := a.Tools.Get(tc.Name)
684 if err != nil {
685 return ToolResult{
686 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
687 IsError: true,
688 }
689 }
690
691 args := tc.Arguments
692
693 // Intercept hook: an extension or other guard can refuse the
694 // call before any side effect happens, OR rewrite the args
695 // seen by the tool. The model sees the reason as the tool
696 // error, learns from it, and (typically) proposes a different
697 // action; rewrites are invisible to the model (they apply only
698 // to the execution).
699 if a.BeforeToolExecute != nil {
700 allowed, reason, modified := a.BeforeToolExecute(tc)
701 if !allowed {
702 if reason == "" {
703 reason = "tool call refused by extension guard"
704 }
705 return ToolResult{
706 Content: []provider.Content{provider.TextBlock{Text: reason}},
707 IsError: true,
708 }
709 }
710 if len(modified) > 0 && json.Valid(modified) {
711 args = modified
712 }
713 }
714
715 if len(args) == 0 {
716 args = json.RawMessage("{}")
717 }
718
719 // Recover panics so a buggy tool does not crash the agent.
720 var res ToolResult
721 func() {
722 defer func() {
723 if r := recover(); r != nil {
724 res = ToolResult{
725 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
726 IsError: true,
727 }
728 }
729 }()
730 out, err := tool.Execute(ctx, args, func(text string) {
731 sink(EvToolProgress{ID: tc.ID, Text: text})
732 })
733 if err != nil {
734 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
735 res = ToolResult{
736 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
737 IsError: true,
738 }
739 return
740 }
741 res = ToolResult{
742 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
743 IsError: true,
744 }
745 return
746 }
747 res = out
748 }()
749 return res
750}
751
752// extractText concatenates all TextBlock content in a message. Used
753// by BeforeAssistantMessage so guards see a single string instead of
754// having to walk provider.Content themselves.
755func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
756 var content []provider.Content
757 for _, c := range msg.Content {
758 tr, ok := c.(provider.ToolResultBlock)
759 if !ok {
760 continue
761 }
762 for _, inner := range tr.Content {
763 switch v := inner.(type) {
764 case provider.TextBlock:
765 // Keep short textual context so the model understands why
766 // the images appeared, but don't duplicate giant read
767 // outputs verbatim.
768 if len(v.Text) > 0 && len(v.Text) <= 500 {
769 content = append(content, v)
770 }
771 case provider.ImageBlock:
772 content = append(content, v)
773 }
774 }
775 }
776 if len(content) == 0 {
777 return provider.Message{}
778 }
779 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
780 content = append([]provider.Content{prefix}, content...)
781 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
782}
783
784func containsString(values []string, want string) bool {
785 for _, value := range values {
786 if value == want {
787 return true
788 }
789 }
790 return false
791}
792
793func extractText(msg provider.Message) string {
794 var out string
795 for _, c := range msg.Content {
796 if tb, ok := c.(provider.TextBlock); ok {
797 if out != "" {
798 out += "\n"
799 }
800 out += tb.Text
801 }
802 }
803 return out
804}
805
806// replaceText returns a copy of msg with every TextBlock replaced by
807// a single TextBlock containing replacement. Non-text content (tool
808// calls, etc.) is preserved in order.
809func replaceText(msg provider.Message, replacement string) provider.Message {
810 out := provider.Message{Role: msg.Role}
811 out.Content = make([]provider.Content, 0, len(msg.Content))
812 replaced := false
813 for _, c := range msg.Content {
814 if _, ok := c.(provider.TextBlock); ok {
815 if !replaced {
816 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
817 replaced = true
818 }
819 continue
820 }
821 out.Content = append(out.Content, c)
822 }
823 if !replaced {
824 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
825 }
826 return out
827}