1package core
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "io"
9 "strings"
10 "sync"
11 "time"
12 "unicode/utf8"
13
14 "github.com/bnema/zut/packages/provider"
15 "github.com/google/uuid"
16)
17
18var (
19 ErrStreamIdleTimeout = errors.New("provider stream idle timeout")
20 ErrMaxSteps = errors.New("agent model step limit exceeded")
21)
22
23const (
24 DefaultStreamIdleTimeout = 5 * time.Minute
25 DefaultStreamMaxRetries = 5
26)
27
28// QueuedMessage is a user prompt waiting for the next safe model-call boundary.
29type QueuedMessage struct {
30 Text string
31 Images []provider.ImageBlock
32 // HostEvent distinguishes host evidence from ordinary input so hosts can
33 // retain it after a failed turn. Explicit cancellation may still discard it.
34 HostEvent bool
35}
36
37type queuedMessage struct {
38 message QueuedMessage
39 accepted time.Time
40}
41
42// requestOpenError marks a failure returned before a provider stream exists.
43// Provider clients own retries for this phase; the agent retry budget is only
44// for reconnecting a stream that opened and then failed.
45type requestOpenError struct{ err error }
46
47func (e *requestOpenError) Error() string { return e.err.Error() }
48func (e *requestOpenError) Unwrap() error { return e.err }
49
50// Agent is a stateful conversation bound to a provider client, a model,
51// and a set of tools.
52type Agent struct {
53 Client provider.Client
54 Model string
55 // System is mutex-protected; use SetSystemPrompt to write it and
56 // PromptConfig to read it.
57 System string
58 // Tools is mutex-protected; use SetTools to write it and PromptConfig
59 // or ToolsSnapshot to read it.
60 Tools Registry
61 MaxSteps int
62 Reasoning string
63 Temperature *float32
64 // FastMode requests the OpenAI fast service tier. The provider boundary
65 // rejects it for providers that do not support that contract.
66 FastMode bool
67
68 // ContextWindow is the effective input context capacity retained from
69 // model resolution. Hosts use it for proactive compaction even when the
70 // model was synthesized for an open-catalog provider.
71 ContextWindow int
72
73 // MaxTokens caps the model's output tokens per turn. Zero leaves
74 // the field unset on the provider request, letting each provider
75 // apply its own default (which can be conservative, e.g. Bedrock
76 // defaults to 4096, truncating long writes/edits). Hosts populate
77 // this from the resolved model's MaxOutput so large single-turn
78 // responses aren't silently cut off with stopReason=length.
79 MaxTokens int
80
81 // BeforeToolExecute, if set, is called immediately before each
82 // tool runs. Returning (allowed=false, reason) short-circuits
83 // the call with an error result containing reason. Optionally,
84 // returning a non-nil modifiedArgs replaces the JSON args the
85 // tool will see, which lets guards redact / augment / patch the
86 // model's request without rewriting the transcript. Empty or
87 // malformed modifiedArgs is ignored.
88 BeforeToolExecute func(call provider.ToolCallBlock) (allowed bool, reason string, modifiedArgs json.RawMessage)
89
90 // BeforeTurn, if set, is called before each turn's model call.
91 // Returning (allowed=false, reason) aborts the turn; reason is
92 // surfaced as an assistant-like status line. Used for rate-
93 // limiting, business-hour gates, and deny-by-default setups.
94 BeforeTurn func(step int) (allowed bool, reason string)
95
96 // BeforeTurnContext is the richer turn-start hook. It receives the
97 // current turn context so blocking interception can honor cancellation.
98 // It may return a bounded, non-transcript context string that is appended
99 // to the provider request's system prompt for this turn only. When set, it
100 // replaces BeforeTurn so a host can combine approval and context
101 // collection into one lifecycle round trip.
102 BeforeTurnContext func(ctx context.Context, step int) (allowed bool, reason, context string)
103
104 // CommitToolResult, if set, commits extension/tool metadata immediately
105 // after each tool call completes. Returning an error converts the result to
106 // an error before it is appended to the transcript or sent to the provider.
107 CommitToolResult func(id string, result ToolResult) error
108 // OnToolResult observes the final result after CommitToolResult succeeds or
109 // converts it to an error.
110 OnToolResult func(id string, result ToolResult)
111
112 // BeforeAssistantMessage, if set, is called after the model's
113 // final assistant message is assembled but before it's appended
114 // to the transcript. Returning (allowed=false) suppresses both
115 // the transcript append and the UI event. A non-empty
116 // replacement rewrites the visible text for the user while
117 // leaving the model's original text in the transcript (so the
118 // model can still see what it said in subsequent turns).
119 BeforeAssistantMessage func(text string) (allowed bool, reason, replacement string)
120
121 // MaxRetries controls reconnection attempts for transient failures after the
122 // provider stream opens. Zero disables this retry layer.
123 // RetryBaseDelay is doubled for each attempt; zero uses 2s.
124 MaxRetries int
125 RetryBaseDelay time.Duration
126 // StreamIdleTimeout cancels an individual provider stream when it emits no
127 // events for this duration. Zero leaves the stream unbounded.
128 StreamIdleTimeout time.Duration
129 streamIdleTimer func(time.Duration) (<-chan time.Time, func(), func())
130
131 // OnEvent, if set, mirrors every AgentEvent the loop emits to
132 // this callback in addition to the per-Prompt sink. Used by the
133 // extension manager to fan events out to subscribed extensions
134 // without each caller having to compose sinks manually.
135 OnEvent func(AgentEvent)
136
137 // OnMessageAppended, if set, fires every time a message is
138 // appended to the in-memory transcript by the agent loop — the
139 // initial user prompt, each finalised assistant message, and
140 // each tool-results message (plus the synthetic OpenAI image
141 // mirror, if any). Hosts wire this to the on-disk session so
142 // that turns are durable as soon as they happen, instead of
143 // only being flushed on a clean exit.
144 OnMessageAppended func(provider.Message)
145
146 // OnUsage, if set, fires after every turn's usage row arrives,
147 // carrying the cumulative usage for the session. Hosts wire
148 // this to the on-disk session so the persisted total stays
149 // current and a crash recovers the right cost figure.
150 OnUsage func(cumulative provider.Usage)
151
152 // OnRetryLifecycle receives sanitized retry audit records. It never
153 // receives provider error text or response bodies.
154 OnRetryLifecycle func(RetryLifecycleRecord)
155
156 // OnTranscriptCompacted, if set, fires after Compact replaces the
157 // in-memory transcript with the synthetic summary plus kept tail.
158 // Hosts wire this to append an explicit compaction checkpoint to
159 // the session log; per-message append hooks do not fire for this
160 // wholesale transcript replacement.
161 OnTranscriptCompacted func(messages []provider.Message)
162
163 mu sync.Mutex
164 messages []provider.Message
165 // rev increments whenever the transcript slice is replaced or a
166 // message is appended. The TUI uses it as a cheap redraw cache key
167 // so editor-only typing doesn't copy/rebuild a long transcript on
168 // every keypress.
169 rev uint64
170 cost CostTracker
171
172 // queued holds user messages submitted while the agent is busy.
173 // The loop appends them as normal user messages at safe
174 // boundaries: before the next model call after a tool batch, or
175 // after a text-only assistant turn finishes. It never interrupts
176 // a running tool or cancels an in-flight provider request.
177 queued []queuedMessage
178 timeContext agentTimeContext
179 hasSessionTimeContext bool
180 // sessionID identifies this conversation thread. cacheSessionID identifies
181 // the root prompt-cache affinity and may be shared by resident child
182 // threads. Hosts bind both before the first provider request; unbound
183 // agents allocate one shared ephemeral value for the process lifetime.
184 sessionID string
185 cacheSessionID string
186 providerRequestStarted bool
187}
188
189// NewAgent returns an Agent with sensible defaults.
190func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
191 return &Agent{
192 Client: client,
193 Model: model,
194 System: system,
195 Tools: tools,
196 MaxSteps: 0, // 0 = unlimited
197 MaxRetries: DefaultStreamMaxRetries,
198 RetryBaseDelay: 2 * time.Second,
199 StreamIdleTimeout: DefaultStreamIdleTimeout,
200 timeContext: newAgentTimeContext(time.Now()),
201 }
202}
203
204// BindSessionID binds a root conversation's durable session as both its
205// thread identity and prompt-cache affinity. Resident children use
206// BindRequestIdentity to retain their own thread while sharing their root's
207// cache affinity.
208func (a *Agent) BindSessionID(id string) error {
209 return a.BindRequestIdentity(id, id)
210}
211
212// BindRequestIdentity assigns the stable root cache and thread identities
213// before the first provider request.
214func (a *Agent) BindRequestIdentity(cacheSessionID, threadID string) error {
215 cacheSessionID = strings.TrimSpace(cacheSessionID)
216 threadID = strings.TrimSpace(threadID)
217 if cacheSessionID == "" || threadID == "" {
218 return errors.New("cache session ID and thread ID are required")
219 }
220 a.mu.Lock()
221 defer a.mu.Unlock()
222 if a.providerRequestStarted {
223 return errors.New("cannot change request identity after the first provider request")
224 }
225 a.cacheSessionID = cacheSessionID
226 a.sessionID = threadID
227 return nil
228}
229
230// SessionID returns the stable conversation-thread identity, allocating an
231// ephemeral identity when persistence is disabled and the host did not bind
232// one during runtime assembly.
233func (a *Agent) SessionID() string {
234 a.mu.Lock()
235 defer a.mu.Unlock()
236 a.ensureRequestIdentityLocked()
237 return a.sessionID
238}
239
240// CacheSessionID returns the stable root cache-affinity identity.
241func (a *Agent) CacheSessionID() string {
242 a.mu.Lock()
243 defer a.mu.Unlock()
244 a.ensureRequestIdentityLocked()
245 return a.cacheSessionID
246}
247
248func (a *Agent) ensureRequestIdentityLocked() {
249 if a.sessionID == "" {
250 a.sessionID = uuid.NewString()
251 }
252 if a.cacheSessionID == "" {
253 a.cacheSessionID = a.sessionID
254 }
255}
256
257func (a *Agent) beginTurn() provider.RequestContext {
258 a.mu.Lock()
259 a.ensureRequestIdentityLocked()
260 context := provider.RequestContext{CacheSessionID: a.cacheSessionID, ThreadID: a.sessionID, TurnID: uuid.NewString()}
261 a.mu.Unlock()
262 return context
263}
264
265func (a *Agent) beginProviderRequest() {
266 a.mu.Lock()
267 a.providerRequestStarted = true
268 a.mu.Unlock()
269}
270
271// SetFastMode changes the service-tier preference for the next model call.
272// It is safe to use while a turn is in flight; the active turn keeps the
273// value it already captured.
274func (a *Agent) SetFastMode(enabled bool) {
275 a.mu.Lock()
276 a.FastMode = enabled
277 a.mu.Unlock()
278}
279
280// FastModeEnabled returns the service-tier preference for a model call.
281func (a *Agent) FastModeEnabled() bool {
282 a.mu.Lock()
283 defer a.mu.Unlock()
284 return a.FastMode
285}
286
287// QueueMessage queues a text and image prompt to be injected as a user message
288// at the next safe boundary of the active agent loop. It is non-blocking in the
289// sense that it never waits for model/tool work; it only takes the transcript
290// mutex briefly. Prompts without text or images are ignored.
291func (a *Agent) QueueMessage(text string, images []provider.ImageBlock) bool {
292 return a.QueuePrompt(QueuedMessage{Text: text, Images: images})
293}
294
295// QueuePrompt queues a prompt with its host provenance at the next safe boundary.
296func (a *Agent) QueuePrompt(prompt QueuedMessage) bool {
297 prompt.Text = strings.TrimSpace(prompt.Text)
298 if prompt.Text == "" && len(prompt.Images) == 0 {
299 return false
300 }
301 message := cloneQueuedMessage(prompt)
302 a.mu.Lock()
303 a.queued = append(a.queued, queuedMessage{message: message, accepted: time.Now()})
304 a.mu.Unlock()
305 return true
306}
307
308// PendingQueuedMessages returns a snapshot of user messages waiting
309// to be injected. Used by hosts to render the visible "sliding in"
310// chips without consuming them.
311func (a *Agent) PendingQueuedMessages() []QueuedMessage {
312 a.mu.Lock()
313 defer a.mu.Unlock()
314 out := make([]QueuedMessage, len(a.queued))
315 for i, queued := range a.queued {
316 out[i] = cloneQueuedMessage(queued.message)
317 }
318 return out
319}
320
321// QueuedMessageCount returns the number of messages waiting to be
322// injected at the next safe boundary.
323func (a *Agent) QueuedMessageCount() int {
324 a.mu.Lock()
325 defer a.mu.Unlock()
326 return len(a.queued)
327}
328
329// PopQueuedMessage removes and returns the most recently queued
330// message. Hosts use this for the slide-back keybinding.
331func (a *Agent) PopQueuedMessage() (QueuedMessage, bool) {
332 a.mu.Lock()
333 defer a.mu.Unlock()
334 n := len(a.queued)
335 if n == 0 {
336 return QueuedMessage{}, false
337 }
338 message := cloneQueuedMessage(a.queued[n-1].message)
339 a.queued = a.queued[:n-1]
340 return message, true
341}
342
343// DrainQueuedMessages discards and returns every queued message.
344// Hosts use this on explicit cancel/clear so stale follow-ups do
345// not run after the user aborted the turn.
346func (a *Agent) DrainQueuedMessages() []QueuedMessage {
347 queued := a.drainQueuedMessages()
348 out := make([]QueuedMessage, len(queued))
349 for i, message := range queued {
350 out[i] = cloneQueuedMessage(message.message)
351 }
352 return out
353}
354
355func cloneQueuedMessage(message QueuedMessage) QueuedMessage {
356 clone := QueuedMessage{Text: message.Text, HostEvent: message.HostEvent}
357 if len(message.Images) == 0 {
358 return clone
359 }
360 clone.Images = make([]provider.ImageBlock, len(message.Images))
361 for i, image := range message.Images {
362 clone.Images[i] = image
363 clone.Images[i].Data = append([]byte(nil), image.Data...)
364 }
365 return clone
366}
367
368func (a *Agent) drainQueuedMessages() []queuedMessage {
369 a.mu.Lock()
370 defer a.mu.Unlock()
371 out := append([]queuedMessage(nil), a.queued...)
372 a.queued = nil
373 return out
374}
375
376func (a *Agent) appendQueuedAsUser(messages []queuedMessage, sink func(AgentEvent)) {
377 for _, queued := range messages {
378 accepted := queued.accepted
379 if accepted.IsZero() {
380 accepted = time.Now()
381 }
382 content := make([]provider.Content, 0, 1+len(queued.message.Images))
383 if queued.message.Text != "" {
384 content = append(content, provider.TextBlock{Text: queued.message.Text})
385 }
386 for _, image := range queued.message.Images {
387 content = append(content, image)
388 }
389 msg := provider.Message{
390 Role: provider.RoleUser,
391 Content: content,
392 Time: accepted,
393 }
394 a.mu.Lock()
395 a.messages = append(a.messages, msg)
396 a.rev++
397 a.mu.Unlock()
398 a.fireMessageAppended(msg)
399 if sink != nil {
400 sink(EvUserMessage{Message: msg})
401 }
402 }
403}
404
405// Messages returns a copy of the current transcript.
406func (a *Agent) Messages() []provider.Message {
407 a.mu.Lock()
408 defer a.mu.Unlock()
409 out := make([]provider.Message, len(a.messages))
410 copy(out, a.messages)
411 return out
412}
413
414// MessageCount returns the current transcript length without copying it.
415func (a *Agent) MessageCount() int {
416 a.mu.Lock()
417 defer a.mu.Unlock()
418 return len(a.messages)
419}
420
421// MessagesFrom returns a copy of the transcript suffix beginning at from and
422// the total transcript length observed under the same lock.
423func (a *Agent) MessagesFrom(from int) ([]provider.Message, int) {
424 a.mu.Lock()
425 defer a.mu.Unlock()
426 if from < 0 {
427 from = 0
428 }
429 if from > len(a.messages) {
430 from = len(a.messages)
431 }
432 out := make([]provider.Message, len(a.messages)-from)
433 copy(out, a.messages[from:])
434 return out, len(a.messages)
435}
436
437// Revision returns a monotonically increasing transcript version.
438// It is cheap to query and changes whenever Messages() would return
439// different transcript content because of append/set operations.
440func (a *Agent) Revision() uint64 {
441 a.mu.Lock()
442 defer a.mu.Unlock()
443 return a.rev
444}
445
446// SetTools swaps the tool registry. Used by /reload-ext to hand
447// the agent a fresh registry after extension subprocesses have been
448// respawned (and their freshly-registered tools merged in).
449func (a *Agent) SetTools(reg Registry) {
450 a.mu.Lock()
451 a.Tools = reg
452 a.mu.Unlock()
453}
454
455// SetSystemPrompt replaces the system prompt used for subsequent turns.
456func (a *Agent) SetSystemPrompt(system string) {
457 a.mu.Lock()
458 a.System = system
459 a.mu.Unlock()
460}
461
462// SetPromptConfig atomically replaces the system prompt and tool registry.
463// It returns the previous registry so the caller can close resources that are
464// no longer owned by the agent.
465func (a *Agent) SetPromptConfig(system string, tools Registry) Registry {
466 a.mu.Lock()
467 defer a.mu.Unlock()
468 oldTools := a.Tools
469 a.System = system
470 a.Tools = tools
471 return oldTools
472}
473
474// PromptConfig returns a consistent snapshot of the system prompt and tools
475// used to construct a provider request.
476func (a *Agent) PromptConfig() (string, Registry) {
477 a.mu.Lock()
478 defer a.mu.Unlock()
479 tools := make(Registry, len(a.Tools))
480 for name, tool := range a.Tools {
481 tools[name] = tool
482 }
483 return a.System, tools
484}
485
486// ToolsSnapshot returns a copy of the current tool registry.
487func (a *Agent) ToolsSnapshot() Registry {
488 a.mu.Lock()
489 defer a.mu.Unlock()
490 tools := make(Registry, len(a.Tools))
491 for name, tool := range a.Tools {
492 tools[name] = tool
493 }
494 return tools
495}
496
497// SetMessages replaces the transcript (used when resuming a session).
498func (a *Agent) SetMessages(msgs []provider.Message) {
499 a.mu.Lock()
500 defer a.mu.Unlock()
501 a.messages = append(a.messages[:0], msgs...)
502 a.rev++
503}
504
505// AppendUserContext adds a user-role message to the transcript without
506// starting a model turn. Hosts use it for context gathered outside the agent
507// loop, such as the output of an explicitly invoked shell command.
508func (a *Agent) AppendUserContext(text string, meta map[string]string) {
509 msg := provider.Message{
510 Role: provider.RoleUser,
511 Content: []provider.Content{provider.TextBlock{Text: text}},
512 Time: time.Now(),
513 Meta: meta,
514 }
515 a.mu.Lock()
516 a.messages = append(a.messages, msg)
517 a.rev++
518 a.mu.Unlock()
519 a.fireMessageAppended(msg)
520}
521
522// Cost returns the cumulative usage.
523func (a *Agent) Cost() provider.Usage {
524 a.mu.Lock()
525 defer a.mu.Unlock()
526 return a.cost.Total
527}
528
529func (a *Agent) addUsage(usage provider.Usage) provider.Usage {
530 a.mu.Lock()
531 defer a.mu.Unlock()
532 return a.cost.Add(usage)
533}
534
535// SeedCost sets the cumulative usage as a baseline before the first
536// turn runs. Used when transferring state from another agent (model
537// or provider switch) so the running cost meter doesn't reset to 0.
538func (a *Agent) SeedCost(u provider.Usage) {
539 a.mu.Lock()
540 defer a.mu.Unlock()
541 a.cost.Seed(u)
542}
543
544// LastTurnUsage returns the per-turn usage of the most recent
545// completed turn. Drives the "context used" gauge in the status bar
546// without waiting for the next turn to land.
547func (a *Agent) LastTurnUsage() provider.Usage {
548 a.mu.Lock()
549 defer a.mu.Unlock()
550 return a.cost.LastTurn
551}
552
553// SeedLastTurnUsage primes the per-turn snapshot. Used on resume so
554// the gauge reflects the prompt size of the last turn in the session
555// file instead of starting at zero.
556func (a *Agent) SeedLastTurnUsage(u provider.Usage) {
557 a.mu.Lock()
558 defer a.mu.Unlock()
559 a.cost.LastTurn = u
560}
561
562// fireMessageAppended invokes OnMessageAppended without holding the
563// agent mutex, so the host's persistence callback can take its own
564// locks without deadlocking the agent loop. Tolerates a nil hook so
565// non-persisting callers (tests, RPC mode) don't have to set it.
566func (a *Agent) fireMessageAppended(m provider.Message) {
567 if a.OnMessageAppended != nil {
568 a.OnMessageAppended(m)
569 }
570}
571
572// Prompt sends a user message and runs the agent loop until the model
573// stops or an error occurs. Events are delivered via sink in order.
574// sink must not block the caller for long; buffer as needed.
575func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
576 if text == "" && len(images) == 0 {
577 return errors.New("prompt is empty")
578 }
579 if sink == nil {
580 sink = func(AgentEvent) {}
581 }
582 sink = a.wrapSink(sink)
583 content := []provider.Content{}
584 if text != "" {
585 content = append(content, provider.TextBlock{Text: text})
586 }
587 for _, img := range images {
588 content = append(content, img)
589 }
590 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
591
592 a.mu.Lock()
593 a.messages = append(a.messages, user)
594 a.rev++
595 a.mu.Unlock()
596 a.fireMessageAppended(user)
597 sink(EvUserMessage{Message: user})
598
599 return a.runLoop(ctx, sink, a.beginTurn())
600}
601
602// Continue runs the agent loop against the existing transcript. Used
603// after appending tool results manually or to retry.
604func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
605 if sink == nil {
606 sink = func(AgentEvent) {}
607 }
608 sink = a.wrapSink(sink)
609 return a.runLoop(ctx, sink, a.beginTurn())
610}
611
612// wrapSink composes the per-call sink with a.OnEvent (if set) so the
613// extension manager (or any other observer) sees every AgentEvent
614// without having to thread itself through every Prompt callsite.
615func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) {
616 if a.OnEvent == nil {
617 return sink
618 }
619 obs := a.OnEvent
620 return func(ev AgentEvent) {
621 obs(ev)
622 sink(ev)
623 }
624}
625
626func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent), requestContext provider.RequestContext) error {
627 for step := 1; a.MaxSteps <= 0 || step <= a.MaxSteps; step++ {
628 // Messages queued while the agent was busy are delivered
629 // before the next model call. This is the safe boundary:
630 // any previous tool batch has already completed and its
631 // results have been appended, but no new provider request has
632 // started yet.
633 if pending := a.drainQueuedMessages(); len(pending) > 0 {
634 a.appendQueuedAsUser(pending, sink)
635 }
636
637 sink(EvTurnStart{Step: step, TurnID: requestContext.TurnID})
638 turnContext := ""
639 if a.BeforeTurnContext != nil {
640 var allowed bool
641 var reason string
642 allowed, reason, turnContext = a.BeforeTurnContext(ctx, step)
643 if !allowed {
644 if reason == "" {
645 reason = "turn blocked by extension guard"
646 }
647 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
648 sink(EvDone{})
649 return nil
650 }
651 } else if a.BeforeTurn != nil {
652 if allowed, reason := a.BeforeTurn(step); !allowed {
653 if reason == "" {
654 reason = "turn blocked by extension guard"
655 }
656 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
657 sink(EvDone{})
658 return nil
659 }
660 }
661
662 a.appendDynamicContext(turnContext)
663
664 var (
665 stop provider.StopReason
666 assistantMsg provider.Message
667 err error
668 )
669 maxAttempts := max(a.MaxRetries+1, 1)
670 for attempt := 0; ; attempt++ {
671 stop, assistantMsg, err = a.oneTurn(ctx, sink, turnContext, requestContext, attempt, maxAttempts)
672 sink(EvTurnEnd{Stop: stop, Err: err})
673 if err == nil {
674 break
675 }
676 retryable := a.canRetryError(err, attempt)
677 var retryCtxErr error
678 if retryable {
679 retryCtxErr = ctx.Err()
680 if retryCtxErr != nil {
681 retryable = false
682 }
683 }
684 reason := classifyRetryReason(err)
685 if !errors.Is(err, context.Canceled) {
686 a.fireRetryLifecycle(RetryLifecycleRecord{
687 Event: RetryLifecycleRequestFailed,
688 Scope: RetryScopeAgent,
689 Attempt: attempt + 1,
690 MaxAttempts: maxAttempts,
691 Reason: reason,
692 Terminal: !retryable,
693 })
694 }
695 if retryCtxErr != nil {
696 return retryCtxErr
697 }
698 if !retryable {
699 break
700 }
701 a.dropLastAssistantMessage()
702 delay := a.retryDelay(attempt)
703 if retryErr := ctx.Err(); retryErr != nil {
704 return retryErr
705 }
706 a.fireRetryLifecycle(RetryLifecycleRecord{
707 Event: RetryLifecycleRetryScheduled,
708 Scope: RetryScopeAgent,
709 Attempt: attempt + 2,
710 MaxAttempts: maxAttempts,
711 Reason: reason,
712 DelayMS: delay.Milliseconds(),
713 })
714 sink(EvRetryScheduled{
715 Scope: RetryScopeAgent,
716 Attempt: attempt + 2,
717 MaxAttempts: maxAttempts,
718 Delay: delay,
719 })
720 if sleepErr := sleepRetry(ctx, delay); sleepErr != nil {
721 return sleepErr
722 }
723 }
724 if err != nil {
725 return err
726 }
727
728 if stop == provider.StopContinue {
729 if err := ctx.Err(); err != nil {
730 sink(EvDone{})
731 return err
732 }
733 continue
734 }
735 if stop == provider.StopToolUse {
736 // Execute each tool call, append a single tool-results message, continue.
737 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
738 a.mu.Lock()
739 a.messages = append(a.messages, toolMsg)
740 a.rev++
741 // OpenAI's chat-completions tool message shape is text-centric.
742 // Vision models reliably consume images when they arrive as user
743 // content, so when a tool result contains images we mirror them
744 // into a synthetic user message immediately after the tool result.
745 // This keeps the transcript self-contained for providers that can
746 // see image blocks in tool messages while making OpenAI vision
747 // models actually receive the image bytes.
748 //
749 // The OpenAI Responses route ("openai-codex") has the same
750 // text-centric tool-output shape: a function_call_output only
751 // carries a string, so images in a tool result never reach the
752 // model. Both providers serialize images correctly when they
753 // arrive as user content, so the mirror covers them both.
754 var imageMirror provider.Message
755 if a.Client != nil && (a.Client.Name() == "openai" || a.Client.Name() == "openai-codex") {
756 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
757 a.messages = append(a.messages, mirror)
758 a.rev++
759 imageMirror = mirror
760 }
761 }
762 a.mu.Unlock()
763 a.fireMessageAppended(toolMsg)
764 if len(imageMirror.Content) > 0 {
765 a.fireMessageAppended(imageMirror)
766 }
767 // If context was cancelled during tool execution, bail out.
768 if err := ctx.Err(); err != nil {
769 sink(EvDone{})
770 return err
771 }
772 _ = hadError
773 continue
774 }
775
776 // If the assistant stopped without tool calls but a message was
777 // queued while it was speaking, loop once more so that message
778 // is appended and answered instead of waiting until a later
779 // top-level prompt.
780 if ctx.Err() == nil && a.QueuedMessageCount() > 0 {
781 continue
782 }
783
784 // Terminal stop (end, length, error, aborted).
785 sink(EvDone{})
786 return nil
787 }
788 if a.MaxSteps > 0 {
789 sink(EvDone{})
790 return fmt.Errorf("%w: %d", ErrMaxSteps, a.MaxSteps)
791 }
792 return nil
793}
794
795func (a *Agent) canRetryError(err error, attempt int) bool {
796 if err == nil || a.MaxRetries <= 0 || attempt >= a.MaxRetries {
797 return false
798 }
799 var openErr *requestOpenError
800 if errors.As(err, &openErr) {
801 return false
802 }
803 // An idle timeout already bounds a full period without provider progress.
804 // Retrying it would restart that deadline for every attempt, allowing one
805 // silent model continuation to keep an agent running for more than thirty
806 // minutes under the default retry policy.
807 if errors.Is(err, ErrStreamIdleTimeout) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
808 return false
809 }
810 msg := strings.ToLower(err.Error())
811 if msg == "" || isNonRetryableProviderLimit(msg) {
812 return false
813 }
814 if isStreamNetworkError(err, msg) {
815 return true
816 }
817 needles := []string{
818 "overloaded", "provider returned error", "rate limit", "ratelimit", "too many requests",
819 "429", "http 429", "500", "http 500", "502", "http 502", "503", "http 503", "504", "http 504",
820 "service unavailable", "server error", "internal error", "ended without", "stream ended before",
821 "did not get a response", "timed out", "timeout", "terminated",
822 // OpenAI's ChatGPT/Codex backend returns this generic message (with a
823 // request ID) for transient server failures and explicitly says
824 // "You can retry your request".
825 "an error occurred while processing your request",
826 // Explicit retry guidance emitted by provider backends (OpenAI
827 // Responses, AWS Bedrock stream exceptions) with varying prefixes.
828 "you can retry your request", "try your request again", "please retry your request",
829 // Capacity messages from the ChatGPT/Codex backend, e.g.
830 // "Our servers are currently overloaded. Please try again later."
831 // The trailing advice also shows up on its own for transient
832 // capacity failures; usage/quota limits are filtered out above by
833 // isNonRetryableProviderLimit before this list is consulted.
834 "servers are currently overloaded", "servers are busy", "try again later",
835 }
836 for _, needle := range needles {
837 if strings.Contains(msg, needle) {
838 return true
839 }
840 }
841 return false
842}
843
844// isStreamNetworkError detects transport errors from an opened provider stream.
845// Keep typed EOF checks because WebSocket readers return io.EOF on a peer close,
846// while provider-specific transports may only preserve a textual error.
847func isStreamNetworkError(err error, msg string) bool {
848 if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
849 return true
850 }
851 for _, needle := range []string{
852 "connection", "network", "unexpected eof", "broken pipe", "socket hang up",
853 "transport failure", "upstream connect",
854 } {
855 if strings.Contains(msg, needle) {
856 return true
857 }
858 }
859 return false
860}
861
862func isNonRetryableProviderLimit(msg string) bool {
863 needles := []string{
864 "usage limit", "monthly usage limit", "freeusagelimit", "gousagelimit",
865 "available balance", "insufficient_quota", "out of budget", "quota exceeded", "billing",
866 }
867 for _, needle := range needles {
868 if strings.Contains(msg, needle) {
869 return true
870 }
871 }
872 return false
873}
874
875func classifyRetryReason(err error) RetryReason {
876 if err == nil {
877 return RetryReasonUnknown
878 }
879 msg := strings.ToLower(err.Error())
880 if isNonRetryableProviderLimit(msg) {
881 return RetryReasonQuota
882 }
883 if strings.Contains(msg, "overload") || strings.Contains(msg, "servers are busy") {
884 return RetryReasonOverload
885 }
886 if strings.Contains(msg, "rate limit") || strings.Contains(msg, "ratelimit") || strings.Contains(msg, "too many requests") || strings.Contains(msg, "429") {
887 return RetryReasonRateLimit
888 }
889 if errors.Is(err, context.DeadlineExceeded) || strings.Contains(msg, "timeout") || strings.Contains(msg, "timed out") {
890 return RetryReasonTimeout
891 }
892 if isStreamNetworkError(err, msg) {
893 return RetryReasonNetwork
894 }
895 for _, needle := range []string{"context window", "context length", "maximum context", "too many tokens"} {
896 if strings.Contains(msg, needle) {
897 return RetryReasonContextWindow
898 }
899 }
900 for _, needle := range []string{"unauthorized", "authentication", "invalid api key", "invalid_api_key", "http 401", "http 403"} {
901 if strings.Contains(msg, needle) {
902 return RetryReasonAuth
903 }
904 }
905 for _, needle := range []string{"server error", "internal error", "service unavailable", "http 500", "http 502", "http 503", "http 504", "provider returned error"} {
906 if strings.Contains(msg, needle) {
907 return RetryReasonServer
908 }
909 }
910 for _, needle := range []string{"bad request", "unsupported", "invalid parameter", "http 400", "http 404", "http 422"} {
911 if strings.Contains(msg, needle) {
912 return RetryReasonClient
913 }
914 }
915 return RetryReasonUnknown
916}
917
918func (a *Agent) fireRetryLifecycle(record RetryLifecycleRecord) {
919 if a.OnRetryLifecycle != nil {
920 a.OnRetryLifecycle(record)
921 }
922}
923
924func (a *Agent) retryDelay(attempt int) time.Duration {
925 base := a.RetryBaseDelay
926 if base <= 0 {
927 base = 2 * time.Second
928 }
929 return base * time.Duration(1<<attempt)
930}
931
932func sleepRetry(ctx context.Context, d time.Duration) error {
933 if d <= 0 {
934 return ctx.Err()
935 }
936 t := time.NewTimer(d)
937 defer t.Stop()
938 select {
939 case <-ctx.Done():
940 return ctx.Err()
941 case <-t.C:
942 return nil
943 }
944}
945
946func (a *Agent) dropLastAssistantMessage() {
947 a.mu.Lock()
948 defer a.mu.Unlock()
949 if n := len(a.messages); n > 0 && a.messages[n-1].Role == provider.RoleAssistant {
950 a.messages = a.messages[:n-1]
951 a.rev++
952 }
953}
954
955// requestLifecycleSink adapts provider-owned pre-stream retries into core
956// events without letting provider concerns leak into mode renderers.
957type requestLifecycleSink struct {
958 sink func(AgentEvent)
959 observe func(RetryLifecycleRecord)
960 provider string
961 model string
962 lastReason RetryReason
963}
964
965func (s *requestLifecycleSink) CacheDiagnostics(diagnostics provider.CacheDiagnostics) {
966 s.sink(EvCacheDiagnostics{
967 Eligible: diagnostics.Eligible,
968 Mode: diagnostics.Mode,
969 Transport: diagnostics.Transport,
970 Continuation: diagnostics.Continuation,
971 })
972}
973
974func (s *requestLifecycleSink) RequestAttempt(attempt, maxAttempts int) {
975 s.sink(EvRequestStarted{
976 Provider: s.provider,
977 Model: s.model,
978 Scope: RetryScopeProvider,
979 Attempt: attempt,
980 MaxAttempts: maxAttempts,
981 })
982}
983
984func (s *requestLifecycleSink) RequestFailed(attempt, maxAttempts int, reason provider.RequestFailureReason, terminal bool) {
985 s.lastReason = retryReasonFromProvider(reason)
986 if s.observe != nil {
987 s.observe(RetryLifecycleRecord{
988 Event: RetryLifecycleRequestFailed,
989 Scope: RetryScopeProvider,
990 Attempt: attempt,
991 MaxAttempts: maxAttempts,
992 Reason: s.lastReason,
993 Terminal: terminal,
994 })
995 }
996}
997
998func (s *requestLifecycleSink) RetryScheduled(attempt, maxAttempts int, delay time.Duration) {
999 s.sink(EvRetryScheduled{
1000 Scope: RetryScopeProvider,
1001 Attempt: attempt,
1002 MaxAttempts: maxAttempts,
1003 Delay: delay,
1004 })
1005 if s.observe != nil {
1006 reason := s.lastReason
1007 if reason == "" {
1008 reason = RetryReasonUnknown
1009 }
1010 s.observe(RetryLifecycleRecord{
1011 Event: RetryLifecycleRetryScheduled,
1012 Scope: RetryScopeProvider,
1013 Attempt: attempt,
1014 MaxAttempts: maxAttempts,
1015 Reason: reason,
1016 DelayMS: delay.Milliseconds(),
1017 })
1018 }
1019 s.lastReason = ""
1020}
1021
1022func retryReasonFromProvider(reason provider.RequestFailureReason) RetryReason {
1023 switch reason {
1024 case provider.RequestFailureOverload:
1025 return RetryReasonOverload
1026 case provider.RequestFailureRateLimit:
1027 return RetryReasonRateLimit
1028 case provider.RequestFailureQuota:
1029 return RetryReasonQuota
1030 case provider.RequestFailureServer:
1031 return RetryReasonServer
1032 case provider.RequestFailureNetwork:
1033 return RetryReasonNetwork
1034 case provider.RequestFailureTimeout:
1035 return RetryReasonTimeout
1036 case provider.RequestFailureClient:
1037 return RetryReasonClient
1038 default:
1039 return RetryReasonUnknown
1040 }
1041}
1042
1043// oneTurn calls the LLM once, forwards events, returns the stop reason
1044// and the assembled assistant message (already appended to the transcript).
1045func newStreamIdleTimer(timeout time.Duration) (<-chan time.Time, func(), func()) {
1046 timer := time.NewTimer(timeout)
1047 stop := func() { timer.Stop() }
1048 reset := func() {
1049 if !timer.Stop() {
1050 select {
1051 case <-timer.C:
1052 default:
1053 }
1054 }
1055 timer.Reset(timeout)
1056 }
1057 return timer.C, stop, reset
1058}
1059
1060func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent), turnContext string, requestContext provider.RequestContext, attempt, maxAttempts int) (provider.StopReason, provider.Message, error) {
1061 fastMode := a.FastModeEnabled()
1062 if err := provider.ValidateFastMode(a.Client.Name(), fastMode); err != nil {
1063 return provider.StopError, provider.Message{}, err
1064 }
1065 system, tools := a.PromptConfig()
1066 // Repair pairs before projecting the copied provider-input view. The
1067 // repair can add stub results for aborted calls; those results must also
1068 // remain in the request so every tool call still has a matching result.
1069 messages := repairToolUseResultPairs(a.Messages())
1070 requestLifecycle := &requestLifecycleSink{
1071 sink: sink,
1072 observe: a.fireRetryLifecycle,
1073 provider: a.Client.Name(),
1074 model: a.Model,
1075 }
1076 req := provider.Request{
1077 Model: a.Model,
1078 System: system,
1079 // Repair any dangling tool_use blocks before sending. A turn
1080 // aborted mid-flight (cancel, connection drop, ECONNREFUSED to a
1081 // dev server, etc.) can leave an assistant tool_use with no
1082 // matching tool_result in the live transcript. The load-time
1083 // repair in OpenSession only runs on restart, so without this the
1084 // next in-process request is rejected by providers like Anthropic
1085 // with "tool_use ids were found without tool_result blocks". The
1086 // repair is pure and a no-op on already-valid transcripts.
1087 Messages: projectProviderMessages(messages),
1088 Tools: tools.Specs(),
1089 Reasoning: a.Reasoning,
1090 FastMode: fastMode,
1091 MaxTokens: a.MaxTokens,
1092 Temperature: a.Temperature,
1093 Lifecycle: requestLifecycle,
1094 Context: requestContext,
1095 }
1096 sink(EvRequestStarted{
1097 Provider: a.Client.Name(),
1098 Model: a.Model,
1099 Scope: RetryScopeAgent,
1100 Attempt: attempt + 1,
1101 MaxAttempts: maxAttempts,
1102 })
1103 streamCtx, cancelStream := context.WithCancel(ctx)
1104 defer cancelStream()
1105 a.beginProviderRequest()
1106 stream, err := a.Client.Stream(streamCtx, req)
1107 if err != nil {
1108 return provider.StopError, provider.Message{}, &requestOpenError{err: err}
1109 }
1110 // A provider may already have queued events when cancellation wins. Keep a
1111 // consumer attached after this turn returns so bounded producer channels
1112 // cannot strand the provider goroutine and its response body.
1113 defer func() { go drainProviderStream(stream) }()
1114
1115 sink(EvAssistantStart{})
1116
1117 var (
1118 stop provider.StopReason
1119 finalErr error
1120 finalMsg provider.Message
1121 )
1122
1123 var idle <-chan time.Time
1124 var resetIdle func()
1125 if a.StreamIdleTimeout > 0 {
1126 newTimer := a.streamIdleTimer
1127 if newTimer == nil {
1128 newTimer = newStreamIdleTimer
1129 }
1130 var stopIdle func()
1131 idle, stopIdle, resetIdle = newTimer(a.StreamIdleTimeout)
1132 defer stopIdle()
1133 }
1134 handleStreamEvent := func(ev provider.Event) (done bool) {
1135 if resetIdle != nil {
1136 resetIdle()
1137 }
1138 switch e := ev.(type) {
1139 case provider.EventStart:
1140 // nothing
1141 case provider.EventTextDelta:
1142 sink(EvTextDelta{Delta: e.Delta})
1143 case provider.EventToolStart:
1144 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
1145 case provider.EventToolArgs:
1146 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
1147 case provider.EventToolEnd:
1148 sink(EvToolUseEnd{ID: e.ID})
1149 case provider.EventUsage:
1150 cum := a.addUsage(e.Usage)
1151 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
1152 if a.OnUsage != nil {
1153 a.OnUsage(cum)
1154 }
1155 case provider.EventDone:
1156 stop = e.Stop
1157 finalErr = e.Err
1158 finalMsg = provider.WithMessageOrigin(e.Message, a.Client.Name(), req.Model)
1159 return true
1160 }
1161 return false
1162 }
1163 for {
1164 // Consume no more than one buffered event before honoring cancellation.
1165 select {
1166 case ev, ok := <-stream:
1167 if !ok || handleStreamEvent(ev) {
1168 goto streamDone
1169 }
1170 default:
1171 }
1172 if err := streamCtx.Err(); err != nil {
1173 return provider.StopError, finalMsg, err
1174 }
1175 select {
1176 case <-streamCtx.Done():
1177 return provider.StopError, finalMsg, streamCtx.Err()
1178 case <-idle:
1179 return provider.StopError, finalMsg, ErrStreamIdleTimeout
1180 case ev, ok := <-stream:
1181 if !ok || handleStreamEvent(ev) {
1182 goto streamDone
1183 }
1184 }
1185 }
1186streamDone:
1187
1188 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
1189 // produce partial content. When the partial message is text only we
1190 // keep whatever was streamed up to the cancel so the user does not
1191 // lose visible work (a cut-off summary is still useful). If the
1192 // partial message already contained tool-call blocks we drop the
1193 // whole thing, because an unmatched tool_use would fail the next
1194 // turn with a tool_result mismatch error.
1195 keep := len(finalMsg.Content) > 0
1196 if stop == provider.StopAborted && keep {
1197 hasToolCall := false
1198 for _, c := range finalMsg.Content {
1199 if _, ok := c.(provider.ToolCallBlock); ok {
1200 hasToolCall = true
1201 break
1202 }
1203 }
1204 if hasToolCall {
1205 keep = false
1206 }
1207 }
1208 if keep {
1209 emit := finalMsg
1210 suppress := false
1211
1212 // BeforeAssistantMessage hook: extensions can suppress or
1213 // rewrite the visible text. The transcript keeps the
1214 // model's original output so the model still sees what it
1215 // said on subsequent turns.
1216 if a.BeforeAssistantMessage != nil {
1217 orig := extractText(finalMsg)
1218 if orig != "" {
1219 allowed, _, replacement := a.BeforeAssistantMessage(orig)
1220 if !allowed {
1221 suppress = true
1222 } else if replacement != "" && replacement != orig {
1223 emit = replaceText(finalMsg, replacement)
1224 }
1225 }
1226 }
1227
1228 a.mu.Lock()
1229 a.messages = append(a.messages, finalMsg)
1230 a.rev++
1231 a.mu.Unlock()
1232 a.fireMessageAppended(finalMsg)
1233 if !suppress {
1234 sink(EvAssistantMessage{Message: emit})
1235 }
1236 // Now surface tool calls as EvToolCall events so UIs can render them
1237 // in order before the tool results arrive.
1238 for _, c := range finalMsg.Content {
1239 if tc, ok := c.(provider.ToolCallBlock); ok {
1240 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
1241 }
1242 }
1243 }
1244
1245 return stop, finalMsg, finalErr
1246}
1247
1248func drainProviderStream(stream <-chan provider.Event) {
1249 for range stream {
1250 }
1251}
1252
1253// ToolResultCommitError supplies a user-safe explanation when a host rejects
1254// a tool result during durable-state commit. Internal storage errors remain in
1255// Err; callers without this type retain the generic failure message.
1256type ToolResultCommitError struct {
1257 Message string
1258 Err error
1259}
1260
1261func (e *ToolResultCommitError) Error() string {
1262 if e == nil || e.Err == nil {
1263 return "tool result commit failed"
1264 }
1265 return e.Err.Error()
1266}
1267
1268func (e *ToolResultCommitError) Unwrap() error {
1269 if e == nil {
1270 return nil
1271 }
1272 return e.Err
1273}
1274
1275func toolResultCommitMessage(err error) string {
1276 var safe *ToolResultCommitError
1277 if errors.As(err, &safe) && strings.TrimSpace(safe.Message) != "" {
1278 return safe.Message
1279 }
1280 return "tool result state could not be persisted"
1281}
1282
1283// executeTools runs every tool call in the assistant message and returns
1284// a single tool-role message carrying all results.
1285func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
1286 var results []provider.Content
1287 var addedTools []string
1288 hadError := false
1289 tools := a.ToolsSnapshot()
1290
1291 for _, c := range msg.Content {
1292 tc, ok := c.(provider.ToolCallBlock)
1293 if !ok {
1294 continue
1295 }
1296 res := a.runOneTool(ctx, tc, tools, sink)
1297 if a.CommitToolResult != nil {
1298 if err := a.CommitToolResult(tc.ID, res); err != nil {
1299 res = ToolResult{
1300 Content: []provider.Content{provider.TextBlock{Text: toolResultCommitMessage(err)}},
1301 IsError: true,
1302 Timing: res.Timing,
1303 }
1304 }
1305 }
1306 if res.IsError {
1307 hadError = true
1308 }
1309 results = append(results, provider.ToolResultBlock{
1310 CallID: tc.ID,
1311 Content: res.Content,
1312 IsError: res.IsError,
1313 Timing: res.Timing,
1314 })
1315 for _, name := range res.ActivateTools {
1316 if _, err := tools.Get(name); err == nil && !containsString(addedTools, name) {
1317 addedTools = append(addedTools, name)
1318 }
1319 }
1320 if a.OnToolResult != nil {
1321 a.OnToolResult(tc.ID, res)
1322 }
1323 if update, ok := res.Details.(PlanUpdate); ok && !res.IsError {
1324 sink(EvPlanUpdate{Update: update})
1325 }
1326 sink(EvToolResult{ID: tc.ID, Result: res})
1327 }
1328
1329 return provider.Message{
1330 Role: provider.RoleTool,
1331 Content: results,
1332 Time: time.Now(),
1333 AddedToolNames: addedTools,
1334 }, hadError
1335}
1336
1337func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, tools Registry, sink func(AgentEvent)) (res ToolResult) {
1338 started := time.Now()
1339 defer func() {
1340 if recover() != nil {
1341 res = ToolResult{
1342 Content: []provider.Content{provider.TextBlock{Text: "tool execution failed"}},
1343 IsError: true,
1344 }
1345 }
1346 res.Timing = &provider.ToolTiming{
1347 StartedAt: started,
1348 CompletedAt: time.Now(),
1349 Duration: time.Since(started),
1350 }
1351 }()
1352
1353 tool, err := tools.Get(tc.Name)
1354 if err != nil {
1355 return ToolResult{
1356 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
1357 IsError: true,
1358 }
1359 }
1360
1361 args := tc.Arguments
1362
1363 // Intercept hook: an extension or other guard can refuse the
1364 // call before any side effect happens, OR rewrite the args
1365 // seen by the tool. The model sees the reason as the tool
1366 // error, learns from it, and (typically) proposes a different
1367 // action; rewrites are invisible to the model (they apply only
1368 // to the execution).
1369 if a.BeforeToolExecute != nil {
1370 allowed, reason, modified := a.BeforeToolExecute(tc)
1371 if !allowed {
1372 if reason == "" {
1373 reason = "tool call refused by extension guard"
1374 }
1375 return ToolResult{
1376 Content: []provider.Content{provider.TextBlock{Text: reason}},
1377 IsError: true,
1378 }
1379 }
1380 if len(modified) > 0 && json.Valid(modified) {
1381 if policy, ok := tool.(ToolArgumentRewritePolicy); ok && !policy.AllowArgumentRewrite() {
1382 return ToolResult{
1383 Content: []provider.Content{provider.TextBlock{Text: "tool arguments cannot be rewritten"}},
1384 IsError: true,
1385 }
1386 }
1387 args = modified
1388 }
1389 }
1390
1391 if len(args) == 0 {
1392 args = json.RawMessage("{}")
1393 }
1394
1395 // Tool panics are recovered by the invocation-level defer above so
1396 // guard, lookup, progress, and Execute panics all receive timing.
1397 sink(EvToolExecutionStarted{ID: tc.ID, Name: tc.Name})
1398 out, err := tool.Execute(ctx, args, func(text string) {
1399 sink(EvToolProgress{ID: tc.ID, Text: text})
1400 })
1401 if err != nil {
1402 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
1403 return ToolResult{
1404 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
1405 IsError: true,
1406 }
1407 }
1408 return ToolResult{
1409 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
1410 IsError: true,
1411 }
1412 }
1413 return out
1414}
1415
1416// extractText concatenates all TextBlock content in a message. Used
1417// by BeforeAssistantMessage so guards see a single string instead of
1418// having to walk provider.Content themselves.
1419func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
1420 var content []provider.Content
1421 for _, c := range msg.Content {
1422 tr, ok := c.(provider.ToolResultBlock)
1423 if !ok {
1424 continue
1425 }
1426 for _, inner := range tr.Content {
1427 switch v := inner.(type) {
1428 case provider.TextBlock:
1429 // Keep short textual context so the model understands why
1430 // the images appeared, but don't duplicate giant read
1431 // outputs verbatim.
1432 if len(v.Text) > 0 && len(v.Text) <= 500 {
1433 content = append(content, v)
1434 }
1435 case provider.ImageBlock:
1436 content = append(content, v)
1437 }
1438 }
1439 }
1440 if len(content) == 0 {
1441 return provider.Message{}
1442 }
1443 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
1444 content = append([]provider.Content{prefix}, content...)
1445 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
1446}
1447
1448const (
1449 maxTurnContextBytes = 16 * 1024
1450 turnContextTruncatedMarker = "\n[extension context truncated]"
1451)
1452
1453func boundedTurnContext(text string) string {
1454 text = strings.TrimSpace(text)
1455 if len(text) <= maxTurnContextBytes {
1456 return text
1457 }
1458 limit := maxTurnContextBytes - len(turnContextTruncatedMarker)
1459 cut := limit
1460 for cut > 0 && !utf8.RuneStart(text[cut]) {
1461 cut--
1462 }
1463 return strings.TrimSpace(text[:cut]) + turnContextTruncatedMarker
1464}
1465
1466func containsString(values []string, want string) bool {
1467 for _, value := range values {
1468 if value == want {
1469 return true
1470 }
1471 }
1472 return false
1473}
1474
1475func extractText(msg provider.Message) string {
1476 var out string
1477 for _, c := range msg.Content {
1478 if tb, ok := c.(provider.TextBlock); ok {
1479 if out != "" {
1480 out += "\n"
1481 }
1482 out += tb.Text
1483 }
1484 }
1485 return out
1486}
1487
1488// replaceText returns a copy of msg with every TextBlock replaced by
1489// a single TextBlock containing replacement. Non-text content (tool
1490// calls, etc.) is preserved in order.
1491func replaceText(msg provider.Message, replacement string) provider.Message {
1492 out := provider.Message{Role: msg.Role}
1493 out.Content = make([]provider.Content, 0, len(msg.Content))
1494 replaced := false
1495 for _, c := range msg.Content {
1496 if _, ok := c.(provider.TextBlock); ok {
1497 if !replaced {
1498 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
1499 replaced = true
1500 }
1501 continue
1502 }
1503 out.Content = append(out.Content, c)
1504 }
1505 if !replaced {
1506 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
1507 }
1508 return out
1509}