1package core
2
3import (
4 "context"
5 "encoding/json"
6 "testing"
7
8 "github.com/patriceckhart/zot/packages/provider"
9)
10
11// recordingTool captures the args it was invoked with so the test
12// can verify the interceptor-rewritten args reached execution.
13type recordingTool struct {
14 lastArgs json.RawMessage
15}
16
17func (r *recordingTool) Name() string { return "echo" }
18func (r *recordingTool) Description() string { return "echoes" }
19func (r *recordingTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
20
21func (r *recordingTool) Execute(_ context.Context, args json.RawMessage, _ func(string)) (ToolResult, error) {
22 r.lastArgs = append(json.RawMessage(nil), args...)
23 return ToolResult{
24 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
25 }, nil
26}
27
28// TestBeforeToolExecuteModifiesArgs verifies that a non-nil
29// modifiedArgs returned from BeforeToolExecute is what the tool
30// actually sees.
31func TestBeforeToolExecuteModifiesArgs(t *testing.T) {
32 rec := &recordingTool{}
33 reg := Registry{"echo": rec}
34 a := NewAgent(nil, "test", "", reg)
35
36 newArgs := json.RawMessage(`{"command":"echo GUARDED: ls"}`)
37 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
38 return true, "", newArgs
39 }
40
41 ctx := context.Background()
42 res := a.runOneTool(ctx, provider.ToolCallBlock{
43 ID: "T1",
44 Name: "echo",
45 Arguments: json.RawMessage(`{"command":"ls"}`),
46 }, func(AgentEvent) {})
47 if res.IsError {
48 t.Fatalf("unexpected error result: %v", res.Content)
49 }
50 if string(rec.lastArgs) != string(newArgs) {
51 t.Errorf("tool saw %s, want %s", string(rec.lastArgs), string(newArgs))
52 }
53}
54
55// TestBeforeToolExecuteInvalidJSONIgnored verifies that returning
56// malformed JSON as modifiedArgs leaves the original args intact
57// (safety: a buggy interceptor can't corrupt the call).
58func TestBeforeToolExecuteInvalidJSONIgnored(t *testing.T) {
59 rec := &recordingTool{}
60 reg := Registry{"echo": rec}
61 a := NewAgent(nil, "test", "", reg)
62
63 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
64 return true, "", json.RawMessage(`{not json`)
65 }
66
67 ctx := context.Background()
68 orig := json.RawMessage(`{"command":"ls"}`)
69 a.runOneTool(ctx, provider.ToolCallBlock{
70 ID: "T1",
71 Name: "echo",
72 Arguments: orig,
73 }, func(AgentEvent) {})
74 if string(rec.lastArgs) != string(orig) {
75 t.Errorf("tool saw %s, want original %s", string(rec.lastArgs), string(orig))
76 }
77}
78
79// TestBeforeToolExecuteBlockSurfacesReason verifies a refusal from
80// the interceptor returns an error ToolResult with the reason text.
81func TestBeforeToolExecuteBlockSurfacesReason(t *testing.T) {
82 rec := &recordingTool{}
83 reg := Registry{"echo": rec}
84 a := NewAgent(nil, "test", "", reg)
85
86 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
87 return false, "nope", nil
88 }
89
90 ctx := context.Background()
91 res := a.runOneTool(ctx, provider.ToolCallBlock{
92 ID: "T1",
93 Name: "echo",
94 Arguments: json.RawMessage(`{"command":"ls"}`),
95 }, func(AgentEvent) {})
96 if !res.IsError {
97 t.Fatal("want error result, got success")
98 }
99 if len(res.Content) == 0 {
100 t.Fatal("no content")
101 }
102 tb, ok := res.Content[0].(provider.TextBlock)
103 if !ok || tb.Text != "nope" {
104 t.Errorf("want reason 'nope', got %v", res.Content[0])
105 }
106 if rec.lastArgs != nil {
107 t.Error("tool ran despite block")
108 }
109}