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
28func TestCallToolUsesPolicyWithoutChangingTranscript(t *testing.T) {
29 rec := &recordingTool{}
30 a := NewAgent(nil, "test", "", Registry{"echo": rec})
31 var events []AgentEvent
32 a.BeforeToolExecute = func(provider.ToolCallBlock) (bool, string, json.RawMessage) {
33 return false, "denied", nil
34 }
35 result := a.CallTool(context.Background(), "extension-call", "echo", json.RawMessage(`{}`), func(ev AgentEvent) {
36 events = append(events, ev)
37 })
38 if !result.IsError || rec.lastArgs != nil || len(a.Messages()) != 0 {
39 t.Fatalf("policy or transcript bypassed: result=%+v, args=%s, messages=%v", result, rec.lastArgs, a.Messages())
40 }
41 if len(events) != 2 {
42 t.Fatalf("expected call/result events, got %v", events)
43 }
44 if _, ok := events[0].(EvToolCall); !ok {
45 t.Fatalf("first event = %T", events[0])
46 }
47 if _, ok := events[1].(EvToolResult); !ok {
48 t.Fatalf("second event = %T", events[1])
49 }
50}
51
52// TestBeforeToolExecuteModifiesArgs verifies that a non-nil
53// modifiedArgs returned from BeforeToolExecute is what the tool
54// actually sees.
55func TestBeforeToolExecuteModifiesArgs(t *testing.T) {
56 rec := &recordingTool{}
57 reg := Registry{"echo": rec}
58 a := NewAgent(nil, "test", "", reg)
59
60 newArgs := json.RawMessage(`{"command":"echo GUARDED: ls"}`)
61 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
62 return true, "", newArgs
63 }
64
65 ctx := context.Background()
66 res := a.runOneTool(ctx, provider.ToolCallBlock{
67 ID: "T1",
68 Name: "echo",
69 Arguments: json.RawMessage(`{"command":"ls"}`),
70 }, func(AgentEvent) {})
71 if res.IsError {
72 t.Fatalf("unexpected error result: %v", res.Content)
73 }
74 if string(rec.lastArgs) != string(newArgs) {
75 t.Errorf("tool saw %s, want %s", string(rec.lastArgs), string(newArgs))
76 }
77}
78
79// TestBeforeToolExecuteInvalidJSONIgnored verifies that returning
80// malformed JSON as modifiedArgs leaves the original args intact
81// (safety: a buggy interceptor can't corrupt the call).
82func TestBeforeToolExecuteInvalidJSONIgnored(t *testing.T) {
83 rec := &recordingTool{}
84 reg := Registry{"echo": rec}
85 a := NewAgent(nil, "test", "", reg)
86
87 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
88 return true, "", json.RawMessage(`{not json`)
89 }
90
91 ctx := context.Background()
92 orig := json.RawMessage(`{"command":"ls"}`)
93 a.runOneTool(ctx, provider.ToolCallBlock{
94 ID: "T1",
95 Name: "echo",
96 Arguments: orig,
97 }, func(AgentEvent) {})
98 if string(rec.lastArgs) != string(orig) {
99 t.Errorf("tool saw %s, want original %s", string(rec.lastArgs), string(orig))
100 }
101}
102
103func TestExecuteToolsSkipsServerToolCalls(t *testing.T) {
104 rec := &recordingTool{}
105 a := NewAgent(nil, "test", "", Registry{"echo": rec})
106 msg, hadError := a.executeTools(context.Background(), provider.Message{
107 Role: provider.RoleAssistant,
108 Content: []provider.Content{
109 provider.ToolCallBlock{ID: "srv", Name: "openrouter:datetime", Server: true},
110 },
111 }, func(AgentEvent) {})
112 if hadError {
113 t.Fatal("server tool produced error")
114 }
115 if len(msg.Content) != 0 {
116 t.Fatalf("server tools must not produce client results, got %+v", msg.Content)
117 }
118 if rec.lastArgs != nil {
119 t.Fatal("client tool executed for a server tool call")
120 }
121}
122
123// TestBeforeToolExecuteBlockSurfacesReason verifies a refusal from
124// the interceptor returns an error ToolResult with the reason text.
125func TestBeforeToolExecuteBlockSurfacesReason(t *testing.T) {
126 rec := &recordingTool{}
127 reg := Registry{"echo": rec}
128 a := NewAgent(nil, "test", "", reg)
129
130 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
131 return false, "nope", nil
132 }
133
134 ctx := context.Background()
135 res := a.runOneTool(ctx, provider.ToolCallBlock{
136 ID: "T1",
137 Name: "echo",
138 Arguments: json.RawMessage(`{"command":"ls"}`),
139 }, func(AgentEvent) {})
140 if !res.IsError {
141 t.Fatal("want error result, got success")
142 }
143 if len(res.Content) == 0 {
144 t.Fatal("no content")
145 }
146 tb, ok := res.Content[0].(provider.TextBlock)
147 if !ok || tb.Text != "nope" {
148 t.Errorf("want reason 'nope', got %v", res.Content[0])
149 }
150 if rec.lastArgs != nil {
151 t.Error("tool ran despite block")
152 }
153}