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 <nicholaivogelfilms@gmail.com>

src/app/api/agent/route.ts

@@ -103,6 +103,7 @@ export async function POST(req: Request): Promise<Response> {
       ...githubTools,
       ...pluginTools,
     },
+    toolChoice: "auto",
     stopWhen: stepCountIs(10),
     onError({ error }) {
       const apiErr = unwrapAPICallError(error)

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,24 +171,37 @@ function friendlyToolName(raw: string): string {
   return TOOL_DISPLAY_NAMES[raw] ?? raw
 }
 
-// shared message + tool rendering for both variants
+interface ChatMessageProps {
-function ChatMessage({
+  readonly msg: UIMessage
+  readonly copiedId: string | null
+  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)
+}
+
+// renders parts in their natural order from the AI SDK
+const ChatMessage = memo(
+  function ChatMessage({
     msg,
     copiedId,
     onCopy,
     onRegenerate,
-}: {
+    isStreaming: msgStreaming = false,
-  readonly msg: {
+  }: ChatMessageProps) {
-    readonly id: string
-    readonly role: string
-    readonly parts: ReadonlyArray<unknown>
-  }
-  readonly copiedId: string | null
-  onCopy: (id: string, text: string) => void
-  onRegenerate: () => void
-}) {
     if (msg.role === "user") {
-    const text = getTextContent(msg.parts)
+      const text = msg.parts
+        .filter(isTextUIPart)
+        .map((p) => p.text)
+        .join("")
       return (
         <Message from="user">
           <MessageContent>{text}</MessageContent>
@@ -185,53 +209,85 @@ function ChatMessage({
       )
     }
 
-  const textParts: string[] = []
+    // walk parts sequentially, flushing text when
-  const toolParts: Array<{
+    // hitting a tool or reasoning part to preserve
-    type: string
+    // interleaving. text flushed before the final
-    state: ToolUIPart["state"]
+    // segment is "thinking" (intermediate chain-of-
-    toolName: string
+    // thought) and rendered muted + collapsible.
-    input: unknown
+    const elements: React.ReactNode[] = []
-    output: unknown
+    let pendingText = ""
-    errorText?: string
+    let allText = ""
-  }> = []
+    let pendingReasoning = ""
+    let reasoningStreaming = false
 
-  for (const part of msg.parts) {
+    let sawToolPart = false
-    const p = part as Record<string, unknown>
+
-    if (p.type === "text" && typeof p.text === "string") {
+    const flushThinking = (
-      textParts.push(p.text)
+      text: string,
-    }
+      idx: number,
-    const pType = p.type as string | undefined
+      streaming = false
-    // handle static (tool-<name>) and dynamic
+    ) => {
-    // (dynamic-tool) tool parts
+      if (!text) return
-    if (
+      elements.push(
-      typeof pType === "string" &&
+        <Reasoning
-      (pType.startsWith("tool-") ||
+          key={`think-${idx}`}
-        pType === "dynamic-tool")
+          isStreaming={streaming}
-    ) {
+          defaultOpen={false}
-      // extract tool name from type field or toolName
+        >
-      const rawName = pType.startsWith("tool-")
+          <ReasoningTrigger />
-        ? pType.slice(5)
+          <ReasoningContent>{text}</ReasoningContent>
-        : ((p.toolName ?? "") as string)
+        </Reasoning>
-      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,
-      })
-    }
     }
 
-  const text = textParts.join("")
+    const flushText = (idx: number, isFinal: boolean) => {
+      if (!pendingText) return
+      if (!isFinal) {
+        // intermediate text before more tools = thinking
+        flushThinking(pendingText, idx)
+      } else {
+        elements.push(
+          <MessageContent key={`text-${idx}`}>
+            <MessageResponse>
+              {pendingText}
+            </MessageResponse>
+          </MessageContent>
+        )
+      }
+      pendingText = ""
+    }
 
-  return (
+    for (let i = 0; i < msg.parts.length; i++) {
-    <Message from="assistant">
+      const part = msg.parts[i]
-      {toolParts.map((tp, i) => (
+
-        <Tool key={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(
+          <Tool key={tp.toolCallId}>
             <ToolHeader
-            title={tp.toolName}
+              title={
+                friendlyToolName(rawName) || "Working"
+              }
               type={tp.type as ToolUIPart["type"]}
               state={tp.state}
             />
@@ -246,16 +302,44 @@ function ChatMessage({
               )}
             </ToolContent>
           </Tool>
-      ))}
+        )
-      {text ? (
+      }
-        <>
+    }
-          <MessageContent>
+
-            <MessageResponse>{text}</MessageResponse>
+    // flush remaining reasoning
-          </MessageContent>
+    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 (
+      <Message from="assistant">
+        {hasContent ? elements : <Loader />}
+        {allText && (
           <Actions>
             <Action
               tooltip="Copy"
-              onClick={() => onCopy(msg.id, text)}
+              onClick={() => onCopy(msg.id, allText)}
             >
               {copiedId === msg.id ? (
                 <Check className="size-4" />
@@ -276,25 +360,23 @@ function ChatMessage({
               <RefreshCcwIcon className="size-4" />
             </Action>
           </Actions>
-        </>
-      ) : (
-        <Loader />
         )}
       </Message>
     )
-}
+  },
-
+  (prev, next) => {
-function getTextContent(
+    if (prev.msg !== next.msg) return false
-  parts: ReadonlyArray<unknown>
+    if (prev.onCopy !== next.onCopy) return false
-): string {
+    if (prev.onRegenerate !== next.onRegenerate)
-  return (parts as ReadonlyArray<{ type: string; text?: string }>)
+      return false
-    .filter(
+    if (prev.isStreaming !== next.isStreaming)
-      (p): p is { type: "text"; text: string } =>
+      return false
-        p.type === "text"
+    const prevCopied = prev.copiedId === prev.msg.id
-    )
+    const nextCopied = next.copiedId === next.msg.id
-    .map((p) => p.text)
+    if (prevCopied !== nextCopied) return false
-    .join("")
+    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) => {
+                onSend={handleIdleSend}
-                  setIsActive(true)
-                  chat.sendMessage({ text })
-                }}
                 className="rounded-2xl"
               />
 
@@ -652,13 +746,19 @@ export function ChatView({ variant }: ChatViewProps) {
             {chat.messages.length > 0 ? (
               <Conversation className="flex-1">
                 <ConversationContent className="mx-auto w-full max-w-3xl">
-                  {chat.messages.map((msg) => (
+                  {chat.messages.map((msg, idx) => (
                     <ChatMessage
                       key={msg.id}
                       msg={msg}
                       copiedId={copiedId}
                       onCopy={handleCopy}
                       onRegenerate={chat.regenerate}
+                      isStreaming={
+                        (chat.status === "streaming" ||
+                          chat.status === "submitted") &&
+                        idx === chat.messages.length - 1 &&
+                        msg.role === "assistant"
+                      }
                     />
                   ))}
                 </ConversationContent>
@@ -698,9 +798,7 @@ export function ChatView({ variant }: ChatViewProps) {
               recorder={recorder}
               status={chat.status}
               isGenerating={chat.isGenerating}
-              onSend={(text) =>
+              onSend={handleActiveSend}
-                chat.sendMessage({ text })
-              }
               onNewChat={chat.messages.length > 0 ? chat.newChat : undefined}
               className="rounded-2xl"
             />
@@ -729,13 +827,19 @@ export function ChatView({ variant }: ChatViewProps) {
               </Suggestions>
             </div>
           ) : (
-            chat.messages.map((msg) => (
+            chat.messages.map((msg, idx) => (
               <ChatMessage
                 key={msg.id}
                 msg={msg}
                 copiedId={copiedId}
                 onCopy={handleCopy}
                 onRegenerate={chat.regenerate}
+                isStreaming={
+                  (chat.status === "streaming" ||
+                    chat.status === "submitted") &&
+                  idx === chat.messages.length - 1 &&
+                  msg.role === "assistant"
+                }
               />
             ))
           )}
@@ -751,9 +855,7 @@ export function ChatView({ variant }: ChatViewProps) {
               recorder={recorder}
               status={chat.status}
               isGenerating={chat.isGenerating}
-              onSend={(text) =>
+              onSend={handleActiveSend}
-                chat.sendMessage({ text })
-              }
               onNewChat={chat.messages.length > 0 ? chat.newChat : undefined}
             />
       </div>

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<typeof Collapsible> & {
   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,7 +153,18 @@ export type ReasoningContentProps = ComponentProps<typeof CollapsibleContent> &
   children: string
 }
 
-export const ReasoningContent = memo(({ className, children, ...props }: ReasoningContentProps) => (
+export const ReasoningContent = memo(({ className, children, ...props }: ReasoningContentProps) => {
+  const { isStreaming } = useReasoning()
+  const scrollRef = useRef<HTMLDivElement>(null)
+
+  // auto-scroll to show latest tokens during streaming
+  useEffect(() => {
+    if (isStreaming && scrollRef.current) {
+      scrollRef.current.scrollTop = scrollRef.current.scrollHeight
+    }
+  }, [children, isStreaming])
+
+  return (
     <CollapsibleContent
       className={cn(
         "mt-4 text-sm",
@@ -162,10 +172,16 @@ export const ReasoningContent = memo(({ className, children, ...props }: Reasoni
         className,
       )}
       {...props}
+    >
+      <div
+        ref={scrollRef}
+        className={cn(isStreaming && "max-h-[3.75rem] overflow-hidden")}
       >
         <Streamdown {...props}>{children}</Streamdown>
+      </div>
     </CollapsibleContent>
-))
+  )
+})
 
 Reasoning.displayName = "Reasoning"
 ReasoningTrigger.displayName = "ReasoningTrigger"

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?"',