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 for attempt := 0; ; attempt++ {
478 stop, assistantMsg, err = a.oneTurn(ctx, sink, turnContext)
479 sink(EvTurnEnd{Stop: stop, Err: err})
480 if err == nil || !a.canRetryError(err, attempt) {
481 break
482 }
483 a.dropLastAssistantMessage()
484 if sleepErr := sleepRetry(ctx, a.retryDelay(attempt)); sleepErr != nil {
485 return sleepErr
486 }
487 }
488 if err != nil {
489 return err
490 }
491
492 if stop == provider.StopToolUse {
493 // Execute each tool call, append a single tool-results message, continue.
494 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
495 a.mu.Lock()
496 a.messages = append(a.messages, toolMsg)
497 a.rev++
498 // OpenAI's chat-completions tool message shape is text-centric.
499 // Vision models reliably consume images when they arrive as user
500 // content, so when a tool result contains images we mirror them
501 // into a synthetic user message immediately after the tool result.
502 // This keeps the transcript self-contained for providers that can
503 // see image blocks in tool messages while making OpenAI vision
504 // models actually receive the image bytes.
505 //
506 // The OpenAI Responses route ("openai-codex") has the same
507 // text-centric tool-output shape: a function_call_output only
508 // carries a string, so images in a tool result never reach the
509 // model. Both providers serialize images correctly when they
510 // arrive as user content, so the mirror covers them both.
511 var imageMirror provider.Message
512 if a.Client != nil && (a.Client.Name() == "openai" || a.Client.Name() == "openai-codex") {
513 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
514 a.messages = append(a.messages, mirror)
515 a.rev++
516 imageMirror = mirror
517 }
518 }
519 a.mu.Unlock()
520 a.fireMessageAppended(toolMsg)
521 if len(imageMirror.Content) > 0 {
522 a.fireMessageAppended(imageMirror)
523 }
524 // If context was cancelled during tool execution, bail out.
525 if err := ctx.Err(); err != nil {
526 sink(EvDone{})
527 return err
528 }
529 _ = hadError
530 continue
531 }
532
533 // If the assistant stopped without tool calls but a message was
534 // queued while it was speaking, loop once more so that message
535 // is appended and answered instead of waiting until a later
536 // top-level prompt.
537 if ctx.Err() == nil && a.QueuedMessageCount() > 0 {
538 continue
539 }
540
541 // Terminal stop (end, length, error, aborted).
542 sink(EvDone{})
543 return nil
544 }
545 if a.MaxSteps > 0 {
546 sink(EvDone{})
547 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
548 }
549 return nil
550}
551
552func (a *Agent) canRetryError(err error, attempt int) bool {
553 if err == nil || a.MaxRetries <= 0 || attempt >= a.MaxRetries {
554 return false
555 }
556 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
557 return false
558 }
559 msg := strings.ToLower(err.Error())
560 if msg == "" || isNonRetryableProviderLimit(msg) {
561 return false
562 }
563 needles := []string{
564 "overloaded", "provider returned error", "rate limit", "ratelimit", "too many requests",
565 "429", "http 429", "500", "http 500", "502", "http 502", "503", "http 503", "504", "http 504",
566 "service unavailable", "server error", "internal error", "network error", "connection error",
567 "connection refused", "connection lost", "fetch failed", "upstream connect", "reset before headers",
568 "socket hang up", "ended without", "stream ended before", "did not get a response", "timed out",
569 "timeout", "terminated", "unexpected eof", "transport failure",
570 // OpenAI's ChatGPT/Codex backend returns this generic message (with a
571 // request ID) for transient server failures and explicitly says
572 // "You can retry your request".
573 "an error occurred while processing your request",
574 // Explicit retry guidance emitted by provider backends (OpenAI
575 // Responses, AWS Bedrock stream exceptions) with varying prefixes.
576 "you can retry your request", "try your request again", "please retry your request",
577 // Capacity messages from the ChatGPT/Codex backend, e.g.
578 // "Our servers are currently overloaded. Please try again later."
579 // The trailing advice also shows up on its own for transient
580 // capacity failures; usage/quota limits are filtered out above by
581 // isNonRetryableProviderLimit before this list is consulted.
582 "servers are currently overloaded", "servers are busy", "try again later",
583 }
584 for _, needle := range needles {
585 if strings.Contains(msg, needle) {
586 return true
587 }
588 }
589 return false
590}
591
592func isNonRetryableProviderLimit(msg string) bool {
593 needles := []string{
594 "usage limit", "monthly usage limit", "freeusagelimit", "gousagelimit",
595 "available balance", "insufficient_quota", "out of budget", "quota exceeded", "billing",
596 }
597 for _, needle := range needles {
598 if strings.Contains(msg, needle) {
599 return true
600 }
601 }
602 return false
603}
604
605func (a *Agent) retryDelay(attempt int) time.Duration {
606 base := a.RetryBaseDelay
607 if base <= 0 {
608 base = 2 * time.Second
609 }
610 return base * time.Duration(1<<attempt)
611}
612
613func sleepRetry(ctx context.Context, d time.Duration) error {
614 if d <= 0 {
615 return ctx.Err()
616 }
617 t := time.NewTimer(d)
618 defer t.Stop()
619 select {
620 case <-ctx.Done():
621 return ctx.Err()
622 case <-t.C:
623 return nil
624 }
625}
626
627func (a *Agent) dropLastAssistantMessage() {
628 a.mu.Lock()
629 defer a.mu.Unlock()
630 if n := len(a.messages); n > 0 && a.messages[n-1].Role == provider.RoleAssistant {
631 a.messages = a.messages[:n-1]
632 a.rev++
633 }
634}
635
636// oneTurn calls the LLM once, forwards events, returns the stop reason
637// and the assembled assistant message (already appended to the transcript).
638func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent), turnContext string) (provider.StopReason, provider.Message, error) {
639 fastMode := a.FastModeEnabled()
640 if err := provider.ValidateFastMode(a.Client.Name(), fastMode); err != nil {
641 return provider.StopError, provider.Message{}, err
642 }
643 system, tools := a.PromptConfig()
644 if contextText := boundedTurnContext(turnContext); contextText != "" {
645 if system != "" {
646 system += "\n\n"
647 }
648 system += "[Extension context for this turn]\n" + contextText
649 }
650 req := provider.Request{
651 Model: a.Model,
652 System: system,
653 // Repair any dangling tool_use blocks before sending. A turn
654 // aborted mid-flight (cancel, connection drop, ECONNREFUSED to a
655 // dev server, etc.) can leave an assistant tool_use with no
656 // matching tool_result in the live transcript. The load-time
657 // repair in OpenSession only runs on restart, so without this the
658 // next in-process request is rejected by providers like Anthropic
659 // with "tool_use ids were found without tool_result blocks". The
660 // repair is pure and a no-op on already-valid transcripts.
661 Messages: repairToolUseResultPairs(a.Messages()),
662 Tools: tools.Specs(),
663 Reasoning: a.Reasoning,
664 FastMode: fastMode,
665 MaxTokens: a.MaxTokens,
666 Temperature: a.Temperature,
667 }
668 stream, err := a.Client.Stream(ctx, req)
669 if err != nil {
670 return provider.StopError, provider.Message{}, err
671 }
672
673 sink(EvAssistantStart{})
674
675 var (
676 stop provider.StopReason
677 finalErr error
678 finalMsg provider.Message
679 )
680
681 for ev := range stream {
682 switch e := ev.(type) {
683 case provider.EventStart:
684 // nothing
685 case provider.EventTextDelta:
686 sink(EvTextDelta{Delta: e.Delta})
687 case provider.EventToolStart:
688 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
689 case provider.EventToolArgs:
690 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
691 case provider.EventToolEnd:
692 sink(EvToolUseEnd{ID: e.ID})
693 case provider.EventUsage:
694 cum := a.cost.Add(e.Usage)
695 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
696 if a.OnUsage != nil {
697 a.OnUsage(cum)
698 }
699 case provider.EventDone:
700 stop = e.Stop
701 finalErr = e.Err
702 finalMsg = e.Message
703 }
704 }
705
706 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
707 // produce partial content. When the partial message is text only we
708 // keep whatever was streamed up to the cancel so the user does not
709 // lose visible work (a cut-off summary is still useful). If the
710 // partial message already contained tool-call blocks we drop the
711 // whole thing, because an unmatched tool_use would fail the next
712 // turn with a tool_result mismatch error.
713 keep := len(finalMsg.Content) > 0
714 if stop == provider.StopAborted && keep {
715 hasToolCall := false
716 for _, c := range finalMsg.Content {
717 if _, ok := c.(provider.ToolCallBlock); ok {
718 hasToolCall = true
719 break
720 }
721 }
722 if hasToolCall {
723 keep = false
724 }
725 }
726 if keep {
727 emit := finalMsg
728 suppress := false
729
730 // BeforeAssistantMessage hook: extensions can suppress or
731 // rewrite the visible text. The transcript keeps the
732 // model's original output so the model still sees what it
733 // said on subsequent turns.
734 if a.BeforeAssistantMessage != nil {
735 orig := extractText(finalMsg)
736 if orig != "" {
737 allowed, _, replacement := a.BeforeAssistantMessage(orig)
738 if !allowed {
739 suppress = true
740 } else if replacement != "" && replacement != orig {
741 emit = replaceText(finalMsg, replacement)
742 }
743 }
744 }
745
746 a.mu.Lock()
747 a.messages = append(a.messages, finalMsg)
748 a.rev++
749 a.mu.Unlock()
750 a.fireMessageAppended(finalMsg)
751 if !suppress {
752 sink(EvAssistantMessage{Message: emit})
753 }
754 // Now surface tool calls as EvToolCall events so UIs can render them
755 // in order before the tool results arrive.
756 for _, c := range finalMsg.Content {
757 if tc, ok := c.(provider.ToolCallBlock); ok {
758 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
759 }
760 }
761 }
762
763 return stop, finalMsg, finalErr
764}
765
766// executeTools runs every tool call in the assistant message and returns
767// a single tool-role message carrying all results.
768func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
769 var results []provider.Content
770 var addedTools []string
771 hadError := false
772 tools := a.ToolsSnapshot()
773
774 for _, c := range msg.Content {
775 tc, ok := c.(provider.ToolCallBlock)
776 if !ok {
777 continue
778 }
779 res := a.runOneTool(ctx, tc, tools, sink)
780 if res.IsError {
781 hadError = true
782 }
783 results = append(results, provider.ToolResultBlock{
784 CallID: tc.ID,
785 Content: res.Content,
786 IsError: res.IsError,
787 })
788 for _, name := range res.ActivateTools {
789 if _, err := tools.Get(name); err == nil && !containsString(addedTools, name) {
790 addedTools = append(addedTools, name)
791 }
792 }
793 if a.OnToolResult != nil {
794 a.OnToolResult(tc.ID, res)
795 }
796 sink(EvToolResult{ID: tc.ID, Result: res})
797 }
798
799 return provider.Message{
800 Role: provider.RoleTool,
801 Content: results,
802 Time: time.Now(),
803 AddedToolNames: addedTools,
804 }, hadError
805}
806
807func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, tools Registry, sink func(AgentEvent)) ToolResult {
808 tool, err := tools.Get(tc.Name)
809 if err != nil {
810 return ToolResult{
811 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
812 IsError: true,
813 }
814 }
815
816 args := tc.Arguments
817
818 // Intercept hook: an extension or other guard can refuse the
819 // call before any side effect happens, OR rewrite the args
820 // seen by the tool. The model sees the reason as the tool
821 // error, learns from it, and (typically) proposes a different
822 // action; rewrites are invisible to the model (they apply only
823 // to the execution).
824 if a.BeforeToolExecute != nil {
825 allowed, reason, modified := a.BeforeToolExecute(tc)
826 if !allowed {
827 if reason == "" {
828 reason = "tool call refused by extension guard"
829 }
830 return ToolResult{
831 Content: []provider.Content{provider.TextBlock{Text: reason}},
832 IsError: true,
833 }
834 }
835 if len(modified) > 0 && json.Valid(modified) {
836 args = modified
837 }
838 }
839
840 if len(args) == 0 {
841 args = json.RawMessage("{}")
842 }
843
844 // Recover panics so a buggy tool does not crash the agent.
845 var res ToolResult
846 func() {
847 defer func() {
848 if r := recover(); r != nil {
849 res = ToolResult{
850 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
851 IsError: true,
852 }
853 }
854 }()
855 out, err := tool.Execute(ctx, args, func(text string) {
856 sink(EvToolProgress{ID: tc.ID, Text: text})
857 })
858 if err != nil {
859 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
860 res = ToolResult{
861 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
862 IsError: true,
863 }
864 return
865 }
866 res = ToolResult{
867 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
868 IsError: true,
869 }
870 return
871 }
872 res = out
873 }()
874 return res
875}
876
877// extractText concatenates all TextBlock content in a message. Used
878// by BeforeAssistantMessage so guards see a single string instead of
879// having to walk provider.Content themselves.
880func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
881 var content []provider.Content
882 for _, c := range msg.Content {
883 tr, ok := c.(provider.ToolResultBlock)
884 if !ok {
885 continue
886 }
887 for _, inner := range tr.Content {
888 switch v := inner.(type) {
889 case provider.TextBlock:
890 // Keep short textual context so the model understands why
891 // the images appeared, but don't duplicate giant read
892 // outputs verbatim.
893 if len(v.Text) > 0 && len(v.Text) <= 500 {
894 content = append(content, v)
895 }
896 case provider.ImageBlock:
897 content = append(content, v)
898 }
899 }
900 }
901 if len(content) == 0 {
902 return provider.Message{}
903 }
904 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
905 content = append([]provider.Content{prefix}, content...)
906 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
907}
908
909const (
910 maxTurnContextBytes = 16 * 1024
911 turnContextTruncatedMarker = "\n[extension context truncated]"
912)
913
914func boundedTurnContext(text string) string {
915 text = strings.TrimSpace(text)
916 if len(text) <= maxTurnContextBytes {
917 return text
918 }
919 limit := maxTurnContextBytes - len(turnContextTruncatedMarker)
920 cut := limit
921 for cut > 0 && !utf8.RuneStart(text[cut]) {
922 cut--
923 }
924 return strings.TrimSpace(text[:cut]) + turnContextTruncatedMarker
925}
926
927func containsString(values []string, want string) bool {
928 for _, value := range values {
929 if value == want {
930 return true
931 }
932 }
933 return false
934}
935
936func extractText(msg provider.Message) string {
937 var out string
938 for _, c := range msg.Content {
939 if tb, ok := c.(provider.TextBlock); ok {
940 if out != "" {
941 out += "\n"
942 }
943 out += tb.Text
944 }
945 }
946 return out
947}
948
949// replaceText returns a copy of msg with every TextBlock replaced by
950// a single TextBlock containing replacement. Non-text content (tool
951// calls, etc.) is preserved in order.
952func replaceText(msg provider.Message, replacement string) provider.Message {
953 out := provider.Message{Role: msg.Role}
954 out.Content = make([]provider.Content, 0, len(msg.Content))
955 replaced := false
956 for _, c := range msg.Content {
957 if _, ok := c.(provider.TextBlock); ok {
958 if !replaced {
959 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
960 replaced = true
961 }
962 continue
963 }
964 out.Content = append(out.Content, c)
965 }
966 if !replaced {
967 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
968 }
969 return out
970}