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