feat(hooks): add lightweight lifecycle hook system

Add a typed lifecycle hook system inspired by OpenClaw, designed for
PicoClaw's ultra-lightweight philosophy. Provides 8 interception points
around the agent loop for observability, content filtering, and guardrails.

Two execution patterns:
- Void hooks (concurrent): message_received, after_tool_call,
  llm_input, llm_output, session_start, session_end
- Modifying hooks (sequential by priority, with cancel):
  message_sending, before_tool_call

Key design choices:
- Zero-cost when unused: all triggers check len==0 and return immediately
- Copy-on-write registration: insertSorted allocates new backing array
  so concurrent readers never race with writers
- Panic recovery in all handler dispatch paths
- sendOutbound wrapper returns cancel status to callers
- MessageTool callback rewired via SetHooks for content filtering

15 tests covering execution, priority ordering, cancel semantics,
concurrency (barrier-based), panic recovery, and error swallowing.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
xj 2026-02-19 02:17:35 -08:00
parent a91a4e5978
commit ebb28f5201
4 changed files with 942 additions and 14 deletions

View file

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

270
pkg/hooks/hooks.go Normal file
View file

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

433
pkg/hooks/hooks_test.go Normal file
View file

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

82
pkg/hooks/types.go Normal file
View file

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