diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 693f2227b..41aa58619 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -20,6 +20,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/providers" "github.com/sipeed/picoclaw/pkg/routing" @@ -38,6 +39,7 @@ type AgentLoop struct { summarizing sync.Map fallback *providers.FallbackChain channelManager *channels.Manager + hooks *hooks.HookRegistry } // processOptions configures how a message is processed @@ -185,7 +187,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { } if !alreadySent { - al.bus.PublishOutbound(bus.OutboundMessage{ + al.sendOutbound(ctx, bus.OutboundMessage{ Channel: msg.Channel, ChatID: msg.ChatID, Content: response, @@ -214,6 +216,56 @@ func (al *AgentLoop) SetChannelManager(cm *channels.Manager) { al.channelManager = cm } +// SetHooks installs a hook registry. Must be called before Run starts. +func (al *AgentLoop) SetHooks(h *hooks.HookRegistry) { + al.hooks = h + + // Rewire MessageTool callbacks to route through sendOutbound for hook interception. + for _, agentID := range al.registry.ListAgentIDs() { + if agent, ok := al.registry.GetAgent(agentID); ok { + if tool, ok := agent.Tools.Get("message"); ok { + if mt, ok := tool.(*tools.MessageTool); ok { + mt.SetSendCallback(func(channel, chatID, content string) error { + if !al.sendOutbound(context.Background(), bus.OutboundMessage{ + Channel: channel, + ChatID: chatID, + Content: content, + }) { + return fmt.Errorf("message canceled by hook") + } + return nil + }) + } + } + } + } +} + +// sendOutbound wraps bus.PublishOutbound with the message_sending hook. +// Returns true if the message was sent, false if canceled by a hook. +func (al *AgentLoop) sendOutbound(ctx context.Context, msg bus.OutboundMessage) bool { + if al.hooks != nil { + event := &hooks.MessageSendingEvent{Channel: msg.Channel, ChatID: msg.ChatID, Content: msg.Content} + al.hooks.TriggerMessageSending(ctx, event) + if event.Cancel { + reason := event.CancelReason + if reason == "" { + reason = "unspecified" + } + logger.WarnCF("hooks", "Outbound message canceled by hook", + map[string]any{ + "channel": msg.Channel, + "chat_id": msg.ChatID, + "reason": reason, + }) + return false + } + msg.Content = event.Content + } + al.bus.PublishOutbound(msg) + return true +} + // RecordLastChannel records the last active channel for this workspace. // This uses the atomic state save mechanism to prevent data loss on crash. func (al *AgentLoop) RecordLastChannel(channel string) error { @@ -283,6 +335,18 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) "session_key": msg.SessionKey, }) + // Fire message_received hook + if al.hooks != nil { + al.hooks.TriggerMessageReceived(ctx, &hooks.MessageReceivedEvent{ + Channel: msg.Channel, + SenderID: msg.SenderID, + ChatID: msg.ChatID, + Content: msg.Content, + Media: msg.Media, + Metadata: msg.Metadata, + }) + } + // Route system messages to processSystemMessage if msg.Channel == "system" { return al.processSystemMessage(ctx, msg) @@ -404,6 +468,18 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt // 1. Update tool contexts al.updateToolContexts(agent, opts.Channel, opts.ChatID) + // Fire session hooks + if al.hooks != nil { + sessionEvt := &hooks.SessionEvent{ + AgentID: agent.ID, + SessionKey: opts.SessionKey, + Channel: opts.Channel, + ChatID: opts.ChatID, + } + al.hooks.TriggerSessionStart(ctx, sessionEvt) + defer al.hooks.TriggerSessionEnd(ctx, sessionEvt) + } + // 2. Build messages (skip history for heartbeat) var history []providers.Message var summary string @@ -443,12 +519,12 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt // 7. Optional: summarization if opts.EnableSummary { - al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID) + al.maybeSummarize(ctx, agent, opts.SessionKey, opts.Channel, opts.ChatID) } // 8. Optional: send response via bus if opts.SendResponse { - al.bus.PublishOutbound(bus.OutboundMessage{ + al.sendOutbound(ctx, bus.OutboundMessage{ Channel: opts.Channel, ChatID: opts.ChatID, Content: finalContent, @@ -545,8 +621,19 @@ func (al *AgentLoop) runLLMIteration( } // Retry loop for context/token errors + llmStart := time.Now() maxRetries := 2 for retry := 0; retry <= maxRetries; retry++ { + // Fire llm_input hook (re-fires after compression so hooks see actual messages) + if al.hooks != nil { + al.hooks.TriggerLLMInput(ctx, &hooks.LLMInputEvent{ + AgentID: agent.ID, + Model: agent.Model, + Messages: messages, + Tools: providerToolDefs, + Iteration: iteration, + }) + } response, err = callLLM() if err == nil { break @@ -565,7 +652,7 @@ func (al *AgentLoop) runLLMIteration( }) if retry == 0 && !constants.IsInternalChannel(opts.Channel) { - al.bus.PublishOutbound(bus.OutboundMessage{ + al.sendOutbound(ctx, bus.OutboundMessage{ Channel: opts.Channel, ChatID: opts.ChatID, Content: "Context window exceeded. Compressing history and retrying...", @@ -584,6 +671,8 @@ func (al *AgentLoop) runLLMIteration( break } + llmDuration := time.Since(llmStart) + if err != nil { logger.ErrorCF("agent", "LLM call failed", map[string]any{ @@ -594,6 +683,18 @@ func (al *AgentLoop) runLLMIteration( return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err) } + // Fire llm_output hook + if al.hooks != nil { + al.hooks.TriggerLLMOutput(ctx, &hooks.LLMOutputEvent{ + AgentID: agent.ID, + Model: agent.Model, + Content: response.Content, + ToolCalls: response.ToolCalls, + Iteration: iteration, + Duration: llmDuration, + }) + } + // Check if no tool calls - we're done if len(response.ToolCalls) == 0 { finalContent = response.Content @@ -684,18 +785,53 @@ func (al *AgentLoop) runLLMIteration( } } - toolResult := agent.Tools.ExecuteWithContext( - ctx, - tc.Name, - tc.Arguments, - opts.Channel, - opts.ChatID, - asyncCallback, - ) + // Fire before_tool_call hook + var toolResult *tools.ToolResult + toolCanceled := false + if al.hooks != nil { + args := tc.Arguments + if args == nil { + args = make(map[string]any) + } + btcEvent := &hooks.BeforeToolCallEvent{ + ToolName: tc.Name, + Args: args, + Channel: opts.Channel, + ChatID: opts.ChatID, + } + al.hooks.TriggerBeforeToolCall(ctx, btcEvent) + if btcEvent.Cancel { + toolCanceled = true + reason := btcEvent.CancelReason + if strings.TrimSpace(reason) == "" { + reason = fmt.Sprintf("tool call %q was canceled by before_tool_call hook", tc.Name) + } + toolResult = tools.ErrorResult(reason) + } + tc.Arguments = btcEvent.Args + } + + toolStart := time.Now() + if !toolCanceled { + toolResult = agent.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, opts.Channel, opts.ChatID, asyncCallback) + } + toolDuration := time.Since(toolStart) + + // Fire after_tool_call hook (fires for both executed and canceled calls) + if al.hooks != nil { + al.hooks.TriggerAfterToolCall(ctx, &hooks.AfterToolCallEvent{ + ToolName: tc.Name, + Args: tc.Arguments, + Channel: opts.Channel, + ChatID: opts.ChatID, + Duration: toolDuration, + Result: toolResult, + }) + } // Send ForUser content to user immediately if not Silent if !toolResult.Silent && toolResult.ForUser != "" && opts.SendResponse { - al.bus.PublishOutbound(bus.OutboundMessage{ + al.sendOutbound(ctx, bus.OutboundMessage{ Channel: opts.Channel, ChatID: opts.ChatID, Content: toolResult.ForUser, @@ -749,7 +885,7 @@ func (al *AgentLoop) updateToolContexts(agent *AgentInstance, channel, chatID st } // maybeSummarize triggers summarization if the session history exceeds thresholds. -func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, chatID string) { +func (al *AgentLoop) maybeSummarize(ctx context.Context, agent *AgentInstance, sessionKey, channel, chatID string) { newHistory := agent.Sessions.GetHistory(sessionKey) tokenEstimate := al.estimateTokens(newHistory) threshold := agent.ContextWindow * 75 / 100 @@ -759,6 +895,13 @@ func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, c if _, loading := al.summarizing.LoadOrStore(summarizeKey, true); !loading { go func() { defer al.summarizing.Delete(summarizeKey) + if !constants.IsInternalChannel(channel) { + al.sendOutbound(ctx, bus.OutboundMessage{ + Channel: channel, + ChatID: chatID, + Content: "Memory threshold reached. Optimizing conversation history...", + }) + } logger.Debug("Memory threshold reached. Optimizing conversation history...") al.summarizeSession(agent, sessionKey) }() diff --git a/pkg/hooks/hooks.go b/pkg/hooks/hooks.go new file mode 100644 index 000000000..24dc8cbff --- /dev/null +++ b/pkg/hooks/hooks.go @@ -0,0 +1,270 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Inspired by and based on nanobot: https://github.com/HKUDS/nanobot +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package hooks + +import ( + "context" + "fmt" + "sync" + + "github.com/sipeed/picoclaw/pkg/logger" +) + +// HookHandler is the callback signature for all hooks. +type HookHandler[T any] func(ctx context.Context, event *T) error + +// HookRegistration tracks a handler with its priority and name. +type HookRegistration[T any] struct { + Handler HookHandler[T] + Priority int // Lower = runs first + Name string +} + +// HookRegistry manages all lifecycle hooks. +type HookRegistry struct { + messageReceived []HookRegistration[MessageReceivedEvent] + messageSending []HookRegistration[MessageSendingEvent] + beforeToolCall []HookRegistration[BeforeToolCallEvent] + afterToolCall []HookRegistration[AfterToolCallEvent] + llmInput []HookRegistration[LLMInputEvent] + llmOutput []HookRegistration[LLMOutputEvent] + sessionStart []HookRegistration[SessionEvent] + sessionEnd []HookRegistration[SessionEvent] + mu sync.RWMutex +} + +// NewHookRegistry creates an empty hook registry. +func NewHookRegistry() *HookRegistry { + return &HookRegistry{} +} + +// insertSorted inserts a registration into a new slice sorted by priority. +// Always allocates a new backing array so concurrent readers of the old slice are safe. +func insertSorted[T any](slice []HookRegistration[T], reg HookRegistration[T]) []HookRegistration[T] { + i := 0 + for i < len(slice) && slice[i].Priority <= reg.Priority { + i++ + } + result := make([]HookRegistration[T], len(slice)+1) + copy(result, slice[:i]) + result[i] = reg + copy(result[i+1:], slice[i:]) + return result +} + +// Registration methods + +func (r *HookRegistry) OnMessageReceived(name string, priority int, handler HookHandler[MessageReceivedEvent]) { + r.mu.Lock() + defer r.mu.Unlock() + r.messageReceived = insertSorted(r.messageReceived, HookRegistration[MessageReceivedEvent]{ + Handler: handler, Priority: priority, Name: name, + }) +} + +func (r *HookRegistry) OnMessageSending(name string, priority int, handler HookHandler[MessageSendingEvent]) { + r.mu.Lock() + defer r.mu.Unlock() + r.messageSending = insertSorted(r.messageSending, HookRegistration[MessageSendingEvent]{ + Handler: handler, Priority: priority, Name: name, + }) +} + +func (r *HookRegistry) OnBeforeToolCall(name string, priority int, handler HookHandler[BeforeToolCallEvent]) { + r.mu.Lock() + defer r.mu.Unlock() + r.beforeToolCall = insertSorted(r.beforeToolCall, HookRegistration[BeforeToolCallEvent]{ + Handler: handler, Priority: priority, Name: name, + }) +} + +func (r *HookRegistry) OnAfterToolCall(name string, priority int, handler HookHandler[AfterToolCallEvent]) { + r.mu.Lock() + defer r.mu.Unlock() + r.afterToolCall = insertSorted(r.afterToolCall, HookRegistration[AfterToolCallEvent]{ + Handler: handler, Priority: priority, Name: name, + }) +} + +func (r *HookRegistry) OnLLMInput(name string, priority int, handler HookHandler[LLMInputEvent]) { + r.mu.Lock() + defer r.mu.Unlock() + r.llmInput = insertSorted(r.llmInput, HookRegistration[LLMInputEvent]{ + Handler: handler, Priority: priority, Name: name, + }) +} + +func (r *HookRegistry) OnLLMOutput(name string, priority int, handler HookHandler[LLMOutputEvent]) { + r.mu.Lock() + defer r.mu.Unlock() + r.llmOutput = insertSorted(r.llmOutput, HookRegistration[LLMOutputEvent]{ + Handler: handler, Priority: priority, Name: name, + }) +} + +func (r *HookRegistry) OnSessionStart(name string, priority int, handler HookHandler[SessionEvent]) { + r.mu.Lock() + defer r.mu.Unlock() + r.sessionStart = insertSorted(r.sessionStart, HookRegistration[SessionEvent]{ + Handler: handler, Priority: priority, Name: name, + }) +} + +func (r *HookRegistry) OnSessionEnd(name string, priority int, handler HookHandler[SessionEvent]) { + r.mu.Lock() + defer r.mu.Unlock() + r.sessionEnd = insertSorted(r.sessionEnd, HookRegistration[SessionEvent]{ + Handler: handler, Priority: priority, Name: name, + }) +} + +// Trigger methods — void hooks + +// triggerVoid runs all handlers concurrently and waits for completion. +// Handlers MUST NOT mutate the event — it is shared across goroutines. +// Errors are logged but do not propagate to the caller. +func triggerVoid[T any](ctx context.Context, hooks []HookRegistration[T], event *T, hookName string) { + if len(hooks) == 0 { + return + } + var wg sync.WaitGroup + for _, h := range hooks { + wg.Add(1) + go func(reg HookRegistration[T]) { + defer wg.Done() + defer func() { + if r := recover(); r != nil { + logger.ErrorCF("hooks", "Hook panic", + map[string]any{ + "hook": hookName, + "handler": reg.Name, + "panic": fmt.Sprintf("%v", r), + }) + } + }() + if err := reg.Handler(ctx, event); err != nil { + logger.WarnCF("hooks", "Hook error", + map[string]any{ + "hook": hookName, + "handler": reg.Name, + "error": err.Error(), + }) + } + }(h) + } + wg.Wait() +} + +// triggerModifying runs handlers sequentially by priority, stopping if Cancel is set. +// The cancelCheck function inspects the event to determine if Cancel was set. +func triggerModifying[T any](ctx context.Context, hooks []HookRegistration[T], event *T, hookName string, cancelCheck func(*T) bool) { + if len(hooks) == 0 { + return + } + for _, h := range hooks { + func() { + defer func() { + if r := recover(); r != nil { + logger.ErrorCF("hooks", "Hook panic", + map[string]any{ + "hook": hookName, + "handler": h.Name, + "panic": fmt.Sprintf("%v", r), + }) + } + }() + if err := h.Handler(ctx, event); err != nil { + logger.WarnCF("hooks", "Hook error", + map[string]any{ + "hook": hookName, + "handler": h.Name, + "error": err.Error(), + }) + } + }() + if cancelCheck(event) { + logger.InfoCF("hooks", "Hook canceled operation", + map[string]any{ + "hook": hookName, + "handler": h.Name, + }) + return + } + } +} + +// TriggerMessageReceived fires all message_received handlers concurrently. +// Handlers must not mutate the event. +func (r *HookRegistry) TriggerMessageReceived(ctx context.Context, event *MessageReceivedEvent) { + r.mu.RLock() + hooks := r.messageReceived + r.mu.RUnlock() + triggerVoid(ctx, hooks, event, "message_received") +} + +func (r *HookRegistry) TriggerMessageSending(ctx context.Context, event *MessageSendingEvent) { + r.mu.RLock() + hooks := r.messageSending + r.mu.RUnlock() + triggerModifying(ctx, hooks, event, "message_sending", func(e *MessageSendingEvent) bool { + return e.Cancel + }) +} + +func (r *HookRegistry) TriggerBeforeToolCall(ctx context.Context, event *BeforeToolCallEvent) { + r.mu.RLock() + hooks := r.beforeToolCall + r.mu.RUnlock() + triggerModifying(ctx, hooks, event, "before_tool_call", func(e *BeforeToolCallEvent) bool { + return e.Cancel + }) +} + +// TriggerAfterToolCall fires all after_tool_call handlers concurrently. +// Handlers must not mutate the event. +func (r *HookRegistry) TriggerAfterToolCall(ctx context.Context, event *AfterToolCallEvent) { + r.mu.RLock() + hooks := r.afterToolCall + r.mu.RUnlock() + triggerVoid(ctx, hooks, event, "after_tool_call") +} + +// TriggerLLMInput fires all llm_input handlers concurrently. +// Handlers must not mutate the event. +func (r *HookRegistry) TriggerLLMInput(ctx context.Context, event *LLMInputEvent) { + r.mu.RLock() + hooks := r.llmInput + r.mu.RUnlock() + triggerVoid(ctx, hooks, event, "llm_input") +} + +// TriggerLLMOutput fires all llm_output handlers concurrently. +// Handlers must not mutate the event. +func (r *HookRegistry) TriggerLLMOutput(ctx context.Context, event *LLMOutputEvent) { + r.mu.RLock() + hooks := r.llmOutput + r.mu.RUnlock() + triggerVoid(ctx, hooks, event, "llm_output") +} + +// TriggerSessionStart fires all session_start handlers concurrently. +// Handlers must not mutate the event. +func (r *HookRegistry) TriggerSessionStart(ctx context.Context, event *SessionEvent) { + r.mu.RLock() + hooks := r.sessionStart + r.mu.RUnlock() + triggerVoid(ctx, hooks, event, "session_start") +} + +// TriggerSessionEnd fires all session_end handlers concurrently. +// Handlers must not mutate the event. +func (r *HookRegistry) TriggerSessionEnd(ctx context.Context, event *SessionEvent) { + r.mu.RLock() + hooks := r.sessionEnd + r.mu.RUnlock() + triggerVoid(ctx, hooks, event, "session_end") +} diff --git a/pkg/hooks/hooks_test.go b/pkg/hooks/hooks_test.go new file mode 100644 index 000000000..eb80c2955 --- /dev/null +++ b/pkg/hooks/hooks_test.go @@ -0,0 +1,433 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Inspired by and based on nanobot: https://github.com/HKUDS/nanobot +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package hooks + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestNewHookRegistry(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + // Triggering all hooks on an empty registry should not panic. + r.TriggerMessageReceived(ctx, &MessageReceivedEvent{Content: "hello"}) + r.TriggerMessageSending(ctx, &MessageSendingEvent{Content: "hello"}) + r.TriggerBeforeToolCall(ctx, &BeforeToolCallEvent{ToolName: "t"}) + r.TriggerAfterToolCall(ctx, &AfterToolCallEvent{ToolName: "t"}) + r.TriggerLLMInput(ctx, &LLMInputEvent{AgentID: "a"}) + r.TriggerLLMOutput(ctx, &LLMOutputEvent{AgentID: "a"}) + r.TriggerSessionStart(ctx, &SessionEvent{AgentID: "a"}) + r.TriggerSessionEnd(ctx, &SessionEvent{AgentID: "a"}) +} + +func TestVoidHookExecution(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var called atomic.Bool + r.OnMessageReceived("test", 0, func(_ context.Context, e *MessageReceivedEvent) error { + called.Store(true) + if e.Content != "ping" { + t.Errorf("Expected content 'ping', got '%s'", e.Content) + } + return nil + }) + + r.TriggerMessageReceived(ctx, &MessageReceivedEvent{Content: "ping"}) + + if !called.Load() { + t.Error("Expected handler to be called") + } +} + +func TestVoidHooksConcurrent(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var count atomic.Int32 + started := make(chan struct{}, 5) + release := make(chan struct{}) + done := make(chan struct{}) + + for i := range 5 { + r.OnMessageReceived("hook-"+string(rune('A'+i)), i, func(_ context.Context, _ *MessageReceivedEvent) error { + started <- struct{}{} + <-release + count.Add(1) + return nil + }) + } + + go func() { + r.TriggerMessageReceived(ctx, &MessageReceivedEvent{Content: "test"}) + close(done) + }() + + // All 5 handlers must reach the barrier concurrently. + for i := range 5 { + select { + case <-started: + case <-time.After(1 * time.Second): + t.Fatalf("timeout waiting for handler %d to start", i+1) + } + } + + // Release all handlers. + close(release) + + select { + case <-done: + case <-time.After(1 * time.Second): + t.Fatal("timeout waiting for handlers to complete") + } + + if count.Load() != 5 { + t.Errorf("Expected 5 handlers called, got %d", count.Load()) + } +} + +func TestModifyingHookPriority(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var mu sync.Mutex + var order []string + + // Register in reverse priority order to verify sorting. + r.OnMessageSending("third", 30, func(_ context.Context, _ *MessageSendingEvent) error { + mu.Lock() + order = append(order, "third") + mu.Unlock() + return nil + }) + r.OnMessageSending("first", 10, func(_ context.Context, _ *MessageSendingEvent) error { + mu.Lock() + order = append(order, "first") + mu.Unlock() + return nil + }) + r.OnMessageSending("second", 20, func(_ context.Context, _ *MessageSendingEvent) error { + mu.Lock() + order = append(order, "second") + mu.Unlock() + return nil + }) + + r.TriggerMessageSending(ctx, &MessageSendingEvent{Content: "hi"}) + + if len(order) != 3 { + t.Fatalf("Expected 3 handlers, got %d", len(order)) + } + if order[0] != "first" || order[1] != "second" || order[2] != "third" { + t.Errorf("Expected [first second third], got %v", order) + } +} + +func TestModifyingHookCancel(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var secondCalled bool + + r.OnMessageSending("canceler", 10, func(_ context.Context, e *MessageSendingEvent) error { + e.Cancel = true + e.CancelReason = "blocked" + return nil + }) + r.OnMessageSending("after-cancel", 20, func(_ context.Context, _ *MessageSendingEvent) error { + secondCalled = true + return nil + }) + + event := &MessageSendingEvent{Content: "hi"} + r.TriggerMessageSending(ctx, event) + + if !event.Cancel { + t.Error("Expected Cancel to be true") + } + if secondCalled { + t.Error("Expected second handler NOT to be called after cancel") + } +} + +func TestBeforeToolCallModification(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + r.OnBeforeToolCall("modifier", 10, func(_ context.Context, e *BeforeToolCallEvent) error { + e.Args["injected"] = "value" + return nil + }) + + event := &BeforeToolCallEvent{ + ToolName: "search", + Args: map[string]any{"query": "test"}, + } + r.TriggerBeforeToolCall(ctx, event) + + if event.Args["injected"] != "value" { + t.Error("Expected injected arg to persist") + } + if event.Args["query"] != "test" { + t.Error("Expected original arg to remain") + } +} + +func TestMessageSendingFilter(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + r.OnMessageSending("rewriter", 10, func(_ context.Context, e *MessageSendingEvent) error { + e.Content = "[filtered] " + e.Content + return nil + }) + + event := &MessageSendingEvent{Content: "hello world"} + r.TriggerMessageSending(ctx, event) + + if event.Content != "[filtered] hello world" { + t.Errorf("Expected '[filtered] hello world', got '%s'", event.Content) + } +} + +func TestZeroCostWhenEmpty(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + // This is primarily a safety/smoke test — no panics, no allocations of note. + for range 100 { + r.TriggerMessageReceived(ctx, &MessageReceivedEvent{}) + r.TriggerMessageSending(ctx, &MessageSendingEvent{}) + r.TriggerBeforeToolCall(ctx, &BeforeToolCallEvent{}) + r.TriggerAfterToolCall(ctx, &AfterToolCallEvent{}) + r.TriggerLLMInput(ctx, &LLMInputEvent{}) + r.TriggerLLMOutput(ctx, &LLMOutputEvent{}) + r.TriggerSessionStart(ctx, &SessionEvent{}) + r.TriggerSessionEnd(ctx, &SessionEvent{}) + } +} + +func TestLLMInputOutput(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var inputCalled, outputCalled atomic.Bool + + r.OnLLMInput("input-hook", 0, func(_ context.Context, e *LLMInputEvent) error { + if e.Model != "gpt-4" { + t.Errorf("Expected model 'gpt-4', got '%s'", e.Model) + } + inputCalled.Store(true) + return nil + }) + + r.OnLLMOutput("output-hook", 0, func(_ context.Context, e *LLMOutputEvent) error { + if e.Content != "response" { + t.Errorf("Expected content 'response', got '%s'", e.Content) + } + outputCalled.Store(true) + return nil + }) + + r.TriggerLLMInput(ctx, &LLMInputEvent{AgentID: "a1", Model: "gpt-4", Iteration: 1}) + r.TriggerLLMOutput(ctx, &LLMOutputEvent{AgentID: "a1", Model: "gpt-4", Content: "response", Iteration: 1}) + + if !inputCalled.Load() { + t.Error("Expected LLM input hook to be called") + } + if !outputCalled.Load() { + t.Error("Expected LLM output hook to be called") + } +} + +func TestSessionStartEnd(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var startCalled, endCalled atomic.Bool + + r.OnSessionStart("start-hook", 0, func(_ context.Context, e *SessionEvent) error { + if e.SessionKey != "sess-1" { + t.Errorf("Expected session key 'sess-1', got '%s'", e.SessionKey) + } + startCalled.Store(true) + return nil + }) + + r.OnSessionEnd("end-hook", 0, func(_ context.Context, e *SessionEvent) error { + if e.SessionKey != "sess-1" { + t.Errorf("Expected session key 'sess-1', got '%s'", e.SessionKey) + } + endCalled.Store(true) + return nil + }) + + event := &SessionEvent{AgentID: "a1", SessionKey: "sess-1", Channel: "test", ChatID: "c1"} + r.TriggerSessionStart(ctx, event) + r.TriggerSessionEnd(ctx, event) + + if !startCalled.Load() { + t.Error("Expected session start hook to be called") + } + if !endCalled.Load() { + t.Error("Expected session end hook to be called") + } +} + +func TestConcurrentRegistrationAndTrigger(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var wg sync.WaitGroup + + // Goroutines registering hooks. + for i := range 10 { + wg.Add(1) + go func() { + defer wg.Done() + r.OnMessageReceived("reg-hook", i, func(_ context.Context, _ *MessageReceivedEvent) error { + return nil + }) + }() + } + + // Goroutines triggering hooks concurrently. + for range 10 { + wg.Add(1) + go func() { + defer wg.Done() + r.TriggerMessageReceived(ctx, &MessageReceivedEvent{Content: "race"}) + }() + } + + wg.Wait() +} + +func TestInsertSorted(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var order []int + + // Register with priorities: 50, 10, 30, 20, 40 + priorities := []int{50, 10, 30, 20, 40} + for _, p := range priorities { + r.OnBeforeToolCall("p-"+string(rune('0'+p)), p, func(_ context.Context, _ *BeforeToolCallEvent) error { + order = append(order, p) + return nil + }) + } + + r.TriggerBeforeToolCall(ctx, &BeforeToolCallEvent{ToolName: "test", Args: map[string]any{}}) + + expected := []int{10, 20, 30, 40, 50} + if len(order) != len(expected) { + t.Fatalf("Expected %d handlers, got %d", len(expected), len(order)) + } + for i, v := range expected { + if order[i] != v { + t.Errorf("Position %d: expected priority %d, got %d", i, v, order[i]) + } + } +} + +func TestAfterToolCallExecution(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + var called bool + var capturedName string + r.OnAfterToolCall("logger", 0, func(_ context.Context, event *AfterToolCallEvent) error { + called = true + capturedName = event.ToolName + return nil + }) + + r.TriggerAfterToolCall(ctx, &AfterToolCallEvent{ + ToolName: "shell", + Args: map[string]any{"cmd": "ls"}, + Channel: "telegram", + ChatID: "123", + }) + + if !called { + t.Error("Expected after_tool_call handler to be called") + } + if capturedName != "shell" { + t.Errorf("Expected ToolName 'shell', got '%s'", capturedName) + } +} + +func TestHandlerErrorsSwallowed(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + // Test void hooks: error in one handler doesn't prevent others from running + var secondCalled bool + r.OnMessageReceived("erroring", 10, func(_ context.Context, _ *MessageReceivedEvent) error { + return fmt.Errorf("handler error") + }) + r.OnMessageReceived("observer", 20, func(_ context.Context, _ *MessageReceivedEvent) error { + secondCalled = true + return nil + }) + + r.TriggerMessageReceived(ctx, &MessageReceivedEvent{Content: "test"}) + if !secondCalled { + t.Error("Expected second void handler to run despite first handler's error") + } + + // Test modifying hooks: error doesn't stop chain (only Cancel does) + var modifySecondCalled bool + r.OnMessageSending("erroring", 10, func(_ context.Context, _ *MessageSendingEvent) error { + return fmt.Errorf("handler error") + }) + r.OnMessageSending("modifier", 20, func(_ context.Context, _ *MessageSendingEvent) error { + modifySecondCalled = true + return nil + }) + + r.TriggerMessageSending(ctx, &MessageSendingEvent{Content: "test"}) + if !modifySecondCalled { + t.Error("Expected second modifying handler to run despite first handler's error") + } +} + +func TestPanicRecovery(t *testing.T) { + r := NewHookRegistry() + ctx := context.Background() + + // Void hook: panic in one handler shouldn't crash, other handlers should still run + var safeHandlerCalled bool + r.OnLLMInput("panicker", 10, func(_ context.Context, _ *LLMInputEvent) error { + panic("boom") + }) + r.OnLLMInput("safe", 10, func(_ context.Context, _ *LLMInputEvent) error { + safeHandlerCalled = true + return nil + }) + + // Should not panic + r.TriggerLLMInput(ctx, &LLMInputEvent{AgentID: "test"}) + if !safeHandlerCalled { + t.Error("Expected safe handler to run despite panicking sibling") + } + + // Modifying hook: panic in handler shouldn't crash + r.OnBeforeToolCall("panicker", 10, func(_ context.Context, _ *BeforeToolCallEvent) error { + panic("boom") + }) + + // Should not panic + r.TriggerBeforeToolCall(ctx, &BeforeToolCallEvent{ToolName: "test"}) +} diff --git a/pkg/hooks/types.go b/pkg/hooks/types.go new file mode 100644 index 000000000..7fd8bcc52 --- /dev/null +++ b/pkg/hooks/types.go @@ -0,0 +1,82 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Inspired by and based on nanobot: https://github.com/HKUDS/nanobot +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package hooks + +import ( + "time" + + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/tools" +) + +// MessageReceivedEvent is fired when an inbound message is consumed from the bus. +type MessageReceivedEvent struct { + Channel string + SenderID string + ChatID string + Content string + Media []string + Metadata map[string]string +} + +// MessageSendingEvent is fired before an outbound message is published. +// Handlers can modify Content or set Cancel to block delivery. +type MessageSendingEvent struct { + Channel string + ChatID string + Content string // Modifiable + Cancel bool + CancelReason string +} + +// BeforeToolCallEvent is fired before a tool is executed. +// Handlers can modify Args, or set Cancel to block execution. +type BeforeToolCallEvent struct { + ToolName string + Args map[string]any // Modifiable + Channel string + ChatID string + Cancel bool + CancelReason string // Message returned to LLM when canceled +} + +// AfterToolCallEvent is fired after a tool completes execution. +type AfterToolCallEvent struct { + ToolName string + Args map[string]any + Channel string + ChatID string + Duration time.Duration + Result *tools.ToolResult +} + +// LLMInputEvent is fired before the LLM provider is called. +type LLMInputEvent struct { + AgentID string + Model string + Messages []providers.Message + Tools []providers.ToolDefinition + Iteration int +} + +// LLMOutputEvent is fired after the LLM provider responds. +type LLMOutputEvent struct { + AgentID string + Model string + Content string + ToolCalls []providers.ToolCall + Iteration int + Duration time.Duration +} + +// SessionEvent is fired at session start and end. +type SessionEvent struct { + AgentID string + SessionKey string + Channel string + ChatID string +}