feat: add policy engine with registry Clone and Remove
New policy.go with ApplyPolicy (allow/deny filtering) and DepthDenyList (leaf agents lose spawn/handoff/list_agents). Add Clone() for shallow registry copy and Remove() for tool unregistration to ToolRegistry.
This commit is contained in:
parent
f3a838fe45
commit
0b251e7ee8
2 changed files with 69 additions and 0 deletions
44
pkg/tools/policy.go
Normal file
44
pkg/tools/policy.go
Normal file
|
|
@ -0,0 +1,44 @@
|
||||||
|
package tools
|
||||||
|
|
||||||
|
// ToolPolicy represents one layer of allow/deny filtering.
|
||||||
|
type ToolPolicy struct {
|
||||||
|
Allow []string
|
||||||
|
Deny []string
|
||||||
|
}
|
||||||
|
|
||||||
|
// ApplyPolicy filters a registry in-place.
|
||||||
|
// Allow (if non-empty): only listed tools survive.
|
||||||
|
// Deny: listed tools removed from whatever remains.
|
||||||
|
func ApplyPolicy(reg *ToolRegistry, policy ToolPolicy) {
|
||||||
|
allowNames := ResolveToolNames(policy.Allow)
|
||||||
|
denyNames := ResolveToolNames(policy.Deny)
|
||||||
|
|
||||||
|
// Allow-list: if non-empty, remove everything not in the allow set
|
||||||
|
if len(allowNames) > 0 {
|
||||||
|
allowSet := make(map[string]struct{}, len(allowNames))
|
||||||
|
for _, name := range allowNames {
|
||||||
|
allowSet[name] = struct{}{}
|
||||||
|
}
|
||||||
|
for _, name := range reg.List() {
|
||||||
|
if _, ok := allowSet[name]; !ok {
|
||||||
|
reg.Remove(name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Deny-list: remove listed tools
|
||||||
|
for _, name := range denyNames {
|
||||||
|
reg.Remove(name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// DepthDenyList returns tools to deny at a given depth.
|
||||||
|
// depth 0: nil (main agent, full access)
|
||||||
|
// depth >= maxDepth: spawn/handoff/list_agents denied (leaf, no further chaining)
|
||||||
|
// between 0 and maxDepth: nil (mid-chain, full access)
|
||||||
|
func DepthDenyList(depth, maxDepth int) []string {
|
||||||
|
if depth >= maxDepth {
|
||||||
|
return []string{"spawn", "handoff", "list_agents"}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
@ -191,6 +191,31 @@ func (r *ToolRegistry) List() []string {
|
||||||
return names
|
return names
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Clone creates a shallow copy of the registry.
|
||||||
|
// Tool instances are shared (not deep-copied), only the map and hooks slice are copied.
|
||||||
|
// This is used for depth-based policy: clone → remove denied tools → pass to RunToolLoop.
|
||||||
|
func (r *ToolRegistry) Clone() *ToolRegistry {
|
||||||
|
r.mu.RLock()
|
||||||
|
defer r.mu.RUnlock()
|
||||||
|
|
||||||
|
cloned := &ToolRegistry{
|
||||||
|
tools: make(map[string]Tool, len(r.tools)),
|
||||||
|
hooks: make([]ToolHook, len(r.hooks)),
|
||||||
|
}
|
||||||
|
for name, tool := range r.tools {
|
||||||
|
cloned.tools[name] = tool
|
||||||
|
}
|
||||||
|
copy(cloned.hooks, r.hooks)
|
||||||
|
return cloned
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove unregisters a tool by name.
|
||||||
|
func (r *ToolRegistry) Remove(name string) {
|
||||||
|
r.mu.Lock()
|
||||||
|
defer r.mu.Unlock()
|
||||||
|
delete(r.tools, name)
|
||||||
|
}
|
||||||
|
|
||||||
// Count returns the number of registered tools.
|
// Count returns the number of registered tools.
|
||||||
func (r *ToolRegistry) Count() int {
|
func (r *ToolRegistry) Count() int {
|
||||||
r.mu.RLock()
|
r.mu.RLock()
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue