1package core
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "testing"
8
9 "github.com/bnema/zut/packages/provider"
10)
11
12// recordingTool captures the args it was invoked with so the test
13// can verify the interceptor-rewritten args reached execution.
14type recordingTool struct {
15 lastArgs json.RawMessage
16}
17
18type resultTool struct {
19 result ToolResult
20 err error
21 panic bool
22}
23
24type immutableResultTool struct{ resultTool }
25
26func (*immutableResultTool) AllowArgumentRewrite() bool { return false }
27
28func (t *resultTool) Name() string { return "result" }
29func (t *resultTool) Description() string { return "result" }
30func (t *resultTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
31func (t *resultTool) Execute(context.Context, json.RawMessage, func(string)) (ToolResult, error) {
32 if t.panic {
33 panic("tool boom")
34 }
35 return t.result, t.err
36}
37
38func (r *recordingTool) Name() string { return "echo" }
39func (r *recordingTool) Description() string { return "echoes" }
40func (r *recordingTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
41
42func (r *recordingTool) Execute(_ context.Context, args json.RawMessage, _ func(string)) (ToolResult, error) {
43 r.lastArgs = append(json.RawMessage(nil), args...)
44 return ToolResult{
45 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
46 }, nil
47}
48
49func TestToolResultPersistenceFailureBecomesToolError(t *testing.T) {
50 agent := NewAgent(nil, "test", "", Registry{
51 "result": &resultTool{result: ToolResult{
52 Content: []provider.Content{provider.TextBlock{Text: "done"}},
53 Details: "durable update",
54 }},
55 })
56 agent.CommitToolResult = func(string, ToolResult) error { return errors.New("disk full") }
57 message := provider.Message{
58 Role: provider.RoleAssistant,
59 Content: []provider.Content{provider.ToolCallBlock{
60 ID: "call-1",
61 Name: "result",
62 Arguments: json.RawMessage(`{}`),
63 }},
64 }
65 var event EvToolResult
66 toolMessage, hadError := agent.executeTools(context.Background(), message, func(ev AgentEvent) {
67 if result, ok := ev.(EvToolResult); ok {
68 event = result
69 }
70 })
71
72 if !hadError {
73 t.Fatal("persistence failure did not mark the tool batch failed")
74 }
75 if !event.Result.IsError || event.Result.Details != nil {
76 t.Fatalf("tool event result = %#v", event.Result)
77 }
78 block, ok := toolMessage.Content[0].(provider.ToolResultBlock)
79 if !ok || !block.IsError {
80 t.Fatalf("tool transcript block = %#v", toolMessage.Content)
81 }
82 if got := block.Content[0].(provider.TextBlock).Text; got != "tool result state could not be persisted" {
83 t.Fatalf("tool error text = %q", got)
84 }
85}
86
87func TestExecuteToolsEmitsPlanUpdateBeforeResult(t *testing.T) {
88 update := PlanUpdate{Plan: []PlanStep{{Step: "Implement", Status: PlanInProgress}}}
89 agent := NewAgent(nil, "test", "", Registry{
90 "result": &resultTool{result: ToolResult{
91 Content: []provider.Content{provider.TextBlock{Text: "Plan updated"}},
92 Details: update,
93 }},
94 })
95 message := provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{
96 provider.ToolCallBlock{ID: "call-1", Name: "result", Arguments: json.RawMessage(`{}`)},
97 }}
98 var events []AgentEvent
99 agent.executeTools(context.Background(), message, func(event AgentEvent) {
100 switch event.(type) {
101 case EvPlanUpdate, EvToolResult:
102 events = append(events, event)
103 }
104 })
105
106 if len(events) != 2 {
107 t.Fatalf("events = %#v, want plan update and tool result", events)
108 }
109 planEvent, ok := events[0].(EvPlanUpdate)
110 if !ok || len(planEvent.Update.Plan) != 1 || planEvent.Update.Plan[0].Step != "Implement" {
111 t.Fatalf("first event = %#v, want plan update", events[0])
112 }
113 if _, ok := events[1].(EvToolResult); !ok {
114 t.Fatalf("second event = %T, want EvToolResult", events[1])
115 }
116}
117
118func TestToolResultCommitErrorUsesSafeMessage(t *testing.T) {
119 agent := NewAgent(nil, "test", "", Registry{
120 "result": &resultTool{result: ToolResult{Content: []provider.Content{provider.TextBlock{Text: "done"}}}},
121 })
122 agent.CommitToolResult = func(string, ToolResult) error {
123 return &ToolResultCommitError{Message: "goal transition rejected: the current goal is still active", Err: errors.New("database path /private/session is unavailable")}
124 }
125 message := provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.ToolCallBlock{ID: "call-1", Name: "result", Arguments: json.RawMessage(`{}`)}}}
126 toolMessage, hadError := agent.executeTools(context.Background(), message, func(AgentEvent) {})
127 if !hadError {
128 t.Fatal("commit error did not mark the tool batch failed")
129 }
130 block := toolMessage.Content[0].(provider.ToolResultBlock)
131 if got := block.Content[0].(provider.TextBlock).Text; got != "goal transition rejected: the current goal is still active" {
132 t.Fatalf("tool error text = %q", got)
133 }
134}
135
136func TestBeforeToolExecuteRejectsRewriteForImmutableTool(t *testing.T) {
137 tool := &immutableResultTool{}
138 agent := NewAgent(nil, "test", "", Registry{"result": tool})
139 agent.BeforeToolExecute = func(provider.ToolCallBlock) (bool, string, json.RawMessage) {
140 return true, "", json.RawMessage(`{"rewritten":true}`)
141 }
142
143 result := agent.runOneTool(context.Background(), provider.ToolCallBlock{
144 ID: "T1", Name: "result", Arguments: json.RawMessage(`{"original":true}`),
145 }, agent.ToolsSnapshot(), func(AgentEvent) {})
146 if !result.IsError {
147 t.Fatalf("result = %#v, want rewrite error", result)
148 }
149 if got := result.Content[0].(provider.TextBlock).Text; got != "tool arguments cannot be rewritten" {
150 t.Fatalf("rewrite error = %q", got)
151 }
152}
153
154// TestBeforeToolExecuteModifiesArgs verifies that a non-nil modifiedArgs
155// returned from BeforeToolExecute is what the tool actually sees.
156func TestBeforeToolExecuteModifiesArgs(t *testing.T) {
157 rec := &recordingTool{}
158 reg := Registry{"echo": rec}
159 a := NewAgent(nil, "test", "", reg)
160
161 newArgs := json.RawMessage(`{"command":"echo GUARDED: ls"}`)
162 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
163 return true, "", newArgs
164 }
165
166 ctx := context.Background()
167 tools := a.ToolsSnapshot()
168 res := a.runOneTool(ctx, provider.ToolCallBlock{
169 ID: "T1",
170 Name: "echo",
171 Arguments: json.RawMessage(`{"command":"ls"}`),
172 }, tools, func(AgentEvent) {})
173 if res.IsError {
174 t.Fatalf("unexpected error result: %v", res.Content)
175 }
176 assertToolTiming(t, res.Timing)
177 if string(rec.lastArgs) != string(newArgs) {
178 t.Errorf("tool saw %s, want %s", string(rec.lastArgs), string(newArgs))
179 }
180}
181
182// TestBeforeToolExecuteInvalidJSONIgnored verifies that returning
183// malformed JSON as modifiedArgs leaves the original args intact
184// (safety: a buggy interceptor can't corrupt the call).
185func TestBeforeToolExecuteInvalidJSONIgnored(t *testing.T) {
186 rec := &recordingTool{}
187 reg := Registry{"echo": rec}
188 a := NewAgent(nil, "test", "", reg)
189
190 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
191 return true, "", json.RawMessage(`{not json`)
192 }
193
194 ctx := context.Background()
195 tools := a.ToolsSnapshot()
196 orig := json.RawMessage(`{"command":"ls"}`)
197 a.runOneTool(ctx, provider.ToolCallBlock{
198 ID: "T1",
199 Name: "echo",
200 Arguments: orig,
201 }, tools, func(AgentEvent) {})
202 if string(rec.lastArgs) != string(orig) {
203 t.Errorf("tool saw %s, want original %s", string(rec.lastArgs), string(orig))
204 }
205}
206
207func TestRunOneToolEmitsExecutionStartedAfterApproval(t *testing.T) {
208 rec := &recordingTool{}
209 a := NewAgent(nil, "test", "", Registry{"echo": rec})
210
211 var started []EvToolExecutionStarted
212 res := a.runOneTool(context.Background(), provider.ToolCallBlock{
213 ID: "T1", Name: "echo", Arguments: json.RawMessage(`{"command":"echo ok"}`),
214 }, a.ToolsSnapshot(), func(ev AgentEvent) {
215 if e, ok := ev.(EvToolExecutionStarted); ok {
216 started = append(started, e)
217 }
218 })
219 if res.IsError {
220 t.Fatalf("runOneTool error = %#v", res)
221 }
222 if len(started) != 1 || started[0].ID != "T1" || started[0].Name != "echo" {
223 t.Fatalf("execution events = %#v", started)
224 }
225}
226
227func TestRunOneToolDoesNotEmitExecutionStartedForMissingOrDeniedTools(t *testing.T) {
228 for _, tc := range []struct {
229 name string
230 tools Registry
231 guard func(provider.ToolCallBlock) (bool, string, json.RawMessage)
232 }{
233 {name: "missing tool", tools: Registry{}},
234 {
235 name: "denied tool",
236 tools: Registry{"echo": &recordingTool{}},
237 guard: func(provider.ToolCallBlock) (bool, string, json.RawMessage) { return false, "denied", nil },
238 },
239 } {
240 t.Run(tc.name, func(t *testing.T) {
241 a := NewAgent(nil, "test", "", tc.tools)
242 a.BeforeToolExecute = tc.guard
243 var started int
244 a.runOneTool(context.Background(), provider.ToolCallBlock{ID: "T1", Name: "echo"}, a.ToolsSnapshot(), func(event AgentEvent) {
245 if _, ok := event.(EvToolExecutionStarted); ok {
246 started++
247 }
248 })
249 if started != 0 {
250 t.Fatalf("execution-started events = %d, want none", started)
251 }
252 })
253 }
254}
255
256// TestBeforeToolExecuteBlockSurfacesReason verifies a refusal from
257// the interceptor returns an error ToolResult with the reason text.
258func TestBeforeToolExecuteBlockSurfacesReason(t *testing.T) {
259 rec := &recordingTool{}
260 reg := Registry{"echo": rec}
261 a := NewAgent(nil, "test", "", reg)
262
263 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
264 return false, "nope", nil
265 }
266
267 ctx := context.Background()
268 tools := a.ToolsSnapshot()
269 res := a.runOneTool(ctx, provider.ToolCallBlock{
270 ID: "T1",
271 Name: "echo",
272 Arguments: json.RawMessage(`{"command":"ls"}`),
273 }, tools, func(AgentEvent) {})
274 if !res.IsError {
275 t.Fatal("want error result, got success")
276 }
277 if len(res.Content) == 0 {
278 t.Fatal("no content")
279 }
280 tb, ok := res.Content[0].(provider.TextBlock)
281 if !ok || tb.Text != "nope" {
282 t.Errorf("want reason 'nope', got %v", res.Content[0])
283 }
284 if rec.lastArgs != nil {
285 t.Error("tool ran despite block")
286 }
287 assertToolTiming(t, res.Timing)
288}
289
290func TestRunOneToolTimingCoversOutcomes(t *testing.T) {
291 cancelled, cancel := context.WithCancel(context.Background())
292 cancel()
293 cases := []struct {
294 name string
295 tools Registry
296 ctx context.Context
297 want string
298 }{
299 {name: "success", tools: Registry{"result": &resultTool{result: ToolResult{Content: []provider.Content{provider.TextBlock{Text: "ok"}}}}}, ctx: context.Background()},
300 {name: "ordinary error", tools: Registry{"result": &resultTool{err: errors.New("ordinary")}}, ctx: context.Background()},
301 {name: "cancellation", tools: Registry{"result": &resultTool{err: context.Canceled}}, ctx: cancelled},
302 {name: "missing", tools: Registry{}, ctx: context.Background()},
303 {name: "panic", tools: Registry{"result": &resultTool{panic: true}}, ctx: context.Background()},
304 }
305 for _, tc := range cases {
306 t.Run(tc.name, func(t *testing.T) {
307 a := NewAgent(nil, "test", "", tc.tools)
308 callName := "result"
309 if tc.name == "missing" {
310 callName = "missing"
311 }
312 res := a.runOneTool(tc.ctx, provider.ToolCallBlock{ID: "timed", Name: callName}, tc.tools, func(AgentEvent) {})
313 assertToolTiming(t, res.Timing)
314 if tc.name == "panic" {
315 text, ok := res.Content[0].(provider.TextBlock)
316 if !ok || text.Text != "tool execution failed" {
317 t.Fatalf("panic result = %#v, want generic tool execution error", res.Content)
318 }
319 }
320 })
321 }
322}
323
324func assertToolTiming(t *testing.T, timing *provider.ToolTiming) {
325 t.Helper()
326 if timing == nil {
327 t.Fatal("tool result has no timing")
328 }
329 if timing.StartedAt.IsZero() || timing.CompletedAt.IsZero() {
330 t.Fatalf("tool timing has zero wall-clock bound: %#v", timing)
331 }
332 if timing.CompletedAt.Before(timing.StartedAt) {
333 t.Fatalf("tool completed before start: %#v", timing)
334 }
335 if timing.Duration < 0 {
336 t.Fatalf("tool duration = %s, want non-negative", timing.Duration)
337 }
338}