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