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