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 agent := NewAgent(nil, "test", "", Registry{
89 "result": &resultTool{result: ToolResult{
90 Content: []provider.Content{provider.TextBlock{Text: "Plan updated"}},
91 Details: PlanOperation{Action: "set", Steps: []PlanStep{{Step: "Implement", Status: PlanInProgress}}},
92 }},
93 })
94 message := provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{
95 provider.ToolCallBlock{ID: "call-1", Name: "result", Arguments: json.RawMessage(`{}`)},
96 }}
97 var events []AgentEvent
98 agent.executeTools(context.Background(), message, func(event AgentEvent) {
99 switch event.(type) {
100 case EvPlanUpdate, EvToolResult:
101 events = append(events, event)
102 }
103 })
104
105 if len(events) != 2 {
106 t.Fatalf("events = %#v, want plan update and tool result", events)
107 }
108 planEvent, ok := events[0].(EvPlanUpdate)
109 if !ok || planEvent.CallID != "call-1" || len(planEvent.Update.Plan) != 1 || planEvent.Update.Plan[0].Step != "Implement" {
110 t.Fatalf("first event = %#v, want plan update", events[0])
111 }
112 if _, ok := events[1].(EvToolResult); !ok {
113 t.Fatalf("second event = %T, want EvToolResult", events[1])
114 }
115}
116
117func TestToolResultCommitErrorUsesSafeMessage(t *testing.T) {
118 agent := NewAgent(nil, "test", "", Registry{
119 "result": &resultTool{result: ToolResult{Content: []provider.Content{provider.TextBlock{Text: "done"}}}},
120 })
121 agent.CommitToolResult = func(string, ToolResult) error {
122 return &ToolResultCommitError{Message: "goal transition rejected: the current goal is still active", Err: errors.New("database path /private/session is unavailable")}
123 }
124 message := provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.ToolCallBlock{ID: "call-1", Name: "result", Arguments: json.RawMessage(`{}`)}}}
125 toolMessage, hadError := agent.executeTools(context.Background(), message, func(AgentEvent) {})
126 if !hadError {
127 t.Fatal("commit error did not mark the tool batch failed")
128 }
129 block := toolMessage.Content[0].(provider.ToolResultBlock)
130 if got := block.Content[0].(provider.TextBlock).Text; got != "goal transition rejected: the current goal is still active" {
131 t.Fatalf("tool error text = %q", got)
132 }
133}
134
135func TestBeforeToolExecuteRejectsRewriteForImmutableTool(t *testing.T) {
136 tool := &immutableResultTool{}
137 agent := NewAgent(nil, "test", "", Registry{"result": tool})
138 agent.BeforeToolExecute = func(provider.ToolCallBlock) (bool, string, json.RawMessage) {
139 return true, "", json.RawMessage(`{"rewritten":true}`)
140 }
141
142 result := agent.runOneTool(context.Background(), provider.ToolCallBlock{
143 ID: "T1", Name: "result", Arguments: json.RawMessage(`{"original":true}`),
144 }, agent.ToolsSnapshot(), func(AgentEvent) {})
145 if !result.IsError {
146 t.Fatalf("result = %#v, want rewrite error", result)
147 }
148 if got := result.Content[0].(provider.TextBlock).Text; got != "tool arguments cannot be rewritten" {
149 t.Fatalf("rewrite error = %q", got)
150 }
151}
152
153// TestBeforeToolExecuteModifiesArgs verifies that a non-nil modifiedArgs
154// returned from BeforeToolExecute is what the tool actually sees.
155func TestBeforeToolExecuteModifiesArgs(t *testing.T) {
156 rec := &recordingTool{}
157 reg := Registry{"echo": rec}
158 a := NewAgent(nil, "test", "", reg)
159
160 newArgs := json.RawMessage(`{"command":"echo GUARDED: ls"}`)
161 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
162 return true, "", newArgs
163 }
164
165 ctx := context.Background()
166 tools := a.ToolsSnapshot()
167 res := a.runOneTool(ctx, provider.ToolCallBlock{
168 ID: "T1",
169 Name: "echo",
170 Arguments: json.RawMessage(`{"command":"ls"}`),
171 }, tools, func(AgentEvent) {})
172 if res.IsError {
173 t.Fatalf("unexpected error result: %v", res.Content)
174 }
175 assertToolTiming(t, res.Timing)
176 if string(rec.lastArgs) != string(newArgs) {
177 t.Errorf("tool saw %s, want %s", string(rec.lastArgs), string(newArgs))
178 }
179}
180
181// TestBeforeToolExecuteInvalidJSONIgnored verifies that returning
182// malformed JSON as modifiedArgs leaves the original args intact
183// (safety: a buggy interceptor can't corrupt the call).
184func TestBeforeToolExecuteInvalidJSONIgnored(t *testing.T) {
185 rec := &recordingTool{}
186 reg := Registry{"echo": rec}
187 a := NewAgent(nil, "test", "", reg)
188
189 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
190 return true, "", json.RawMessage(`{not json`)
191 }
192
193 ctx := context.Background()
194 tools := a.ToolsSnapshot()
195 orig := json.RawMessage(`{"command":"ls"}`)
196 a.runOneTool(ctx, provider.ToolCallBlock{
197 ID: "T1",
198 Name: "echo",
199 Arguments: orig,
200 }, tools, func(AgentEvent) {})
201 if string(rec.lastArgs) != string(orig) {
202 t.Errorf("tool saw %s, want original %s", string(rec.lastArgs), string(orig))
203 }
204}
205
206func TestRunOneToolEmitsExecutionStartedAfterApproval(t *testing.T) {
207 rec := &recordingTool{}
208 a := NewAgent(nil, "test", "", Registry{"echo": rec})
209
210 var started []EvToolExecutionStarted
211 res := a.runOneTool(context.Background(), provider.ToolCallBlock{
212 ID: "T1", Name: "echo", Arguments: json.RawMessage(`{"command":"echo ok"}`),
213 }, a.ToolsSnapshot(), func(ev AgentEvent) {
214 if e, ok := ev.(EvToolExecutionStarted); ok {
215 started = append(started, e)
216 }
217 })
218 if res.IsError {
219 t.Fatalf("runOneTool error = %#v", res)
220 }
221 if len(started) != 1 || started[0].ID != "T1" || started[0].Name != "echo" {
222 t.Fatalf("execution events = %#v", started)
223 }
224}
225
226func TestRunOneToolDoesNotEmitExecutionStartedForMissingOrDeniedTools(t *testing.T) {
227 for _, tc := range []struct {
228 name string
229 tools Registry
230 guard func(provider.ToolCallBlock) (bool, string, json.RawMessage)
231 }{
232 {name: "missing tool", tools: Registry{}},
233 {
234 name: "denied tool",
235 tools: Registry{"echo": &recordingTool{}},
236 guard: func(provider.ToolCallBlock) (bool, string, json.RawMessage) { return false, "denied", nil },
237 },
238 } {
239 t.Run(tc.name, func(t *testing.T) {
240 a := NewAgent(nil, "test", "", tc.tools)
241 a.BeforeToolExecute = tc.guard
242 var started int
243 a.runOneTool(context.Background(), provider.ToolCallBlock{ID: "T1", Name: "echo"}, a.ToolsSnapshot(), func(event AgentEvent) {
244 if _, ok := event.(EvToolExecutionStarted); ok {
245 started++
246 }
247 })
248 if started != 0 {
249 t.Fatalf("execution-started events = %d, want none", started)
250 }
251 })
252 }
253}
254
255// TestBeforeToolExecuteBlockSurfacesReason verifies a refusal from
256// the interceptor returns an error ToolResult with the reason text.
257func TestBeforeToolExecuteBlockSurfacesReason(t *testing.T) {
258 rec := &recordingTool{}
259 reg := Registry{"echo": rec}
260 a := NewAgent(nil, "test", "", reg)
261
262 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
263 return false, "nope", nil
264 }
265
266 ctx := context.Background()
267 tools := a.ToolsSnapshot()
268 res := a.runOneTool(ctx, provider.ToolCallBlock{
269 ID: "T1",
270 Name: "echo",
271 Arguments: json.RawMessage(`{"command":"ls"}`),
272 }, tools, func(AgentEvent) {})
273 if !res.IsError {
274 t.Fatal("want error result, got success")
275 }
276 if len(res.Content) == 0 {
277 t.Fatal("no content")
278 }
279 tb, ok := res.Content[0].(provider.TextBlock)
280 if !ok || tb.Text != "nope" {
281 t.Errorf("want reason 'nope', got %v", res.Content[0])
282 }
283 if rec.lastArgs != nil {
284 t.Error("tool ran despite block")
285 }
286 assertToolTiming(t, res.Timing)
287}
288
289func TestRunOneToolTimingCoversOutcomes(t *testing.T) {
290 cancelled, cancel := context.WithCancel(context.Background())
291 cancel()
292 cases := []struct {
293 name string
294 tools Registry
295 ctx context.Context
296 want string
297 }{
298 {name: "success", tools: Registry{"result": &resultTool{result: ToolResult{Content: []provider.Content{provider.TextBlock{Text: "ok"}}}}}, ctx: context.Background()},
299 {name: "ordinary error", tools: Registry{"result": &resultTool{err: errors.New("ordinary")}}, ctx: context.Background()},
300 {name: "cancellation", tools: Registry{"result": &resultTool{err: context.Canceled}}, ctx: cancelled},
301 {name: "missing", tools: Registry{}, ctx: context.Background()},
302 {name: "panic", tools: Registry{"result": &resultTool{panic: true}}, ctx: context.Background()},
303 }
304 for _, tc := range cases {
305 t.Run(tc.name, func(t *testing.T) {
306 a := NewAgent(nil, "test", "", tc.tools)
307 callName := "result"
308 if tc.name == "missing" {
309 callName = "missing"
310 }
311 res := a.runOneTool(tc.ctx, provider.ToolCallBlock{ID: "timed", Name: callName}, tc.tools, func(AgentEvent) {})
312 assertToolTiming(t, res.Timing)
313 if tc.name == "panic" {
314 text, ok := res.Content[0].(provider.TextBlock)
315 if !ok || text.Text != "tool execution failed" {
316 t.Fatalf("panic result = %#v, want generic tool execution error", res.Content)
317 }
318 }
319 })
320 }
321}
322
323func assertToolTiming(t *testing.T, timing *provider.ToolTiming) {
324 t.Helper()
325 if timing == nil {
326 t.Fatal("tool result has no timing")
327 }
328 if timing.StartedAt.IsZero() || timing.CompletedAt.IsZero() {
329 t.Fatalf("tool timing has zero wall-clock bound: %#v", timing)
330 }
331 if timing.CompletedAt.Before(timing.StartedAt) {
332 t.Fatalf("tool completed before start: %#v", timing)
333 }
334 if timing.Duration < 0 {
335 t.Fatalf("tool duration = %s, want non-negative", timing.Duration)
336 }
337}