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