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