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

This commit is contained in:
xj 2026-02-22 19:17:01 -08:00
parent f84f55f80f
commit db436aa5ca
4 changed files with 62 additions and 39 deletions

View file

@ -122,7 +122,7 @@ func registerSharedTools(
// Message tool // Message tool
messageTool := tools.NewMessageTool() 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{ msgBus.PublishOutbound(bus.OutboundMessage{
Channel: channel, Channel: channel,
ChatID: chatID, ChatID: chatID,
@ -219,9 +219,9 @@ func (al *AgentLoop) SetChannelManager(cm *channels.Manager) {
} }
// SetHooks installs a hook registry. Must be called before Run starts. // 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() { 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 al.hooks = h
@ -230,8 +230,8 @@ func (al *AgentLoop) SetHooks(h *hooks.HookRegistry) {
if agent, ok := al.registry.GetAgent(agentID); ok { if agent, ok := al.registry.GetAgent(agentID); ok {
if tool, ok := agent.Tools.Get("message"); 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,
@ -247,17 +247,17 @@ func (al *AgentLoop) SetHooks(h *hooks.HookRegistry) {
} }
} }
} }
return nil
} }
// SetPluginManager installs a plugin manager and routes its hook registry into the loop. // SetPluginManager installs a plugin manager and routes its hook registry into the loop.
// Must be called before Run starts. // Must be called before Run starts.
func (al *AgentLoop) SetPluginManager(pm *plugin.Manager) { func (al *AgentLoop) SetPluginManager(pm *plugin.Manager) error {
al.pluginManager = pm al.pluginManager = pm
if pm == nil { 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. // 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 { if err := pm.RegisterAll(plugins...); err != nil {
return err return err
} }
al.SetPluginManager(pm) return al.SetPluginManager(pm)
return nil
} }
// sendOutbound wraps bus.PublishOutbound with the message_sending hook. // sendOutbound wraps bus.PublishOutbound with the message_sending hook.

View file

@ -56,7 +56,9 @@ func TestSetPluginManagerInstallsHookRegistry(t *testing.T) {
t.Fatalf("register plugin: %v", err) 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 { if al.pluginManager == nil {
t.Fatal("expected plugin manager to be set") 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-*") tmpDir, err := os.MkdirTemp("", "agent-plugin-test-*")
if err != nil { if err != nil {
t.Fatalf("Failed to create temp dir: %v", err) t.Fatalf("Failed to create temp dir: %v", err)
@ -100,11 +102,7 @@ func TestSetHooksPanicsWhenRunning(t *testing.T) {
al := NewAgentLoop(cfg, msgBus, &mockProvider{}) al := NewAgentLoop(cfg, msgBus, &mockProvider{})
al.running.Store(true) 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())
} }

View file

@ -5,10 +5,10 @@ import (
"fmt" "fmt"
) )
type SendCallback func(channel, chatID, content string) error type SendCallbackWithContext func(ctx context.Context, channel, chatID, content string) error
type MessageTool struct { type MessageTool struct {
sendCallback SendCallback sendCallback SendCallbackWithContext
defaultChannel string defaultChannel string
defaultChatID string defaultChatID string
sentInRound bool // Tracks whether a message was sent in the current processing round 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 return t.sentInRound
} }
func (t *MessageTool) SetSendCallback(callback SendCallback) { func (t *MessageTool) SetSendCallback(callback SendCallbackWithContext) {
t.sendCallback = callback 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} 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{ return &ToolResult{
ForLLM: fmt.Sprintf("sending message: %v", err), ForLLM: fmt.Sprintf("sending message: %v", err),
IsError: true, IsError: true,

View file

@ -11,7 +11,7 @@ func TestMessageTool_Execute_Success(t *testing.T) {
tool.SetContext("test-channel", "test-chat-id") tool.SetContext("test-channel", "test-chat-id")
var sentChannel, sentChatID, sentContent string 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 sentChannel = channel
sentChatID = chatID sentChatID = chatID
sentContent = content sentContent = content
@ -63,7 +63,7 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) {
tool.SetContext("default-channel", "default-chat-id") tool.SetContext("default-channel", "default-chat-id")
var sentChannel, sentChatID string var sentChannel, sentChatID string
tool.SetSendCallback(func(channel, chatID, content string) error { tool.SetSendCallback(func(_ context.Context, channel, chatID, content string) error {
sentChannel = channel sentChannel = channel
sentChatID = chatID sentChatID = chatID
return nil 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) { func TestMessageTool_Execute_SendFailure(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
tool.SetContext("test-channel", "test-chat-id") tool.SetContext("test-channel", "test-chat-id")
sendErr := errors.New("network error") 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 return sendErr
}) })
@ -153,7 +179,7 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
// No SetContext called, so defaultChannel and defaultChatID are empty // 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 return nil
}) })