diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 0a24cc8e3..0e8b9cce1 100644 --- a/pkg/agent/loop.go +++ b/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 { - 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) + if mt, ok := tool.(*tools.MessageTool); ok { + mt.SetSendCallback(func(ctx context.Context, channel, chatID, content string) error { + if sent, reason := al.sendOutbound(ctx, bus.OutboundMessage{ + Channel: channel, + ChatID: chatID, + Content: content, + }); !sent { + if strings.TrimSpace(reason) == "" { + reason = "unspecified" } - 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 + return al.SetHooks(nil) } - 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 nil + return al.SetPluginManager(pm) } // sendOutbound wraps bus.PublishOutbound with the message_sending hook. diff --git a/pkg/agent/plugin_test.go b/pkg/agent/plugin_test.go index c02cf6eaa..e208282b3 100644 --- a/pkg/agent/plugin_test.go +++ b/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 r := recover(); r == nil { - t.Fatal("expected panic when calling SetHooks while running") - } - }() - - al.SetHooks(hooks.NewHookRegistry()) + if err := al.SetHooks(hooks.NewHookRegistry()); err == nil { + t.Fatal("expected error when calling SetHooks while running") + } } diff --git a/pkg/tools/message.go b/pkg/tools/message.go index 15ef4ff73..6391ee8ab 100644 --- a/pkg/tools/message.go +++ b/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, diff --git a/pkg/tools/message_test.go b/pkg/tools/message_test.go index 717c1117b..a111e7a83 100644 --- a/pkg/tools/message_test.go +++ b/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 })