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 {
|
||||
tools map[string]Tool
|
||||
hooks []ToolHook
|
||||
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) {
|
||||
r.mu.Lock()
|
||||
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()
|
||||
result := tool.Execute(ctx, args)
|
||||
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
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue