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