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// CallTool executes a host-initiated tool without adding a model tool-use
732// message to the transcript. The caller supplies a unique ID and event sink.
733// It shares the model tool path's guards, confirmation hook and error handling.
734func (a *Agent) CallTool(ctx context.Context, id, name string, args json.RawMessage, sink func(AgentEvent)) ToolResult {
735 if sink == nil {
736 sink = func(AgentEvent) {}
737 }
738 sink(EvToolCall{ID: id, Name: name, Args: args})
739 return a.runOneTool(ctx, provider.ToolCallBlock{ID: id, Name: name, Arguments: args}, sink)
740}
741
742// executeTools runs every tool call in the assistant message and returns
743// a single tool-role message carrying all results.
744func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
745 var results []provider.Content
746 var addedTools []string
747 hadError := false
748
749 for _, c := range msg.Content {
750 tc, ok := c.(provider.ToolCallBlock)
751 if !ok || tc.Server {
752 continue
753 }
754 res := a.runOneTool(ctx, tc, sink)
755 if res.IsError {
756 hadError = true
757 }
758 results = append(results, provider.ToolResultBlock{
759 CallID: tc.ID,
760 Content: res.Content,
761 IsError: res.IsError,
762 })
763 for _, name := range res.ActivateTools {
764 if _, err := a.Tools.Get(name); err == nil && !containsString(addedTools, name) {
765 addedTools = append(addedTools, name)
766 }
767 }
768 }
769
770 return provider.Message{
771 Role: provider.RoleTool,
772 Content: results,
773 Time: time.Now(),
774 AddedToolNames: addedTools,
775 }, hadError
776}
777
778func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) (result ToolResult) {
779 args := tc.Arguments
780 status := ""
781 executed := false
782 defer func() {
783 if status == "" {
784 status = "completed"
785 if result.IsError {
786 status = "failed"
787 switch result.Status {
788 case "cancelled", "timed_out", "blocked":
789 status = result.Status
790 }
791 }
792 }
793 if sink != nil {
794 sink(EvToolResult{ID: tc.ID, Name: tc.Name, Args: args, Status: status, Executed: executed, Result: result})
795 }
796 }()
797 if err := ctx.Err(); err != nil {
798 status = executionErrorStatus(err)
799 return ToolResult{IsError: true, Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}}}
800 }
801 tool, err := a.Tools.Get(tc.Name)
802 if err != nil {
803 return ToolResult{
804 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
805 IsError: true,
806 }
807 }
808
809 // Intercept hook: an extension or other guard can refuse the
810 // call before any side effect happens, OR rewrite the args
811 // seen by the tool. The model sees the reason as the tool
812 // error, learns from it, and (typically) proposes a different
813 // action; rewrites are invisible to the model (they apply only
814 // to the execution).
815 before := a.BeforeToolExecute
816 if a.BeforeToolExecuteContext != nil {
817 before = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
818 return a.BeforeToolExecuteContext(ctx, call)
819 }
820 }
821 if before != nil {
822 allowed, reason, modified := before(tc)
823 if len(modified) > 0 && json.Valid(modified) {
824 args = modified
825 }
826 if !allowed {
827 status = "blocked"
828 if ctx.Err() != nil {
829 status = executionErrorStatus(ctx.Err())
830 }
831 if reason == "" {
832 reason = "tool call refused by extension guard"
833 }
834 return ToolResult{
835 Content: []provider.Content{provider.TextBlock{Text: reason}},
836 IsError: true,
837 }
838 }
839 }
840
841 if err := ctx.Err(); err != nil {
842 status = executionErrorStatus(err)
843 return ToolResult{IsError: true, Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}}}
844 }
845 if len(args) == 0 {
846 args = json.RawMessage("{}")
847 }
848
849 // Recover panics so a buggy tool does not crash the agent.
850 var res ToolResult
851 func() {
852 defer func() {
853 if r := recover(); r != nil {
854 res = ToolResult{
855 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
856 IsError: true,
857 }
858 }
859 }()
860 executed = true
861 out, err := tool.Execute(ctx, args, func(text string) {
862 sink(EvToolProgress{ID: tc.ID, Text: text})
863 })
864 if err != nil {
865 status = executionErrorStatus(err)
866 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
867 res = ToolResult{
868 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
869 IsError: true,
870 }
871 return
872 }
873 res = ToolResult{
874 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
875 IsError: true,
876 }
877 return
878 }
879 res = out
880 }()
881 return res
882}
883
884func executionErrorStatus(err error) string {
885 if errors.Is(err, context.DeadlineExceeded) {
886 return "timed_out"
887 }
888 if errors.Is(err, context.Canceled) {
889 return "cancelled"
890 }
891 var policyErr *ToolPolicyError
892 if errors.As(err, &policyErr) {
893 return "blocked"
894 }
895 return "failed"
896}
897
898// extractText concatenates all TextBlock content in a message. Used
899// by BeforeAssistantMessage so guards see a single string instead of
900// having to walk provider.Content themselves.
901func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
902 var content []provider.Content
903 for _, c := range msg.Content {
904 tr, ok := c.(provider.ToolResultBlock)
905 if !ok {
906 continue
907 }
908 for _, inner := range tr.Content {
909 switch v := inner.(type) {
910 case provider.TextBlock:
911 // Keep short textual context so the model understands why
912 // the images appeared, but don't duplicate giant read
913 // outputs verbatim.
914 if len(v.Text) > 0 && len(v.Text) <= 500 {
915 content = append(content, v)
916 }
917 case provider.ImageBlock:
918 content = append(content, v)
919 }
920 }
921 }
922 if len(content) == 0 {
923 return provider.Message{}
924 }
925 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
926 content = append([]provider.Content{prefix}, content...)
927 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
928}
929
930func containsString(values []string, want string) bool {
931 for _, value := range values {
932 if value == want {
933 return true
934 }
935 }
936 return false
937}
938
939func extractText(msg provider.Message) string {
940 var out string
941 for _, c := range msg.Content {
942 if tb, ok := c.(provider.TextBlock); ok {
943 if out != "" {
944 out += "\n"
945 }
946 out += tb.Text
947 }
948 }
949 return out
950}
951
952// replaceText returns a copy of msg with every TextBlock replaced by
953// a single TextBlock containing replacement. Non-text content (tool
954// calls, etc.) is preserved in order.
955func replaceText(msg provider.Message, replacement string) provider.Message {
956 out := provider.Message{Role: msg.Role}
957 out.Content = make([]provider.Content, 0, len(msg.Content))
958 replaced := false
959 for _, c := range msg.Content {
960 if _, ok := c.(provider.TextBlock); ok {
961 if !replaced {
962 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
963 replaced = true
964 }
965 continue
966 }
967 out.Content = append(out.Content, c)
968 }
969 if !replaced {
970 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
971 }
972 return out
973}