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,102 +171,175 @@ 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
-  msg,
-  copiedId,
-  onCopy,
-  onRegenerate,
-}: {
-  readonly msg: {
-    readonly id: string
-    readonly role: string
-    readonly parts: ReadonlyArray<unknown>
-  }
   readonly copiedId: string | null
-  onCopy: (id: string, text: string) => void
+  readonly onCopy: (id: string, text: string) => void
-  onRegenerate: () => void
+  readonly onRegenerate: () => void
-}) {
+  readonly isStreaming?: boolean
-  if (msg.role === "user") {
+}
-    const text = getTextContent(msg.parts)
+
-    return (
+type AnyToolPart = ToolUIPart | DynamicToolUIPart
-      <Message from="user">
+
-        <MessageContent>{text}</MessageContent>
+function extractToolName(part: AnyToolPart): string {
-      </Message>
+  if (part.type === "dynamic-tool") {
-    )
+    return part.toolName ?? ""
   }
+  return part.type.slice(5)
+}
 
-  const textParts: string[] = []
+// renders parts in their natural order from the AI SDK
-  const toolParts: Array<{
+const ChatMessage = memo(
-    type: string
+  function ChatMessage({
-    state: ToolUIPart["state"]
+    msg,
-    toolName: string
+    copiedId,
-    input: unknown
+    onCopy,
-    output: unknown
+    onRegenerate,
-    errorText?: string
+    isStreaming: msgStreaming = false,
-  }> = []
+  }: ChatMessageProps) {
-
+    if (msg.role === "user") {
-  for (const part of msg.parts) {
+      const text = msg.parts
-    const p = part as Record<string, unknown>
+        .filter(isTextUIPart)
-    if (p.type === "text" && typeof p.text === "string") {
+        .map((p) => p.text)
-      textParts.push(p.text)
+        .join("")
+      return (
+        <Message from="user">
+          <MessageContent>{text}</MessageContent>
+        </Message>
+      )
     }
-    const pType = p.type as string | undefined
+
-    // handle static (tool-<name>) and dynamic
+    // walk parts sequentially, flushing text when
-    // (dynamic-tool) tool parts
+    // hitting a tool or reasoning part to preserve
-    if (
+    // interleaving. text flushed before the final
-      typeof pType === "string" &&
+    // segment is "thinking" (intermediate chain-of-
-      (pType.startsWith("tool-") ||
+    // thought) and rendered muted + collapsible.
-        pType === "dynamic-tool")
+    const elements: React.ReactNode[] = []
-    ) {
+    let pendingText = ""
-      // extract tool name from type field or toolName
+    let allText = ""
-      const rawName = pType.startsWith("tool-")
+    let pendingReasoning = ""
-        ? pType.slice(5)
+    let reasoningStreaming = false
-        : ((p.toolName ?? "") as string)
+
-      toolParts.push({
+    let sawToolPart = false
-        type: pType,
+
-        state: p.state as ToolUIPart["state"],
+    const flushThinking = (
-        toolName:
+      text: string,
-          friendlyToolName(rawName) || "Working",
+      idx: number,
-        input: p.input,
+      streaming = false
-        output: p.output,
+    ) => {
-        errorText: p.errorText as string | undefined,
+      if (!text) return
-      })
+      elements.push(
+        <Reasoning
+          key={`think-${idx}`}
+          isStreaming={streaming}
+          defaultOpen={false}
+        >
+          <ReasoningTrigger />
+          <ReasoningContent>{text}</ReasoningContent>
+        </Reasoning>
+      )
     }
-  }
 
-  const text = textParts.join("")
+    const flushText = (idx: number, isFinal: boolean) => {
-
+      if (!pendingText) return
-  return (
+      if (!isFinal) {
-    <Message from="assistant">
+        // intermediate text before more tools = thinking
-      {toolParts.map((tp, i) => (
+        flushThinking(pendingText, idx)
-        <Tool key={i}>
+      } else {
-          <ToolHeader
+        elements.push(
-            title={tp.toolName}
+          <MessageContent key={`text-${idx}`}>
-            type={tp.type as ToolUIPart["type"]}
+            <MessageResponse>
-            state={tp.state}
+              {pendingText}
-          />
+            </MessageResponse>
-          <ToolContent>
-            <ToolInput input={tp.input} />
-            {(tp.state === "output-available" ||
-              tp.state === "output-error") && (
-              <ToolOutput
-                output={tp.output}
-                errorText={tp.errorText}
-              />
-            )}
-          </ToolContent>
-        </Tool>
-      ))}
-      {text ? (
-        <>
-          <MessageContent>
-            <MessageResponse>{text}</MessageResponse>
           </MessageContent>
+        )
+      }
+      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(
+          <Tool key={tp.toolCallId}>
+            <ToolHeader
+              title={
+                friendlyToolName(rawName) || "Working"
+              }
+              type={tp.type as ToolUIPart["type"]}
+              state={tp.state}
+            />
+            <ToolContent>
+              <ToolInput input={tp.input} />
+              {(tp.state === "output-available" ||
+                tp.state === "output-error") && (
+                <ToolOutput
+                  output={tp.output}
+                  errorText={tp.errorText}
+                />
+              )}
+            </ToolContent>
+          </Tool>
+        )
+      }
+    }
+
+    // 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 (
+      <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>
-        </>
+        )}
-      ) : (
+      </Message>
-        <Loader />
-      )}
-    </Message>
-  )
-}
-
-function getTextContent(
-  parts: ReadonlyArray<unknown>
-): 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) => {
+                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,18 +153,35 @@ export type ReasoningContentProps = ComponentProps<typeof CollapsibleContent> &
   children: string
 }
 
-export const ReasoningContent = memo(({ className, children, ...props }: ReasoningContentProps) => (
+export const ReasoningContent = memo(({ className, children, ...props }: ReasoningContentProps) => {
-  <CollapsibleContent
+  const { isStreaming } = useReasoning()
-    className={cn(
+  const scrollRef = useRef<HTMLDivElement>(null)
-      "mt-4 text-sm",
+
-      "data-[state=closed]:fade-out-0 data-[state=closed]:slide-out-to-top-2 data-[state=open]:slide-in-from-top-2 text-muted-foreground outline-none data-[state=closed]:animate-out data-[state=open]:animate-in",
+  // auto-scroll to show latest tokens during streaming
-      className,
+  useEffect(() => {
-    )}
+    if (isStreaming && scrollRef.current) {
-    {...props}
+      scrollRef.current.scrollTop = scrollRef.current.scrollHeight
-  >
+    }
-    <Streamdown {...props}>{children}</Streamdown>
+  }, [children, isStreaming])
-  </CollapsibleContent>
+
-))
+  return (
+    <CollapsibleContent
+      className={cn(
+        "mt-4 text-sm",
+        "data-[state=closed]:fade-out-0 data-[state=closed]:slide-out-to-top-2 data-[state=open]:slide-in-from-top-2 text-muted-foreground outline-none data-[state=closed]:animate-out data-[state=open]:animate-in",
+        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?"',