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