feat: add ToolHook interface for before/after execution interception
Introduce ToolHook with BeforeExecute/AfterExecute lifecycle methods in ToolRegistry. Hooks enable policy enforcement and loop detection pipelines without modifying individual tools. BeforeExecute can block execution; AfterExecute always runs for observability.
This commit is contained in:
parent
f73c153578
commit
42d73bde4d
2 changed files with 52 additions and 0 deletions
17
pkg/tools/hooks.go
Normal file
17
pkg/tools/hooks.go
Normal file
|
|
@ -0,0 +1,17 @@
|
||||||
|
package tools
|
||||||
|
|
||||||
|
import "context"
|
||||||
|
|
||||||
|
// ToolHook allows intercepting tool execution for policy enforcement,
|
||||||
|
// loop detection, logging, or other cross-cutting concerns.
|
||||||
|
//
|
||||||
|
// Hooks are called by the ToolRegistry around tool execution:
|
||||||
|
// - BeforeExecute: called before the tool runs. Return non-nil error to block execution.
|
||||||
|
// - AfterExecute: called after the tool completes (even on error). Cannot block.
|
||||||
|
//
|
||||||
|
// Multiple hooks are executed in registration order. If any BeforeExecute returns
|
||||||
|
// an error, subsequent hooks and the tool itself are skipped.
|
||||||
|
type ToolHook interface {
|
||||||
|
BeforeExecute(ctx context.Context, toolName string, args map[string]interface{}) error
|
||||||
|
AfterExecute(ctx context.Context, toolName string, args map[string]interface{}, result *ToolResult)
|
||||||
|
}
|
||||||
|
|
@ -12,6 +12,7 @@ import (
|
||||||
|
|
||||||
type ToolRegistry struct {
|
type ToolRegistry struct {
|
||||||
tools map[string]Tool
|
tools map[string]Tool
|
||||||
|
hooks []ToolHook
|
||||||
mu sync.RWMutex
|
mu sync.RWMutex
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -21,6 +22,14 @@ func NewToolRegistry() *ToolRegistry {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// AddHook registers a hook that will be called around tool executions.
|
||||||
|
// Hooks are called in registration order.
|
||||||
|
func (r *ToolRegistry) AddHook(hook ToolHook) {
|
||||||
|
r.mu.Lock()
|
||||||
|
defer r.mu.Unlock()
|
||||||
|
r.hooks = append(r.hooks, hook)
|
||||||
|
}
|
||||||
|
|
||||||
func (r *ToolRegistry) Register(tool Tool) {
|
func (r *ToolRegistry) Register(tool Tool) {
|
||||||
r.mu.Lock()
|
r.mu.Lock()
|
||||||
defer r.mu.Unlock()
|
defer r.mu.Unlock()
|
||||||
|
|
@ -71,6 +80,27 @@ func (r *ToolRegistry) ExecuteWithContext(ctx context.Context, name string, args
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Run BeforeExecute hooks — any hook returning error blocks execution.
|
||||||
|
r.mu.RLock()
|
||||||
|
hooks := r.hooks
|
||||||
|
r.mu.RUnlock()
|
||||||
|
|
||||||
|
for _, hook := range hooks {
|
||||||
|
if err := hook.BeforeExecute(ctx, name, args); err != nil {
|
||||||
|
logger.WarnCF("tool", "Hook blocked tool execution",
|
||||||
|
map[string]interface{}{
|
||||||
|
"tool": name,
|
||||||
|
"error": err.Error(),
|
||||||
|
})
|
||||||
|
blocked := ErrorResult(fmt.Sprintf("tool %q blocked by hook: %s", name, err.Error()))
|
||||||
|
// Still run AfterExecute hooks for observability
|
||||||
|
for _, h := range hooks {
|
||||||
|
h.AfterExecute(ctx, name, args, blocked)
|
||||||
|
}
|
||||||
|
return blocked
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
start := time.Now()
|
start := time.Now()
|
||||||
result := tool.Execute(ctx, args)
|
result := tool.Execute(ctx, args)
|
||||||
duration := time.Since(start)
|
duration := time.Since(start)
|
||||||
|
|
@ -98,6 +128,11 @@ func (r *ToolRegistry) ExecuteWithContext(ctx context.Context, name string, args
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Run AfterExecute hooks for observability
|
||||||
|
for _, hook := range hooks {
|
||||||
|
hook.AfterExecute(ctx, name, args, result)
|
||||||
|
}
|
||||||
|
|
||||||
return result
|
return result
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue