fix(agent): make hook setup recoverable and preserve request context

pkg/agent/loop.go

@@ -122,7 +122,7 @@ func registerSharedTools(
 
 		// Message tool
 		messageTool := tools.NewMessageTool()
-		messageTool.SetSendCallback(func(channel, chatID, content string) error {
+		messageTool.SetSendCallback(func(_ context.Context, channel, chatID, content string) error {
 			msgBus.PublishOutbound(bus.OutboundMessage{
 				Channel: channel,
 				ChatID:  chatID,
@@ -219,9 +219,9 @@ 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) {
+func (al *AgentLoop) SetHooks(h *hooks.HookRegistry) error {
 	if al.running.Load() {
-		panic("SetHooks must be called before Run starts")
+		return fmt.Errorf("SetHooks must be called before Run starts")
 	}
 	al.hooks = h
 
@@ -229,35 +229,35 @@ func (al *AgentLoop) SetHooks(h *hooks.HookRegistry) {
 	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 {
+				if mt, ok := tool.(*tools.MessageTool); ok {
-						mt.SetSendCallback(func(channel, chatID, content string) error {
+					mt.SetSendCallback(func(ctx context.Context, channel, chatID, content string) error {
-							if sent, reason := al.sendOutbound(context.Background(), bus.OutboundMessage{
+						if sent, reason := al.sendOutbound(ctx, bus.OutboundMessage{
-								Channel: channel,
+							Channel: channel,
-								ChatID:  chatID,
+							ChatID:  chatID,
-								Content: content,
+							Content: content,
-							}); !sent {
+						}); !sent {
-								if strings.TrimSpace(reason) == "" {
+							if strings.TrimSpace(reason) == "" {
-									reason = "unspecified"
+								reason = "unspecified"
-								}
-								return fmt.Errorf("message canceled by hook: %s", reason)
 							}
-							return nil
+							return fmt.Errorf("message canceled by hook: %s", reason)
-						})
+						}
+						return nil
+					})
 				}
 			}
 		}
 	}
+	return nil
 }
 
 // SetPluginManager installs a plugin manager and routes its hook registry into the loop.
 // Must be called before Run starts.
-func (al *AgentLoop) SetPluginManager(pm *plugin.Manager) {
+func (al *AgentLoop) SetPluginManager(pm *plugin.Manager) error {
 	al.pluginManager = pm
 	if pm == nil {
-		al.SetHooks(nil)
+		return al.SetHooks(nil)
-		return
 	}
-	al.SetHooks(pm.HookRegistry())
+	return al.SetHooks(pm.HookRegistry())
 }
 
 // EnablePlugins is a convenience helper to build and install a plugin manager.
@@ -266,8 +266,7 @@ func (al *AgentLoop) EnablePlugins(plugins ...plugin.Plugin) error {
 	if err := pm.RegisterAll(plugins...); err != nil {
 		return err
 	}
-	al.SetPluginManager(pm)
+	return al.SetPluginManager(pm)
-	return nil
 }
 
 // sendOutbound wraps bus.PublishOutbound with the message_sending hook.

pkg/agent/plugin_test.go

@@ -56,7 +56,9 @@ func TestSetPluginManagerInstallsHookRegistry(t *testing.T) {
 		t.Fatalf("register plugin: %v", err)
 	}
 
-	al.SetPluginManager(pm)
+	if err := al.SetPluginManager(pm); err != nil {
+		t.Fatalf("SetPluginManager: %v", err)
+	}
 
 	if al.pluginManager == nil {
 		t.Fatal("expected plugin manager to be set")
@@ -78,7 +80,7 @@ func TestSetPluginManagerInstallsHookRegistry(t *testing.T) {
 	}
 }
 
-func TestSetHooksPanicsWhenRunning(t *testing.T) {
+func TestSetHooksReturnsErrorWhenRunning(t *testing.T) {
 	tmpDir, err := os.MkdirTemp("", "agent-plugin-test-*")
 	if err != nil {
 		t.Fatalf("Failed to create temp dir: %v", err)
@@ -100,11 +102,7 @@ func TestSetHooksPanicsWhenRunning(t *testing.T) {
 	al := NewAgentLoop(cfg, msgBus, &mockProvider{})
 	al.running.Store(true)
 
-	defer func() {
+	if err := al.SetHooks(hooks.NewHookRegistry()); err == nil {
-		if r := recover(); r == nil {
+		t.Fatal("expected error when calling SetHooks while running")
-			t.Fatal("expected panic when calling SetHooks while running")
+	}
-		}
-	}()
-
-	al.SetHooks(hooks.NewHookRegistry())
 }

pkg/tools/message.go

@@ -5,10 +5,10 @@ import (
 	"fmt"
 )
 
-type SendCallback func(channel, chatID, content string) error
+type SendCallbackWithContext func(ctx context.Context, channel, chatID, content string) error
 
 type MessageTool struct {
-	sendCallback   SendCallback
+	sendCallback   SendCallbackWithContext
 	defaultChannel string
 	defaultChatID  string
 	sentInRound    bool // Tracks whether a message was sent in the current processing round
@@ -58,7 +58,7 @@ func (t *MessageTool) HasSentInRound() bool {
 	return t.sentInRound
 }
 
-func (t *MessageTool) SetSendCallback(callback SendCallback) {
+func (t *MessageTool) SetSendCallback(callback SendCallbackWithContext) {
 	t.sendCallback = callback
 }
 
@@ -86,7 +86,7 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *ToolRes
 		return &ToolResult{ForLLM: "Message sending not configured", IsError: true}
 	}
 
-	if err := t.sendCallback(channel, chatID, content); err != nil {
+	if err := t.sendCallback(ctx, channel, chatID, content); err != nil {
 		return &ToolResult{
 			ForLLM:  fmt.Sprintf("sending message: %v", err),
 			IsError: true,

pkg/tools/message_test.go

@@ -11,7 +11,7 @@ func TestMessageTool_Execute_Success(t *testing.T) {
 	tool.SetContext("test-channel", "test-chat-id")
 
 	var sentChannel, sentChatID, sentContent string
-	tool.SetSendCallback(func(channel, chatID, content string) error {
+	tool.SetSendCallback(func(_ context.Context, channel, chatID, content string) error {
 		sentChannel = channel
 		sentChatID = chatID
 		sentContent = content
@@ -63,7 +63,7 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) {
 	tool.SetContext("default-channel", "default-chat-id")
 
 	var sentChannel, sentChatID string
-	tool.SetSendCallback(func(channel, chatID, content string) error {
+	tool.SetSendCallback(func(_ context.Context, channel, chatID, content string) error {
 		sentChannel = channel
 		sentChatID = chatID
 		return nil
@@ -94,12 +94,38 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) {
 	}
 }
 
+func TestMessageTool_Execute_PropagatesContext(t *testing.T) {
+	tool := NewMessageTool()
+	tool.SetContext("test-channel", "test-chat-id")
+
+	type keyType string
+	const key keyType = "k"
+	ctx := context.WithValue(context.Background(), key, "v")
+
+	seen := ""
+	tool.SetSendCallback(func(cbCtx context.Context, channel, chatID, content string) error {
+		val, _ := cbCtx.Value(key).(string)
+		seen = val
+		return nil
+	})
+
+	result := tool.Execute(ctx, map[string]any{
+		"content": "context test",
+	})
+	if result.IsError {
+		t.Fatalf("unexpected error: %v", result.ForLLM)
+	}
+	if seen != "v" {
+		t.Fatalf("expected propagated context value 'v', got %q", seen)
+	}
+}
+
 func TestMessageTool_Execute_SendFailure(t *testing.T) {
 	tool := NewMessageTool()
 	tool.SetContext("test-channel", "test-chat-id")
 
 	sendErr := errors.New("network error")
-	tool.SetSendCallback(func(channel, chatID, content string) error {
+	tool.SetSendCallback(func(_ context.Context, channel, chatID, content string) error {
 		return sendErr
 	})
 
@@ -153,7 +179,7 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) {
 	tool := NewMessageTool()
 	// No SetContext called, so defaultChannel and defaultChatID are empty
 
-	tool.SetSendCallback(func(channel, chatID, content string) error {
+	tool.SetSendCallback(func(_ context.Context, channel, chatID, content string) error {
 		return nil
 	})