refactor(ai-elements): wrap live message bubble (#525)

* refactor(ai-elements): migrate live message bubble

* fix(ai-elements): preserve ordered message ownership

* fix(ai-elements): preserve reasoning and falsy tool outputs

---------

Co-authored-by: hshum <hshum@users.noreply.github.com>

Author: HomenShum

src/components/ai-elements/message.tsx

@@ -12,23 +12,27 @@ import {
   TooltipTrigger,
 } from "@/components/ai-ui/tooltip";
 import { cn } from "@/lib/utils";
-import { cjk } from "@streamdown/cjk";
-import { code } from "@streamdown/code";
-import { math } from "@streamdown/math";
-import { mermaid } from "@streamdown/mermaid";
 import type { UIMessage } from "ai";
 import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
 import type { ComponentProps, HTMLAttributes, ReactElement } from "react";
 import {
   createContext,
+  lazy,
   memo,
+  Suspense,
   useCallback,
   useContext,
   useEffect,
   useMemo,
   useState,
 } from "react";
-import { Streamdown } from "streamdown";
+import type { Streamdown } from "streamdown";
+
+const StreamdownRenderer = lazy(() =>
+  import("./streamdown-renderer").then((module) => ({
+    default: module.StreamdownRenderer,
+  }))
+);
 
 export type MessageProps = HTMLAttributes<HTMLDivElement> & {
   from: UIMessage["role"];
@@ -321,18 +325,25 @@ export const MessageBranchPage = ({
 
 export type MessageResponseProps = ComponentProps<typeof Streamdown>;
 
-const streamdownPlugins = { cjk, code, math, mermaid };
-
 export const MessageResponse = memo(
-  ({ className, ...props }: MessageResponseProps) => (
-    <Streamdown
-      className={cn(
-        "size-full [&>*:first-child]:mt-0 [&>*:last-child]:mb-0",
-        className
-      )}
-      plugins={streamdownPlugins}
-      {...props}
-    />
+  ({ children, className, ...props }: MessageResponseProps) => (
+    <Suspense
+      fallback={
+        <div className={cn("size-full whitespace-pre-wrap", className)}>
+          {children}
+        </div>
+      }
+    >
+      <StreamdownRenderer
+        className={cn(
+          "size-full [&>*:first-child]:mt-0 [&>*:last-child]:mb-0",
+          className
+        )}
+        {...props}
+      >
+        {children}
+      </StreamdownRenderer>
+    </Suspense>
   ),
   (prevProps, nextProps) =>
     prevProps.children === nextProps.children &&

src/components/ai-elements/reasoning.tsx

@@ -7,26 +7,29 @@ import {
   CollapsibleTrigger,
 } from "@/components/ai-ui/collapsible";
 import { cn } from "@/lib/utils";
-import { cjk } from "@streamdown/cjk";
-import { code } from "@streamdown/code";
-import { math } from "@streamdown/math";
-import { mermaid } from "@streamdown/mermaid";
 import { BrainIcon, ChevronDownIcon } from "lucide-react";
 import type { ComponentProps, ReactNode } from "react";
 import {
   createContext,
+  lazy,
   memo,
+  Suspense,
   useCallback,
   useContext,
   useEffect,
   useMemo,
   useRef,
   useState,
 } from "react";
-import { Streamdown } from "streamdown";
 
 import { Shimmer } from "./shimmer";
 
+const ReasoningStreamdown = lazy(() =>
+  import("./streamdown-renderer").then((module) => ({
+    default: module.StreamdownRenderer,
+  }))
+);
+
 interface ReasoningContextValue {
   isStreaming: boolean;
   isOpen: boolean;
@@ -204,8 +207,6 @@ export type ReasoningContentProps = ComponentProps<
   children: string;
 };
 
-const streamdownPlugins = { cjk, code, math, mermaid };
-
 export const ReasoningContent = memo(
   ({ className, children, ...props }: ReasoningContentProps) => (
     <CollapsibleContent
@@ -216,7 +217,15 @@ export const ReasoningContent = memo(
       )}
       {...props}
     >
-      <Streamdown plugins={streamdownPlugins}>{children}</Streamdown>
+      <Suspense
+        fallback={
+          <div className="whitespace-pre-wrap text-muted-foreground">
+            {children}
+          </div>
+        }
+      >
+        <ReasoningStreamdown>{children}</ReasoningStreamdown>
+      </Suspense>
     </CollapsibleContent>
   )
 );

src/components/ai-elements/streamdown-renderer.tsx

@@ -0,0 +1,16 @@
+"use client";
+
+import { cjk } from "@streamdown/cjk";
+import { code } from "@streamdown/code";
+import { math } from "@streamdown/math";
+import { mermaid } from "@streamdown/mermaid";
+import type { ComponentProps } from "react";
+import { Streamdown } from "streamdown";
+
+const streamdownPlugins = { cjk, code, math, mermaid };
+
+export type StreamdownRendererProps = ComponentProps<typeof Streamdown>;
+
+export const StreamdownRenderer = (props: StreamdownRendererProps) => (
+  <Streamdown plugins={streamdownPlugins} {...props} />
+);

src/components/ai-elements/tool.tsx

@@ -157,18 +157,21 @@ export const ToolOutput = ({
   errorText,
   ...props
 }: ToolOutputProps) => {
-  if (!(output || errorText)) {
+  const hasOutput = output !== undefined && output !== null;
+  if (!hasOutput && !errorText) {
     return null;
   }
 
-  let Output = <div>{output as ReactNode}</div>;
+  let Output = hasOutput ? <div>{output as ReactNode}</div> : null;
 
-  if (typeof output === "object" && !isValidElement(output)) {
+  if (hasOutput && typeof output === "object" && !isValidElement(output)) {
     Output = (
       <DeferredCodeBlock code={JSON.stringify(output, null, 2)} />
     );
-  } else if (typeof output === "string") {
+  } else if (hasOutput && typeof output === "string") {
     Output = <DeferredCodeBlock code={output} />;
+  } else if (hasOutput) {
+    Output = <DeferredCodeBlock code={String(output)} />;
   }
 
   return (