diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 263eeb4dd..1613d7a41 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -23,6 +23,7 @@ import ( "github.com/sipeed/picoclaw/pkg/channels" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" + "github.com/sipeed/picoclaw/pkg/hooks" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/mcp" "github.com/sipeed/picoclaw/pkg/media" @@ -46,6 +47,7 @@ type AgentLoop struct { channelManager *channels.Manager mediaStore media.MediaStore transcriber voice.Transcriber + hookManager *hooks.HookManager } // processOptions configures how a message is processed @@ -77,6 +79,27 @@ func NewAgentLoop( cooldown := providers.NewCooldownTracker() fallbackChain := providers.NewFallbackChain(cooldown) + // Initialize hook manager from config + var hookManager *hooks.HookManager + hookRules := convertHooksConfig(cfg.Hooks) + if len(hookRules) > 0 { + hookManager = hooks.NewHookManager(hookRules) + logger.InfoCF("agent", "Hook manager initialized", + map[string]any{ + "pre_message": len(hookRules[hooks.PreMessage]), + "post_message": len(hookRules[hooks.PostMessage]), + "pre_tool": len(hookRules[hooks.PreToolUse]), + "post_tool": len(hookRules[hooks.PostToolUse]), + }) + + // Inject hook manager into all agent tool registries + for _, agentID := range registry.ListAgentIDs() { + if agent, ok := registry.GetAgent(agentID); ok { + agent.Tools.SetHookManager(hookManager) + } + } + } + // Create state manager using default agent's workspace for channel recording defaultAgent := registry.GetDefaultAgent() var stateManager *state.Manager @@ -91,9 +114,32 @@ func NewAgentLoop( state: stateManager, summarizing: sync.Map{}, fallback: fallbackChain, + hookManager: hookManager, } } +// convertHooksConfig converts config.HooksConfig to the hooks package's rule map. +func convertHooksConfig(cfg config.HooksConfig) map[hooks.Event][]hooks.HookRule { + rules := make(map[hooks.Event][]hooks.HookRule) + + convert := func(event hooks.Event, cfgRules []config.HookRuleConfig) { + for _, r := range cfgRules { + rules[event] = append(rules[event], hooks.HookRule{ + Matcher: r.Matcher, + Command: r.Command, + InjectOutput: r.InjectOutput, + }) + } + } + + convert(hooks.PreMessage, cfg.PreMessage) + convert(hooks.PostMessage, cfg.PostMessage) + convert(hooks.PreToolUse, cfg.PreToolUse) + convert(hooks.PostToolUse, cfg.PostToolUse) + + return rules +} + // registerSharedTools registers tools that are shared across all agents (web, message, spawn). func registerSharedTools( cfg *config.Config, @@ -659,7 +705,20 @@ func (al *AgentLoop) runAgentLoop( } } - // 1. Build messages (skip history for heartbeat) + // 1. PreMessage hook: inject context before building messages + userMessage := opts.UserMessage + if al.hookManager != nil && al.hookManager.HasHooks(hooks.PreMessage) { + injected := al.hookManager.CollectInjectedOutput(ctx, hooks.PreMessage, hooks.HookPayload{ + Message: opts.UserMessage, + Channel: opts.Channel, + ChatID: opts.ChatID, + }) + if injected != "" { + userMessage = userMessage + "\n\n[Hook Context]\n" + injected + } + } + + // 2. Build messages (skip history for heartbeat) var history []providers.Message var summary string if !opts.NoHistory { @@ -669,7 +728,7 @@ func (al *AgentLoop) runAgentLoop( messages := agent.ContextBuilder.BuildMessages( history, summary, - opts.UserMessage, + userMessage, opts.Media, opts.Channel, opts.ChatID, @@ -679,10 +738,10 @@ func (al *AgentLoop) runAgentLoop( maxMediaSize := al.cfg.Agents.Defaults.GetMaxMediaSize() messages = resolveMediaRefs(messages, al.mediaStore, maxMediaSize) - // 2. Save user message to session + // 3. Save user message to session agent.Sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage) - // 3. Run LLM iteration loop + // 4. Run LLM iteration loop finalContent, iteration, err := al.runLLMIteration(ctx, agent, messages, opts) if err != nil { return "", err @@ -691,21 +750,30 @@ func (al *AgentLoop) runAgentLoop( // If last tool had ForUser content and we already sent it, we might not need to send final response // This is controlled by the tool's Silent flag and ForUser content - // 4. Handle empty response + // 5. Handle empty response if finalContent == "" { finalContent = opts.DefaultResponse } - // 5. Save final assistant message to session + // 6. Save final assistant message to session agent.Sessions.AddMessage(opts.SessionKey, "assistant", finalContent) agent.Sessions.Save(opts.SessionKey) - // 6. Optional: summarization + // 7. PostMessage hook: fire after response is saved + if al.hookManager != nil && al.hookManager.HasHooks(hooks.PostMessage) { + al.hookManager.Trigger(ctx, hooks.PostMessage, hooks.HookPayload{ + Message: finalContent, + Channel: opts.Channel, + ChatID: opts.ChatID, + }) + } + + // 8. Optional: summarization if opts.EnableSummary { al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID) } - // 7. Optional: send response via bus + // 9. Optional: send response via bus if opts.SendResponse { al.bus.PublishOutbound(ctx, bus.OutboundMessage{ Channel: opts.Channel, @@ -714,7 +782,7 @@ func (al *AgentLoop) runAgentLoop( }) } - // 8. Log response + // 10. Log response responsePreview := utils.Truncate(finalContent, 120) logger.InfoCF("agent", fmt.Sprintf("Response: %s", responsePreview), map[string]any{ diff --git a/pkg/config/config.go b/pkg/config/config.go index 3cfebf5e8..b7f5dbe6f 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -58,6 +58,7 @@ type Config struct { Tools ToolsConfig `json:"tools"` Heartbeat HeartbeatConfig `json:"heartbeat"` Devices DevicesConfig `json:"devices"` + Hooks HooksConfig `json:"hooks,omitempty"` } // MarshalJSON implements custom JSON marshaling for Config @@ -67,6 +68,7 @@ func (c Config) MarshalJSON() ([]byte, error) { aux := &struct { Providers *ProvidersConfig `json:"providers,omitempty"` Session *SessionConfig `json:"session,omitempty"` + Hooks *HooksConfig `json:"hooks,omitempty"` *Alias }{ Alias: (*Alias)(&c), @@ -82,6 +84,11 @@ func (c Config) MarshalJSON() ([]byte, error) { aux.Session = &c.Session } + // Only include hooks if configured + if !c.Hooks.IsEmpty() { + aux.Hooks = &c.Hooks + } + return json.Marshal(aux) } @@ -588,6 +595,29 @@ type MediaCleanupConfig struct { Interval int `json:"interval_minutes" env:"PICOCLAW_MEDIA_CLEANUP_INTERVAL"` } +// HookRuleConfig defines a single user-configurable hook. +type HookRuleConfig struct { + Matcher string `json:"matcher"` // Tool name filter; empty = match all + Command string `json:"command"` // Shell command to execute + InjectOutput bool `json:"inject_output"` // Append stdout to result/context +} + +// HooksConfig defines user-configurable hooks at key lifecycle points. +// Each event maps to a list of hook rules that are executed sequentially. +// Default is empty (no hooks). See pkg/hooks for event documentation. +type HooksConfig struct { + PreMessage []HookRuleConfig `json:"PreMessage,omitempty"` + PostMessage []HookRuleConfig `json:"PostMessage,omitempty"` + PreToolUse []HookRuleConfig `json:"PreToolUse,omitempty"` + PostToolUse []HookRuleConfig `json:"PostToolUse,omitempty"` +} + +// IsEmpty returns true when no hook rules are configured. +func (h HooksConfig) IsEmpty() bool { + return len(h.PreMessage) == 0 && len(h.PostMessage) == 0 && + len(h.PreToolUse) == 0 && len(h.PostToolUse) == 0 +} + type ToolsConfig struct { AllowReadPaths []string `json:"allow_read_paths" env:"PICOCLAW_TOOLS_ALLOW_READ_PATHS"` AllowWritePaths []string `json:"allow_write_paths" env:"PICOCLAW_TOOLS_ALLOW_WRITE_PATHS"` diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 10ebc7c90..83a045404 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -479,3 +479,78 @@ func TestDefaultConfig_WorkspacePath_WithPicoclawHome(t *testing.T) { t.Errorf("Workspace path with PICOCLAW_HOME = %q, want %q", cfg.Agents.Defaults.Workspace, want) } } + +func TestConfig_HooksConfig_Parse(t *testing.T) { + jsonData := `{ + "hooks": { + "PostToolUse": [ + { + "matcher": "exec", + "command": "~/.picoclaw/scripts/error-detector.sh", + "inject_output": true + } + ], + "PreMessage": [ + { + "matcher": "", + "command": "~/.picoclaw/scripts/activator.sh", + "inject_output": true + } + ] + } + }` + + cfg := DefaultConfig() + if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { + t.Fatalf("unmarshal: %v", err) + } + + if len(cfg.Hooks.PostToolUse) != 1 { + t.Fatalf("expected 1 PostToolUse hook, got %d", len(cfg.Hooks.PostToolUse)) + } + hook := cfg.Hooks.PostToolUse[0] + if hook.Matcher != "exec" { + t.Errorf("expected matcher 'exec', got %q", hook.Matcher) + } + if hook.Command != "~/.picoclaw/scripts/error-detector.sh" { + t.Errorf("expected command path, got %q", hook.Command) + } + if !hook.InjectOutput { + t.Error("expected inject_output true") + } + + if len(cfg.Hooks.PreMessage) != 1 { + t.Fatalf("expected 1 PreMessage hook, got %d", len(cfg.Hooks.PreMessage)) + } + preHook := cfg.Hooks.PreMessage[0] + if preHook.Matcher != "" { + t.Errorf("expected empty matcher, got %q", preHook.Matcher) + } +} + +func TestConfig_HooksConfig_EmptyByDefault(t *testing.T) { + cfg := DefaultConfig() + if len(cfg.Hooks.PreMessage) != 0 { + t.Error("expected no PreMessage hooks by default") + } + if len(cfg.Hooks.PostMessage) != 0 { + t.Error("expected no PostMessage hooks by default") + } + if len(cfg.Hooks.PreToolUse) != 0 { + t.Error("expected no PreToolUse hooks by default") + } + if len(cfg.Hooks.PostToolUse) != 0 { + t.Error("expected no PostToolUse hooks by default") + } +} + +func TestConfig_HooksConfig_OmittedFromJSON(t *testing.T) { + cfg := DefaultConfig() + data, err := json.Marshal(cfg) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if strings.Contains(string(data), `"hooks"`) { + t.Error("expected hooks to be omitted from JSON when empty") + } +} diff --git a/pkg/hooks/executor.go b/pkg/hooks/executor.go new file mode 100644 index 000000000..dcf1e76e6 --- /dev/null +++ b/pkg/hooks/executor.go @@ -0,0 +1,123 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package hooks + +import ( + "bytes" + "context" + "fmt" + "os" + "os/exec" + "runtime" + "strings" + "time" +) + +const ( + // DefaultTimeout is the maximum time a hook script can run. + DefaultTimeout = 30 * time.Second + + // MaxOutputBytes is the maximum bytes captured from hook stdout. + MaxOutputBytes = 64 * 1024 // 64KB +) + +// Executor runs hook commands via os/exec. +type Executor struct { + Timeout time.Duration + MaxOutputBytes int +} + +// NewExecutor creates an Executor with default settings. +func NewExecutor() *Executor { + return &Executor{ + Timeout: DefaultTimeout, + MaxOutputBytes: MaxOutputBytes, + } +} + +// Run executes a command with the given stdin data and extra environment variables. +// It returns the stdout output (truncated to MaxOutputBytes) and any error. +// The command inherits the current process environment plus the extra env vars. +func (e *Executor) Run(ctx context.Context, command string, stdinData []byte, extraEnv []string) HookResult { + if strings.TrimSpace(command) == "" { + return HookResult{Err: fmt.Errorf("empty hook command")} + } + + // Expand ~ in command path + command = expandHome(command) + + timeout := e.Timeout + if timeout <= 0 { + timeout = DefaultTimeout + } + + cmdCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + var cmd *exec.Cmd + if runtime.GOOS == "windows" { + cmd = exec.CommandContext(cmdCtx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command) + } else { + cmd = exec.CommandContext(cmdCtx, "sh", "-c", command) + } + + // Inherit current env + hook-specific vars + cmd.Env = append(os.Environ(), extraEnv...) + + // Pass payload via stdin + if len(stdinData) > 0 { + cmd.Stdin = bytes.NewReader(stdinData) + } + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + + output := stdout.String() + maxOutput := e.MaxOutputBytes + if maxOutput <= 0 { + maxOutput = MaxOutputBytes + } + if len(output) > maxOutput { + output = output[:maxOutput] + } + + if err != nil { + if cmdCtx.Err() == context.DeadlineExceeded { + return HookResult{ + Output: output, + Err: fmt.Errorf("hook timed out after %v: %s", timeout, command), + } + } + stderrStr := stderr.String() + if len(stderrStr) > 1024 { + stderrStr = stderrStr[:1024] + } + return HookResult{ + Output: output, + Err: fmt.Errorf("hook failed: %w (stderr: %s)", err, stderrStr), + } + } + + return HookResult{Output: output} +} + +// expandHome replaces leading ~ with the user's home directory. +func expandHome(path string) string { + if path == "" || path[0] != '~' { + return path + } + home, err := os.UserHomeDir() + if err != nil { + return path + } + if len(path) > 1 && path[1] == '/' { + return home + path[1:] + } + return home +} diff --git a/pkg/hooks/hooks.go b/pkg/hooks/hooks.go new file mode 100644 index 000000000..65d9a4bbc --- /dev/null +++ b/pkg/hooks/hooks.go @@ -0,0 +1,268 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package hooks + +import ( + "context" + "encoding/json" + "strings" + "sync" + + "github.com/sipeed/picoclaw/pkg/logger" +) + +// Event represents a hook trigger point in the agent lifecycle. +type Event string + +const ( + // PreMessage fires before the agent processes a user message. + // inject_output appends hook stdout to user message context. + PreMessage Event = "PreMessage" + + // PostMessage fires after the agent sends its final response. + // Typically used for logging/analytics (fire-and-forget). + PostMessage Event = "PostMessage" + + // PreToolUse fires before a tool's Execute method is called. + // Typically used for validation/logging (fire-and-forget). + PreToolUse Event = "PreToolUse" + + // PostToolUse fires after a tool's Execute method completes. + // inject_output appends hook stdout to result.ForLLM. + PostToolUse Event = "PostToolUse" +) + +// HookRule defines a single hook configuration entry. +type HookRule struct { + // Matcher filters by tool name (for tool hooks). + // Empty string matches all tools. + // Supports exact match ("exec"), wildcard ("*"), and prefix glob ("mcp_*"). + Matcher string `json:"matcher"` + + // Command is the shell command to execute. + // Supports ~ expansion for home directory. + Command string `json:"command"` + + // InjectOutput controls whether hook stdout is appended to the result. + // For PostToolUse: appended to result.ForLLM. + // For PreMessage: appended to user message context. + InjectOutput bool `json:"inject_output"` +} + +// HookPayload is the JSON payload passed to hook scripts via stdin. +type HookPayload struct { + Event Event `json:"event"` + ToolName string `json:"tool_name,omitempty"` + ToolArgs map[string]any `json:"tool_args,omitempty"` + ToolOutput string `json:"tool_output,omitempty"` + ToolError bool `json:"tool_error,omitempty"` + Channel string `json:"channel,omitempty"` + ChatID string `json:"chat_id,omitempty"` + Message string `json:"message,omitempty"` +} + +// HookResult contains the output from a hook execution. +type HookResult struct { + Output string // stdout from the hook script + Err error // execution error, if any +} + +// HookManager manages hook rules and triggers hook execution. +// It is safe for concurrent use. +type HookManager struct { + rules map[Event][]HookRule + executor *Executor + mu sync.RWMutex +} + +// NewHookManager creates a HookManager from configuration. +// If rules is nil or empty, the manager is effectively a no-op. +func NewHookManager(rules map[Event][]HookRule) *HookManager { + if rules == nil { + rules = make(map[Event][]HookRule) + } + return &HookManager{ + rules: rules, + executor: NewExecutor(), + } +} + +// HasHooks returns true if any hooks are configured for the given event. +// This is a fast-path check to avoid unnecessary work when no hooks exist. +func (hm *HookManager) HasHooks(event Event) bool { + hm.mu.RLock() + defer hm.mu.RUnlock() + return len(hm.rules[event]) > 0 +} + +// Trigger fires all matching hooks for the given event and payload. +// It returns a slice of HookResults (one per matching rule). +// Non-matching rules (by matcher) are skipped. +// Hooks are executed sequentially in configuration order to ensure +// deterministic output ordering. +func (hm *HookManager) Trigger(ctx context.Context, event Event, payload HookPayload) []HookResult { + hm.mu.RLock() + rules := hm.rules[event] + hm.mu.RUnlock() + + if len(rules) == 0 { + return nil + } + + payload.Event = event + + // Build env vars from payload + env := buildEnvVars(payload) + + // Serialize payload to JSON for stdin + stdinData, err := json.Marshal(payload) + if err != nil { + logger.ErrorCF("hooks", "Failed to marshal hook payload", + map[string]any{"event": string(event), "error": err.Error()}) + return nil + } + + var results []HookResult + for _, rule := range rules { + // Matcher filtering: empty matcher matches everything; + // non-empty must match the tool name. + if rule.Matcher != "" && !matchesToolName(rule.Matcher, payload.ToolName) { + continue + } + + logger.InfoCF("hooks", "Triggering hook", + map[string]any{ + "event": string(event), + "command": rule.Command, + "matcher": rule.Matcher, + "tool": payload.ToolName, + }) + + result := hm.executor.Run(ctx, rule.Command, stdinData, env) + if result.Err != nil { + logger.WarnCF("hooks", "Hook execution failed", + map[string]any{ + "event": string(event), + "command": rule.Command, + "error": result.Err.Error(), + }) + } else { + logger.DebugCF("hooks", "Hook completed", + map[string]any{ + "event": string(event), + "command": rule.Command, + "output_len": len(result.Output), + }) + } + + results = append(results, result) + } + + return results +} + +// CollectInjectedOutput triggers all hooks for the given event and returns +// concatenated stdout from hooks that have inject_output=true and succeeded. +// Hooks with inject_output=false are still executed for side effects. +func (hm *HookManager) CollectInjectedOutput( + ctx context.Context, + event Event, + payload HookPayload, +) string { + hm.mu.RLock() + rules := hm.rules[event] + hm.mu.RUnlock() + + if len(rules) == 0 { + return "" + } + + payload.Event = event + env := buildEnvVars(payload) + stdinData, err := json.Marshal(payload) + if err != nil { + return "" + } + + var parts []string + for _, rule := range rules { + if rule.Matcher != "" && !matchesToolName(rule.Matcher, payload.ToolName) { + continue + } + + logger.InfoCF("hooks", "Triggering hook", + map[string]any{ + "event": string(event), + "command": rule.Command, + "inject_output": rule.InjectOutput, + }) + + result := hm.executor.Run(ctx, rule.Command, stdinData, env) + + if result.Err != nil { + logger.WarnCF("hooks", "Hook execution failed", + map[string]any{ + "event": string(event), + "command": rule.Command, + "error": result.Err.Error(), + }) + continue + } + + if rule.InjectOutput && strings.TrimSpace(result.Output) != "" { + parts = append(parts, strings.TrimSpace(result.Output)) + } + } + + return strings.Join(parts, "\n") +} + +// matchesToolName checks if the matcher matches the tool name. +// Supports exact match, wildcard "*" (matches everything), +// and prefix glob with trailing "*" (e.g. "mcp_*" matches "mcp_github"). +func matchesToolName(matcher, toolName string) bool { + if matcher == "*" { + return true + } + if strings.HasSuffix(matcher, "*") { + prefix := strings.TrimSuffix(matcher, "*") + return strings.HasPrefix(toolName, prefix) + } + return matcher == toolName +} + +// buildEnvVars constructs the environment variable slice for hook scripts. +// Large values (like tool output) are truncated to prevent E2BIG errors; +// the full payload is available via stdin JSON. +func buildEnvVars(payload HookPayload) []string { + env := []string{ + "PICOCLAW_HOOK_EVENT=" + string(payload.Event), + } + if payload.ToolName != "" { + env = append(env, "PICOCLAW_TOOL_NAME="+payload.ToolName) + } + if payload.ToolOutput != "" { + // Truncate tool output in env var to prevent arg-list-too-long errors. + // Full output is available via stdin JSON. + output := payload.ToolOutput + if len(output) > 8192 { + output = output[:8192] + } + env = append(env, "PICOCLAW_TOOL_OUTPUT="+output) + } + if payload.ToolError { + env = append(env, "PICOCLAW_TOOL_ERROR=true") + } else { + env = append(env, "PICOCLAW_TOOL_ERROR=false") + } + if payload.Channel != "" { + env = append(env, "PICOCLAW_CHANNEL="+payload.Channel) + } + if payload.ChatID != "" { + env = append(env, "PICOCLAW_CHAT_ID="+payload.ChatID) + } + return env +} diff --git a/pkg/hooks/hooks_test.go b/pkg/hooks/hooks_test.go new file mode 100644 index 000000000..b035a884b --- /dev/null +++ b/pkg/hooks/hooks_test.go @@ -0,0 +1,375 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package hooks + +import ( + "context" + "strings" + "testing" + "time" +) + +func TestMatchesToolName_ExactMatch(t *testing.T) { + if !matchesToolName("exec", "exec") { + t.Error("expected exact match") + } + if matchesToolName("exec", "read_file") { + t.Error("expected no match for different name") + } +} + +func TestMatchesToolName_Wildcard(t *testing.T) { + if !matchesToolName("*", "exec") { + t.Error("expected wildcard to match anything") + } + if !matchesToolName("*", "") { + t.Error("expected wildcard to match empty") + } +} + +func TestMatchesToolName_PrefixGlob(t *testing.T) { + if !matchesToolName("mcp_*", "mcp_github") { + t.Error("expected prefix glob to match mcp_github") + } + if matchesToolName("mcp_*", "exec") { + t.Error("expected prefix glob to not match exec") + } +} + +func TestMatchesToolName_EmptyMatcher(t *testing.T) { + // Empty matcher is handled by the caller (Trigger), not matchesToolName. + // matchesToolName("", "exec") does exact match against "", which is false. + if matchesToolName("", "exec") { + t.Error("empty matcher should not match via matchesToolName") + } +} + +func TestBuildEnvVars_Complete(t *testing.T) { + payload := HookPayload{ + Event: PostToolUse, + ToolName: "exec", + ToolOutput: "hello", + ToolError: true, + Channel: "telegram", + ChatID: "chat-42", + } + + env := buildEnvVars(payload) + envMap := make(map[string]string) + for _, e := range env { + parts := strings.SplitN(e, "=", 2) + if len(parts) == 2 { + envMap[parts[0]] = parts[1] + } + } + + if envMap["PICOCLAW_HOOK_EVENT"] != "PostToolUse" { + t.Errorf("expected event PostToolUse, got %q", envMap["PICOCLAW_HOOK_EVENT"]) + } + if envMap["PICOCLAW_TOOL_NAME"] != "exec" { + t.Errorf("expected tool_name exec, got %q", envMap["PICOCLAW_TOOL_NAME"]) + } + if envMap["PICOCLAW_TOOL_OUTPUT"] != "hello" { + t.Errorf("expected tool_output hello, got %q", envMap["PICOCLAW_TOOL_OUTPUT"]) + } + if envMap["PICOCLAW_TOOL_ERROR"] != "true" { + t.Errorf("expected tool_error true, got %q", envMap["PICOCLAW_TOOL_ERROR"]) + } + if envMap["PICOCLAW_CHANNEL"] != "telegram" { + t.Errorf("expected channel telegram, got %q", envMap["PICOCLAW_CHANNEL"]) + } + if envMap["PICOCLAW_CHAT_ID"] != "chat-42" { + t.Errorf("expected chat_id chat-42, got %q", envMap["PICOCLAW_CHAT_ID"]) + } +} + +func TestBuildEnvVars_TruncatesLargeOutput(t *testing.T) { + bigOutput := strings.Repeat("x", 10000) + payload := HookPayload{ + ToolOutput: bigOutput, + } + env := buildEnvVars(payload) + for _, e := range env { + if strings.HasPrefix(e, "PICOCLAW_TOOL_OUTPUT=") { + val := strings.TrimPrefix(e, "PICOCLAW_TOOL_OUTPUT=") + if len(val) > 8192 { + t.Errorf("expected output truncated to 8192, got %d", len(val)) + } + } + } +} + +func TestBuildEnvVars_ErrorFalse(t *testing.T) { + payload := HookPayload{ + ToolError: false, + } + env := buildEnvVars(payload) + for _, e := range env { + if strings.HasPrefix(e, "PICOCLAW_TOOL_ERROR=") { + val := strings.TrimPrefix(e, "PICOCLAW_TOOL_ERROR=") + if val != "false" { + t.Errorf("expected tool_error false, got %q", val) + } + } + } +} + +func TestNewHookManager_NilRules(t *testing.T) { + hm := NewHookManager(nil) + if hm.HasHooks(PreMessage) { + t.Error("expected no hooks for nil rules") + } +} + +func TestNewHookManager_EmptyRules(t *testing.T) { + hm := NewHookManager(map[Event][]HookRule{}) + if hm.HasHooks(PostToolUse) { + t.Error("expected no hooks for empty rules") + } +} + +func TestHookManager_HasHooks(t *testing.T) { + rules := map[Event][]HookRule{ + PostToolUse: { + {Matcher: "exec", Command: "echo test"}, + }, + } + hm := NewHookManager(rules) + + if !hm.HasHooks(PostToolUse) { + t.Error("expected hooks for PostToolUse") + } + if hm.HasHooks(PreMessage) { + t.Error("expected no hooks for PreMessage") + } +} + +func TestHookManager_Trigger_EchoCommand(t *testing.T) { + rules := map[Event][]HookRule{ + PostToolUse: { + {Matcher: "", Command: "echo hook-output"}, + }, + } + hm := NewHookManager(rules) + + results := hm.Trigger(context.Background(), PostToolUse, HookPayload{ + ToolName: "exec", + }) + + 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) != "hook-output" { + t.Errorf("expected output 'hook-output', got %q", results[0].Output) + } +} + +func TestHookManager_Trigger_MatcherFilters(t *testing.T) { + rules := map[Event][]HookRule{ + PostToolUse: { + {Matcher: "read_file", Command: "echo should-not-run"}, + }, + } + hm := NewHookManager(rules) + + results := hm.Trigger(context.Background(), PostToolUse, HookPayload{ + ToolName: "exec", + }) + + if len(results) != 0 { + t.Errorf("expected 0 results (matcher filtered), got %d", len(results)) + } +} + +func TestHookManager_Trigger_Timeout(t *testing.T) { + rules := map[Event][]HookRule{ + PreMessage: { + {Command: "sleep 10"}, + }, + } + hm := NewHookManager(rules) + hm.executor.Timeout = 100 * time.Millisecond + + results := hm.Trigger(context.Background(), PreMessage, HookPayload{}) + + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + if results[0].Err == nil { + t.Error("expected timeout error") + } + if !strings.Contains(results[0].Err.Error(), "timed out") { + t.Errorf("expected timeout error message, got %q", results[0].Err.Error()) + } +} + +func TestHookManager_CollectInjectedOutput(t *testing.T) { + rules := map[Event][]HookRule{ + PostToolUse: { + {Matcher: "exec", Command: "echo injected-context", InjectOutput: true}, + {Matcher: "exec", Command: "echo not-injected", InjectOutput: false}, + }, + } + hm := NewHookManager(rules) + + output := hm.CollectInjectedOutput(context.Background(), PostToolUse, HookPayload{ + ToolName: "exec", + }) + + trimmed := strings.TrimSpace(output) + if trimmed != "injected-context" { + t.Errorf("expected only inject_output=true content, got %q", trimmed) + } +} + +func TestHookManager_Trigger_StdinPayload(t *testing.T) { + // Verify that the hook receives JSON payload via stdin + rules := map[Event][]HookRule{ + PostToolUse: { + {Matcher: "", Command: "cat", InjectOutput: true}, + }, + } + hm := NewHookManager(rules) + + output := hm.CollectInjectedOutput(context.Background(), PostToolUse, HookPayload{ + ToolName: "exec", + ToolOutput: "test-output", + Channel: "telegram", + }) + + if !strings.Contains(output, `"tool_name":"exec"`) { + t.Errorf("expected stdin to contain tool_name, got %q", output) + } + if !strings.Contains(output, `"channel":"telegram"`) { + t.Errorf("expected stdin to contain channel, got %q", output) + } +} + +func TestHookManager_Trigger_FailedCommand(t *testing.T) { + rules := map[Event][]HookRule{ + PreMessage: { + {Command: "exit 1"}, + }, + } + hm := NewHookManager(rules) + + results := hm.Trigger(context.Background(), PreMessage, HookPayload{}) + + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + if results[0].Err == nil { + t.Error("expected error for exit 1") + } +} + +func TestHookManager_Trigger_NoRulesForEvent(t *testing.T) { + rules := map[Event][]HookRule{ + PostToolUse: {{Command: "echo test"}}, + } + hm := NewHookManager(rules) + + results := hm.Trigger(context.Background(), PreMessage, HookPayload{}) + if results != nil { + t.Errorf("expected nil results for unregistered event, got %d", len(results)) + } +} + +func TestHookManager_Trigger_MultipleRules(t *testing.T) { + rules := map[Event][]HookRule{ + PostToolUse: { + {Matcher: "", Command: "echo first"}, + {Matcher: "", Command: "echo second"}, + }, + } + hm := NewHookManager(rules) + + results := hm.Trigger(context.Background(), PostToolUse, HookPayload{ + ToolName: "exec", + }) + + if len(results) != 2 { + t.Fatalf("expected 2 results, got %d", len(results)) + } + if strings.TrimSpace(results[0].Output) != "first" { + t.Errorf("expected first result 'first', got %q", results[0].Output) + } + if strings.TrimSpace(results[1].Output) != "second" { + t.Errorf("expected second result 'second', got %q", results[1].Output) + } +} + +func TestHookManager_CollectInjectedOutput_MultipleInject(t *testing.T) { + rules := map[Event][]HookRule{ + PreMessage: { + {Command: "echo line1", InjectOutput: true}, + {Command: "echo line2", InjectOutput: true}, + }, + } + hm := NewHookManager(rules) + + output := hm.CollectInjectedOutput(context.Background(), PreMessage, HookPayload{}) + + if !strings.Contains(output, "line1") || !strings.Contains(output, "line2") { + t.Errorf("expected both lines in output, got %q", output) + } +} + +func TestExpandHome(t *testing.T) { + result := expandHome("/absolute/path") + if result != "/absolute/path" { + t.Errorf("expected absolute path unchanged, got %q", result) + } + + result = expandHome("") + if result != "" { + t.Errorf("expected empty string unchanged, got %q", result) + } + + // Verify ~ expansion doesn't panic + result = expandHome("~/test") + if strings.HasPrefix(result, "~") { + t.Log("Home dir expansion may not work in test env, skipping") + } +} + +func TestExecutor_MaxOutputTruncation(t *testing.T) { + e := NewExecutor() + e.MaxOutputBytes = 10 + + result := e.Run(context.Background(), "echo 'this is a long output that should be truncated'", nil, nil) + if result.Err != nil { + t.Fatalf("expected no error, got %v", result.Err) + } + if len(result.Output) > 10 { + t.Errorf("expected output truncated to 10 bytes, got %d", len(result.Output)) + } +} + +func TestExecutor_EmptyCommand(t *testing.T) { + e := NewExecutor() + result := e.Run(context.Background(), "", nil, nil) + if result.Err == nil { + t.Error("expected error for empty command") + } + if !strings.Contains(result.Err.Error(), "empty hook command") { + t.Errorf("expected empty command error, got %q", result.Err.Error()) + } +} + +func TestExecutor_ExtraEnvVars(t *testing.T) { + e := NewExecutor() + result := e.Run(context.Background(), "echo $TEST_HOOK_VAR", nil, []string{"TEST_HOOK_VAR=hello-hook"}) + if result.Err != nil { + t.Fatalf("expected no error, got %v", result.Err) + } + if strings.TrimSpace(result.Output) != "hello-hook" { + t.Errorf("expected 'hello-hook', got %q", result.Output) + } +} diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index ca8436c67..9ef62e6c2 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -7,12 +7,14 @@ import ( "sync" "time" + "github.com/sipeed/picoclaw/pkg/hooks" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" ) type ToolRegistry struct { tools map[string]Tool + hooks *hooks.HookManager mu sync.RWMutex } @@ -22,6 +24,13 @@ func NewToolRegistry() *ToolRegistry { } } +// SetHookManager injects a HookManager for pre/post tool execution hooks. +func (r *ToolRegistry) SetHookManager(hm *hooks.HookManager) { + r.mu.Lock() + defer r.mu.Unlock() + r.hooks = hm +} + func (r *ToolRegistry) Register(tool Tool) { r.mu.Lock() defer r.mu.Unlock() @@ -74,6 +83,16 @@ func (r *ToolRegistry) ExecuteWithContext( // Always inject — tools validate what they require. ctx = WithToolContext(ctx, channel, chatID) + // Pre-tool hook + if r.hooks != nil && r.hooks.HasHooks(hooks.PreToolUse) { + r.hooks.Trigger(ctx, hooks.PreToolUse, hooks.HookPayload{ + ToolName: name, + ToolArgs: args, + Channel: channel, + ChatID: chatID, + }) + } + // If tool implements AsyncExecutor and callback is provided, use ExecuteAsync. // The callback is a call parameter, not mutable state on the tool instance. var result *ToolResult @@ -89,6 +108,21 @@ func (r *ToolRegistry) ExecuteWithContext( } duration := time.Since(start) + // Post-tool hook + if r.hooks != nil && r.hooks.HasHooks(hooks.PostToolUse) { + injected := r.hooks.CollectInjectedOutput(ctx, hooks.PostToolUse, hooks.HookPayload{ + ToolName: name, + ToolArgs: args, + ToolOutput: result.ForLLM, + ToolError: result.IsError, + Channel: channel, + ChatID: chatID, + }) + if injected != "" { + result.ForLLM = result.ForLLM + "\n\n[Hook Output]\n" + injected + } + } + // Log based on result type if result.IsError { logger.ErrorCF("tool", "Tool execution failed",