diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 005423882..0a24cc8e3 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -220,23 +220,29 @@ func (al *AgentLoop) SetChannelManager(cm *channels.Manager) { // SetHooks installs a hook registry. Must be called before Run starts. func (al *AgentLoop) SetHooks(h *hooks.HookRegistry) { + if al.running.Load() { + panic("SetHooks must be called before Run starts") + } 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 - }) + if mt, ok := tool.(*tools.MessageTool); ok { + mt.SetSendCallback(func(channel, chatID, content string) error { + if sent, reason := al.sendOutbound(context.Background(), bus.OutboundMessage{ + Channel: channel, + ChatID: chatID, + Content: content, + }); !sent { + if strings.TrimSpace(reason) == "" { + reason = "unspecified" + } + return fmt.Errorf("message canceled by hook: %s", reason) + } + return nil + }) } } } @@ -265,8 +271,8 @@ func (al *AgentLoop) EnablePlugins(plugins ...plugin.Plugin) error { } // 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 { +// Returns whether the message was sent and, if canceled, the cancel reason. +func (al *AgentLoop) sendOutbound(ctx context.Context, msg bus.OutboundMessage) (bool, string) { if al.hooks != nil { event := &hooks.MessageSendingEvent{Channel: msg.Channel, ChatID: msg.ChatID, Content: msg.Content} al.hooks.TriggerMessageSending(ctx, event) @@ -281,12 +287,12 @@ func (al *AgentLoop) sendOutbound(ctx context.Context, msg bus.OutboundMessage) "chat_id": msg.ChatID, "reason": reason, }) - return false + return false, reason } msg.Content = event.Content } al.bus.PublishOutbound(msg) - return true + return true, "" } // RecordLastChannel records the last active channel for this workspace. diff --git a/pkg/agent/plugin_test.go b/pkg/agent/plugin_test.go index f872364c7..7bbcd3624 100644 --- a/pkg/agent/plugin_test.go +++ b/pkg/agent/plugin_test.go @@ -61,7 +61,7 @@ func TestSetPluginManagerInstallsHookRegistry(t *testing.T) { t.Fatal("expected agent loop hooks to use plugin manager registry") } - sent := al.sendOutbound(context.Background(), bus.OutboundMessage{ + sent, reason := al.sendOutbound(context.Background(), bus.OutboundMessage{ Channel: "cli", ChatID: "direct", Content: "hello", @@ -69,5 +69,38 @@ func TestSetPluginManagerInstallsHookRegistry(t *testing.T) { if sent { t.Fatal("expected outbound message to be blocked by plugin") } + if reason == "" { + t.Fatal("expected cancel reason to be propagated") + } } +func TestSetHooksPanicsWhenRunning(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-plugin-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, &mockProvider{}) + al.running.Store(true) + + defer func() { + if r := recover(); r == nil { + t.Fatal("expected panic when calling SetHooks while running") + } + }() + + al.SetHooks(hooks.NewHookRegistry()) +} diff --git a/pkg/hooks/hooks_test.go b/pkg/hooks/hooks_test.go index eb80c2955..427ec0d07 100644 --- a/pkg/hooks/hooks_test.go +++ b/pkg/hooks/hooks_test.go @@ -293,12 +293,12 @@ func TestConcurrentRegistrationAndTrigger(t *testing.T) { // Goroutines registering hooks. for i := range 10 { wg.Add(1) - go func() { + go func(idx int) { defer wg.Done() - r.OnMessageReceived("reg-hook", i, func(_ context.Context, _ *MessageReceivedEvent) error { + r.OnMessageReceived(fmt.Sprintf("reg-hook-%d", idx), idx, func(_ context.Context, _ *MessageReceivedEvent) error { return nil }) - }() + }(i) } // Goroutines triggering hooks concurrently. @@ -322,7 +322,7 @@ func TestInsertSorted(t *testing.T) { // 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 { + r.OnBeforeToolCall(fmt.Sprintf("p-%d", p), p, func(_ context.Context, _ *BeforeToolCallEvent) error { order = append(order, p) return nil })