From c44bd6138cf3ce272c3810cf9a864b8417f55f33 Mon Sep 17 00:00:00 2001 From: LC <64722907+lc6464@users.noreply.github.com> Date: Tue, 28 Apr 2026 10:17:12 +0800 Subject: [PATCH] refactor(pico): unify message kind handling of tool_calls and thought (#2680) * refactor(pico): unify message kind handling of tool_calls and thought * fix(pico): add legacy compatibility for thought payload in Send method Co-authored-by: Copilot --------- Co-authored-by: Copilot --- pkg/channels/pico/client_test.go | 53 +++++++++++++++++-- pkg/channels/pico/pico.go | 13 +++-- pkg/channels/pico/pico_test.go | 45 +++++++++++++++- pkg/channels/pico/protocol.go | 12 ++++- .../features/chat/assistant-message-state.ts | 18 +++++-- 5 files changed, 126 insertions(+), 15 deletions(-) diff --git a/pkg/channels/pico/client_test.go b/pkg/channels/pico/client_test.go index 5ee028bae..2b167e457 100644 --- a/pkg/channels/pico/client_test.go +++ b/pkg/channels/pico/client_test.go @@ -333,18 +333,33 @@ func TestIsThoughtPayload(t *testing.T) { want bool }{ { - name: "explicit thought bool", + name: "explicit thought kind", + payload: map[string]any{PayloadKeyKind: MessageKindThought}, + want: true, + }, + { + name: "thought kind ignores case and whitespace", + payload: map[string]any{PayloadKeyKind: " ThOuGhT "}, + want: true, + }, + { + name: "legacy thought bool remains supported for inbound compatibility", payload: map[string]any{PayloadKeyThought: true}, want: true, }, { - name: "thought false", + name: "legacy thought false", payload: map[string]any{PayloadKeyThought: false}, want: false, }, { - name: "thought string ignored", - payload: map[string]any{PayloadKeyThought: "true"}, + name: "tool calls kind", + payload: map[string]any{PayloadKeyKind: MessageKindToolCalls}, + want: false, + }, + { + name: "non-string kind ignored", + payload: map[string]any{PayloadKeyKind: true}, want: false, }, { @@ -380,7 +395,7 @@ func TestPicoClientChannel_HandleServerMessage_IgnoresThought(t *testing.T) { Type: TypeMessageCreate, Payload: map[string]any{ PayloadKeyContent: "internal reasoning", - PayloadKeyThought: true, + PayloadKeyKind: MessageKindThought, }, }) @@ -390,3 +405,31 @@ func TestPicoClientChannel_HandleServerMessage_IgnoresThought(t *testing.T) { case <-time.After(150 * time.Millisecond): } } + +func TestPicoClientChannel_HandleServerMessage_IgnoresLegacyThoughtBool(t *testing.T) { + mb := bus.NewMessageBus() + bc := &config.Channel{Type: config.ChannelPicoClient, Enabled: true} + ch, err := NewPicoClientChannel(bc, &config.PicoClientSettings{ + URL: "ws://localhost:8080/ws", + }, mb) + if err != nil { + t.Fatalf("NewPicoClientChannel() error = %v", err) + } + + ch.ctx = context.Background() + pc := &picoConn{sessionID: "sess-thought-legacy"} + + ch.handleServerMessage(pc, PicoMessage{ + Type: TypeMessageCreate, + Payload: map[string]any{ + PayloadKeyContent: "legacy internal reasoning", + PayloadKeyThought: true, + }, + }) + + select { + case msg := <-mb.InboundChan(): + t.Fatalf("expected no inbound publish for legacy thought payload, got %+v", msg) + case <-time.After(150 * time.Millisecond): + } +} diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index 9bd8a5b5d..d1de8f4d5 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -323,10 +323,18 @@ func (c *PicoChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri payload := map[string]any{ PayloadKeyContent: content, - PayloadKeyThought: isThought, "message_id": msgID, } - if isToolCalls { + switch { + case isThought: + payload[PayloadKeyKind] = MessageKindThought + + // This field is kept solely for compatibility with legacy pico clients that + // do not yet support the newer "kind" field. + // DO NOT use it for any purpose other than legacy client compatibility. + payload[PayloadKeyThought] = true + + case isToolCalls: payload[PayloadKeyKind] = MessageKindToolCalls if toolCalls, ok := picoToolCallsPayload(msg); ok { payload[PayloadKeyToolCalls] = toolCalls @@ -457,7 +465,6 @@ func (c *PicoChannel) SendPlaceholder(ctx context.Context, chatID string) (strin msgID := uuid.New().String() outMsg := newMessage(TypeMessageCreate, map[string]any{ PayloadKeyContent: text, - PayloadKeyThought: false, "message_id": msgID, }) diff --git a/pkg/channels/pico/pico_test.go b/pkg/channels/pico/pico_test.go index 22ed5451a..bbe73a222 100644 --- a/pkg/channels/pico/pico_test.go +++ b/pkg/channels/pico/pico_test.go @@ -131,8 +131,8 @@ func TestSend_ThoughtMessageDoesNotFinalizeTrackedToolFeedback(t *testing.T) { if got := payload[PayloadKeyContent]; got != "thinking trace" { t.Fatalf("thought content = %#v, want %q", got, "thinking trace") } - if got := payload[PayloadKeyThought]; got != true { - t.Fatalf("thought flag = %#v, want true", got) + if got := payload[PayloadKeyKind]; got != MessageKindThought { + t.Fatalf("thought kind = %#v, want %q", got, MessageKindThought) } if got := payload["message_id"]; got == "msg-progress" || got == nil || got == "" { t.Fatalf("thought message_id = %#v, want new non-progress id", got) @@ -193,6 +193,47 @@ func TestSend_ThoughtMessageDoesNotFinalizeTrackedToolFeedback(t *testing.T) { } } +func TestSendPlaceholder_EmitsNormalMessageWithoutKind(t *testing.T) { + ch := newTestPicoChannel(t) + ch.bc.Placeholder.Enabled = true + + if err := ch.Start(context.Background()); err != nil { + t.Fatalf("Start() error = %v", err) + } + defer ch.Stop(context.Background()) + + clientConn, received, cleanup := newTestPicoWebSocket(t) + defer cleanup() + ch.addConnForTest(&picoConn{id: "conn-1", conn: clientConn, sessionID: "sess-1"}) + + msgID, err := ch.SendPlaceholder(context.Background(), "pico:sess-1") + if err != nil { + t.Fatalf("SendPlaceholder() error = %v", err) + } + if msgID == "" { + t.Fatal("expected placeholder message id") + } + + select { + case msg := <-received: + if msg.Type != TypeMessageCreate { + t.Fatalf("placeholder message type = %q, want %q", msg.Type, TypeMessageCreate) + } + payload := msg.Payload + if got := payload["message_id"]; got != msgID { + t.Fatalf("placeholder message_id = %#v, want %q", got, msgID) + } + if got := payload[PayloadKeyContent]; got != "Thinking..." { + t.Fatalf("placeholder content = %#v, want %q", got, "Thinking...") + } + if got, ok := payload[PayloadKeyKind]; ok { + t.Fatalf("placeholder kind = %#v, want absent", got) + } + case <-time.After(time.Second): + t.Fatal("expected placeholder message to be delivered") + } +} + func TestCreateAndAddConnection_RespectsMaxConnectionsConcurrently(t *testing.T) { ch := newTestPicoChannel(t) diff --git a/pkg/channels/pico/protocol.go b/pkg/channels/pico/protocol.go index 46e8fa3ee..6e3a5ca89 100644 --- a/pkg/channels/pico/protocol.go +++ b/pkg/channels/pico/protocol.go @@ -1,6 +1,9 @@ package pico -import "time" +import ( + "strings" + "time" +) // Protocol message types. const ( @@ -47,6 +50,13 @@ func newMessage(msgType string, payload map[string]any) PicoMessage { } func isThoughtPayload(payload map[string]any) bool { + kind, _ := payload[PayloadKeyKind].(string) + if strings.EqualFold(strings.TrimSpace(kind), MessageKindThought) { + return true + } + + // Keep pico_client inbound-compatible with legacy servers that still send + // the pre-kind boolean thought marker. thought, _ := payload[PayloadKeyThought].(bool) return thought } diff --git a/web/frontend/src/features/chat/assistant-message-state.ts b/web/frontend/src/features/chat/assistant-message-state.ts index 54eb6163f..c26b2665c 100644 --- a/web/frontend/src/features/chat/assistant-message-state.ts +++ b/web/frontend/src/features/chat/assistant-message-state.ts @@ -19,14 +19,25 @@ export interface AssistantMessageUpdateState { toolCalls?: AssistantToolCalls } +function normalizeAssistantMessageKind( + payload: Record, +): string | undefined { + if (typeof payload.kind !== "string") { + return undefined + } + const kind = payload.kind.trim().toLowerCase() + return kind || undefined +} + function parseAssistantMessageKind( payload: Record, toolCalls?: AssistantToolCalls, ): AssistantMessageKind { - if (payload.thought === true) { + const kind = normalizeAssistantMessageKind(payload) + if (kind === "thought") { return "thought" } - if (payload.kind === "tool_calls" || toolCalls) { + if (kind === "tool_calls" || toolCalls) { return "tool_calls" } return "normal" @@ -36,8 +47,7 @@ function hasExplicitAssistantKindPayload( payload: Record, ): boolean { return ( - typeof payload.thought === "boolean" || - payload.kind === "tool_calls" || + normalizeAssistantMessageKind(payload) !== undefined || payload.tool_calls !== undefined ) }