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 return a.promptWithPrelude(ctx, text, images, nil, "", sink)
349}
350
351// PromptWithTool runs a host-requested tool after the user prompt and before
352// the first model turn. The call and result are persisted as a matched pair.
353// A cancelled execution still records its result, but does not call the model.
354func (a *Agent) PromptWithTool(ctx context.Context, text string, call provider.ToolCallBlock, origin string, sink func(AgentEvent)) error {
355 if strings.TrimSpace(text) == "" || call.ID == "" || call.Name == "" || !json.Valid(call.Arguments) || !strings.HasPrefix(strings.TrimSpace(string(call.Arguments)), "{") {
356 return fmt.Errorf("tool prompt requires text, tool ID, name and JSON object arguments")
357 }
358 return a.promptWithPrelude(ctx, text, nil, &call, origin, sink)
359}
360
361func (a *Agent) promptWithPrelude(ctx context.Context, text string, images []provider.ImageBlock, call *provider.ToolCallBlock, origin string, sink func(AgentEvent)) error {
362 if sink == nil {
363 sink = func(AgentEvent) {}
364 }
365 sink = a.wrapSink(sink)
366 if a.OnEvent != nil {
367 a.OnEvent(EvPromptSubmit{Text: text, ImageCount: len(images)})
368 }
369 content := []provider.Content{}
370 if text != "" {
371 content = append(content, provider.TextBlock{Text: text})
372 }
373 for _, img := range images {
374 content = append(content, img)
375 }
376 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
377
378 a.mu.Lock()
379 a.messages = append(a.messages, user)
380 a.rev++
381 a.mu.Unlock()
382 a.fireMessageAppended(user)
383 sink(EvUserMessage{Message: user})
384
385 if call != nil {
386 if err := ctx.Err(); err != nil {
387 return err
388 }
389 sink(EvToolCall{ID: call.ID, Name: call.Name, Args: call.Arguments})
390 result := a.runOneTool(ctx, *call, sink)
391 assistant := provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{*call}, Time: time.Now(), Meta: map[string]string{"origin_extension": origin, "synthetic_tool_call": "true"}}
392 tool := provider.Message{Role: provider.RoleTool, Content: []provider.Content{provider.ToolResultBlock{CallID: call.ID, Content: result.Content, IsError: result.IsError}}, Time: time.Now()}
393 for _, name := range result.ActivateTools {
394 if _, err := a.Tools.Get(name); err == nil && !containsString(tool.AddedToolNames, name) {
395 tool.AddedToolNames = append(tool.AddedToolNames, name)
396 }
397 }
398 var mirror provider.Message
399 if a.Client != nil && (a.Client.Name() == "openai" || a.Client.Name() == "openai-codex") {
400 mirror = mirrorToolImagesAsUser(tool)
401 }
402 a.mu.Lock()
403 a.messages = append(a.messages, assistant, tool)
404 a.rev += 2
405 if len(mirror.Content) > 0 {
406 a.messages = append(a.messages, mirror)
407 a.rev++
408 }
409 a.mu.Unlock()
410 a.fireMessageAppended(assistant)
411 a.fireMessageAppended(tool)
412 if len(mirror.Content) > 0 {
413 a.fireMessageAppended(mirror)
414 }
415 if err := ctx.Err(); err != nil {
416 return err
417 }
418 }
419 return a.runLoop(ctx, sink)
420}
421
422// Continue runs the agent loop against the existing transcript. Used
423// after appending tool results manually or to retry.
424func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
425 if sink == nil {
426 sink = func(AgentEvent) {}
427 }
428 sink = a.wrapSink(sink)
429 return a.runLoop(ctx, sink)
430}
431
432// wrapSink composes the per-call sink with a.OnEvent (if set) so the
433// extension manager (or any other observer) sees every AgentEvent
434// without having to thread itself through every Prompt callsite.
435func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) {
436 if a.OnEvent == nil {
437 return sink
438 }
439 obs := a.OnEvent
440 return func(ev AgentEvent) {
441 obs(ev)
442 sink(ev)
443 }
444}
445
446func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
447 for step := 1; a.MaxSteps <= 0 || step <= a.MaxSteps; step++ {
448 // Preparation is cached across steps and user messages. Only an
449 // explicit prompt/model/session change invokes BeforeStart again.
450 if err := a.prepareStart(ctx); err != nil {
451 sink(EvDone{})
452 return err
453 }
454 // Messages queued while the agent was busy are delivered
455 // before the next model call. This is the safe boundary:
456 // any previous tool batch has already completed and its
457 // results have been appended, but no new provider request has
458 // started yet.
459 if pending := a.drainQueuedMessages(); len(pending) > 0 {
460 a.appendQueuedAsUser(pending, sink)
461 }
462
463 sink(EvTurnStart{Step: step})
464 if a.BeforeTurn != nil {
465 if allowed, reason := a.BeforeTurn(step); !allowed {
466 if reason == "" {
467 reason = "turn blocked by extension guard"
468 }
469 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
470 sink(EvDone{})
471 return nil
472 }
473 }
474
475 var (
476 stop provider.StopReason
477 assistantMsg provider.Message
478 err error
479 )
480 for attempt := 0; ; attempt++ {
481 stop, assistantMsg, err = a.oneTurn(ctx, sink)
482 sink(EvTurnEnd{Stop: stop, Err: err})
483 if err == nil || !a.canRetryError(err, attempt) {
484 break
485 }
486 a.dropLastAssistantMessage()
487 if sleepErr := sleepRetry(ctx, a.retryDelay(attempt)); sleepErr != nil {
488 return sleepErr
489 }
490 }
491 if err != nil {
492 return err
493 }
494
495 if stop == provider.StopToolUse {
496 // Execute each client tool call, append a single tool-results message, continue.
497 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
498 if len(toolMsg.Content) == 0 {
499 // Provider-executed (server) tools need no client results.
500 continue
501 }
502 a.mu.Lock()
503 a.messages = append(a.messages, toolMsg)
504 a.rev++
505 // OpenAI's chat-completions tool message shape is text-centric.
506 // Vision models reliably consume images when they arrive as user
507 // content, so when a tool result contains images we mirror them
508 // into a synthetic user message immediately after the tool result.
509 // This keeps the transcript self-contained for providers that can
510 // see image blocks in tool messages while making OpenAI vision
511 // models actually receive the image bytes.
512 //
513 // The OpenAI Responses route ("openai-codex") has the same
514 // text-centric tool-output shape: a function_call_output only
515 // carries a string, so images in a tool result never reach the
516 // model. Both providers serialize images correctly when they
517 // arrive as user content, so the mirror covers them both.
518 var imageMirror provider.Message
519 if a.Client != nil && (a.Client.Name() == "openai" || a.Client.Name() == "openai-codex") {
520 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
521 a.messages = append(a.messages, mirror)
522 a.rev++
523 imageMirror = mirror
524 }
525 }
526 a.mu.Unlock()
527 a.fireMessageAppended(toolMsg)
528 if len(imageMirror.Content) > 0 {
529 a.fireMessageAppended(imageMirror)
530 }
531 // If context was cancelled during tool execution, bail out.
532 if err := ctx.Err(); err != nil {
533 sink(EvDone{})
534 return err
535 }
536 _ = hadError
537 continue
538 }
539
540 // If the assistant stopped without tool calls but a message was
541 // queued while it was speaking, loop once more so that message
542 // is appended and answered instead of waiting until a later
543 // top-level prompt.
544 if ctx.Err() == nil && a.QueuedMessageCount() > 0 {
545 continue
546 }
547
548 // Terminal stop (end, length, error, aborted).
549 sink(EvDone{})
550 return nil
551 }
552 if a.MaxSteps > 0 {
553 sink(EvDone{})
554 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
555 }
556 return nil
557}
558
559func (a *Agent) canRetryError(err error, attempt int) bool {
560 if err == nil || a.MaxRetries <= 0 || attempt >= a.MaxRetries {
561 return false
562 }
563 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
564 return false
565 }
566 msg := strings.ToLower(err.Error())
567 if msg == "" || isNonRetryableProviderLimit(msg) {
568 return false
569 }
570 needles := []string{
571 "overloaded", "provider returned error", "rate limit", "ratelimit", "too many requests",
572 "429", "http 429", "500", "http 500", "502", "http 502", "503", "http 503", "504", "http 504",
573 "service unavailable", "server error", "internal error", "network error", "connection error",
574 "connection refused", "connection lost", "fetch failed", "upstream connect", "reset before headers",
575 "socket hang up", "ended without", "stream ended before", "did not get a response", "timed out",
576 "timeout", "terminated", "unexpected eof", "transport failure",
577 // OpenAI's ChatGPT/Codex backend returns this generic message (with a
578 // request ID) for transient server failures and explicitly says
579 // "You can retry your request".
580 "an error occurred while processing your request",
581 // Explicit retry guidance emitted by provider backends (OpenAI
582 // Responses, AWS Bedrock stream exceptions) with varying prefixes.
583 "you can retry your request", "try your request again", "please retry your request",
584 // Capacity messages from the ChatGPT/Codex backend, e.g.
585 // "Our servers are currently overloaded. Please try again later."
586 // The trailing advice also shows up on its own for transient
587 // capacity failures; usage/quota limits are filtered out above by
588 // isNonRetryableProviderLimit before this list is consulted.
589 "servers are currently overloaded", "servers are busy", "try again later",
590 }
591 for _, needle := range needles {
592 if strings.Contains(msg, needle) {
593 return true
594 }
595 }
596 return false
597}
598
599func isNonRetryableProviderLimit(msg string) bool {
600 needles := []string{
601 "usage limit", "monthly usage limit", "freeusagelimit", "gousagelimit",
602 "available balance", "insufficient_quota", "out of budget", "quota exceeded", "billing",
603 }
604 for _, needle := range needles {
605 if strings.Contains(msg, needle) {
606 return true
607 }
608 }
609 return false
610}
611
612func (a *Agent) retryDelay(attempt int) time.Duration {
613 base := a.RetryBaseDelay
614 if base <= 0 {
615 base = 2 * time.Second
616 }
617 return base * time.Duration(1<<attempt)
618}
619
620func sleepRetry(ctx context.Context, d time.Duration) error {
621 if d <= 0 {
622 return ctx.Err()
623 }
624 t := time.NewTimer(d)
625 defer t.Stop()
626 select {
627 case <-ctx.Done():
628 return ctx.Err()
629 case <-t.C:
630 return nil
631 }
632}
633
634func (a *Agent) dropLastAssistantMessage() {
635 a.mu.Lock()
636 defer a.mu.Unlock()
637 if n := len(a.messages); n > 0 && a.messages[n-1].Role == provider.RoleAssistant {
638 a.messages = a.messages[:n-1]
639 a.rev++
640 }
641}
642
643// oneTurn calls the LLM once, forwards events, returns the stop reason
644// and the assembled assistant message (already appended to the transcript).
645func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
646 var req provider.Request
647 for {
648 if err := a.prepareStart(ctx); err != nil {
649 return provider.StopError, provider.Message{}, err
650 }
651 a.mu.Lock()
652 // A reset can also arrive after runLoop's preparation, for example
653 // while BeforeTurn waits. Validate and snapshot under the same lock
654 // so this request never receives an unprepared replacement prompt.
655 if a.BeforeStart != nil && !a.startCurrentLocked() {
656 a.mu.Unlock()
657 continue
658 }
659 req = provider.Request{
660 Model: a.Model,
661 System: a.System,
662 // Repair any dangling tool_use blocks before sending. A turn
663 // aborted mid-flight (cancel, connection drop, ECONNREFUSED to a
664 // dev server, etc.) can leave an assistant tool_use with no
665 // matching tool_result in the live transcript. The load-time
666 // repair in OpenSession only runs on restart, so without this the
667 // next in-process request is rejected by providers like Anthropic
668 // with "tool_use ids were found without tool_result blocks". The
669 // repair is pure and a no-op on already-valid transcripts.
670 Messages: repairToolUseResultPairs(append([]provider.Message(nil), a.messages...)),
671 Tools: a.Tools.Specs(),
672 Reasoning: a.Reasoning,
673 MaxTokens: a.MaxTokens,
674 Temperature: a.Temperature,
675 SessionID: a.SessionID,
676 MaxToolCalls: a.MaxToolCalls,
677 }
678 a.mu.Unlock()
679 break
680 }
681 stream, err := a.Client.Stream(ctx, req)
682 if err != nil {
683 return provider.StopError, provider.Message{}, err
684 }
685
686 sink(EvAssistantStart{})
687
688 var (
689 stop provider.StopReason
690 finalErr error
691 finalMsg provider.Message
692 )
693
694 for ev := range stream {
695 switch e := ev.(type) {
696 case provider.EventStart:
697 // nothing
698 case provider.EventTextDelta:
699 sink(EvTextDelta{Delta: e.Delta})
700 case provider.EventToolStart:
701 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
702 case provider.EventToolArgs:
703 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
704 case provider.EventToolEnd:
705 sink(EvToolUseEnd{ID: e.ID})
706 case provider.EventUsage:
707 cum := a.cost.Add(e.Usage)
708 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
709 if a.OnUsage != nil {
710 a.OnUsage(cum)
711 }
712 case provider.EventDone:
713 stop = e.Stop
714 finalErr = e.Err
715 finalMsg = e.Message
716 }
717 }
718
719 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
720 // produce partial content. When the partial message is text only we
721 // keep whatever was streamed up to the cancel so the user does not
722 // lose visible work (a cut-off summary is still useful). If the
723 // partial message already contained tool-call blocks we drop the
724 // whole thing, because an unmatched tool_use would fail the next
725 // turn with a tool_result mismatch error.
726 keep := len(finalMsg.Content) > 0
727 if stop == provider.StopAborted && keep {
728 hasToolCall := false
729 for _, c := range finalMsg.Content {
730 if _, ok := c.(provider.ToolCallBlock); ok {
731 hasToolCall = true
732 break
733 }
734 }
735 if hasToolCall {
736 keep = false
737 }
738 }
739 if keep {
740 emit := finalMsg
741 suppress := false
742
743 // BeforeAssistantMessage hook: extensions can suppress or
744 // rewrite the visible text. The transcript keeps the
745 // model's original output so the model still sees what it
746 // said on subsequent turns.
747 if a.BeforeAssistantMessage != nil {
748 orig := extractText(finalMsg)
749 if orig != "" {
750 allowed, _, replacement := a.BeforeAssistantMessage(orig)
751 if !allowed {
752 suppress = true
753 } else if replacement != "" && replacement != orig {
754 emit = replaceText(finalMsg, replacement)
755 }
756 }
757 }
758
759 a.mu.Lock()
760 a.messages = append(a.messages, finalMsg)
761 a.rev++
762 a.mu.Unlock()
763 a.fireMessageAppended(finalMsg)
764 if !suppress {
765 sink(EvAssistantMessage{Message: emit})
766 }
767 // Now surface tool calls as EvToolCall events so UIs can render them
768 // in order before the tool results arrive.
769 for _, c := range finalMsg.Content {
770 if tc, ok := c.(provider.ToolCallBlock); ok {
771 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
772 }
773 }
774 }
775
776 return stop, finalMsg, finalErr
777}
778
779// CallTool executes a host-initiated tool without adding a model tool-use
780// message to the transcript. The caller supplies a unique ID and event sink.
781// It shares the model tool path's guards, confirmation hook and error handling.
782func (a *Agent) CallTool(ctx context.Context, id, name string, args json.RawMessage, sink func(AgentEvent)) ToolResult {
783 if sink == nil {
784 sink = func(AgentEvent) {}
785 }
786 sink(EvToolCall{ID: id, Name: name, Args: args})
787 return a.runOneTool(ctx, provider.ToolCallBlock{ID: id, Name: name, Arguments: args}, sink)
788}
789
790// executeTools runs every tool call in the assistant message and returns
791// a single tool-role message carrying all results.
792func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
793 var results []provider.Content
794 var addedTools []string
795 hadError := false
796
797 for _, c := range msg.Content {
798 tc, ok := c.(provider.ToolCallBlock)
799 if !ok || tc.Server {
800 continue
801 }
802 res := a.runOneTool(ctx, tc, sink)
803 if res.IsError {
804 hadError = true
805 }
806 results = append(results, provider.ToolResultBlock{
807 CallID: tc.ID,
808 Content: res.Content,
809 IsError: res.IsError,
810 })
811 for _, name := range res.ActivateTools {
812 if _, err := a.Tools.Get(name); err == nil && !containsString(addedTools, name) {
813 addedTools = append(addedTools, name)
814 }
815 }
816 }
817
818 return provider.Message{
819 Role: provider.RoleTool,
820 Content: results,
821 Time: time.Now(),
822 AddedToolNames: addedTools,
823 }, hadError
824}
825
826func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) (result ToolResult) {
827 args := tc.Arguments
828 status := ""
829 executed := false
830 defer func() {
831 if status == "" {
832 status = "completed"
833 if result.IsError {
834 status = "failed"
835 switch result.Status {
836 case "cancelled", "timed_out", "blocked":
837 status = result.Status
838 }
839 }
840 }
841 if sink != nil {
842 sink(EvToolResult{ID: tc.ID, Name: tc.Name, Args: args, Status: status, Executed: executed, Result: result})
843 }
844 }()
845 if err := ctx.Err(); err != nil {
846 status = executionErrorStatus(err)
847 return ToolResult{IsError: true, Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}}}
848 }
849 tool, err := a.Tools.Get(tc.Name)
850 if err != nil {
851 return ToolResult{
852 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
853 IsError: true,
854 }
855 }
856
857 // Intercept hook: an extension or other guard can refuse the
858 // call before any side effect happens, OR rewrite the args
859 // seen by the tool. The model sees the reason as the tool
860 // error, learns from it, and (typically) proposes a different
861 // action; rewrites are invisible to the model (they apply only
862 // to the execution).
863 before := a.BeforeToolExecute
864 if a.BeforeToolExecuteContext != nil {
865 before = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
866 return a.BeforeToolExecuteContext(ctx, call)
867 }
868 }
869 if before != nil {
870 allowed, reason, modified := before(tc)
871 if len(modified) > 0 && json.Valid(modified) {
872 args = modified
873 }
874 if !allowed {
875 status = "blocked"
876 if ctx.Err() != nil {
877 status = executionErrorStatus(ctx.Err())
878 }
879 if reason == "" {
880 reason = "tool call refused by extension guard"
881 }
882 return ToolResult{
883 Content: []provider.Content{provider.TextBlock{Text: reason}},
884 IsError: true,
885 }
886 }
887 }
888
889 if err := ctx.Err(); err != nil {
890 status = executionErrorStatus(err)
891 return ToolResult{IsError: true, Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}}}
892 }
893 if len(args) == 0 {
894 args = json.RawMessage("{}")
895 }
896
897 // Recover panics so a buggy tool does not crash the agent.
898 var res ToolResult
899 func() {
900 defer func() {
901 if r := recover(); r != nil {
902 res = ToolResult{
903 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
904 IsError: true,
905 }
906 }
907 }()
908 executed = true
909 out, err := tool.Execute(ctx, args, func(text string) {
910 sink(EvToolProgress{ID: tc.ID, Text: text})
911 })
912 if err != nil {
913 status = executionErrorStatus(err)
914 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
915 res = ToolResult{
916 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
917 IsError: true,
918 }
919 return
920 }
921 res = ToolResult{
922 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
923 IsError: true,
924 }
925 return
926 }
927 res = out
928 }()
929 return res
930}
931
932func executionErrorStatus(err error) string {
933 if errors.Is(err, context.DeadlineExceeded) {
934 return "timed_out"
935 }
936 if errors.Is(err, context.Canceled) {
937 return "cancelled"
938 }
939 var policyErr *ToolPolicyError
940 if errors.As(err, &policyErr) {
941 return "blocked"
942 }
943 return "failed"
944}
945
946// extractText concatenates all TextBlock content in a message. Used
947// by BeforeAssistantMessage so guards see a single string instead of
948// having to walk provider.Content themselves.
949func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
950 var content []provider.Content
951 for _, c := range msg.Content {
952 tr, ok := c.(provider.ToolResultBlock)
953 if !ok {
954 continue
955 }
956 for _, inner := range tr.Content {
957 switch v := inner.(type) {
958 case provider.TextBlock:
959 // Keep short textual context so the model understands why
960 // the images appeared, but don't duplicate giant read
961 // outputs verbatim.
962 if len(v.Text) > 0 && len(v.Text) <= 500 {
963 content = append(content, v)
964 }
965 case provider.ImageBlock:
966 content = append(content, v)
967 }
968 }
969 }
970 if len(content) == 0 {
971 return provider.Message{}
972 }
973 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
974 content = append([]provider.Content{prefix}, content...)
975 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
976}
977
978func containsString(values []string, want string) bool {
979 for _, value := range values {
980 if value == want {
981 return true
982 }
983 }
984 return false
985}
986
987func extractText(msg provider.Message) string {
988 var out string
989 for _, c := range msg.Content {
990 if tb, ok := c.(provider.TextBlock); ok {
991 if out != "" {
992 out += "\n"
993 }
994 out += tb.Text
995 }
996 }
997 return out
998}
999
1000// replaceText returns a copy of msg with every TextBlock replaced by
1001// a single TextBlock containing replacement. Non-text content (tool
1002// calls, etc.) is preserved in order.
1003func replaceText(msg provider.Message, replacement string) provider.Message {
1004 out := provider.Message{Role: msg.Role}
1005 out.Content = make([]provider.Content, 0, len(msg.Content))
1006 replaced := false
1007 for _, c := range msg.Content {
1008 if _, ok := c.(provider.TextBlock); ok {
1009 if !replaced {
1010 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
1011 replaced = true
1012 }
1013 continue
1014 }
1015 out.Content = append(out.Content, c)
1016 }
1017 if !replaced {
1018 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
1019 }
1020 return out
1021}