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