1package core
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "sync"
9 "time"
10
11 "github.com/patriceckhart/zot/internal/provider"
12)
13
14// Agent is a stateful conversation bound to a provider client, a model,
15// and a set of tools.
16type Agent struct {
17 Client provider.Client
18 Model string
19 System string
20 Tools Registry
21 MaxSteps int
22 Reasoning string
23
24 // BeforeToolExecute, if set, is called immediately before each
25 // tool runs. Returning (allowed=false, reason) short-circuits
26 // the call with an error result containing reason. Optionally,
27 // returning a non-nil modifiedArgs replaces the JSON args the
28 // tool will see, which lets guards redact / augment / patch the
29 // model's request without rewriting the transcript. Empty or
30 // malformed modifiedArgs is ignored.
31 BeforeToolExecute func(call provider.ToolCallBlock) (allowed bool, reason string, modifiedArgs json.RawMessage)
32
33 // BeforeTurn, if set, is called before each turn's model call.
34 // Returning (allowed=false, reason) aborts the turn; reason is
35 // surfaced as an assistant-like status line. Used for rate-
36 // limiting, business-hour gates, and deny-by-default setups.
37 BeforeTurn func(step int) (allowed bool, reason string)
38
39 // BeforeAssistantMessage, if set, is called after the model's
40 // final assistant message is assembled but before it's appended
41 // to the transcript. Returning (allowed=false) suppresses both
42 // the transcript append and the UI event. A non-empty
43 // replacement rewrites the visible text for the user while
44 // leaving the model's original text in the transcript (so the
45 // model can still see what it said in subsequent turns).
46 BeforeAssistantMessage func(text string) (allowed bool, reason, replacement string)
47
48 // OnEvent, if set, mirrors every AgentEvent the loop emits to
49 // this callback in addition to the per-Prompt sink. Used by the
50 // extension manager to fan events out to subscribed extensions
51 // without each caller having to compose sinks manually.
52 OnEvent func(AgentEvent)
53
54 // OnMessageAppended, if set, fires every time a message is
55 // appended to the in-memory transcript by the agent loop — the
56 // initial user prompt, each finalised assistant message, and
57 // each tool-results message (plus the synthetic OpenAI image
58 // mirror, if any). Hosts wire this to the on-disk session so
59 // that turns are durable as soon as they happen, instead of
60 // only being flushed on a clean exit.
61 OnMessageAppended func(provider.Message)
62
63 // OnUsage, if set, fires after every turn's usage row arrives,
64 // carrying the cumulative usage for the session. Hosts wire
65 // this to the on-disk session so the persisted total stays
66 // current and a crash recovers the right cost figure.
67 OnUsage func(cumulative provider.Usage)
68
69 // OnTranscriptCompacted, if set, fires after Compact replaces the
70 // in-memory transcript with the synthetic summary plus kept tail.
71 // Hosts wire this to append an explicit compaction checkpoint to
72 // the session log; per-message append hooks do not fire for this
73 // wholesale transcript replacement.
74 OnTranscriptCompacted func(messages []provider.Message)
75
76 mu sync.Mutex
77 messages []provider.Message
78 // rev increments whenever the transcript slice is replaced or a
79 // message is appended. The TUI uses it as a cheap redraw cache key
80 // so editor-only typing doesn't copy/rebuild a long transcript on
81 // every keypress.
82 rev uint64
83 cost CostTracker
84}
85
86// NewAgent returns an Agent with sensible defaults.
87func NewAgent(client provider.Client, model, system string, tools Registry) *Agent {
88 return &Agent{
89 Client: client,
90 Model: model,
91 System: system,
92 Tools: tools,
93 MaxSteps: 0, // 0 = unlimited
94 }
95}
96
97// Messages returns a copy of the current transcript.
98func (a *Agent) Messages() []provider.Message {
99 a.mu.Lock()
100 defer a.mu.Unlock()
101 out := make([]provider.Message, len(a.messages))
102 copy(out, a.messages)
103 return out
104}
105
106// Revision returns a monotonically increasing transcript version.
107// It is cheap to query and changes whenever Messages() would return
108// different transcript content because of append/set operations.
109func (a *Agent) Revision() uint64 {
110 a.mu.Lock()
111 defer a.mu.Unlock()
112 return a.rev
113}
114
115// SetTools swaps the tool registry. Used by /reload-ext to hand
116// the agent a fresh registry after extension subprocesses have been
117// respawned (and their freshly-registered tools merged in).
118func (a *Agent) SetTools(reg Registry) {
119 a.mu.Lock()
120 a.Tools = reg
121 a.mu.Unlock()
122}
123
124// SetMessages replaces the transcript (used when resuming a session).
125func (a *Agent) SetMessages(msgs []provider.Message) {
126 a.mu.Lock()
127 defer a.mu.Unlock()
128 a.messages = append(a.messages[:0], msgs...)
129 a.rev++
130}
131
132// Cost returns the cumulative usage.
133func (a *Agent) Cost() provider.Usage {
134 a.mu.Lock()
135 defer a.mu.Unlock()
136 return a.cost.Total
137}
138
139// SeedCost sets the cumulative usage as a baseline before the first
140// turn runs. Used when transferring state from another agent (model
141// or provider switch) so the running cost meter doesn't reset to 0.
142func (a *Agent) SeedCost(u provider.Usage) {
143 a.mu.Lock()
144 defer a.mu.Unlock()
145 a.cost.Total = u
146}
147
148// fireMessageAppended invokes OnMessageAppended without holding the
149// agent mutex, so the host's persistence callback can take its own
150// locks without deadlocking the agent loop. Tolerates a nil hook so
151// non-persisting callers (tests, RPC mode) don't have to set it.
152func (a *Agent) fireMessageAppended(m provider.Message) {
153 if a.OnMessageAppended != nil {
154 a.OnMessageAppended(m)
155 }
156}
157
158// Prompt sends a user message and runs the agent loop until the model
159// stops or an error occurs. Events are delivered via sink in order.
160// sink must not block the caller for long; buffer as needed.
161func (a *Agent) Prompt(ctx context.Context, text string, images []provider.ImageBlock, sink func(AgentEvent)) error {
162 if sink == nil {
163 sink = func(AgentEvent) {}
164 }
165 sink = a.wrapSink(sink)
166 content := []provider.Content{}
167 if text != "" {
168 content = append(content, provider.TextBlock{Text: text})
169 }
170 for _, img := range images {
171 content = append(content, img)
172 }
173 user := provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
174
175 a.mu.Lock()
176 a.messages = append(a.messages, user)
177 a.rev++
178 a.mu.Unlock()
179 a.fireMessageAppended(user)
180 sink(EvUserMessage{Message: user})
181
182 return a.runLoop(ctx, sink)
183}
184
185// Continue runs the agent loop against the existing transcript. Used
186// after appending tool results manually or to retry.
187func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error {
188 if sink == nil {
189 sink = func(AgentEvent) {}
190 }
191 sink = a.wrapSink(sink)
192 return a.runLoop(ctx, sink)
193}
194
195// wrapSink composes the per-call sink with a.OnEvent (if set) so the
196// extension manager (or any other observer) sees every AgentEvent
197// without having to thread itself through every Prompt callsite.
198func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) {
199 if a.OnEvent == nil {
200 return sink
201 }
202 obs := a.OnEvent
203 return func(ev AgentEvent) {
204 obs(ev)
205 sink(ev)
206 }
207}
208
209func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error {
210 for step := 1; a.MaxSteps <= 0 || step <= a.MaxSteps; step++ {
211 sink(EvTurnStart{Step: step})
212 if a.BeforeTurn != nil {
213 if allowed, reason := a.BeforeTurn(step); !allowed {
214 if reason == "" {
215 reason = "turn blocked by extension guard"
216 }
217 sink(EvTurnEnd{Stop: provider.StopError, Err: fmt.Errorf("%s", reason)})
218 sink(EvDone{})
219 return nil
220 }
221 }
222 stop, assistantMsg, err := a.oneTurn(ctx, sink)
223 sink(EvTurnEnd{Stop: stop, Err: err})
224 if err != nil {
225 return err
226 }
227
228 if stop == provider.StopToolUse {
229 // Execute each tool call, append a single tool-results message, continue.
230 toolMsg, hadError := a.executeTools(ctx, assistantMsg, sink)
231 a.mu.Lock()
232 a.messages = append(a.messages, toolMsg)
233 a.rev++
234 // OpenAI's chat-completions tool message shape is text-centric.
235 // Vision models reliably consume images when they arrive as user
236 // content, so when a tool result contains images we mirror them
237 // into a synthetic user message immediately after the tool result.
238 // This keeps the transcript self-contained for providers that can
239 // see image blocks in tool messages while making OpenAI vision
240 // models actually receive the image bytes.
241 var imageMirror provider.Message
242 if a.Client != nil && a.Client.Name() == "openai" {
243 if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 {
244 a.messages = append(a.messages, mirror)
245 a.rev++
246 imageMirror = mirror
247 }
248 }
249 a.mu.Unlock()
250 a.fireMessageAppended(toolMsg)
251 if len(imageMirror.Content) > 0 {
252 a.fireMessageAppended(imageMirror)
253 }
254 // If context was cancelled during tool execution, bail out.
255 if err := ctx.Err(); err != nil {
256 sink(EvDone{})
257 return err
258 }
259 _ = hadError
260 continue
261 }
262
263 // Terminal stop (end, length, error, aborted).
264 sink(EvDone{})
265 return nil
266 }
267 if a.MaxSteps > 0 {
268 sink(EvDone{})
269 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
270 }
271 return nil
272}
273
274// oneTurn calls the LLM once, forwards events, returns the stop reason
275// and the assembled assistant message (already appended to the transcript).
276func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
277 req := provider.Request{
278 Model: a.Model,
279 System: a.System,
280 Messages: a.Messages(),
281 Tools: a.Tools.Specs(),
282 Reasoning: a.Reasoning,
283 }
284 stream, err := a.Client.Stream(ctx, req)
285 if err != nil {
286 return provider.StopError, provider.Message{}, err
287 }
288
289 sink(EvAssistantStart{})
290
291 var (
292 stop provider.StopReason
293 finalErr error
294 finalMsg provider.Message
295 )
296
297 for ev := range stream {
298 switch e := ev.(type) {
299 case provider.EventStart:
300 // nothing
301 case provider.EventTextDelta:
302 sink(EvTextDelta{Delta: e.Delta})
303 case provider.EventToolStart:
304 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
305 case provider.EventToolArgs:
306 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
307 case provider.EventToolEnd:
308 sink(EvToolUseEnd{ID: e.ID})
309 case provider.EventUsage:
310 cum := a.cost.Add(e.Usage)
311 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
312 if a.OnUsage != nil {
313 a.OnUsage(cum)
314 }
315 case provider.EventDone:
316 stop = e.Stop
317 finalErr = e.Err
318 finalMsg = e.Message
319 }
320 }
321
322 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
323 // produce partial content. When the partial message is text only we
324 // keep whatever was streamed up to the cancel so the user does not
325 // lose visible work (a cut-off summary is still useful). If the
326 // partial message already contained tool-call blocks we drop the
327 // whole thing, because an unmatched tool_use would fail the next
328 // turn with a tool_result mismatch error.
329 keep := len(finalMsg.Content) > 0
330 if stop == provider.StopAborted && keep {
331 hasToolCall := false
332 for _, c := range finalMsg.Content {
333 if _, ok := c.(provider.ToolCallBlock); ok {
334 hasToolCall = true
335 break
336 }
337 }
338 if hasToolCall {
339 keep = false
340 }
341 }
342 if keep {
343 emit := finalMsg
344 suppress := false
345
346 // BeforeAssistantMessage hook: extensions can suppress or
347 // rewrite the visible text. The transcript keeps the
348 // model's original output so the model still sees what it
349 // said on subsequent turns.
350 if a.BeforeAssistantMessage != nil {
351 orig := extractText(finalMsg)
352 if orig != "" {
353 allowed, _, replacement := a.BeforeAssistantMessage(orig)
354 if !allowed {
355 suppress = true
356 } else if replacement != "" && replacement != orig {
357 emit = replaceText(finalMsg, replacement)
358 }
359 }
360 }
361
362 a.mu.Lock()
363 a.messages = append(a.messages, finalMsg)
364 a.rev++
365 a.mu.Unlock()
366 a.fireMessageAppended(finalMsg)
367 if !suppress {
368 sink(EvAssistantMessage{Message: emit})
369 }
370 // Now surface tool calls as EvToolCall events so UIs can render them
371 // in order before the tool results arrive.
372 for _, c := range finalMsg.Content {
373 if tc, ok := c.(provider.ToolCallBlock); ok {
374 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
375 }
376 }
377 }
378
379 return stop, finalMsg, finalErr
380}
381
382// executeTools runs every tool call in the assistant message and returns
383// a single tool-role message carrying all results.
384func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
385 var results []provider.Content
386 hadError := false
387
388 for _, c := range msg.Content {
389 tc, ok := c.(provider.ToolCallBlock)
390 if !ok {
391 continue
392 }
393 res := a.runOneTool(ctx, tc, sink)
394 if res.IsError {
395 hadError = true
396 }
397 results = append(results, provider.ToolResultBlock{
398 CallID: tc.ID,
399 Content: res.Content,
400 IsError: res.IsError,
401 })
402 sink(EvToolResult{ID: tc.ID, Result: res})
403 }
404
405 return provider.Message{
406 Role: provider.RoleTool,
407 Content: results,
408 Time: time.Now(),
409 }, hadError
410}
411
412func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
413 tool, err := a.Tools.Get(tc.Name)
414 if err != nil {
415 return ToolResult{
416 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
417 IsError: true,
418 }
419 }
420
421 args := tc.Arguments
422
423 // Intercept hook: an extension or other guard can refuse the
424 // call before any side effect happens, OR rewrite the args
425 // seen by the tool. The model sees the reason as the tool
426 // error, learns from it, and (typically) proposes a different
427 // action; rewrites are invisible to the model (they apply only
428 // to the execution).
429 if a.BeforeToolExecute != nil {
430 allowed, reason, modified := a.BeforeToolExecute(tc)
431 if !allowed {
432 if reason == "" {
433 reason = "tool call refused by extension guard"
434 }
435 return ToolResult{
436 Content: []provider.Content{provider.TextBlock{Text: reason}},
437 IsError: true,
438 }
439 }
440 if len(modified) > 0 && json.Valid(modified) {
441 args = modified
442 }
443 }
444
445 if len(args) == 0 {
446 args = json.RawMessage("{}")
447 }
448
449 // Recover panics so a buggy tool does not crash the agent.
450 var res ToolResult
451 func() {
452 defer func() {
453 if r := recover(); r != nil {
454 res = ToolResult{
455 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
456 IsError: true,
457 }
458 }
459 }()
460 out, err := tool.Execute(ctx, args, func(text string) {
461 sink(EvToolProgress{ID: tc.ID, Text: text})
462 })
463 if err != nil {
464 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
465 res = ToolResult{
466 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
467 IsError: true,
468 }
469 return
470 }
471 res = ToolResult{
472 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
473 IsError: true,
474 }
475 return
476 }
477 res = out
478 }()
479 return res
480}
481
482// extractText concatenates all TextBlock content in a message. Used
483// by BeforeAssistantMessage so guards see a single string instead of
484// having to walk provider.Content themselves.
485func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
486 var content []provider.Content
487 for _, c := range msg.Content {
488 tr, ok := c.(provider.ToolResultBlock)
489 if !ok {
490 continue
491 }
492 for _, inner := range tr.Content {
493 switch v := inner.(type) {
494 case provider.TextBlock:
495 // Keep short textual context so the model understands why
496 // the images appeared, but don't duplicate giant read
497 // outputs verbatim.
498 if len(v.Text) > 0 && len(v.Text) <= 500 {
499 content = append(content, v)
500 }
501 case provider.ImageBlock:
502 content = append(content, v)
503 }
504 }
505 }
506 if len(content) == 0 {
507 return provider.Message{}
508 }
509 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
510 content = append([]provider.Content{prefix}, content...)
511 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
512}
513
514func extractText(msg provider.Message) string {
515 var out string
516 for _, c := range msg.Content {
517 if tb, ok := c.(provider.TextBlock); ok {
518 if out != "" {
519 out += "\n"
520 }
521 out += tb.Text
522 }
523 }
524 return out
525}
526
527// replaceText returns a copy of msg with every TextBlock replaced by
528// a single TextBlock containing replacement. Non-text content (tool
529// calls, etc.) is preserved in order.
530func replaceText(msg provider.Message, replacement string) provider.Message {
531 out := provider.Message{Role: msg.Role}
532 out.Content = make([]provider.Content, 0, len(msg.Content))
533 replaced := false
534 for _, c := range msg.Content {
535 if _, ok := c.(provider.TextBlock); ok {
536 if !replaced {
537 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
538 replaced = true
539 }
540 continue
541 }
542 out.Content = append(out.Content, c)
543 }
544 if !replaced {
545 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
546 }
547 return out
548}