From a25c8a26bcf2f8bbdc523fb9894660bfb3c1aac7 Mon Sep 17 00:00:00 2001 From: Nicholai Date: Sat, 7 Feb 2026 13:28:32 -0700 Subject: [PATCH] fix(agent): improve chat message rendering (#56) Refactor ChatMessage to use AI SDK type guards and render parts in natural order. Collapse reasoning by default so thinking tokens don't flood the screen. Cap reasoning content height during streaming. Add tool workflow guidance to system prompt. Co-authored-by: Nicholai --- src/app/api/agent/route.ts | 1 + src/components/agent/chat-view.tsx | 340 +++++++++++++++++++---------- src/components/ai/reasoning.tsx | 50 +++-- src/lib/agent/system-prompt.ts | 36 +-- 4 files changed, 275 insertions(+), 152 deletions(-) diff --git a/src/app/api/agent/route.ts b/src/app/api/agent/route.ts index 406a9e7..a73aa8a 100755 --- a/src/app/api/agent/route.ts +++ b/src/app/api/agent/route.ts @@ -103,6 +103,7 @@ export async function POST(req: Request): Promise { ...githubTools, ...pluginTools, }, + toolChoice: "auto", stopWhen: stepCountIs(10), onError({ error }) { const apiErr = unwrapAPICallError(error) diff --git a/src/components/agent/chat-view.tsx b/src/components/agent/chat-view.tsx index 55417cb..725cf33 100755 --- a/src/components/agent/chat-view.tsx +++ b/src/components/agent/chat-view.tsx @@ -1,6 +1,6 @@ "use client" -import { useState, useCallback, useRef, useEffect } from "react" +import { useState, useCallback, useRef, useEffect, memo } from "react" import { CopyIcon, ThumbsUpIcon, @@ -20,9 +20,20 @@ import { IconAlertCircle, IconEye, } from "@tabler/icons-react" -import type { ToolUIPart } from "ai" +import { + isTextUIPart, + isToolUIPart, + isReasoningUIPart, + type UIMessage, + type ToolUIPart, + type DynamicToolUIPart, +} from "ai" import { cn } from "@/lib/utils" -import { Button } from "@/components/ui/button" +import { + Reasoning, + ReasoningTrigger, + ReasoningContent, +} from "@/components/ai/reasoning" import { Conversation, ConversationContent, @@ -160,102 +171,175 @@ function friendlyToolName(raw: string): string { return TOOL_DISPLAY_NAMES[raw] ?? raw } -// shared message + tool rendering for both variants -function ChatMessage({ - msg, - copiedId, - onCopy, - onRegenerate, -}: { - readonly msg: { - readonly id: string - readonly role: string - readonly parts: ReadonlyArray - } +interface ChatMessageProps { + readonly msg: UIMessage readonly copiedId: string | null - onCopy: (id: string, text: string) => void - onRegenerate: () => void -}) { - if (msg.role === "user") { - const text = getTextContent(msg.parts) - return ( - - {text} - - ) + readonly onCopy: (id: string, text: string) => void + readonly onRegenerate: () => void + readonly isStreaming?: boolean +} + +type AnyToolPart = ToolUIPart | DynamicToolUIPart + +function extractToolName(part: AnyToolPart): string { + if (part.type === "dynamic-tool") { + return part.toolName ?? "" } + return part.type.slice(5) +} - const textParts: string[] = [] - const toolParts: Array<{ - type: string - state: ToolUIPart["state"] - toolName: string - input: unknown - output: unknown - errorText?: string - }> = [] - - for (const part of msg.parts) { - const p = part as Record - if (p.type === "text" && typeof p.text === "string") { - textParts.push(p.text) +// renders parts in their natural order from the AI SDK +const ChatMessage = memo( + function ChatMessage({ + msg, + copiedId, + onCopy, + onRegenerate, + isStreaming: msgStreaming = false, + }: ChatMessageProps) { + if (msg.role === "user") { + const text = msg.parts + .filter(isTextUIPart) + .map((p) => p.text) + .join("") + return ( + + {text} + + ) } - const pType = p.type as string | undefined - // handle static (tool-) and dynamic - // (dynamic-tool) tool parts - if ( - typeof pType === "string" && - (pType.startsWith("tool-") || - pType === "dynamic-tool") - ) { - // extract tool name from type field or toolName - const rawName = pType.startsWith("tool-") - ? pType.slice(5) - : ((p.toolName ?? "") as string) - toolParts.push({ - type: pType, - state: p.state as ToolUIPart["state"], - toolName: - friendlyToolName(rawName) || "Working", - input: p.input, - output: p.output, - errorText: p.errorText as string | undefined, - }) + + // walk parts sequentially, flushing text when + // hitting a tool or reasoning part to preserve + // interleaving. text flushed before the final + // segment is "thinking" (intermediate chain-of- + // thought) and rendered muted + collapsible. + const elements: React.ReactNode[] = [] + let pendingText = "" + let allText = "" + let pendingReasoning = "" + let reasoningStreaming = false + + let sawToolPart = false + + const flushThinking = ( + text: string, + idx: number, + streaming = false + ) => { + if (!text) return + elements.push( + + + {text} + + ) } - } - const text = textParts.join("") - - return ( - - {toolParts.map((tp, i) => ( - - - - - {(tp.state === "output-available" || - tp.state === "output-error") && ( - - )} - - - ))} - {text ? ( - <> - - {text} + const flushText = (idx: number, isFinal: boolean) => { + if (!pendingText) return + if (!isFinal) { + // intermediate text before more tools = thinking + flushThinking(pendingText, idx) + } else { + elements.push( + + + {pendingText} + + ) + } + pendingText = "" + } + + for (let i = 0; i < msg.parts.length; i++) { + const part = msg.parts[i] + + if (isReasoningUIPart(part)) { + pendingReasoning += part.text + reasoningStreaming = part.state === "streaming" + continue + } + + if (isTextUIPart(part)) { + pendingText += part.text + allText += part.text + continue + } + + if (isToolUIPart(part)) { + sawToolPart = true + // flush reasoning accumulated before this tool + flushThinking(pendingReasoning, i, reasoningStreaming) + pendingReasoning = "" + reasoningStreaming = false + // flush text as thinking (not final) + flushText(i, false) + const tp = part as AnyToolPart + const rawName = extractToolName(tp) + elements.push( + + + + + {(tp.state === "output-available" || + tp.state === "output-error") && ( + + )} + + + ) + } + } + + // flush remaining reasoning + flushThinking( + pendingReasoning, + msg.parts.length, + reasoningStreaming + ) + + // while streaming, if no tool calls have arrived yet + // and text is substantial, it's likely chain-of-thought + // that'll be reclassified as thinking once tools come in. + // render it collapsed so it doesn't flood the screen. + const COT_THRESHOLD = 500 + if ( + msgStreaming && + !sawToolPart && + pendingText.length > COT_THRESHOLD + ) { + flushThinking(pendingText, msg.parts.length, true) + pendingText = "" + } + + // flush remaining text as the final response + flushText(msg.parts.length, true) + + const hasContent = elements.length > 0 + + return ( + + {hasContent ? elements : } + {allText && ( onCopy(msg.id, text)} + onClick={() => onCopy(msg.id, allText)} > {copiedId === msg.id ? ( @@ -276,25 +360,23 @@ function ChatMessage({ - - ) : ( - - )} - - ) -} - -function getTextContent( - parts: ReadonlyArray -): string { - return (parts as ReadonlyArray<{ type: string; text?: string }>) - .filter( - (p): p is { type: "text"; text: string } => - p.type === "text" + )} + ) - .map((p) => p.text) - .join("") -} + }, + (prev, next) => { + if (prev.msg !== next.msg) return false + if (prev.onCopy !== next.onCopy) return false + if (prev.onRegenerate !== next.onRegenerate) + return false + if (prev.isStreaming !== next.isStreaming) + return false + const prevCopied = prev.copiedId === prev.msg.id + const nextCopied = next.copiedId === next.msg.id + if (prevCopied !== nextCopied) return false + return true + } +) function ChatInput({ textareaRef, @@ -533,6 +615,21 @@ export function ChatView({ variant }: ChatViewProps) { [isPage, chat.sendMessage] ) + const handleIdleSend = useCallback( + (text: string) => { + setIsActive(true) + chat.sendMessage({ text }) + }, + [chat.sendMessage] + ) + + const handleActiveSend = useCallback( + (text: string) => { + chat.sendMessage({ text }) + }, + [chat.sendMessage] + ) + const suggestions = isPage ? DASHBOARD_SUGGESTIONS : getSuggestionsForPath(chat.pathname) @@ -593,10 +690,7 @@ export function ChatView({ variant }: ChatViewProps) { recorder={recorder} status={chat.status} isGenerating={chat.isGenerating} - onSend={(text) => { - setIsActive(true) - chat.sendMessage({ text }) - }} + onSend={handleIdleSend} className="rounded-2xl" /> @@ -652,13 +746,19 @@ export function ChatView({ variant }: ChatViewProps) { {chat.messages.length > 0 ? ( - {chat.messages.map((msg) => ( + {chat.messages.map((msg, idx) => ( ))} @@ -698,9 +798,7 @@ export function ChatView({ variant }: ChatViewProps) { recorder={recorder} status={chat.status} isGenerating={chat.isGenerating} - onSend={(text) => - chat.sendMessage({ text }) - } + onSend={handleActiveSend} onNewChat={chat.messages.length > 0 ? chat.newChat : undefined} className="rounded-2xl" /> @@ -729,13 +827,19 @@ export function ChatView({ variant }: ChatViewProps) { ) : ( - chat.messages.map((msg) => ( + chat.messages.map((msg, idx) => ( )) )} @@ -751,9 +855,7 @@ export function ChatView({ variant }: ChatViewProps) { recorder={recorder} status={chat.status} isGenerating={chat.isGenerating} - onSend={(text) => - chat.sendMessage({ text }) - } + onSend={handleActiveSend} onNewChat={chat.messages.length > 0 ? chat.newChat : undefined} /> diff --git a/src/components/ai/reasoning.tsx b/src/components/ai/reasoning.tsx index 1499bd9..d1c7f44 100755 --- a/src/components/ai/reasoning.tsx +++ b/src/components/ai/reasoning.tsx @@ -3,7 +3,7 @@ import { useControllableState } from "@radix-ui/react-use-controllable-state" import { BrainIcon, ChevronDownIcon } from "lucide-react" import type { ComponentProps, ReactNode } from "react" -import { createContext, memo, useContext, useEffect, useState } from "react" +import { createContext, memo, useContext, useEffect, useRef, useState } from "react" import { Streamdown } from "streamdown" import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible" import { cn } from "@/lib/utils" @@ -34,7 +34,7 @@ export type ReasoningProps = ComponentProps & { duration?: number } -const AUTO_CLOSE_DELAY = 1000 +const CLOSE_GRACE_MS = 300 const MS_IN_S = 1000 export const Reasoning = memo( @@ -73,14 +73,13 @@ export const Reasoning = memo( } }, [isStreaming, startTime, setDuration]) - // Auto-open when streaming starts, auto-close when streaming ends (once only) + // close gracefully once streaming finishes useEffect(() => { if (defaultOpen && !isStreaming && isOpen && !hasAutoClosed) { - // Add a small delay before closing to allow user to see the content const timer = setTimeout(() => { setIsOpen(false) setHasAutoClosed(true) - }, AUTO_CLOSE_DELAY) + }, CLOSE_GRACE_MS) return () => clearTimeout(timer) } @@ -154,18 +153,35 @@ export type ReasoningContentProps = ComponentProps & children: string } -export const ReasoningContent = memo(({ className, children, ...props }: ReasoningContentProps) => ( - - {children} - -)) +export const ReasoningContent = memo(({ className, children, ...props }: ReasoningContentProps) => { + const { isStreaming } = useReasoning() + const scrollRef = useRef(null) + + // auto-scroll to show latest tokens during streaming + useEffect(() => { + if (isStreaming && scrollRef.current) { + scrollRef.current.scrollTop = scrollRef.current.scrollHeight + } + }, [children, isStreaming]) + + return ( + +
+ {children} +
+ + ) +}) Reasoning.displayName = "Reasoning" ReasoningTrigger.displayName = "ReasoningTrigger" diff --git a/src/lib/agent/system-prompt.ts b/src/lib/agent/system-prompt.ts index f0103df..407e8c7 100755 --- a/src/lib/agent/system-prompt.ts +++ b/src/lib/agent/system-prompt.ts @@ -5,22 +5,6 @@ import type { PromptSection } from "@/lib/agent/plugins/types" type PromptMode = "full" | "minimal" | "none" -type ToolCategory = - | "data" - | "navigation" - | "ui" - | "memory" - | "github" - | "skills" - | "feedback" - -interface ToolMeta { - readonly name: string - readonly summary: string - readonly category: ToolCategory - readonly adminOnly?: true -} - interface DashboardSummary { readonly id: string readonly name: string @@ -38,6 +22,22 @@ interface PromptContext { readonly mode?: PromptMode } +type ToolCategory = + | "data" + | "navigation" + | "ui" + | "memory" + | "github" + | "skills" + | "feedback" + +interface ToolMeta { + readonly name: string + readonly summary: string + readonly category: ToolCategory + readonly adminOnly?: true +} + interface DerivedState { readonly mode: PromptMode readonly page: string @@ -665,6 +665,10 @@ function buildGuidelines( return [ ...core, + "- Tool workflow: data requests -> queryData immediately. " + + "Navigation -> navigateTo, brief confirmation. " + + "Dashboards -> queryData first, then generateUI. " + + "Memories -> save proactively with rememberContext.", '- "How\'s development going?" means fetch repo_stats and ' + 'recent commits right now, not "Would you like to see ' + 'commits or PRs?"',