streamdown

packages/ui/@/components/ai-elements/message.tsx

@@ -1,7 +1,12 @@
 "use client";
 
 import type { UIMessage } from "ai";
-import type { ComponentProps, HTMLAttributes, ReactElement } from "react";
+import type {
+  ComponentProps,
+  HTMLAttributes,
+  ReactElement,
+  ReactNode,
+} from "react";
 
 import { Button } from "@/components/ui/button";
 import {
@@ -25,6 +30,7 @@ import {
   useMemo,
   useState,
 } from "react";
+import { Streamdown } from "streamdown";
 
 export type MessageProps = HTMLAttributes<HTMLDivElement> & {
   from: UIMessage["role"];
@@ -33,7 +39,7 @@ export type MessageProps = HTMLAttributes<HTMLDivElement> & {
 export const Message = ({ className, from, ...props }: MessageProps) => (
   <div
     className={cn(
-      "group flex w-full max-w-[95%] flex-col gap-2",
+      "group flex w-full max-w-4xl flex-col gap-2",
       from === "user" ? "is-user ml-auto justify-end" : "is-assistant",
       className
     )}
@@ -297,16 +303,60 @@ export const MessageBranchPage = ({
 
 export type MessageResponseProps = ComponentProps<"div">;
 
+function extractText(children: ReactNode): string | null {
+  if (children == null || typeof children === "boolean") {
+    return "";
+  }
+
+  if (
+    typeof children === "string" ||
+    typeof children === "number" ||
+    typeof children === "bigint"
+  ) {
+    return String(children);
+  }
+
+  if (Array.isArray(children)) {
+    let result = "";
+    for (const child of children) {
+      const part = extractText(child);
+      if (part === null) {
+        return null;
+      }
+      result += part;
+    }
+    return result;
+  }
+
+  return null;
+}
+
 export const MessageResponse = memo(
-  ({ className, ...props }: MessageResponseProps) => (
-    <div
-      className={cn(
-        "size-full whitespace-pre-wrap break-words [&>*:first-child]:mt-0 [&>*:last-child]:mb-0",
-        className
-      )}
-      {...props}
-    />
-  ),
+  ({ className, children, ...props }: MessageResponseProps) => {
+    const text = extractText(children);
+
+    return (
+      <div
+        className={cn(
+          "size-full break-words [&>*:first-child]:mt-0 [&>*:last-child]:mb-0",
+          className
+        )}
+        {...props}
+      >
+        {text !== null ? (
+          <Streamdown
+            className="space-y-3 whitespace-pre-wrap"
+            mode="streaming"
+            parseIncompleteMarkdown={true}
+          >
+            {text}
+          </Streamdown>
+        ) : (
+          children
+        )}
+      </div>
+    );
+  },
   (prevProps, nextProps) => prevProps.children === nextProps.children
 );
 

packages/ui/@/components/ai-elements/prompt-input.tsx

@@ -6,7 +6,7 @@ import { Button } from "@/components/ui/button";
 import { Textarea } from "@/components/ui/textarea";
 import { cn } from "@/lib/utils";
 import { ArrowUpIcon } from "lucide-react";
-import { forwardRef } from "react";
+import { forwardRef, useCallback, useLayoutEffect, useRef } from "react";
 
 export type PromptInputProps = ComponentProps<"form">;
 
@@ -27,35 +27,76 @@ PromptInput.displayName = "PromptInput";
 
 export type PromptInputTextareaProps = ComponentProps<typeof Textarea>;
 
+const PROMPT_INPUT_TEXTAREA_MAX_HEIGHT = 220;
+
+function autoResizeTextarea(
+  textarea: HTMLTextAreaElement,
+  maxHeight = PROMPT_INPUT_TEXTAREA_MAX_HEIGHT
+): void {
+  textarea.style.height = "auto";
+  const nextHeight = Math.min(textarea.scrollHeight, maxHeight);
+  textarea.style.height = `${nextHeight}px`;
+  textarea.style.overflowY =
+    textarea.scrollHeight > maxHeight ? "auto" : "hidden";
+}
+
 export const PromptInputTextarea = forwardRef<
   HTMLTextAreaElement,
   PromptInputTextareaProps
->(({ className, onKeyDown, ...props }, ref) => (
-  <Textarea
-    className={cn(
-      "min-h-[88px] max-h-[220px] resize-none border-0 bg-transparent p-2 shadow-none focus-visible:ring-0",
-      className
-    )}
-    onKeyDown={(event) => {
-      onKeyDown?.(event);
-      if (event.defaultPrevented) {
-        return;
-      }
+>(({ className, onInput, onKeyDown, rows, value, ...props }, ref) => {
+  const innerRef = useRef<HTMLTextAreaElement | null>(null);
 
-      const isComposing = event.nativeEvent.isComposing;
-      if (isComposing) {
-        return;
+  const assignRef = useCallback(
+    (node: HTMLTextAreaElement | null) => {
+      innerRef.current = node;
+      if (typeof ref === "function") {
+        ref(node);
+      } else if (ref) {
+        ref.current = node;
       }
+    },
+    [ref]
+  );
 
-      if (event.key === "Enter" && !event.shiftKey) {
-        event.preventDefault();
-        event.currentTarget.form?.requestSubmit();
-      }
-    }}
-    ref={ref}
-    {...props}
-  />
-));
+  useLayoutEffect(() => {
+    if (innerRef.current) {
+      autoResizeTextarea(innerRef.current);
+    }
+  }, [value]);
+
+  return (
+    <Textarea
+      className={cn(
+        "min-h-0 max-h-[220px] resize-none overflow-y-hidden border-0 bg-transparent p-2 shadow-none focus-visible:ring-0",
+        className
+      )}
+      onInput={(event) => {
+        autoResizeTextarea(event.currentTarget);
+        onInput?.(event);
+      }}
+      onKeyDown={(event) => {
+        onKeyDown?.(event);
+        if (event.defaultPrevented) {
+          return;
+        }
+
+        const isComposing = event.nativeEvent.isComposing;
+        if (isComposing) {
+          return;
+        }
+
+        if (event.key === "Enter" && !event.shiftKey) {
+          event.preventDefault();
+          event.currentTarget.form?.requestSubmit();
+        }
+      }}
+      ref={assignRef}
+      rows={rows ?? 1}
+      value={value}
+      {...props}
+    />
+  );
+});
 
 PromptInputTextarea.displayName = "PromptInputTextarea";
 
@@ -82,4 +123,3 @@ export const PromptInputSubmit = ({
     {children ?? <ArrowUpIcon className="size-4" />}
   </Button>
 );
-

packages/ui/package.json

@@ -45,6 +45,7 @@
     "react-day-picker": "^8.9.1",
     "react-hook-form": "^7.48.2",
     "react-remove-scroll": "^2.5.7",
+    "streamdown": "^2.1.0",
     "tailwind-merge": "^2.0.0",
     "tailwindcss-animate": "^1.0.7",
     "uniqolor": "1.1.1",