1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "testing"
8
9 "github.com/bnema/zut/packages/agent/tools"
10 "github.com/bnema/zut/packages/core"
11 "github.com/bnema/zut/packages/provider"
12)
13
14type recordingWebSearchTool struct {
15 executions int
16}
17
18func (t *recordingWebSearchTool) Name() string { return "web_search" }
19func (t *recordingWebSearchTool) Description() string { return "test web search" }
20func (t *recordingWebSearchTool) Schema() json.RawMessage {
21 return json.RawMessage(`{"type":"object"}`)
22}
23func (t *recordingWebSearchTool) Preview(context.Context, json.RawMessage) (core.ToolResult, error) {
24 return core.ToolResult{}, nil
25}
26func (t *recordingWebSearchTool) Execute(context.Context, json.RawMessage, func(string)) (core.ToolResult, error) {
27 t.executions++
28 return core.ToolResult{}, nil
29}
30
31type executeOnlyWebSearchTool struct {
32 executions int
33}
34
35func (t *executeOnlyWebSearchTool) Name() string { return "web_search" }
36func (t *executeOnlyWebSearchTool) Description() string { return "test web search" }
37func (t *executeOnlyWebSearchTool) Schema() json.RawMessage {
38 return json.RawMessage(`{"type":"object"}`)
39}
40func (t *executeOnlyWebSearchTool) Execute(context.Context, json.RawMessage, func(string)) (core.ToolResult, error) {
41 t.executions++
42 return core.ToolResult{}, nil
43}
44
45func TestWebSearchSessionGuardPreservesPreviewCapability(t *testing.T) {
46 guard := &webSearchSessionGuard{}
47 previewable := &recordingWebSearchTool{}
48 registry := guard.wrapRegistry(core.Registry{"web_search": previewable})
49 wrapped := registry["web_search"]
50 previewer, ok := wrapped.(core.ToolPreviewer)
51 if !ok {
52 t.Fatalf("preview-capable tool wrapper = %T, want core.ToolPreviewer", wrapped)
53 }
54 if _, err := previewer.Preview(context.Background(), json.RawMessage(`{}`)); !errors.Is(err, errWebSearchSessionRevoked) {
55 t.Fatalf("revoked Preview error = %v, want session revocation", err)
56 }
57 if got := guard.wrapRegistry(registry)["web_search"]; got != wrapped {
58 t.Fatalf("re-wrapped preview tool = %T, want existing wrapper", got)
59 }
60
61 executeOnly := &executeOnlyWebSearchTool{}
62 registry = guard.wrapRegistry(core.Registry{"web_search": executeOnly})
63 if _, ok := registry["web_search"].(core.ToolPreviewer); ok {
64 t.Fatalf("execute-only tool wrapper = %T, unexpectedly advertises core.ToolPreviewer", registry["web_search"])
65 }
66}
67
68func TestWebSearchSessionGuardRevokesSnapshottedToolBeforeExecute(t *testing.T) {
69 guard := &webSearchSessionGuard{}
70 underlying := &recordingWebSearchTool{}
71 registry := guard.wrapRegistry(core.Registry{"web_search": underlying})
72 guard.setAvailable(true)
73
74 // This reference models core's per-execution registry snapshot.
75 snapshotted := registry["web_search"]
76 guard.setAvailable(false)
77 delete(registry, "web_search")
78
79 if _, err := snapshotted.Execute(context.Background(), json.RawMessage(`{}`), nil); !errors.Is(err, errWebSearchSessionRevoked) {
80 t.Fatalf("snapshotted Execute error = %v, want revocation", err)
81 }
82 if underlying.executions != 0 {
83 t.Fatalf("revoked underlying tool executed %d times", underlying.executions)
84 }
85}
86
87type snapshottedWebSearchClient struct {
88 calls int
89}
90
91func (c *snapshottedWebSearchClient) Name() string { return "snapshot-test" }
92
93func (c *snapshottedWebSearchClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
94 c.calls++
95 events := make(chan provider.Event, 2)
96 if c.calls == 1 {
97 events <- provider.EventStart{Provider: c.Name(), Model: req.Model}
98 events <- provider.EventDone{Stop: provider.StopToolUse, Message: provider.Message{
99 Role: provider.RoleAssistant,
100 Content: []provider.Content{
101 provider.ToolCallBlock{ID: "revoke", Name: "revoke_web_search", Arguments: json.RawMessage(`{}`)},
102 provider.ToolCallBlock{ID: "search", Name: "web_search", Arguments: json.RawMessage(`{}`)},
103 },
104 }}
105 } else {
106 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
107 Role: provider.RoleAssistant,
108 Content: []provider.Content{provider.TextBlock{Text: "done"}},
109 }}
110 }
111 close(events)
112 return events, nil
113}
114
115type webSearchRevokerTool struct {
116 guard *webSearchSessionGuard
117 agent *core.Agent
118}
119
120func (t *webSearchRevokerTool) Name() string { return "revoke_web_search" }
121func (t *webSearchRevokerTool) Description() string { return "revoke web search" }
122func (t *webSearchRevokerTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
123func (t *webSearchRevokerTool) Execute(context.Context, json.RawMessage, func(string)) (core.ToolResult, error) {
124 t.guard.setAvailable(false)
125 registry := t.agent.ToolsSnapshot()
126 delete(registry, "web_search")
127 t.agent.SetTools(registry)
128 return core.ToolResult{}, nil
129}
130
131func TestWebSearchSessionGuardBlocksCoreSnapshotWithoutConfirmation(t *testing.T) {
132 guard := &webSearchSessionGuard{}
133 webSearch := &recordingWebSearchTool{}
134 revoker := &webSearchRevokerTool{guard: guard}
135 registry := guard.wrapRegistry(core.NewRegistry(revoker, webSearch))
136 guard.setAvailable(true)
137 agent := core.NewAgent(&snapshottedWebSearchClient{}, "model", "", registry)
138 revoker.agent = agent
139
140 // BeforeToolExecute is deliberately nil: this exercises the yolo/no-
141 // confirmation path after core has snapshotted both tools for the batch.
142 if err := agent.Prompt(context.Background(), "search", nil, nil); err != nil {
143 t.Fatalf("Prompt: %v", err)
144 }
145 if webSearch.executions != 0 {
146 t.Fatalf("snapshotted web_search executed %d times after revocation", webSearch.executions)
147 }
148}
149
150func TestWebSearchSessionGuardRefreshDoesNotReviveOldGeneration(t *testing.T) {
151 guard := &webSearchSessionGuard{}
152 oldUnderlying := &recordingWebSearchTool{}
153 oldRegistry := guard.wrapRegistry(core.Registry{"web_search": oldUnderlying})
154 guard.setAvailable(true)
155 oldSnapshot := oldRegistry["web_search"]
156
157 guard.setAvailable(false)
158 freshUnderlying := &recordingWebSearchTool{}
159 freshRegistry := guard.wrapRegistry(core.Registry{"web_search": freshUnderlying})
160 guard.setAvailable(true)
161
162 if _, err := oldSnapshot.Execute(context.Background(), json.RawMessage(`{}`), nil); !errors.Is(err, errWebSearchSessionRevoked) {
163 t.Fatalf("old generation Execute error = %v, want revocation", err)
164 }
165 if _, err := freshRegistry["web_search"].Execute(context.Background(), json.RawMessage(`{}`), nil); err != nil {
166 t.Fatalf("fresh generation Execute error: %v", err)
167 }
168 if oldUnderlying.executions != 0 || freshUnderlying.executions != 1 {
169 t.Fatalf("execution counts old=%d fresh=%d, want 0 and 1", oldUnderlying.executions, freshUnderlying.executions)
170 }
171}
172
173func TestWebSearchToolAllowedForSessionHonorsInvocationCeilings(t *testing.T) {
174 cases := []struct {
175 name string
176 args Args
177 wantAllowed bool
178 wantExplicit bool
179 }{
180 {name: "default", args: Args{}, wantAllowed: true},
181 {name: "no tools", args: Args{NoTools: true}, wantAllowed: false},
182 {name: "permission set", args: Args{PermissionSet: &tools.PermissionSet{}}, wantAllowed: false},
183 {name: "explicit web search", args: Args{ToolsSet: true, Tools: []string{"web_search"}}, wantAllowed: true, wantExplicit: true},
184 {name: "programmatic explicit web search", args: Args{Tools: []string{"web_search"}}, wantAllowed: true, wantExplicit: true},
185 {name: "no tools overrides explicit web search", args: Args{NoTools: true, ToolsSet: true, Tools: []string{"web_search"}}, wantAllowed: false},
186 {name: "permission set overrides explicit web search", args: Args{PermissionSet: &tools.PermissionSet{}, ToolsSet: true, Tools: []string{"web_search"}}, wantAllowed: false},
187 {name: "explicit tools exclude web search", args: Args{ToolsSet: true, Tools: []string{"read"}}, wantAllowed: false},
188 }
189 for _, tc := range cases {
190 t.Run(tc.name, func(t *testing.T) {
191 if got := webSearchToolAllowedForSession(tc.args); got != tc.wantAllowed {
192 t.Fatalf("allowed = %v, want %v", got, tc.wantAllowed)
193 }
194 if got := webSearchExplicitlyEnabledForSession(tc.args); got != tc.wantExplicit {
195 t.Fatalf("explicit override = %v, want %v", got, tc.wantExplicit)
196 }
197 })
198 }
199}