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: 50,
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; 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 sink(EvDone{})
211 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
212}
213
214// oneTurn calls the LLM once, forwards events, returns the stop reason
215// and the assembled assistant message (already appended to the transcript).
216func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
217 req := provider.Request{
218 Model: a.Model,
219 System: a.System,
220 Messages: a.Messages(),
221 Tools: a.Tools.Specs(),
222 Reasoning: a.Reasoning,
223 }
224 stream, err := a.Client.Stream(ctx, req)
225 if err != nil {
226 return provider.StopError, provider.Message{}, err
227 }
228
229 sink(EvAssistantStart{})
230
231 var (
232 stop provider.StopReason
233 finalErr error
234 finalMsg provider.Message
235 )
236
237 for ev := range stream {
238 switch e := ev.(type) {
239 case provider.EventStart:
240 // nothing
241 case provider.EventTextDelta:
242 sink(EvTextDelta{Delta: e.Delta})
243 case provider.EventToolStart:
244 sink(EvToolUseStart{ID: e.ID, Name: e.Name})
245 case provider.EventToolArgs:
246 sink(EvToolUseArgs{ID: e.ID, Delta: e.Delta})
247 case provider.EventToolEnd:
248 sink(EvToolUseEnd{ID: e.ID})
249 case provider.EventUsage:
250 cum := a.cost.Add(e.Usage)
251 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
252 case provider.EventDone:
253 stop = e.Stop
254 finalErr = e.Err
255 finalMsg = e.Message
256 }
257 }
258
259 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
260 // produce partial content. When the partial message is text only we
261 // keep whatever was streamed up to the cancel so the user does not
262 // lose visible work (a cut-off summary is still useful). If the
263 // partial message already contained tool-call blocks we drop the
264 // whole thing, because an unmatched tool_use would fail the next
265 // turn with a tool_result mismatch error.
266 keep := len(finalMsg.Content) > 0
267 if stop == provider.StopAborted && keep {
268 hasToolCall := false
269 for _, c := range finalMsg.Content {
270 if _, ok := c.(provider.ToolCallBlock); ok {
271 hasToolCall = true
272 break
273 }
274 }
275 if hasToolCall {
276 keep = false
277 }
278 }
279 if keep {
280 emit := finalMsg
281 suppress := false
282
283 // BeforeAssistantMessage hook: extensions can suppress or
284 // rewrite the visible text. The transcript keeps the
285 // model's original output so the model still sees what it
286 // said on subsequent turns.
287 if a.BeforeAssistantMessage != nil {
288 orig := extractText(finalMsg)
289 if orig != "" {
290 allowed, _, replacement := a.BeforeAssistantMessage(orig)
291 if !allowed {
292 suppress = true
293 } else if replacement != "" && replacement != orig {
294 emit = replaceText(finalMsg, replacement)
295 }
296 }
297 }
298
299 a.mu.Lock()
300 a.messages = append(a.messages, finalMsg)
301 a.mu.Unlock()
302 if !suppress {
303 sink(EvAssistantMessage{Message: emit})
304 }
305 // Now surface tool calls as EvToolCall events so UIs can render them
306 // in order before the tool results arrive.
307 for _, c := range finalMsg.Content {
308 if tc, ok := c.(provider.ToolCallBlock); ok {
309 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
310 }
311 }
312 }
313
314 return stop, finalMsg, finalErr
315}
316
317// executeTools runs every tool call in the assistant message and returns
318// a single tool-role message carrying all results.
319func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
320 var results []provider.Content
321 hadError := false
322
323 for _, c := range msg.Content {
324 tc, ok := c.(provider.ToolCallBlock)
325 if !ok {
326 continue
327 }
328 res := a.runOneTool(ctx, tc, sink)
329 if res.IsError {
330 hadError = true
331 }
332 results = append(results, provider.ToolResultBlock{
333 CallID: tc.ID,
334 Content: res.Content,
335 IsError: res.IsError,
336 })
337 sink(EvToolResult{ID: tc.ID, Result: res})
338 }
339
340 return provider.Message{
341 Role: provider.RoleTool,
342 Content: results,
343 Time: time.Now(),
344 }, hadError
345}
346
347func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
348 tool, err := a.Tools.Get(tc.Name)
349 if err != nil {
350 return ToolResult{
351 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
352 IsError: true,
353 }
354 }
355
356 args := tc.Arguments
357
358 // Intercept hook: an extension or other guard can refuse the
359 // call before any side effect happens, OR rewrite the args
360 // seen by the tool. The model sees the reason as the tool
361 // error, learns from it, and (typically) proposes a different
362 // action; rewrites are invisible to the model (they apply only
363 // to the execution).
364 if a.BeforeToolExecute != nil {
365 allowed, reason, modified := a.BeforeToolExecute(tc)
366 if !allowed {
367 if reason == "" {
368 reason = "tool call refused by extension guard"
369 }
370 return ToolResult{
371 Content: []provider.Content{provider.TextBlock{Text: reason}},
372 IsError: true,
373 }
374 }
375 if len(modified) > 0 && json.Valid(modified) {
376 args = modified
377 }
378 }
379
380 if len(args) == 0 {
381 args = json.RawMessage("{}")
382 }
383
384 // Recover panics so a buggy tool does not crash the agent.
385 var res ToolResult
386 func() {
387 defer func() {
388 if r := recover(); r != nil {
389 res = ToolResult{
390 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
391 IsError: true,
392 }
393 }
394 }()
395 out, err := tool.Execute(ctx, args, func(text string) {
396 sink(EvToolProgress{ID: tc.ID, Text: text})
397 })
398 if err != nil {
399 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
400 res = ToolResult{
401 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
402 IsError: true,
403 }
404 return
405 }
406 res = ToolResult{
407 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
408 IsError: true,
409 }
410 return
411 }
412 res = out
413 }()
414 return res
415}
416
417// extractText concatenates all TextBlock content in a message. Used
418// by BeforeAssistantMessage so guards see a single string instead of
419// having to walk provider.Content themselves.
420func mirrorToolImagesAsUser(msg provider.Message) provider.Message {
421 var content []provider.Content
422 for _, c := range msg.Content {
423 tr, ok := c.(provider.ToolResultBlock)
424 if !ok {
425 continue
426 }
427 for _, inner := range tr.Content {
428 switch v := inner.(type) {
429 case provider.TextBlock:
430 // Keep short textual context so the model understands why
431 // the images appeared, but don't duplicate giant read
432 // outputs verbatim.
433 if len(v.Text) > 0 && len(v.Text) <= 500 {
434 content = append(content, v)
435 }
436 case provider.ImageBlock:
437 content = append(content, v)
438 }
439 }
440 }
441 if len(content) == 0 {
442 return provider.Message{}
443 }
444 prefix := provider.TextBlock{Text: "Tool output included the following image content:"}
445 content = append([]provider.Content{prefix}, content...)
446 return provider.Message{Role: provider.RoleUser, Content: content, Time: time.Now()}
447}
448
449func extractText(msg provider.Message) string {
450 var out string
451 for _, c := range msg.Content {
452 if tb, ok := c.(provider.TextBlock); ok {
453 if out != "" {
454 out += "\n"
455 }
456 out += tb.Text
457 }
458 }
459 return out
460}
461
462// replaceText returns a copy of msg with every TextBlock replaced by
463// a single TextBlock containing replacement. Non-text content (tool
464// calls, etc.) is preserved in order.
465func replaceText(msg provider.Message, replacement string) provider.Message {
466 out := provider.Message{Role: msg.Role}
467 out.Content = make([]provider.Content, 0, len(msg.Content))
468 replaced := false
469 for _, c := range msg.Content {
470 if _, ok := c.(provider.TextBlock); ok {
471 if !replaced {
472 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
473 replaced = true
474 }
475 continue
476 }
477 out.Content = append(out.Content, c)
478 }
479 if !replaced {
480 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
481 }
482 return out
483}