diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index ac230aa86..03fdfec82 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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 { diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index a67c8d040..56ea000c8 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -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") diff --git a/pkg/channels/pico/client.go b/pkg/channels/pico/client.go index b4bfd09e5..bf3e38cf4 100644 --- a/pkg/channels/pico/client.go +++ b/pkg/channels/pico/client.go @@ -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) diff --git a/pkg/channels/pico/client_test.go b/pkg/channels/pico/client_test.go index b40606647..732589432 100644 --- a/pkg/channels/pico/client_test.go +++ b/pkg/channels/pico/client_test.go @@ -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): + } +} diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index e22da1ba1..6525c2d4a 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -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 { diff --git a/pkg/channels/pico/protocol.go b/pkg/channels/pico/protocol.go index 3f8ba8643..ecdc2d140 100644 --- a/pkg/channels/pico/protocol.go +++ b/pkg/channels/pico/protocol.go @@ -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, diff --git a/pkg/providers/antigravity_provider.go b/pkg/providers/antigravity_provider.go index 8a1890212..b5ab847d5 100644 --- a/pkg/providers/antigravity_provider.go +++ b/pkg/providers/antigravity_provider.go @@ -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 } diff --git a/pkg/providers/antigravity_provider_test.go b/pkg/providers/antigravity_provider_test.go index 238765321..9155e2d56 100644 --- a/pkg/providers/antigravity_provider_test.go +++ b/pkg/providers/antigravity_provider_test.go @@ -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) + } +} diff --git a/web/frontend/src/components/chat/assistant-message.tsx b/web/frontend/src/components/chat/assistant-message.tsx index 9966226b2..418516172 100644 --- a/web/frontend/src/components/chat/assistant-message.tsx +++ b/web/frontend/src/components/chat/assistant-message.tsx @@ -1,5 +1,6 @@ -import { IconCheck, IconCopy } from "@tabler/icons-react" +import { IconBrain, IconCheck, IconCopy } from "@tabler/icons-react" import { useState } from "react" +import { useTranslation } from "react-i18next" import ReactMarkdown from "react-markdown" import rehypeRaw from "rehype-raw" import rehypeSanitize from "rehype-sanitize" @@ -7,16 +8,20 @@ import remarkGfm from "remark-gfm" import { Button } from "@/components/ui/button" import { formatMessageTime } from "@/hooks/use-pico-chat" +import { cn } from "@/lib/utils" interface AssistantMessageProps { content: string + isThought?: boolean timestamp?: string | number } export function AssistantMessage({ content, + isThought = false, timestamp = "", }: AssistantMessageProps) { + const { t } = useTranslation() const [isCopied, setIsCopied] = useState(false) const formattedTimestamp = timestamp !== "" ? formatMessageTime(timestamp) : "" @@ -33,6 +38,12 @@ export function AssistantMessage({