opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

43
pkg/agent/loop.go
View file

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

16
pkg/agent/plugin_test.go
View file

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

8
pkg/tools/message.go
View file

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

34
pkg/tools/message_test.go
View file

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