fix(pico): separate thought and normal messages

This commit is contained in:
lc6464
2026-04-10 20:23:12 +08:00
parent 748ac58dd1
commit c8bac699fe
15 changed files with 300 additions and 24 deletions
+48 -7
View File
@@ -105,6 +105,8 @@ const (
toolLimitResponse = "I've reached `max_tool_iterations` without a final response. Increase `max_tool_iterations` in config.json if this task needs more tool steps."
handledToolResponseSummary = "Requested output delivered via tool attachment."
sessionKeyAgentPrefix = "agent:"
metadataKeyMessageKind = "message_kind"
messageKindThought = "thought"
metadataKeyAccountID = "account_id"
metadataKeyGuildID = "guild_id"
metadataKeyTeamID = "team_id"
@@ -1622,6 +1624,41 @@ func (al *AgentLoop) targetReasoningChannelID(channelName string) (chatID string
return ""
}
func (al *AgentLoop) publishPicoReasoning(ctx context.Context, reasoningContent, chatID string) {
if reasoningContent == "" || chatID == "" {
return
}
if ctx.Err() != nil {
return
}
pubCtx, pubCancel := context.WithTimeout(ctx, 5*time.Second)
defer pubCancel()
if err := al.bus.PublishOutbound(pubCtx, bus.OutboundMessage{
Channel: "pico",
ChatID: chatID,
Content: reasoningContent,
Metadata: map[string]string{
metadataKeyMessageKind: messageKindThought,
},
}); err != nil {
if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) ||
errors.Is(err, bus.ErrBusClosed) {
logger.DebugCF("agent", "Pico reasoning publish skipped (timeout/cancel)", map[string]any{
"channel": "pico",
"error": err.Error(),
})
} else {
logger.WarnCF("agent", "Failed to publish pico reasoning (best-effort)", map[string]any{
"channel": "pico",
"error": err.Error(),
})
}
}
}
func (al *AgentLoop) handleReasoning(
ctx context.Context,
reasoningContent, channelName, channelID string,
@@ -2223,12 +2260,16 @@ turnLoop:
if reasoningContent == "" {
reasoningContent = response.ReasoningContent
}
go al.handleReasoning(
turnCtx,
reasoningContent,
ts.channel,
al.targetReasoningChannelID(ts.channel),
)
if ts.channel == "pico" {
al.publishPicoReasoning(turnCtx, reasoningContent, ts.chatID)
} else {
go al.handleReasoning(
turnCtx,
reasoningContent,
ts.channel,
al.targetReasoningChannelID(ts.channel),
)
}
al.emitEvent(
EventKindLLMResponse,
ts.eventMeta("runTurn", "turn.llm.response"),
@@ -2277,7 +2318,7 @@ turnLoop:
if len(response.ToolCalls) == 0 || gracefulTerminal {
responseContent := response.Content
if responseContent == "" && response.ReasoningContent != "" {
if responseContent == "" && response.ReasoningContent != "" && ts.channel != "pico" {
responseContent = response.ReasoningContent
}
if steerMsgs := al.dequeueSteeringMessagesForScope(ts.sessionKey); len(steerMsgs) > 0 {
+56
View File
@@ -2660,6 +2660,62 @@ func TestProcessMessage_PublishesReasoningContentToReasoningChannel(t *testing.T
}
}
func TestProcessMessage_PicoPublishesReasoningAsThoughtMessage(t *testing.T) {
tmpDir := t.TempDir()
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Workspace: tmpDir,
ModelName: "test-model",
MaxTokens: 4096,
MaxToolIterations: 10,
},
},
}
msgBus := bus.NewMessageBus()
provider := &reasoningContentProvider{
response: "final answer",
reasoningContent: "thinking trace",
}
al := NewAgentLoop(cfg, msgBus, provider)
response, err := al.processMessage(context.Background(), bus.InboundMessage{
Channel: "pico",
SenderID: "user1",
ChatID: "pico:test-session",
Content: "hello",
})
if err != nil {
t.Fatalf("processMessage() error = %v", err)
}
if response != "final answer" {
t.Fatalf("processMessage() response = %q, want %q", response, "final answer")
}
var thoughtMsg *bus.OutboundMessage
deadline := time.After(3 * time.Second)
for thoughtMsg == nil {
select {
case outbound := <-msgBus.OutboundChan():
msg := outbound
if msg.Content == "thinking trace" {
thoughtMsg = &msg
}
case <-deadline:
t.Fatal("expected thought outbound message for pico")
}
}
if thoughtMsg.Channel != "pico" || thoughtMsg.ChatID != "pico:test-session" {
t.Fatalf("thought message route = %s/%s, want pico/pico:test-session", thoughtMsg.Channel, thoughtMsg.ChatID)
}
if thoughtMsg.Metadata[metadataKeyMessageKind] != messageKindThought {
t.Fatalf("thought metadata kind = %q, want %q", thoughtMsg.Metadata[metadataKeyMessageKind], messageKindThought)
}
}
func TestProcessHeartbeat_DoesNotPublishToolFeedback(t *testing.T) {
tmpDir := t.TempDir()
heartbeatFile := filepath.Join(tmpDir, "heartbeat-task.txt")
+6 -2
View File
@@ -242,7 +242,11 @@ func (c *PicoClientChannel) handleInbound(pc *picoConn, msg PicoMessage) {
}
func (c *PicoClientChannel) handleServerMessage(pc *picoConn, msg PicoMessage) {
content, _ := msg.Payload["content"].(string)
if isThoughtPayload(msg.Payload) {
return
}
content, _ := msg.Payload[PayloadKeyContent].(string)
if strings.TrimSpace(content) == "" {
return
}
@@ -285,7 +289,7 @@ func (c *PicoClientChannel) Send(ctx context.Context, msg bus.OutboundMessage) (
}
outMsg := newMessage(TypeMessageSend, map[string]any{
"content": msg.Content,
PayloadKeyContent: msg.Content,
})
outMsg.SessionID = strings.TrimPrefix(msg.ChatID, "pico_client:")
return nil, pc.writeJSON(outMsg)
+64
View File
@@ -316,3 +316,67 @@ func TestPicoChannel_HandleMessageSend_AllowsMediaOnly(t *testing.T) {
t.Fatal("timed out waiting for inbound media message")
}
}
func TestIsThoughtPayload(t *testing.T) {
tests := []struct {
name string
payload map[string]any
want bool
}{
{
name: "explicit thought bool",
payload: map[string]any{PayloadKeyThought: true},
want: true,
},
{
name: "thought false",
payload: map[string]any{PayloadKeyThought: false},
want: false,
},
{
name: "thought string ignored",
payload: map[string]any{PayloadKeyThought: "true"},
want: false,
},
{
name: "default normal",
payload: map[string]any{PayloadKeyContent: "hello"},
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := isThoughtPayload(tt.payload); got != tt.want {
t.Fatalf("isThoughtPayload() = %v, want %v", got, tt.want)
}
})
}
}
func TestPicoClientChannel_HandleServerMessage_IgnoresThought(t *testing.T) {
mb := bus.NewMessageBus()
ch, err := NewPicoClientChannel(config.PicoClientConfig{
URL: "ws://localhost:8080/ws",
}, mb)
if err != nil {
t.Fatalf("NewPicoClientChannel() error = %v", err)
}
ch.ctx = context.Background()
pc := &picoConn{sessionID: "sess-thought"}
ch.handleServerMessage(pc, PicoMessage{
Type: TypeMessageCreate,
Payload: map[string]any{
PayloadKeyContent: "internal reasoning",
PayloadKeyThought: true,
},
})
select {
case msg := <-mb.InboundChan():
t.Fatalf("expected no inbound publish for thought payload, got %+v", msg)
case <-time.After(150 * time.Millisecond):
}
}
+13 -3
View File
@@ -39,6 +39,13 @@ var allowedInlineImageMIMETypes = map[string]struct{}{
"image/bmp": {},
}
func outboundMessageIsThought(metadata map[string]string) bool {
if len(metadata) == 0 {
return false
}
return strings.EqualFold(strings.TrimSpace(metadata["message_kind"]), MessageKindThought)
}
// writeJSON sends a JSON message to the connection with write locking.
func (pc *picoConn) writeJSON(v any) error {
if pc.closed.Load() {
@@ -247,9 +254,11 @@ func (c *PicoChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri
if !c.IsRunning() {
return nil, channels.ErrNotRunning
}
isThought := outboundMessageIsThought(msg.Metadata)
outMsg := newMessage(TypeMessageCreate, map[string]any{
"content": msg.Content,
PayloadKeyContent: msg.Content,
PayloadKeyThought: isThought,
})
return nil, c.broadcastToSession(msg.ChatID, outMsg)
@@ -288,8 +297,9 @@ func (c *PicoChannel) SendPlaceholder(ctx context.Context, chatID string) (strin
msgID := uuid.New().String()
outMsg := newMessage(TypeMessageCreate, map[string]any{
"content": text,
"message_id": msgID,
PayloadKeyContent: text,
PayloadKeyThought: false,
"message_id": msgID,
})
if err := c.broadcastToSession(chatID, outMsg); err != nil {
+10
View File
@@ -19,6 +19,11 @@ const (
TypePong = "pong"
PicoTokenPrefix = "pico-"
PayloadKeyContent = "content"
PayloadKeyThought = "thought"
MessageKindThought = "thought"
)
// PicoMessage is the wire format for all Pico Protocol messages.
@@ -39,6 +44,11 @@ func newMessage(msgType string, payload map[string]any) PicoMessage {
}
}
func isThoughtPayload(payload map[string]any) bool {
thought, _ := payload[PayloadKeyThought].(bool)
return thought
}
func newErrorWithPayload(code, message string, extra map[string]any) PicoMessage {
payload := map[string]any{
"code": code,
+12 -5
View File
@@ -389,6 +389,7 @@ type antigravityJSONResponse struct {
Content struct {
Parts []struct {
Text string `json:"text,omitempty"`
Thought bool `json:"thought,omitempty"`
ThoughtSignature string `json:"thoughtSignature,omitempty"`
ThoughtSignatureSnake string `json:"thought_signature,omitempty"`
FunctionCall *antigravityFunctionCall `json:"functionCall,omitempty"`
@@ -406,6 +407,7 @@ type antigravityJSONResponse struct {
func (p *AntigravityProvider) parseSSEResponse(body string) (*LLMResponse, error) {
var contentParts []string
var reasoningParts []string
var toolCalls []ToolCall
var usage *UsageInfo
var finishReason string
@@ -433,7 +435,11 @@ func (p *AntigravityProvider) parseSSEResponse(body string) (*LLMResponse, error
for _, candidate := range resp.Candidates {
for _, part := range candidate.Content.Parts {
if part.Text != "" {
contentParts = append(contentParts, part.Text)
if part.Thought {
reasoningParts = append(reasoningParts, part.Text)
} else {
contentParts = append(contentParts, part.Text)
}
}
if part.FunctionCall != nil {
argumentsJSON, _ := json.Marshal(part.FunctionCall.Args)
@@ -475,10 +481,11 @@ func (p *AntigravityProvider) parseSSEResponse(body string) (*LLMResponse, error
}
return &LLMResponse{
Content: strings.Join(contentParts, ""),
ToolCalls: toolCalls,
FinishReason: mappedFinish,
Usage: usage,
Content: strings.Join(contentParts, ""),
ReasoningContent: strings.Join(reasoningParts, ""),
ToolCalls: toolCalls,
FinishReason: mappedFinish,
Usage: usage,
}, nil
}
@@ -54,3 +54,27 @@ func TestResolveToolResponseNameInfersNameFromGeneratedCallID(t *testing.T) {
t.Fatalf("expected inferred tool name search_docs, got %q", got)
}
}
func TestParseSSEResponse_SplitsThoughtAndVisibleContent(t *testing.T) {
p := &AntigravityProvider{}
body := "data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hidden reasoning\",\"thought\":true},{\"text\":\"visible answer\"}],\"role\":\"model\"},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":8,\"candidatesTokenCount\":17,\"totalTokenCount\":216}}}\n" +
"data: [DONE]\n"
resp, err := p.parseSSEResponse(body)
if err != nil {
t.Fatalf("parseSSEResponse() error = %v", err)
}
if resp.Content != "visible answer" {
t.Fatalf("Content = %q, want %q", resp.Content, "visible answer")
}
if resp.ReasoningContent != "hidden reasoning" {
t.Fatalf("ReasoningContent = %q, want %q", resp.ReasoningContent, "hidden reasoning")
}
if resp.FinishReason != "stop" {
t.Fatalf("FinishReason = %q, want %q", resp.FinishReason, "stop")
}
if resp.Usage == nil || resp.Usage.TotalTokens != 216 {
t.Fatalf("Usage.TotalTokens = %v, want %d", resp.Usage, 216)
}
}