web_search_guard.go RawHistoryFolder
1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "sync/atomic"
8
9 "github.com/bnema/zut/packages/core"
10)
11
12var errWebSearchSessionRevoked = errors.New("web_search is unavailable in this session")
13
14// webSearchSessionGuard is shared by every web_search tool instance built for
15// one interactive session. Registry refreshes replace the advertised tool, but
16// an older tool already snapshotted by core still consults this guard at its
17// execution boundary.
18type webSearchSessionGuard struct {
19 available atomic.Bool
20 generation atomic.Uint64
21}
22
23func (g *webSearchSessionGuard) setAvailable(available bool) {
24 if g == nil {
25 return
26 }
27 if !available {
28 // Increment only on an actual allow-to-deny transition. A wrapper built
29 // while denied then committed by a later successful refresh belongs to
30 // the new generation; wrappers from before revocation never revive.
31 if g.available.Swap(false) {
32 g.generation.Add(1)
33 }
34 return
35 }
36 g.available.Store(true)
37}
38
39func (g *webSearchSessionGuard) wrapRegistry(reg core.Registry) core.Registry {
40 if g == nil || reg == nil {
41 return reg
42 }
43 tool, ok := reg["web_search"]
44 if !ok {
45 return reg
46 }
47 if webSearchToolGuard(tool) == g {
48 return reg
49 }
50 guarded := &guardedWebSearchTool{
51 Tool: tool,
52 guard: g,
53 generation: g.generation.Load(),
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
63 return reg
64}
65
66// guardedWebSearchTool deliberately wraps only web_search. All extension
67// interception, confirmation, and permission behavior remains in the normal
68// core path; this final check only closes the stale-registry execution window.
69type guardedWebSearchTool struct {
70 core.Tool
71 guard *webSearchSessionGuard
72 generation uint64
73}
74
75func 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
86func (t *guardedWebSearchTool) available() bool {
87 return t != nil && t.guard != nil && t.guard.available.Load() && t.generation == t.guard.generation.Load()
88}
89
90func (t *guardedWebSearchTool) Execute(ctx context.Context, args json.RawMessage, progress func(string)) (core.ToolResult, error) {
91 if !t.available() {
92 return core.ToolResult{}, errWebSearchSessionRevoked
93 }
94 return t.Tool.Execute(ctx, args, progress)
95}
96
97// guardedWebSearchPreviewTool preserves preview support only when the wrapped
98// tool already implements core.ToolPreviewer.
99type guardedWebSearchPreviewTool struct {
100 *guardedWebSearchTool
101 previewer core.ToolPreviewer
102}
103
104func (t *guardedWebSearchPreviewTool) Preview(ctx context.Context, args json.RawMessage) (core.ToolResult, error) {
105 if !t.available() {
106 return core.ToolResult{}, errWebSearchSessionRevoked
107 }
108 return t.previewer.Preview(ctx, args)
109}