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