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