From 9982ee29a88cbb8790ab74958b91e5127b347511 Mon Sep 17 00:00:00 2001
From: lc6464 <64722907+lc6464@users.noreply.github.com>
Date: Thu, 9 Apr 2026 22:59:36 +0800
Subject: [PATCH] fix(pico): avoid duplicate final websocket message

---
 pkg/agent/loop.go      | 10 +++-------
 pkg/agent/loop_test.go | 30 ++++++++++++++++++++++--------
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 8aa71a168..431376be3 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -2253,16 +2253,12 @@ turnLoop:
 		}
 		logger.DebugCF("agent", "LLM response", llmResponseFields)
 
-		if al.bus != nil && ts.channel == "pico" {
-			liveContent := response.Content
-			if liveContent == "" && len(response.ToolCalls) == 0 && response.ReasoningContent != "" {
-				liveContent = response.ReasoningContent
-			}
-			if strings.TrimSpace(liveContent) != "" {
+		if al.bus != nil && ts.channel == "pico" && len(response.ToolCalls) > 0 {
+			if strings.TrimSpace(response.Content) != "" {
 				al.bus.PublishOutbound(turnCtx, bus.OutboundMessage{
 					Channel: ts.channel,
 					ChatID:  ts.chatID,
-					Content: liveContent,
+					Content: response.Content,
 				})
 			}
 		}
diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go
index 7c10e11aa..371f91d89 100644
--- a/pkg/agent/loop_test.go
+++ b/pkg/agent/loop_test.go
@@ -2766,7 +2766,7 @@ func TestProcessMessage_PublishesToolFeedbackWhenEnabled(t *testing.T) {
 	}
 }
 
-func TestProcessMessage_PicoPublishesAssistantContentDuringToolCalls(t *testing.T) {
+func TestRun_PicoPublishesAssistantContentDuringToolCallsWithoutFinalDuplicate(t *testing.T) {
 	tmpDir := t.TempDir()
 
 	cfg := &config.Config{
@@ -2790,17 +2790,21 @@ func TestProcessMessage_PicoPublishesAssistantContentDuringToolCalls(t *testing.
 	}
 	agent.Tools.Register(&toolLimitTestTool{})
 
-	response, err := al.processMessage(context.Background(), bus.InboundMessage{
+	runCtx, runCancel := context.WithCancel(context.Background())
+	defer runCancel()
+
+	runDone := make(chan error, 1)
+	go func() {
+		runDone <- al.Run(runCtx)
+	}()
+
+	if err := msgBus.PublishInbound(context.Background(), bus.InboundMessage{
 		Channel:  "pico",
 		SenderID: "user-1",
 		ChatID:   "session-1",
 		Content:  "run with tools",
-	})
-	if err != nil {
-		t.Fatalf("processMessage() error = %v", err)
-	}
-	if response != "final model text" {
-		t.Fatalf("processMessage() response = %q, want %q", response, "final model text")
+	}); err != nil {
+		t.Fatalf("PublishInbound() error = %v", err)
 	}
 
 	outputs := make([]string, 0, 2)
@@ -2821,6 +2825,16 @@ func TestProcessMessage_PicoPublishesAssistantContentDuringToolCalls(t *testing.
 		t.Fatalf("second outbound content = %q, want %q", outputs[1], "final model text")
 	}
 
+	runCancel()
+	select {
+	case err := <-runDone:
+		if err != nil {
+			t.Fatalf("Run() error = %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for Run() to exit")
+	}
+
 	select {
 	case outbound := <-msgBus.OutboundChan():
 		if outbound.Content == "final model text" {