fix(agent): address web search review findings
- commit
- c4de20dcc57394948fc59b906a59437c6c92967c
- author
- brice · · 2026-08-06 20:26 UTC
- parent
- 66d3695
- signature
- GPG Verified
CD88F9E564E0102646FB6077A1A07C7C7C3E349D· arch-gpg Verified for the committer email
11 files · +125 -47
M docs/web-search.md +2 -2
| @@ -57,8 +57,8 @@ Only direct HTTP(S) result URLs are returned. Titles are limited to 300 Unicode | |||
| 57 | 57 | ||
| 58 | 58 | ## Transcript and integrations | |
| 59 | 59 | ||
| 60 | - | The query is stored in the normal tool-call arguments. Returned titles, URLs, snippets, and errors are normal tool-result content: they can be stored in session transcripts and exports and exposed through JSON and RPC events/transcript retrieval. SDK event consumers see the same tool result only when the SDK explicitly enables `web_search`. Extensions can intercept the normal tool call, but cannot claim the reserved `web_search` name or replace the built-in. | |
| 60 | + | The query is stored in the normal tool-call arguments. Returned titles, URLs, snippets, and errors are normal tool-result content: they can be stored in session transcripts and exports and exposed through JSON and RPC events/transcript retrieval. SDK event consumers see the same tool result only when the SDK explicitly enables `web_search`. Extensions can intercept the normal tool call, but cannot claim the reserved `web_search` name or replace the built-in. Do not submit credentials, secrets, or sensitive local paths in a query: arbitrary query text is retained in those transcript and export surfaces. | |
| 61 | 61 | ||
| 62 | - | The untrusted-content label is not a privacy boundary. Raw HTML, response headers, cookies, credentials, proxy configuration, and local paths are not returned or persisted by this tool. A configured environment proxy may observe the request, but proxy settings are never displayed or included in transcripts, JSON/RPC events, or extension content. | |
| 62 | + | The untrusted-content label is not a privacy boundary. Raw HTML, response headers, cookies, proxy configuration, and internal request metadata are not returned or persisted by this tool. A configured environment proxy may observe the request, but proxy settings are never displayed or included in transcripts, JSON/RPC events, or extension content. | |
| 63 | 63 | ||
| 64 | 64 | Fetching a result URL is a separate future capability requiring its own SSRF, redirect, content-type, size-limit, and network-policy design; `web_search` does not fetch it. | |
| @@ -57,8 +57,8 @@ Only direct HTTP(S) result URLs are returned. Titles are limited to 300 Unicode | |||||
| 57 | 57 | ||||
| 58 | ## Transcript and integrations | 58 | ## Transcript and integrations | ||
| 59 | 59 | ||||
| 60 | - | The query is stored in the normal tool-call arguments. Returned titles, URLs, snippets, and errors are normal tool-result content: they can be stored in session transcripts and exports and exposed through JSON and RPC events/transcript retrieval. SDK event consumers see the same tool result only when the SDK explicitly enables `web_search`. Extensions can intercept the normal tool call, but cannot claim the reserved `web_search` name or replace the built-in. | 60 | + | The query is stored in the normal tool-call arguments. Returned titles, URLs, snippets, and errors are normal tool-result content: they can be stored in session transcripts and exports and exposed through JSON and RPC events/transcript retrieval. SDK event consumers see the same tool result only when the SDK explicitly enables `web_search`. Extensions can intercept the normal tool call, but cannot claim the reserved `web_search` name or replace the built-in. Do not submit credentials, secrets, or sensitive local paths in a query: arbitrary query text is retained in those transcript and export surfaces. |
| 61 | 61 | ||||
| 62 | - | The untrusted-content label is not a privacy boundary. Raw HTML, response headers, cookies, credentials, proxy configuration, and local paths are not returned or persisted by this tool. A configured environment proxy may observe the request, but proxy settings are never displayed or included in transcripts, JSON/RPC events, or extension content. | 62 | + | The untrusted-content label is not a privacy boundary. Raw HTML, response headers, cookies, proxy configuration, and internal request metadata are not returned or persisted by this tool. A configured environment proxy may observe the request, but proxy settings are never displayed or included in transcripts, JSON/RPC events, or extension content. |
| 63 | 63 | ||||
| 64 | Fetching a result URL is a separate future capability requiring its own SSRF, redirect, content-type, size-limit, and network-policy design; `web_search` does not fetch it. | 64 | Fetching a result URL is a separate future capability requiring its own SSRF, redirect, content-type, size-limit, and network-policy design; `web_search` does not fetch it. | ||
M docs/zutfiles.md +1 -1
| @@ -440,7 +440,7 @@ Only declare commands the agent genuinely needs. `requirements.bin` checks that | |||
| 440 | 440 | ||
| 441 | 441 | ## Network and environment permissions | |
| 442 | 442 | ||
| 443 | - | The manifest shape includes `permissions.net.allow` and `permissions.env.read`, but the local runtime does not enforce them yet. To avoid presenting unenforced declarations as security controls, it rejects manifests containing any network host or environment variable. `permissions.net` is rejected rather than ignored; it does not authorize the built-in `web_search` tool. Packaged agents receive no native web search until a future enforced network-permission design exists. | |
| 443 | + | The manifest shape includes `permissions.net.allow` and `permissions.env.read`, but the local runtime does not enforce them yet. To avoid presenting unenforced declarations as security controls, it rejects manifests containing any network host or environment variable. `permissions.net` is rejected rather than ignored; it does not authorize the built-in `web_search` tool. Packaged agents receive no native web search until a future runtime design enforces network permissions. | |
| 444 | 444 | ||
| 445 | 445 | Use empty arrays or omit these sections: | |
| 446 | 446 | ||
| @@ -440,7 +440,7 @@ Only declare commands the agent genuinely needs. `requirements.bin` checks that | |||||
| 440 | 440 | ||||
| 441 | ## Network and environment permissions | 441 | ## Network and environment permissions | ||
| 442 | 442 | ||||
| 443 | - | The manifest shape includes `permissions.net.allow` and `permissions.env.read`, but the local runtime does not enforce them yet. To avoid presenting unenforced declarations as security controls, it rejects manifests containing any network host or environment variable. `permissions.net` is rejected rather than ignored; it does not authorize the built-in `web_search` tool. Packaged agents receive no native web search until a future enforced network-permission design exists. | 443 | + | The manifest shape includes `permissions.net.allow` and `permissions.env.read`, but the local runtime does not enforce them yet. To avoid presenting unenforced declarations as security controls, it rejects manifests containing any network host or environment variable. `permissions.net` is rejected rather than ignored; it does not authorize the built-in `web_search` tool. Packaged agents receive no native web search until a future runtime design enforces network permissions. |
| 444 | 444 | ||||
| 445 | Use empty arrays or omit these sections: | 445 | Use empty arrays or omit these sections: | ||
| 446 | 446 | ||||
M go.mod +3 -3
| @@ -9,13 +9,13 @@ require ( | |||
| 9 | 9 | github.com/mattn/go-runewidth v0.0.16 | |
| 10 | 10 | github.com/sahilm/fuzzy v0.1.1 | |
| 11 | 11 | golang.org/x/image v0.43.0 | |
| 12 | - | golang.org/x/net v0.30.0 | |
| 12 | + | golang.org/x/net v0.57.0 | |
| 13 | 13 | ) | |
| 14 | 14 | ||
| 15 | 15 | require ( | |
| 16 | 16 | github.com/dlclark/regexp2 v1.11.5 // indirect | |
| 17 | 17 | github.com/kylelemons/godebug v1.1.0 // indirect | |
| 18 | 18 | github.com/rivo/uniseg v0.2.0 // indirect | |
| 19 | - | golang.org/x/sys v0.26.0 | |
| 20 | - | golang.org/x/term v0.25.0 | |
| 19 | + | golang.org/x/sys v0.47.0 | |
| 20 | + | golang.org/x/term v0.45.0 | |
| 21 | 21 | ) | |
| @@ -9,13 +9,13 @@ require ( | |||||
| 9 | github.com/mattn/go-runewidth v0.0.16 | 9 | github.com/mattn/go-runewidth v0.0.16 | ||
| 10 | github.com/sahilm/fuzzy v0.1.1 | 10 | github.com/sahilm/fuzzy v0.1.1 | ||
| 11 | golang.org/x/image v0.43.0 | 11 | golang.org/x/image v0.43.0 | ||
| 12 | - | golang.org/x/net v0.30.0 | 12 | + | golang.org/x/net v0.57.0 |
| 13 | ) | 13 | ) | ||
| 14 | 14 | ||||
| 15 | require ( | 15 | require ( | ||
| 16 | github.com/dlclark/regexp2 v1.11.5 // indirect | 16 | github.com/dlclark/regexp2 v1.11.5 // indirect | ||
| 17 | github.com/kylelemons/godebug v1.1.0 // indirect | 17 | github.com/kylelemons/godebug v1.1.0 // indirect | ||
| 18 | github.com/rivo/uniseg v0.2.0 // indirect | 18 | github.com/rivo/uniseg v0.2.0 // indirect | ||
| 19 | - | golang.org/x/sys v0.26.0 | 19 | + | golang.org/x/sys v0.47.0 |
| 20 | - | golang.org/x/term v0.25.0 | 20 | + | golang.org/x/term v0.45.0 |
| 21 | ) | 21 | ) | ||
M go.sum +6 -6 generated file: collapsed
M packages/agent/cli.go +10 -4
| @@ -1478,7 +1478,7 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 1478 | 1478 | return reg | |
| 1479 | 1479 | } | |
| 1480 | 1480 | prepareInteractiveRegistry := func(reg core.Registry) core.Registry { | |
| 1481 | - | injectSubagentTools(reg) | |
| 1481 | + | reg = injectSubagentTools(reg) | |
| 1482 | 1482 | return webSearchGuard.wrapRegistry(reg) | |
| 1483 | 1483 | } | |
| 1484 | 1484 | prepareResolvedInteractiveRegistry := func(reg core.Registry, policy subagents.WebSearchPolicy) core.Registry { | |
| @@ -2269,7 +2269,9 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 2269 | 2269 | current := liveInteractiveAgent(iv, ag) | |
| 2270 | 2270 | webSearchPolicy, err := refreshAgentToolsAndPrompt(args, sharedSandbox, extToolAdapter, current, prepareInteractiveRegistry, iv) | |
| 2271 | 2271 | if err != nil { | |
| 2272 | - | setWebSearchAvailable(false) | |
| 2272 | + | if !errors.Is(err, errInteractiveAgentChanged) { | |
| 2273 | + | setWebSearchAvailable(false) | |
| 2274 | + | } | |
| 2273 | 2275 | return err | |
| 2274 | 2276 | } | |
| 2275 | 2277 | setWebSearchAvailable(webSearchPolicy == subagents.WebSearchAllow) | |
| @@ -2278,8 +2280,12 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 2278 | 2280 | SetWebSearchAvailable: setWebSearchAvailable, | |
| 2279 | 2281 | RefreshPrompt: func() error { | |
| 2280 | 2282 | current := liveInteractiveAgent(iv, ag) | |
| 2281 | - | _, err := refreshAgentToolsAndPrompt(args, sharedSandbox, extToolAdapter, current, prepareInteractiveRegistry, iv) | |
| 2282 | - | return err | |
| 2283 | + | webSearchPolicy, err := refreshAgentToolsAndPrompt(args, sharedSandbox, extToolAdapter, current, prepareInteractiveRegistry, iv) | |
| 2284 | + | if err != nil { | |
| 2285 | + | return err | |
| 2286 | + | } | |
| 2287 | + | setWebSearchAvailable(webSearchPolicy == subagents.WebSearchAllow) | |
| 2288 | + | return nil | |
| 2283 | 2289 | }, | |
| 2284 | 2290 | AuthManager: mgr, | |
| 2285 | 2291 | LlamaCPPConfig: ResolveLlamaCPPConfig, | |
| @@ -1478,7 +1478,7 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 1478 | return reg | 1478 | return reg | ||
| 1479 | } | 1479 | } | ||
| 1480 | prepareInteractiveRegistry := func(reg core.Registry) core.Registry { | 1480 | prepareInteractiveRegistry := func(reg core.Registry) core.Registry { | ||
| 1481 | - | injectSubagentTools(reg) | 1481 | + | reg = injectSubagentTools(reg) |
| 1482 | return webSearchGuard.wrapRegistry(reg) | 1482 | return webSearchGuard.wrapRegistry(reg) | ||
| 1483 | } | 1483 | } | ||
| 1484 | prepareResolvedInteractiveRegistry := func(reg core.Registry, policy subagents.WebSearchPolicy) core.Registry { | 1484 | prepareResolvedInteractiveRegistry := func(reg core.Registry, policy subagents.WebSearchPolicy) core.Registry { | ||
| @@ -2269,7 +2269,9 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 2269 | current := liveInteractiveAgent(iv, ag) | 2269 | current := liveInteractiveAgent(iv, ag) | ||
| 2270 | webSearchPolicy, err := refreshAgentToolsAndPrompt(args, sharedSandbox, extToolAdapter, current, prepareInteractiveRegistry, iv) | 2270 | webSearchPolicy, err := refreshAgentToolsAndPrompt(args, sharedSandbox, extToolAdapter, current, prepareInteractiveRegistry, iv) | ||
| 2271 | if err != nil { | 2271 | if err != nil { | ||
| 2272 | - | setWebSearchAvailable(false) | 2272 | + | if !errors.Is(err, errInteractiveAgentChanged) { |
| 2273 | + | setWebSearchAvailable(false) | |||
| 2274 | + | } | |||
| 2273 | return err | 2275 | return err | ||
| 2274 | } | 2276 | } | ||
| 2275 | setWebSearchAvailable(webSearchPolicy == subagents.WebSearchAllow) | 2277 | setWebSearchAvailable(webSearchPolicy == subagents.WebSearchAllow) | ||
| @@ -2278,8 +2280,12 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 2278 | SetWebSearchAvailable: setWebSearchAvailable, | 2280 | SetWebSearchAvailable: setWebSearchAvailable, | ||
| 2279 | RefreshPrompt: func() error { | 2281 | RefreshPrompt: func() error { | ||
| 2280 | current := liveInteractiveAgent(iv, ag) | 2282 | current := liveInteractiveAgent(iv, ag) | ||
| 2281 | - | _, err := refreshAgentToolsAndPrompt(args, sharedSandbox, extToolAdapter, current, prepareInteractiveRegistry, iv) | 2283 | + | webSearchPolicy, err := refreshAgentToolsAndPrompt(args, sharedSandbox, extToolAdapter, current, prepareInteractiveRegistry, iv) |
| 2282 | - | return err | 2284 | + | if err != nil { |
| 2285 | + | return err | |||
| 2286 | + | } | |||
| 2287 | + | setWebSearchAvailable(webSearchPolicy == subagents.WebSearchAllow) | |||
| 2288 | + | return nil | |||
| 2283 | }, | 2289 | }, | ||
| 2284 | AuthManager: mgr, | 2290 | AuthManager: mgr, | ||
| 2285 | LlamaCPPConfig: ResolveLlamaCPPConfig, | 2291 | LlamaCPPConfig: ResolveLlamaCPPConfig, | ||
M packages/agent/modes/interactive.go +6 -1
| @@ -203,7 +203,10 @@ type InteractiveConfig struct { | |||
| 203 | 203 | RefreshTools func() error | |
| 204 | 204 | ||
| 205 | 205 | // SetWebSearchAvailable updates session-wide execution and generic-child | |
| 206 | - | // ceilings after a live registry commit or fail-closed revocation. | |
| 206 | + | // ceilings after a live registry commit or fail-closed revocation. It can | |
| 207 | + | // run while agentMu and i.mu are held, so callbacks must not re-enter | |
| 208 | + | // Interactive or wait on either lock. The CLI callback only updates its | |
| 209 | + | // web-search guard and supervisor policy. | |
| 207 | 210 | SetWebSearchAvailable func(bool) | |
| 208 | 211 | ||
| 209 | 212 | // RefreshPrompt re-resolves the complete live prompt and tool registry | |
| @@ -4739,6 +4742,8 @@ func (i *Interactive) webSearchUnavailableHint() string { | |||
| 4739 | 4742 | func (i *Interactive) setWebSearchAvailable(available bool) { | |
| 4740 | 4743 | // Advance before publishing the new gate/child ceiling. A resolver that | |
| 4741 | 4744 | // captured the prior generation can no longer commit after this transition. | |
| 4745 | + | // Callers may hold agentMu and i.mu; the callback must complete without | |
| 4746 | + | // calling back into Interactive or requiring either lock. | |
| 4742 | 4747 | i.webSearchPolicyGeneration.Add(1) | |
| 4743 | 4748 | if i.cfg.SetWebSearchAvailable != nil { | |
| 4744 | 4749 | i.cfg.SetWebSearchAvailable(available) | |
| @@ -203,7 +203,10 @@ type InteractiveConfig struct { | |||||
| 203 | RefreshTools func() error | 203 | RefreshTools func() error | ||
| 204 | 204 | ||||
| 205 | // SetWebSearchAvailable updates session-wide execution and generic-child | 205 | // SetWebSearchAvailable updates session-wide execution and generic-child | ||
| 206 | - | // ceilings after a live registry commit or fail-closed revocation. | 206 | + | // ceilings after a live registry commit or fail-closed revocation. It can |
| 207 | + | // run while agentMu and i.mu are held, so callbacks must not re-enter | |||
| 208 | + | // Interactive or wait on either lock. The CLI callback only updates its | |||
| 209 | + | // web-search guard and supervisor policy. | |||
| 207 | SetWebSearchAvailable func(bool) | 210 | SetWebSearchAvailable func(bool) | ||
| 208 | 211 | ||||
| 209 | // RefreshPrompt re-resolves the complete live prompt and tool registry | 212 | // RefreshPrompt re-resolves the complete live prompt and tool registry | ||
| @@ -4739,6 +4742,8 @@ func (i *Interactive) webSearchUnavailableHint() string { | |||||
| 4739 | func (i *Interactive) setWebSearchAvailable(available bool) { | 4742 | func (i *Interactive) setWebSearchAvailable(available bool) { | ||
| 4740 | // Advance before publishing the new gate/child ceiling. A resolver that | 4743 | // Advance before publishing the new gate/child ceiling. A resolver that | ||
| 4741 | // captured the prior generation can no longer commit after this transition. | 4744 | // captured the prior generation can no longer commit after this transition. | ||
| 4745 | + | // Callers may hold agentMu and i.mu; the callback must complete without | |||
| 4746 | + | // calling back into Interactive or requiring either lock. | |||
| 4742 | i.webSearchPolicyGeneration.Add(1) | 4747 | i.webSearchPolicyGeneration.Add(1) | ||
| 4743 | if i.cfg.SetWebSearchAvailable != nil { | 4748 | if i.cfg.SetWebSearchAvailable != nil { | ||
| 4744 | i.cfg.SetWebSearchAvailable(available) | 4749 | i.cfg.SetWebSearchAvailable(available) | ||
M packages/agent/subagents/persist.go +1 -1
| @@ -782,7 +782,7 @@ func (f *Supervisor) resume(ctx context.Context, id string, resuming bool) (*Age | |||
| 782 | 782 | SessionPath: existing.SessionPath, | |
| 783 | 783 | } | |
| 784 | 784 | ||
| 785 | - | if m.Subagent == "" && len(m.Tools) == 0 && len(f.cfg.Policy.AllowedTools) > 0 { | |
| 785 | + | if len(m.Tools) == 0 && len(f.cfg.Policy.AllowedTools) > 0 { | |
| 786 | 786 | m.Tools = append([]string(nil), f.cfg.Policy.AllowedTools...) | |
| 787 | 787 | } | |
| 788 | 788 | if err := f.validateSpawnOptions(SpawnRequest{MaxTurns: m.MaxTurns, Tools: m.Tools}); err != nil { | |
| @@ -782,7 +782,7 @@ func (f *Supervisor) resume(ctx context.Context, id string, resuming bool) (*Age | |||||
| 782 | SessionPath: existing.SessionPath, | 782 | SessionPath: existing.SessionPath, | ||
| 783 | } | 783 | } | ||
| 784 | 784 | ||||
| 785 | - | if m.Subagent == "" && len(m.Tools) == 0 && len(f.cfg.Policy.AllowedTools) > 0 { | 785 | + | if len(m.Tools) == 0 && len(f.cfg.Policy.AllowedTools) > 0 { |
| 786 | m.Tools = append([]string(nil), f.cfg.Policy.AllowedTools...) | 786 | m.Tools = append([]string(nil), f.cfg.Policy.AllowedTools...) | ||
| 787 | } | 787 | } | ||
| 788 | if err := f.validateSpawnOptions(SpawnRequest{MaxTurns: m.MaxTurns, Tools: m.Tools}); err != nil { | 788 | if err := f.validateSpawnOptions(SpawnRequest{MaxTurns: m.MaxTurns, Tools: m.Tools}); err != nil { | ||
M packages/agent/subagents/persist_test.go +2 -0
| @@ -1260,6 +1260,7 @@ func TestWebSearchPolicyPersistsThroughReloadResumeAndChildArgv(t *testing.T) { | |||
| 1260 | 1260 | } | |
| 1261 | 1261 | ||
| 1262 | 1262 | first := newSupervisor(WebSearchAllow) | |
| 1263 | + | t.Cleanup(first.StopAll) | |
| 1263 | 1264 | spawned, err := first.Spawn(context.Background(), "web policy") | |
| 1264 | 1265 | if err != nil { | |
| 1265 | 1266 | t.Fatalf("spawn: %v", err) | |
| @@ -1282,6 +1283,7 @@ func TestWebSearchPolicyPersistsThroughReloadResumeAndChildArgv(t *testing.T) { | |||
| 1282 | 1283 | } | |
| 1283 | 1284 | ||
| 1284 | 1285 | second := newSupervisor(WebSearchDeny) | |
| 1286 | + | t.Cleanup(second.StopAll) | |
| 1285 | 1287 | if loaded, errs := second.Reload(); loaded != 1 || len(errs) != 0 { | |
| 1286 | 1288 | t.Fatalf("reload loaded=%d errs=%v; want one clean load", loaded, errs) | |
| 1287 | 1289 | } | |
| @@ -1260,6 +1260,7 @@ func TestWebSearchPolicyPersistsThroughReloadResumeAndChildArgv(t *testing.T) { | |||||
| 1260 | } | 1260 | } | ||
| 1261 | 1261 | ||||
| 1262 | first := newSupervisor(WebSearchAllow) | 1262 | first := newSupervisor(WebSearchAllow) | ||
| 1263 | + | t.Cleanup(first.StopAll) | |||
| 1263 | spawned, err := first.Spawn(context.Background(), "web policy") | 1264 | spawned, err := first.Spawn(context.Background(), "web policy") | ||
| 1264 | if err != nil { | 1265 | if err != nil { | ||
| 1265 | t.Fatalf("spawn: %v", err) | 1266 | t.Fatalf("spawn: %v", err) | ||
| @@ -1282,6 +1283,7 @@ func TestWebSearchPolicyPersistsThroughReloadResumeAndChildArgv(t *testing.T) { | |||||
| 1282 | } | 1283 | } | ||
| 1283 | 1284 | ||||
| 1284 | second := newSupervisor(WebSearchDeny) | 1285 | second := newSupervisor(WebSearchDeny) | ||
| 1286 | + | t.Cleanup(second.StopAll) | |||
| 1285 | if loaded, errs := second.Reload(); loaded != 1 || len(errs) != 0 { | 1287 | if loaded, errs := second.Reload(); loaded != 1 || len(errs) != 0 { | ||
| 1286 | t.Fatalf("reload loaded=%d errs=%v; want one clean load", loaded, errs) | 1288 | t.Fatalf("reload loaded=%d errs=%v; want one clean load", loaded, errs) | ||
| 1287 | } | 1289 | } | ||
M packages/agent/tools/web_search_test.go +25 -19
| @@ -8,7 +8,9 @@ import ( | |||
| 8 | 8 | "net" | |
| 9 | 9 | "net/http" | |
| 10 | 10 | "net/http/httptest" | |
| 11 | + | "net/url" | |
| 11 | 12 | "strings" | |
| 13 | + | "sync/atomic" | |
| 12 | 14 | "testing" | |
| 13 | 15 | "time" | |
| 14 | 16 | ||
| @@ -75,9 +77,9 @@ const webSearchFixture = `<!doctype html><html><body> | |||
| 75 | 77 | </body></html>` | |
| 76 | 78 | ||
| 77 | 79 | func TestWebSearchRequestAndResultContract(t *testing.T) { | |
| 78 | - | var calls int | |
| 80 | + | var calls atomic.Int64 | |
| 79 | 81 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 80 | - | calls++ | |
| 82 | + | calls.Add(1) | |
| 81 | 83 | if r.Method != http.MethodGet { | |
| 82 | 84 | t.Errorf("method = %s, want GET", r.Method) | |
| 83 | 85 | } | |
| @@ -107,8 +109,8 @@ func TestWebSearchRequestAndResultContract(t *testing.T) { | |||
| 107 | 109 | if isError { | |
| 108 | 110 | t.Fatalf("unexpected error: %s", text) | |
| 109 | 111 | } | |
| 110 | - | if calls != 1 { | |
| 111 | - | t.Fatalf("calls = %d, want 1", calls) | |
| 112 | + | if got := calls.Load(); got != 1 { | |
| 113 | + | t.Fatalf("calls = %d, want 1", got) | |
| 112 | 114 | } | |
| 113 | 115 | if !strings.HasPrefix(text, webSearchDisclosure+"\n\n[1] One result\n https://example.com/one\n First snippet.") { | |
| 114 | 116 | t.Fatalf("unexpected result text:\n%s", text) | |
| @@ -141,14 +143,10 @@ func TestWebSearchEndpointIsFixed(t *testing.T) { | |||
| 141 | 143 | } | |
| 142 | 144 | } | |
| 143 | 145 | ||
| 144 | - | func TestWebSearchUsesEnvironmentProxy(t *testing.T) { | |
| 145 | - | for _, name := range []string{"HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "all_proxy", "no_proxy"} { | |
| 146 | - | t.Setenv(name, "") | |
| 147 | - | } | |
| 148 | - | ||
| 149 | - | var proxyCalls int | |
| 146 | + | func TestWebSearchUsesConfiguredProxyTransport(t *testing.T) { | |
| 147 | + | var proxyCalls atomic.Int64 | |
| 150 | 148 | proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 151 | - | proxyCalls++ | |
| 149 | + | proxyCalls.Add(1) | |
| 152 | 150 | if r.URL.Scheme != "http" || r.URL.Host != "web-search.invalid" { | |
| 153 | 151 | t.Errorf("proxy request URL = %q, want absolute web-search.invalid URL", r.URL) | |
| 154 | 152 | } | |
| @@ -162,15 +160,23 @@ func TestWebSearchUsesEnvironmentProxy(t *testing.T) { | |||
| 162 | 160 | _, _ = w.Write([]byte(webSearchFixture)) | |
| 163 | 161 | })) | |
| 164 | 162 | defer proxy.Close() | |
| 165 | - | t.Setenv("HTTP_PROXY", proxy.URL) | |
| 166 | 163 | ||
| 167 | - | tool := &WebSearchTool{endpointOverride: "http://web-search.invalid/html/"} | |
| 164 | + | proxyURL, err := url.Parse(proxy.URL) | |
| 165 | + | if err != nil { | |
| 166 | + | t.Fatal(err) | |
| 167 | + | } | |
| 168 | + | transport := &http.Transport{Proxy: http.ProxyURL(proxyURL)} | |
| 169 | + | t.Cleanup(transport.CloseIdleConnections) | |
| 170 | + | tool := &WebSearchTool{ | |
| 171 | + | endpointOverride: "http://web-search.invalid/html/", | |
| 172 | + | transport: transport, | |
| 173 | + | } | |
| 168 | 174 | text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "environment"}) | |
| 169 | 175 | if isError { | |
| 170 | - | t.Fatalf("unexpected error through environment proxy: %s", text) | |
| 176 | + | t.Fatalf("unexpected error through configured proxy: %s", text) | |
| 171 | 177 | } | |
| 172 | - | if proxyCalls != 1 { | |
| 173 | - | t.Fatalf("proxy calls = %d, want 1", proxyCalls) | |
| 178 | + | if got := proxyCalls.Load(); got != 1 { | |
| 179 | + | t.Fatalf("proxy calls = %d, want 1", got) | |
| 174 | 180 | } | |
| 175 | 181 | } | |
| 176 | 182 | ||
| @@ -333,10 +339,10 @@ func TestWebSearchDoesNotUseCookies(t *testing.T) { | |||
| 333 | 339 | } | |
| 334 | 340 | ||
| 335 | 341 | func TestWebSearchRejectsRedirectAndOversizedResponse(t *testing.T) { | |
| 336 | - | var followed bool | |
| 342 | + | var followed atomic.Bool | |
| 337 | 343 | redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 338 | 344 | if r.URL.Path == "/next" { | |
| 339 | - | followed = true | |
| 345 | + | followed.Store(true) | |
| 340 | 346 | } | |
| 341 | 347 | http.Redirect(w, r, "/next", http.StatusFound) | |
| 342 | 348 | })) | |
| @@ -346,7 +352,7 @@ func TestWebSearchRejectsRedirectAndOversizedResponse(t *testing.T) { | |||
| 346 | 352 | if !isError || text != "web search: DuckDuckGo request failed" { | |
| 347 | 353 | t.Fatalf("redirect result = (%q, %v)", text, isError) | |
| 348 | 354 | } | |
| 349 | - | if followed { | |
| 355 | + | if followed.Load() { | |
| 350 | 356 | t.Fatal("redirect was followed") | |
| 351 | 357 | } | |
| 352 | 358 | ||
| @@ -8,7 +8,9 @@ import ( | |||||
| 8 | "net" | 8 | "net" | ||
| 9 | "net/http" | 9 | "net/http" | ||
| 10 | "net/http/httptest" | 10 | "net/http/httptest" | ||
| 11 | + | "net/url" | |||
| 11 | "strings" | 12 | "strings" | ||
| 13 | + | "sync/atomic" | |||
| 12 | "testing" | 14 | "testing" | ||
| 13 | "time" | 15 | "time" | ||
| 14 | 16 | ||||
| @@ -75,9 +77,9 @@ const webSearchFixture = `<!doctype html><html><body> | |||||
| 75 | </body></html>` | 77 | </body></html>` | ||
| 76 | 78 | ||||
| 77 | func TestWebSearchRequestAndResultContract(t *testing.T) { | 79 | func TestWebSearchRequestAndResultContract(t *testing.T) { | ||
| 78 | - | var calls int | 80 | + | var calls atomic.Int64 |
| 79 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | 81 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | ||
| 80 | - | calls++ | 82 | + | calls.Add(1) |
| 81 | if r.Method != http.MethodGet { | 83 | if r.Method != http.MethodGet { | ||
| 82 | t.Errorf("method = %s, want GET", r.Method) | 84 | t.Errorf("method = %s, want GET", r.Method) | ||
| 83 | } | 85 | } | ||
| @@ -107,8 +109,8 @@ func TestWebSearchRequestAndResultContract(t *testing.T) { | |||||
| 107 | if isError { | 109 | if isError { | ||
| 108 | t.Fatalf("unexpected error: %s", text) | 110 | t.Fatalf("unexpected error: %s", text) | ||
| 109 | } | 111 | } | ||
| 110 | - | if calls != 1 { | 112 | + | if got := calls.Load(); got != 1 { |
| 111 | - | t.Fatalf("calls = %d, want 1", calls) | 113 | + | t.Fatalf("calls = %d, want 1", got) |
| 112 | } | 114 | } | ||
| 113 | if !strings.HasPrefix(text, webSearchDisclosure+"\n\n[1] One result\n https://example.com/one\n First snippet.") { | 115 | if !strings.HasPrefix(text, webSearchDisclosure+"\n\n[1] One result\n https://example.com/one\n First snippet.") { | ||
| 114 | t.Fatalf("unexpected result text:\n%s", text) | 116 | t.Fatalf("unexpected result text:\n%s", text) | ||
| @@ -141,14 +143,10 @@ func TestWebSearchEndpointIsFixed(t *testing.T) { | |||||
| 141 | } | 143 | } | ||
| 142 | } | 144 | } | ||
| 143 | 145 | ||||
| 144 | - | func TestWebSearchUsesEnvironmentProxy(t *testing.T) { | 146 | + | func TestWebSearchUsesConfiguredProxyTransport(t *testing.T) { |
| 145 | - | for _, name := range []string{"HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "all_proxy", "no_proxy"} { | 147 | + | var proxyCalls atomic.Int64 |
| 146 | - | t.Setenv(name, "") | |||
| 147 | - | } | |||
| 148 | - | ||||
| 149 | - | var proxyCalls int | |||
| 150 | proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | 148 | proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | ||
| 151 | - | proxyCalls++ | 149 | + | proxyCalls.Add(1) |
| 152 | if r.URL.Scheme != "http" || r.URL.Host != "web-search.invalid" { | 150 | if r.URL.Scheme != "http" || r.URL.Host != "web-search.invalid" { | ||
| 153 | t.Errorf("proxy request URL = %q, want absolute web-search.invalid URL", r.URL) | 151 | t.Errorf("proxy request URL = %q, want absolute web-search.invalid URL", r.URL) | ||
| 154 | } | 152 | } | ||
| @@ -162,15 +160,23 @@ func TestWebSearchUsesEnvironmentProxy(t *testing.T) { | |||||
| 162 | _, _ = w.Write([]byte(webSearchFixture)) | 160 | _, _ = w.Write([]byte(webSearchFixture)) | ||
| 163 | })) | 161 | })) | ||
| 164 | defer proxy.Close() | 162 | defer proxy.Close() | ||
| 165 | - | t.Setenv("HTTP_PROXY", proxy.URL) | |||
| 166 | 163 | ||||
| 167 | - | tool := &WebSearchTool{endpointOverride: "http://web-search.invalid/html/"} | 164 | + | proxyURL, err := url.Parse(proxy.URL) |
| 165 | + | if err != nil { | |||
| 166 | + | t.Fatal(err) | |||
| 167 | + | } | |||
| 168 | + | transport := &http.Transport{Proxy: http.ProxyURL(proxyURL)} | |||
| 169 | + | t.Cleanup(transport.CloseIdleConnections) | |||
| 170 | + | tool := &WebSearchTool{ | |||
| 171 | + | endpointOverride: "http://web-search.invalid/html/", | |||
| 172 | + | transport: transport, | |||
| 173 | + | } | |||
| 168 | text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "environment"}) | 174 | text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "environment"}) | ||
| 169 | if isError { | 175 | if isError { | ||
| 170 | - | t.Fatalf("unexpected error through environment proxy: %s", text) | 176 | + | t.Fatalf("unexpected error through configured proxy: %s", text) |
| 171 | } | 177 | } | ||
| 172 | - | if proxyCalls != 1 { | 178 | + | if got := proxyCalls.Load(); got != 1 { |
| 173 | - | t.Fatalf("proxy calls = %d, want 1", proxyCalls) | 179 | + | t.Fatalf("proxy calls = %d, want 1", got) |
| 174 | } | 180 | } | ||
| 175 | } | 181 | } | ||
| 176 | 182 | ||||
| @@ -333,10 +339,10 @@ func TestWebSearchDoesNotUseCookies(t *testing.T) { | |||||
| 333 | } | 339 | } | ||
| 334 | 340 | ||||
| 335 | func TestWebSearchRejectsRedirectAndOversizedResponse(t *testing.T) { | 341 | func TestWebSearchRejectsRedirectAndOversizedResponse(t *testing.T) { | ||
| 336 | - | var followed bool | 342 | + | var followed atomic.Bool |
| 337 | redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | 343 | redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | ||
| 338 | if r.URL.Path == "/next" { | 344 | if r.URL.Path == "/next" { | ||
| 339 | - | followed = true | 345 | + | followed.Store(true) |
| 340 | } | 346 | } | ||
| 341 | http.Redirect(w, r, "/next", http.StatusFound) | 347 | http.Redirect(w, r, "/next", http.StatusFound) | ||
| 342 | })) | 348 | })) | ||
| @@ -346,7 +352,7 @@ func TestWebSearchRejectsRedirectAndOversizedResponse(t *testing.T) { | |||||
| 346 | if !isError || text != "web search: DuckDuckGo request failed" { | 352 | if !isError || text != "web search: DuckDuckGo request failed" { | ||
| 347 | t.Fatalf("redirect result = (%q, %v)", text, isError) | 353 | t.Fatalf("redirect result = (%q, %v)", text, isError) | ||
| 348 | } | 354 | } | ||
| 349 | - | if followed { | 355 | + | if followed.Load() { |
| 350 | t.Fatal("redirect was followed") | 356 | t.Fatal("redirect was followed") | ||
| 351 | } | 357 | } | ||
| 352 | 358 | ||||
M packages/agent/web_search_guard.go +32 -10
| @@ -44,14 +44,22 @@ func (g *webSearchSessionGuard) wrapRegistry(reg core.Registry) core.Registry { | |||
| 44 | 44 | if !ok { | |
| 45 | 45 | return reg | |
| 46 | 46 | } | |
| 47 | - | if guarded, ok := tool.(*guardedWebSearchTool); ok && guarded.guard == g { | |
| 47 | + | if webSearchToolGuard(tool) == g { | |
| 48 | 48 | return reg | |
| 49 | 49 | } | |
| 50 | - | reg["web_search"] = &guardedWebSearchTool{ | |
| 50 | + | guarded := &guardedWebSearchTool{ | |
| 51 | 51 | Tool: tool, | |
| 52 | 52 | guard: g, | |
| 53 | 53 | generation: g.generation.Load(), | |
| 54 | 54 | } | |
| 55 | + | if previewer, ok := tool.(core.ToolPreviewer); ok { | |
| 56 | + | reg["web_search"] = &guardedWebSearchPreviewTool{ | |
| 57 | + | guardedWebSearchTool: guarded, | |
| 58 | + | previewer: previewer, | |
| 59 | + | } | |
| 60 | + | return reg | |
| 61 | + | } | |
| 62 | + | reg["web_search"] = guarded | |
| 55 | 63 | return reg | |
| 56 | 64 | } | |
| 57 | 65 | ||
| @@ -64,24 +72,38 @@ type guardedWebSearchTool struct { | |||
| 64 | 72 | generation uint64 | |
| 65 | 73 | } | |
| 66 | 74 | ||
| 75 | + | func webSearchToolGuard(tool core.Tool) *webSearchSessionGuard { | |
| 76 | + | switch tool := tool.(type) { | |
| 77 | + | case *guardedWebSearchTool: | |
| 78 | + | return tool.guard | |
| 79 | + | case *guardedWebSearchPreviewTool: | |
| 80 | + | return tool.guard | |
| 81 | + | default: | |
| 82 | + | return nil | |
| 83 | + | } | |
| 84 | + | } | |
| 85 | + | ||
| 67 | 86 | func (t *guardedWebSearchTool) available() bool { | |
| 68 | 87 | return t != nil && t.guard != nil && t.guard.available.Load() && t.generation == t.guard.generation.Load() | |
| 69 | 88 | } | |
| 70 | 89 | ||
| 71 | - | func (t *guardedWebSearchTool) Preview(ctx context.Context, args json.RawMessage) (core.ToolResult, error) { | |
| 90 | + | func (t *guardedWebSearchTool) Execute(ctx context.Context, args json.RawMessage, progress func(string)) (core.ToolResult, error) { | |
| 72 | 91 | if !t.available() { | |
| 73 | 92 | return core.ToolResult{}, errWebSearchSessionRevoked | |
| 74 | 93 | } | |
| 75 | - | previewer, ok := t.Tool.(core.ToolPreviewer) | |
| 76 | - | if !ok { | |
| 77 | - | return core.ToolResult{}, errors.New("web_search preview is unavailable") | |
| 78 | - | } | |
| 79 | - | return previewer.Preview(ctx, args) | |
| 94 | + | return t.Tool.Execute(ctx, args, progress) | |
| 80 | 95 | } | |
| 81 | 96 | ||
| 82 | - | func (t *guardedWebSearchTool) Execute(ctx context.Context, args json.RawMessage, progress func(string)) (core.ToolResult, error) { | |
| 97 | + | // guardedWebSearchPreviewTool preserves preview support only when the wrapped | |
| 98 | + | // tool already implements core.ToolPreviewer. | |
| 99 | + | type guardedWebSearchPreviewTool struct { | |
| 100 | + | *guardedWebSearchTool | |
| 101 | + | previewer core.ToolPreviewer | |
| 102 | + | } | |
| 103 | + | ||
| 104 | + | func (t *guardedWebSearchPreviewTool) Preview(ctx context.Context, args json.RawMessage) (core.ToolResult, error) { | |
| 83 | 105 | if !t.available() { | |
| 84 | 106 | return core.ToolResult{}, errWebSearchSessionRevoked | |
| 85 | 107 | } | |
| 86 | - | return t.Tool.Execute(ctx, args, progress) | |
| 108 | + | return t.previewer.Preview(ctx, args) | |
| 87 | 109 | } | |
| @@ -44,14 +44,22 @@ func (g *webSearchSessionGuard) wrapRegistry(reg core.Registry) core.Registry { | |||||
| 44 | if !ok { | 44 | if !ok { | ||
| 45 | return reg | 45 | return reg | ||
| 46 | } | 46 | } | ||
| 47 | - | if guarded, ok := tool.(*guardedWebSearchTool); ok && guarded.guard == g { | 47 | + | if webSearchToolGuard(tool) == g { |
| 48 | return reg | 48 | return reg | ||
| 49 | } | 49 | } | ||
| 50 | - | reg["web_search"] = &guardedWebSearchTool{ | 50 | + | guarded := &guardedWebSearchTool{ |
| 51 | Tool: tool, | 51 | Tool: tool, | ||
| 52 | guard: g, | 52 | guard: g, | ||
| 53 | generation: g.generation.Load(), | 53 | generation: g.generation.Load(), | ||
| 54 | } | 54 | } | ||
| 55 | + | if previewer, ok := tool.(core.ToolPreviewer); ok { | |||
| 56 | + | reg["web_search"] = &guardedWebSearchPreviewTool{ | |||
| 57 | + | guardedWebSearchTool: guarded, | |||
| 58 | + | previewer: previewer, | |||
| 59 | + | } | |||
| 60 | + | return reg | |||
| 61 | + | } | |||
| 62 | + | reg["web_search"] = guarded | |||
| 55 | return reg | 63 | return reg | ||
| 56 | } | 64 | } | ||
| 57 | 65 | ||||
| @@ -64,24 +72,38 @@ type guardedWebSearchTool struct { | |||||
| 64 | generation uint64 | 72 | generation uint64 | ||
| 65 | } | 73 | } | ||
| 66 | 74 | ||||
| 75 | + | func webSearchToolGuard(tool core.Tool) *webSearchSessionGuard { | |||
| 76 | + | switch tool := tool.(type) { | |||
| 77 | + | case *guardedWebSearchTool: | |||
| 78 | + | return tool.guard | |||
| 79 | + | case *guardedWebSearchPreviewTool: | |||
| 80 | + | return tool.guard | |||
| 81 | + | default: | |||
| 82 | + | return nil | |||
| 83 | + | } | |||
| 84 | + | } | |||
| 85 | + | ||||
| 67 | func (t *guardedWebSearchTool) available() bool { | 86 | func (t *guardedWebSearchTool) available() bool { | ||
| 68 | return t != nil && t.guard != nil && t.guard.available.Load() && t.generation == t.guard.generation.Load() | 87 | return t != nil && t.guard != nil && t.guard.available.Load() && t.generation == t.guard.generation.Load() | ||
| 69 | } | 88 | } | ||
| 70 | 89 | ||||
| 71 | - | func (t *guardedWebSearchTool) Preview(ctx context.Context, args json.RawMessage) (core.ToolResult, error) { | 90 | + | func (t *guardedWebSearchTool) Execute(ctx context.Context, args json.RawMessage, progress func(string)) (core.ToolResult, error) { |
| 72 | if !t.available() { | 91 | if !t.available() { | ||
| 73 | return core.ToolResult{}, errWebSearchSessionRevoked | 92 | return core.ToolResult{}, errWebSearchSessionRevoked | ||
| 74 | } | 93 | } | ||
| 75 | - | previewer, ok := t.Tool.(core.ToolPreviewer) | 94 | + | return t.Tool.Execute(ctx, args, progress) |
| 76 | - | if !ok { | |||
| 77 | - | return core.ToolResult{}, errors.New("web_search preview is unavailable") | |||
| 78 | - | } | |||
| 79 | - | return previewer.Preview(ctx, args) | |||
| 80 | } | 95 | } | ||
| 81 | 96 | ||||
| 82 | - | func (t *guardedWebSearchTool) Execute(ctx context.Context, args json.RawMessage, progress func(string)) (core.ToolResult, error) { | 97 | + | // guardedWebSearchPreviewTool preserves preview support only when the wrapped |
| 98 | + | // tool already implements core.ToolPreviewer. | |||
| 99 | + | type guardedWebSearchPreviewTool struct { | |||
| 100 | + | *guardedWebSearchTool | |||
| 101 | + | previewer core.ToolPreviewer | |||
| 102 | + | } | |||
| 103 | + | ||||
| 104 | + | func (t *guardedWebSearchPreviewTool) Preview(ctx context.Context, args json.RawMessage) (core.ToolResult, error) { | |||
| 83 | if !t.available() { | 105 | if !t.available() { | ||
| 84 | return core.ToolResult{}, errWebSearchSessionRevoked | 106 | return core.ToolResult{}, errWebSearchSessionRevoked | ||
| 85 | } | 107 | } | ||
| 86 | - | return t.Tool.Execute(ctx, args, progress) | 108 | + | return t.previewer.Preview(ctx, args) |
| 87 | } | 109 | } | ||
M packages/agent/web_search_guard_test.go +37 -0
| @@ -28,6 +28,43 @@ func (t *recordingWebSearchTool) Execute(context.Context, json.RawMessage, func( | |||
| 28 | 28 | return core.ToolResult{}, nil | |
| 29 | 29 | } | |
| 30 | 30 | ||
| 31 | + | type executeOnlyWebSearchTool struct { | |
| 32 | + | executions int | |
| 33 | + | } | |
| 34 | + | ||
| 35 | + | func (t *executeOnlyWebSearchTool) Name() string { return "web_search" } | |
| 36 | + | func (t *executeOnlyWebSearchTool) Description() string { return "test web search" } | |
| 37 | + | func (t *executeOnlyWebSearchTool) Schema() json.RawMessage { | |
| 38 | + | return json.RawMessage(`{"type":"object"}`) | |
| 39 | + | } | |
| 40 | + | func (t *executeOnlyWebSearchTool) Execute(context.Context, json.RawMessage, func(string)) (core.ToolResult, error) { | |
| 41 | + | t.executions++ | |
| 42 | + | return core.ToolResult{}, nil | |
| 43 | + | } | |
| 44 | + | ||
| 45 | + | func 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 | + | ||
| 31 | 68 | func TestWebSearchSessionGuardRevokesSnapshottedToolBeforeExecute(t *testing.T) { | |
| 32 | 69 | guard := &webSearchSessionGuard{} | |
| 33 | 70 | underlying := &recordingWebSearchTool{} | |
| @@ -28,6 +28,43 @@ func (t *recordingWebSearchTool) Execute(context.Context, json.RawMessage, func( | |||||
| 28 | return core.ToolResult{}, nil | 28 | return core.ToolResult{}, nil | ||
| 29 | } | 29 | } | ||
| 30 | 30 | ||||
| 31 | + | type executeOnlyWebSearchTool struct { | |||
| 32 | + | executions int | |||
| 33 | + | } | |||
| 34 | + | ||||
| 35 | + | func (t *executeOnlyWebSearchTool) Name() string { return "web_search" } | |||
| 36 | + | func (t *executeOnlyWebSearchTool) Description() string { return "test web search" } | |||
| 37 | + | func (t *executeOnlyWebSearchTool) Schema() json.RawMessage { | |||
| 38 | + | return json.RawMessage(`{"type":"object"}`) | |||
| 39 | + | } | |||
| 40 | + | func (t *executeOnlyWebSearchTool) Execute(context.Context, json.RawMessage, func(string)) (core.ToolResult, error) { | |||
| 41 | + | t.executions++ | |||
| 42 | + | return core.ToolResult{}, nil | |||
| 43 | + | } | |||
| 44 | + | ||||
| 45 | + | func 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 | + | ||||
| 31 | func TestWebSearchSessionGuardRevokesSnapshottedToolBeforeExecute(t *testing.T) { | 68 | func TestWebSearchSessionGuardRevokesSnapshottedToolBeforeExecute(t *testing.T) { | ||
| 32 | guard := &webSearchSessionGuard{} | 69 | guard := &webSearchSessionGuard{} | ||
| 33 | underlying := &recordingWebSearchTool{} | 70 | underlying := &recordingWebSearchTool{} | ||