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:
Leandro Barbosa 2026-02-18 15:04:48 -03:00
parent f73c153578
commit 42d73bde4d
2 changed files with 52 additions and 0 deletions

17
pkg/tools/hooks.go Normal file
View 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)
}

View file

@ -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
} }