diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 5749149c1..521bea865 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -257,11 +257,25 @@ func (al *AgentLoop) Run(ctx context.Context) error { } continued, continueErr := al.drainQueuedSteeringContinuations(ctx, target) if continueErr != nil { - al.maybePublishError(ctx, m.Channel, m.ChatID, sessionKey, continueErr) + al.maybePublishErrorWithPolicy( + ctx, + m.Channel, + m.ChatID, + sessionKey, + continueErr, + finalResponseAlwaysPublish, + ) return } if continued != "" { - al.PublishResponseIfNeeded(ctx, target.Channel, target.ChatID, target.SessionKey, continued) + al.publishResponseIfNeededWithPolicy( + ctx, + target.Channel, + target.ChatID, + target.SessionKey, + continued, + finalResponseAlwaysPublish, + ) } return } diff --git a/pkg/agent/agent_outbound.go b/pkg/agent/agent_outbound.go index f4a01adfd..7285bce97 100644 --- a/pkg/agent/agent_outbound.go +++ b/pkg/agent/agent_outbound.go @@ -17,40 +17,88 @@ import ( "github.com/sipeed/picoclaw/pkg/utils" ) +type finalResponseDeliveryPolicy uint8 + +const ( + finalResponseSuppressIfMessageToolSent finalResponseDeliveryPolicy = iota + finalResponseAlwaysPublish +) + func (al *AgentLoop) maybePublishError(ctx context.Context, channel, chatID, sessionKey string, err error) bool { + return al.maybePublishErrorWithPolicy( + ctx, + channel, + chatID, + sessionKey, + err, + finalResponseSuppressIfMessageToolSent, + ) +} + +func (al *AgentLoop) maybePublishErrorWithPolicy( + ctx context.Context, + channel, chatID, sessionKey string, + err error, + policy finalResponseDeliveryPolicy, +) bool { if errors.Is(err, context.Canceled) { return false } - al.PublishResponseIfNeeded(ctx, channel, chatID, sessionKey, fmt.Sprintf("Error processing message: %v", err)) + al.publishResponseIfNeededWithPolicy( + ctx, + channel, + chatID, + sessionKey, + fmt.Sprintf("Error processing message: %v", err), + policy, + ) return true } -func (al *AgentLoop) publishResponseOrError( +func (al *AgentLoop) publishResponseOrErrorWithPolicy( ctx context.Context, channel, chatID, sessionKey string, response string, err error, + policy finalResponseDeliveryPolicy, ) { if err != nil { - if !al.maybePublishError(ctx, channel, chatID, sessionKey, err) { + if !al.maybePublishErrorWithPolicy(ctx, channel, chatID, sessionKey, err, policy) { return } response = "" } - al.PublishResponseIfNeeded(ctx, channel, chatID, sessionKey, response) + al.publishResponseIfNeededWithPolicy(ctx, channel, chatID, sessionKey, response, policy) } func (al *AgentLoop) PublishResponseIfNeeded(ctx context.Context, channel, chatID, sessionKey, response string) { + al.publishResponseIfNeededWithPolicy( + ctx, + channel, + chatID, + sessionKey, + response, + finalResponseSuppressIfMessageToolSent, + ) +} + +func (al *AgentLoop) publishResponseIfNeededWithPolicy( + ctx context.Context, + channel, chatID, sessionKey, response string, + policy finalResponseDeliveryPolicy, +) { if response == "" { return } alreadySentToSameChat := false - defaultAgent := al.GetRegistry().GetDefaultAgent() - if defaultAgent != nil { - if tool, ok := defaultAgent.Tools.Get("message"); ok { - if mt, ok := tool.(*tools.MessageTool); ok { - alreadySentToSameChat = mt.HasSentTo(sessionKey, channel, chatID) + if policy == finalResponseSuppressIfMessageToolSent { + defaultAgent := al.GetRegistry().GetDefaultAgent() + if defaultAgent != nil { + if tool, ok := defaultAgent.Tools.Get("message"); ok { + if mt, ok := tool.(*tools.MessageTool); ok { + alreadySentToSameChat = mt.HasSentTo(sessionKey, channel, chatID) + } } } } diff --git a/pkg/agent/agent_steering.go b/pkg/agent/agent_steering.go index 9b136e7cd..21800c1cf 100644 --- a/pkg/agent/agent_steering.go +++ b/pkg/agent/agent_steering.go @@ -15,7 +15,15 @@ func (al *AgentLoop) processMessageSync(ctx context.Context, msg bus.InboundMess } response, err := al.processMessage(ctx, msg) - al.publishResponseOrError(ctx, msg.Channel, msg.ChatID, msg.SessionKey, response, err) + al.publishResponseOrErrorWithPolicy( + ctx, + msg.Channel, + msg.ChatID, + msg.SessionKey, + response, + err, + finalResponseAlwaysPublish, + ) } func (al *AgentLoop) runTurnWithSteering(ctx context.Context, initialMsg bus.InboundMessage) { @@ -58,7 +66,14 @@ func (al *AgentLoop) runTurnWithSteering(ctx context.Context, initialMsg bus.Inb // Publish final response if finalResponse != "" { - al.PublishResponseIfNeeded(ctx, target.Channel, target.ChatID, target.SessionKey, finalResponse) + al.publishResponseIfNeededWithPolicy( + ctx, + target.Channel, + target.ChatID, + target.SessionKey, + finalResponse, + finalResponseAlwaysPublish, + ) } } diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index 7a869ec94..e73b59369 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -194,6 +194,105 @@ func newTestAgentLoop( return al, cfg, msgBus, provider, func() { os.RemoveAll(tmpDir) } } +func TestPublishResponseIfNeededWithPolicy_AlwaysPublishesFinalAfterMessageTool(t *testing.T) { + al, _, msgBus, _, cleanup := newTestAgentLoop(t) + defer cleanup() + + agent := al.registry.GetDefaultAgent() + if agent == nil { + t.Fatal("expected default agent") + } + rawTool, ok := agent.Tools.Get("message") + if !ok { + mt := tools.NewMessageTool() + agent.Tools.Register(mt) + rawTool = mt + } + mt, ok := rawTool.(*tools.MessageTool) + if !ok { + t.Fatalf("message tool type = %T", rawTool) + } + mt.SetSendCallback(func(ctx context.Context, channel, chatID, content, replyToMessageID string) error { + return nil + }) + + ctx := tools.WithToolSessionContext(context.Background(), routing.DefaultAgentID, "session-msg-1", nil) + res := mt.Execute(ctx, map[string]any{ + "content": "working on it", + "channel": "telegram", + "chat_id": "-100123", + }) + if res == nil || res.IsError { + t.Fatalf("message tool execute failed: %+v", res) + } + + al.publishResponseIfNeededWithPolicy( + context.Background(), + "telegram", + "-100123", + "session-msg-1", + "final result", + finalResponseAlwaysPublish, + ) + + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Content != "final result" { + t.Fatalf("outbound content = %q, want final result", outbound.Content) + } + case <-time.After(2 * time.Second): + t.Fatal("expected outbound response") + } +} + +func TestPublishResponseIfNeededWithPolicy_SuppressesWhenMessageToolAlreadySent(t *testing.T) { + al, _, msgBus, _, cleanup := newTestAgentLoop(t) + defer cleanup() + + agent := al.registry.GetDefaultAgent() + if agent == nil { + t.Fatal("expected default agent") + } + rawTool, ok := agent.Tools.Get("message") + if !ok { + mt := tools.NewMessageTool() + agent.Tools.Register(mt) + rawTool = mt + } + mt, ok := rawTool.(*tools.MessageTool) + if !ok { + t.Fatalf("message tool type = %T", rawTool) + } + mt.SetSendCallback(func(ctx context.Context, channel, chatID, content, replyToMessageID string) error { + return nil + }) + + ctx := tools.WithToolSessionContext(context.Background(), routing.DefaultAgentID, "session-msg-2", nil) + res := mt.Execute(ctx, map[string]any{ + "content": "working on it", + "channel": "telegram", + "chat_id": "-100123", + }) + if res == nil || res.IsError { + t.Fatalf("message tool execute failed: %+v", res) + } + + al.publishResponseIfNeededWithPolicy( + context.Background(), + "telegram", + "-100123", + "session-msg-2", + "final result", + finalResponseSuppressIfMessageToolSent, + ) + + select { + case outbound := <-msgBus.OutboundChan(): + t.Fatalf("unexpected outbound response: %+v", outbound) + case <-time.After(150 * time.Millisecond): + } +} + func TestNewAgentLoop_RegistersWebSearchTool(t *testing.T) { cfg := config.DefaultConfig() cfg.Agents.Defaults.Workspace = t.TempDir()