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 a.mu.Unlock()
185 // If context was cancelled during tool execution, bail out.
186 if err := ctx.Err(); err != nil {
187 sink(EvDone{})
188 return err
189 }
190 _ = hadError
191 continue
192 }
193
194 // Terminal stop (end, length, error, aborted).
195 sink(EvDone{})
196 return nil
197 }
198 sink(EvDone{})
199 return fmt.Errorf("max steps (%d) exceeded", a.MaxSteps)
200}
201
202// oneTurn calls the LLM once, forwards events, returns the stop reason
203// and the assembled assistant message (already appended to the transcript).
204func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.StopReason, provider.Message, error) {
205 req := provider.Request{
206 Model: a.Model,
207 System: a.System,
208 Messages: a.Messages(),
209 Tools: a.Tools.Specs(),
210 Reasoning: a.Reasoning,
211 }
212 stream, err := a.Client.Stream(ctx, req)
213 if err != nil {
214 return provider.StopError, provider.Message{}, err
215 }
216
217 sink(EvAssistantStart{})
218
219 var (
220 stop provider.StopReason
221 finalErr error
222 finalMsg provider.Message
223 )
224
225 for ev := range stream {
226 switch e := ev.(type) {
227 case provider.EventStart:
228 // nothing
229 case provider.EventTextDelta:
230 sink(EvTextDelta{Delta: e.Delta})
231 case provider.EventToolStart, provider.EventToolArgs, provider.EventToolEnd:
232 // The provider emits these to support partial-JSON UIs.
233 // We surface the complete tool call once per block at EvToolCall
234 // below (after EventDone), so for now we drop them.
235 case provider.EventUsage:
236 cum := a.cost.Add(e.Usage)
237 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
238 case provider.EventDone:
239 stop = e.Stop
240 finalErr = e.Err
241 finalMsg = e.Message
242 }
243 }
244
245 // Append assistant message to transcript. Aborted turns (Esc / Ctrl+C)
246 // produce partial, mid-sentence content that would confuse subsequent
247 // turns if it stayed in the transcript, drop it instead.
248 if len(finalMsg.Content) > 0 && stop != provider.StopAborted {
249 emit := finalMsg
250 suppress := false
251
252 // BeforeAssistantMessage hook: extensions can suppress or
253 // rewrite the visible text. The transcript keeps the
254 // model's original output so the model still sees what it
255 // said on subsequent turns.
256 if a.BeforeAssistantMessage != nil {
257 orig := extractText(finalMsg)
258 if orig != "" {
259 allowed, _, replacement := a.BeforeAssistantMessage(orig)
260 if !allowed {
261 suppress = true
262 } else if replacement != "" && replacement != orig {
263 emit = replaceText(finalMsg, replacement)
264 }
265 }
266 }
267
268 a.mu.Lock()
269 a.messages = append(a.messages, finalMsg)
270 a.mu.Unlock()
271 if !suppress {
272 sink(EvAssistantMessage{Message: emit})
273 }
274 // Now surface tool calls as EvToolCall events so UIs can render them
275 // in order before the tool results arrive.
276 for _, c := range finalMsg.Content {
277 if tc, ok := c.(provider.ToolCallBlock); ok {
278 sink(EvToolCall{ID: tc.ID, Name: tc.Name, Args: tc.Arguments})
279 }
280 }
281 }
282
283 return stop, finalMsg, finalErr
284}
285
286// executeTools runs every tool call in the assistant message and returns
287// a single tool-role message carrying all results.
288func (a *Agent) executeTools(ctx context.Context, msg provider.Message, sink func(AgentEvent)) (provider.Message, bool) {
289 var results []provider.Content
290 hadError := false
291
292 for _, c := range msg.Content {
293 tc, ok := c.(provider.ToolCallBlock)
294 if !ok {
295 continue
296 }
297 res := a.runOneTool(ctx, tc, sink)
298 if res.IsError {
299 hadError = true
300 }
301 results = append(results, provider.ToolResultBlock{
302 CallID: tc.ID,
303 Content: res.Content,
304 IsError: res.IsError,
305 })
306 sink(EvToolResult{ID: tc.ID, Result: res})
307 }
308
309 return provider.Message{
310 Role: provider.RoleTool,
311 Content: results,
312 Time: time.Now(),
313 }, hadError
314}
315
316func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink func(AgentEvent)) ToolResult {
317 tool, err := a.Tools.Get(tc.Name)
318 if err != nil {
319 return ToolResult{
320 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
321 IsError: true,
322 }
323 }
324
325 args := tc.Arguments
326
327 // Intercept hook: an extension or other guard can refuse the
328 // call before any side effect happens, OR rewrite the args
329 // seen by the tool. The model sees the reason as the tool
330 // error, learns from it, and (typically) proposes a different
331 // action; rewrites are invisible to the model (they apply only
332 // to the execution).
333 if a.BeforeToolExecute != nil {
334 allowed, reason, modified := a.BeforeToolExecute(tc)
335 if !allowed {
336 if reason == "" {
337 reason = "tool call refused by extension guard"
338 }
339 return ToolResult{
340 Content: []provider.Content{provider.TextBlock{Text: reason}},
341 IsError: true,
342 }
343 }
344 if len(modified) > 0 && json.Valid(modified) {
345 args = modified
346 }
347 }
348
349 if len(args) == 0 {
350 args = json.RawMessage("{}")
351 }
352
353 // Recover panics so a buggy tool does not crash the agent.
354 var res ToolResult
355 func() {
356 defer func() {
357 if r := recover(); r != nil {
358 res = ToolResult{
359 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("panic: %v", r)}},
360 IsError: true,
361 }
362 }
363 }()
364 out, err := tool.Execute(ctx, args, func(text string) {
365 sink(EvToolProgress{ID: tc.ID, Text: text})
366 })
367 if err != nil {
368 if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
369 res = ToolResult{
370 Content: []provider.Content{provider.TextBlock{Text: "aborted: " + err.Error()}},
371 IsError: true,
372 }
373 return
374 }
375 res = ToolResult{
376 Content: []provider.Content{provider.TextBlock{Text: err.Error()}},
377 IsError: true,
378 }
379 return
380 }
381 res = out
382 }()
383 return res
384}
385
386// extractText concatenates all TextBlock content in a message. Used
387// by BeforeAssistantMessage so guards see a single string instead of
388// having to walk provider.Content themselves.
389func extractText(msg provider.Message) string {
390 var out string
391 for _, c := range msg.Content {
392 if tb, ok := c.(provider.TextBlock); ok {
393 if out != "" {
394 out += "\n"
395 }
396 out += tb.Text
397 }
398 }
399 return out
400}
401
402// replaceText returns a copy of msg with every TextBlock replaced by
403// a single TextBlock containing replacement. Non-text content (tool
404// calls, etc.) is preserved in order.
405func replaceText(msg provider.Message, replacement string) provider.Message {
406 out := provider.Message{Role: msg.Role}
407 out.Content = make([]provider.Content, 0, len(msg.Content))
408 replaced := false
409 for _, c := range msg.Content {
410 if _, ok := c.(provider.TextBlock); ok {
411 if !replaced {
412 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
413 replaced = true
414 }
415 continue
416 }
417 out.Content = append(out.Content, c)
418 }
419 if !replaced {
420 out.Content = append(out.Content, provider.TextBlock{Text: replacement})
421 }
422 return out
423}