From 75b351d37978c88c62b1a5a12960528f6825a435 Mon Sep 17 00:00:00 2001
From: Simon Farshid <simon.farshid@outlook.com>
Date: Tue, 2 Jul 2024 13:06:02 -0700
Subject: [PATCH] feat: styles revamp (#374)

---
 .../registry/assistant-ui/assistant-modal.tsx |  21 +-
 .../registry/assistant-ui/full/thread.tsx     | 227 ++++++++++--------
 .../registry/assistant-ui/markdown-text.tsx   |  45 ++--
 .../registry/assistant-ui/thread.tsx          | 103 ++++----
 4 files changed, 212 insertions(+), 184 deletions(-)

diff --git a/packages/shadcn-registry/registry/assistant-ui/assistant-modal.tsx b/packages/shadcn-registry/registry/assistant-ui/assistant-modal.tsx
index 3766463ae..f8ddb2cbf 100644
--- a/packages/shadcn-registry/registry/assistant-ui/assistant-modal.tsx
+++ b/packages/shadcn-registry/registry/assistant-ui/assistant-modal.tsx
@@ -9,7 +9,6 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-import { cn } from "@/lib/utils";
 import { TooltipProvider } from "@radix-ui/react-tooltip";
 import { type FC, forwardRef } from "react";
 import { AssistantModalPrimitive } from "@assistant-ui/react";
@@ -17,16 +16,14 @@ import { AssistantModalPrimitive } from "@assistant-ui/react";
 export const AssistantModal: FC = () => {
   return (
     <AssistantModalPrimitive.Root>
-      <AssistantModalPrimitive.Anchor className="fixed bottom-4 right-4 size-12">
+      <AssistantModalPrimitive.Anchor className="fixed bottom-4 right-4 size-11">
         <AssistantModalPrimitive.Trigger asChild>
           <FloatingAssistantButton />
         </AssistantModalPrimitive.Trigger>
       </AssistantModalPrimitive.Anchor>
       <AssistantModalPrimitive.Content
         sideOffset={16}
-        className={
-          "data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out data-[state=open]:zoom-in data-[state=open]:slide-in-from-bottom-1/2 data-[state=open]:slide-in-from-right-1/2 data-[state=closed]:slide-out-to-bottom-1/2 data-[state=closed]:slide-out-to-right-1/2 bg-popover text-popover-foreground data-[state=closed]:animate-out data-[state=open]:animate-in z-50 h-[500px] w-[400px] rounded-xl border p-0 shadow-md outline-none"
-        }
+        className="bg-popover text-popover-foreground data-[state=closed]:animate-out data-[state=open]:animate-in data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out data-[state=open]:zoom-in data-[state=open]:slide-in-from-bottom-1/2 data-[state=open]:slide-in-from-right-1/2 data-[state=closed]:slide-out-to-bottom-1/2 data-[state=closed]:slide-out-to-right-1/2 z-50 h-[500px] w-[400px] overflow-clip rounded-xl border p-0 shadow-md outline-none [&>div]:bg-inherit"
       >
         <Thread />
       </AssistantModalPrimitive.Content>
@@ -54,19 +51,13 @@ const FloatingAssistantButton = forwardRef<
             ref={ref}
           >
             <BotIcon
-              className={cn(
-                "absolute size-6 transition-all",
-                state === "open" && "rotate-90 scale-0",
-                state === "closed" && "rotate-0 scale-100",
-              )}
+              data-state={state}
+              className="absolute size-6 transition-all data-[state=closed]:rotate-0 data-[state=open]:rotate-90 data-[state=closed]:scale-100 data-[state=open]:scale-0"
             />
 
             <ChevronDownIcon
-              className={cn(
-                "absolute size-6 transition-all",
-                state === "open" && "rotate-0 scale-100",
-                state === "closed" && "-rotate-90 scale-0",
-              )}
+              data-state={state}
+              className="absolute size-6 transition-all data-[state=closed]:-rotate-90 data-[state=open]:rotate-0 data-[state=closed]:scale-0 data-[state=open]:scale-100"
             />
             <span className="sr-only">{tooltip}</span>
           </Button>
diff --git a/packages/shadcn-registry/registry/assistant-ui/full/thread.tsx b/packages/shadcn-registry/registry/assistant-ui/full/thread.tsx
index 1682e2b54..99271d598 100644
--- a/packages/shadcn-registry/registry/assistant-ui/full/thread.tsx
+++ b/packages/shadcn-registry/registry/assistant-ui/full/thread.tsx
@@ -3,6 +3,7 @@
 import {
   ActionBarPrimitive,
   BranchPickerPrimitive,
+  BranchPickerPrimitiveRootProps,
   ComposerPrimitive,
   MessagePrimitive,
   ThreadPrimitive,
@@ -33,11 +34,9 @@ import { MarkdownText } from "@/components/ui/assistant-ui/markdown-text";
 export const Thread: FC = () => {
   return (
     <TooltipProvider>
-      <ThreadPrimitive.Root className="flex h-full flex-col items-center pb-4">
-        <ThreadPrimitive.Viewport className="flex w-full flex-grow flex-col items-center overflow-y-scroll scroll-smooth px-4 pt-12">
-          <ThreadPrimitive.Empty>
-            <ThreadEmpty />
-          </ThreadPrimitive.Empty>
+      <ThreadPrimitive.Root className="bg-background h-full">
+        <ThreadPrimitive.Viewport className="flex h-full flex-col items-center overflow-y-scroll scroll-smooth px-4 pt-8">
+          <ThreadWelcome />
 
           <ThreadPrimitive.Messages
             components={{
@@ -47,58 +46,72 @@ export const Thread: FC = () => {
             }}
           />
 
-          <ThreadScrollToBottom />
+          <div className="sticky bottom-0 mt-4 flex w-full max-w-2xl flex-grow flex-col items-center justify-end rounded-t-lg bg-inherit pb-4">
+            <ThreadScrollToBottom />
+            <Composer />
+          </div>
         </ThreadPrimitive.Viewport>
-
-        <Composer />
       </ThreadPrimitive.Root>
     </TooltipProvider>
   );
 };
 
-const ThreadEmpty: FC = () => {
+const ThreadScrollToBottom: FC = () => {
   return (
-    <div className="flex flex-grow flex-col items-center justify-center">
-      <Avatar>
-        <AvatarFallback>C</AvatarFallback>
-      </Avatar>
-      <p className="mt-4 text-xl">How can I help you today?</p>
-    </div>
+    <ThreadPrimitive.ScrollToBottom asChild>
+      <TooltipIconButton
+        tooltip="Scroll to bottom"
+        variant="outline"
+        className="absolute -top-8 rounded-full disabled:invisible"
+      >
+        <ArrowDownIcon />
+      </TooltipIconButton>
+    </ThreadPrimitive.ScrollToBottom>
   );
 };
 
-const ThreadScrollToBottom: FC = () => {
+const ThreadWelcome: FC = () => {
   return (
-    <div className="sticky bottom-0">
-      <ThreadPrimitive.ScrollToBottom asChild>
-        <IconButton
-          tooltip="Scroll to bottom"
-          variant="outline"
-          className="absolute -top-10 rounded-full disabled:invisible"
-        >
-          <ArrowDownIcon className="size-4" />
-        </IconButton>
-      </ThreadPrimitive.ScrollToBottom>
-    </div>
+    <ThreadPrimitive.Empty>
+      <div className="flex flex-grow basis-full flex-col items-center justify-center">
+        <Avatar>
+          <AvatarFallback>C</AvatarFallback>
+        </Avatar>
+        <p className="mt-4 text-lg font-medium">How can I help you today?</p>
+      </div>
+    </ThreadPrimitive.Empty>
   );
 };
 
 const Composer: FC = () => {
   return (
-    <ComposerPrimitive.Root className="flex w-[calc(100%-32px)] max-w-[42rem] items-end rounded-lg border p-0.5 transition-shadow focus-within:shadow-sm">
+    <ComposerPrimitive.Root className="relative flex w-full items-end rounded-lg border transition-shadow focus-within:shadow-sm">
       <ComposerPrimitive.Input
         autoFocus
         placeholder="Write a message..."
-        className="placeholder:text-foreground/50 h-12 max-h-40 flex-grow resize-none bg-transparent p-3.5 text-sm outline-none"
+        rows={1}
+        className="placeholder:text-muted-foreground size-full max-h-40 resize-none bg-transparent p-4 pr-12 text-sm outline-none"
       />
       <ThreadPrimitive.If running={false}>
-        <ComposerPrimitive.Send className="bg-foreground m-2 flex h-8 w-8 items-center justify-center rounded-md text-2xl font-bold shadow transition-opacity disabled:opacity-10">
-          <SendHorizonalIcon className="text-background size-4" />
+        <ComposerPrimitive.Send asChild>
+          <TooltipIconButton
+            tooltip="Send"
+            variant="default"
+            className="absolute bottom-0 right-0 m-2.5 size-8 p-2 transition-opacity"
+          >
+            <SendHorizonalIcon />
+          </TooltipIconButton>
         </ComposerPrimitive.Send>
       </ThreadPrimitive.If>
       <ThreadPrimitive.If running>
-        <ComposerPrimitive.Cancel className="border-foreground m-3.5 flex size-5 items-center justify-center rounded-full border-2">
-          <div className="bg-foreground size-2 rounded-[1px]" />
+        <ComposerPrimitive.Cancel asChild>
+          <TooltipIconButton
+            tooltip="Cancel"
+            variant="outline"
+            className="absolute bottom-0 right-0 m-2.5 size-8 p-2 transition-opacity"
+          >
+            <CircleStopIcon />
+          </TooltipIconButton>
         </ComposerPrimitive.Cancel>
       </ThreadPrimitive.If>
     </ComposerPrimitive.Root>
@@ -107,42 +120,38 @@ const Composer: FC = () => {
 
 const UserMessage: FC = () => {
   return (
-    <MessagePrimitive.Root className="relative mb-6 flex w-full max-w-2xl flex-col items-end gap-2 pl-24">
-      <div className="relative mr-1 flex items-start gap-3">
-        <ActionBarPrimitive.Root
-          hideWhenRunning
-          autohide="not-last"
-          className="mt-2"
-        >
-          <ActionBarPrimitive.Edit asChild>
-            <IconButton tooltip="Edit">
-              <PencilIcon className="size-4" />
-            </IconButton>
-          </ActionBarPrimitive.Edit>
-        </ActionBarPrimitive.Root>
-
-        <div className="bg-foreground/5 text-foreground max-w-xl break-words rounded-3xl px-5 py-2.5">
-          <MessagePrimitive.Content />
-        </div>
+    <MessagePrimitive.Root className="my-4 grid w-full max-w-2xl auto-rows-auto grid-cols-[minmax(72px,1fr)_auto] gap-y-2">
+      <ActionBarPrimitive.Root
+        hideWhenRunning
+        autohide="not-last"
+        className="col-start-1 mr-3 mt-2.5 flex flex-col items-end"
+      >
+        <ActionBarPrimitive.Edit asChild>
+          <TooltipIconButton tooltip="Edit">
+            <PencilIcon />
+          </TooltipIconButton>
+        </ActionBarPrimitive.Edit>
+      </ActionBarPrimitive.Root>
+
+      <div className="bg-muted text-foreground col-start-2 row-start-1 max-w-xl break-words rounded-3xl px-5 py-2.5">
+        <MessagePrimitive.Content />
       </div>
 
-      <BranchPicker />
+      <BranchPicker className="col-span-full col-start-1 row-start-2 -mr-1 justify-end" />
     </MessagePrimitive.Root>
   );
 };
 
 const EditComposer: FC = () => {
   return (
-    <ComposerPrimitive.Root className="bg-foreground/5 mb-4 flex w-full max-w-2xl flex-col gap-2 rounded-xl">
+    <ComposerPrimitive.Root className="bg-muted my-4 flex w-full max-w-2xl flex-col gap-2 rounded-xl">
       <ComposerPrimitive.Input className="text-foreground flex h-8 w-full resize-none bg-transparent p-5 pb-0 outline-none" />
 
       <div className="mx-3 mb-3 flex items-center justify-center gap-2 self-end">
         <ComposerPrimitive.Cancel asChild>
-          <Button variant="secondary" className="bg-transparent">
-            Cancel
-          </Button>
+          <Button variant="ghost">Cancel</Button>
         </ComposerPrimitive.Cancel>
-        <ComposerPrimitive.Send>
+        <ComposerPrimitive.Send asChild>
           <Button>Send</Button>
         </ComposerPrimitive.Send>
       </div>
@@ -152,72 +161,76 @@ const EditComposer: FC = () => {
 
 const AssistantMessage: FC = () => {
   return (
-    <MessagePrimitive.Root className="relative mb-6 flex w-full max-w-2xl gap-3">
-      <Avatar>
+    <MessagePrimitive.Root className="relative my-4 grid w-full max-w-2xl grid-cols-[auto_auto_1fr] grid-rows-[auto_1fr]">
+      <Avatar className="col-start-1 row-span-full row-start-1 mr-4">
         <AvatarFallback>A</AvatarFallback>
       </Avatar>
 
-      <div className="mt-2 flex-grow">
-        <MessagePrimitive.InProgress className="bg-foreground inline-block size-3 animate-pulse rounded-full" />
-        <div className="text-foreground max-w-xl break-words">
-          <MessagePrimitive.Content components={{ Text: MarkdownText }} />
-        </div>
+      <div className="text-foreground col-span-2 col-start-2 row-start-1 my-1.5 max-w-xl break-words leading-7">
+        <MessagePrimitive.Content components={{ Text: MarkdownText }} />
+      </div>
 
-        <div className="flex pt-2">
-          <BranchPicker />
+      <BranchPicker className="col-start-2 row-start-2 -ml-2 mr-2" />
 
-          <ActionBarPrimitive.Root
-            hideWhenRunning
-            autohide="not-last"
-            autohideFloat="single-branch"
-            className="data-[floating]:bg-background z-50 flex items-center gap-1 rounded-lg data-[floating]:absolute data-[floating]:border data-[floating]:p-1"
-          >
-            <ActionBarPrimitive.Copy asChild>
-              <IconButton tooltip="Copy">
-                <MessagePrimitive.If copied>
-                  <CheckIcon className="size-4" />
-                </MessagePrimitive.If>
-                <MessagePrimitive.If copied={false}>
-                  <CopyIcon className="size-4" />
-                </MessagePrimitive.If>
-              </IconButton>
-            </ActionBarPrimitive.Copy>
-            <ActionBarPrimitive.Reload asChild>
-              <IconButton tooltip="Refresh">
-                <RefreshCwIcon className="size-4" />
-              </IconButton>
-            </ActionBarPrimitive.Reload>
-          </ActionBarPrimitive.Root>
-        </div>
-      </div>
+      <ActionBarPrimitive.Root
+        hideWhenRunning
+        autohide="not-last"
+        autohideFloat="single-branch"
+        className="text-muted-foreground data-[floating]:bg-background col-start-3 row-start-2 -ml-1 flex gap-1 rounded-lg data-[floating]:absolute data-[floating]:border-2 data-[floating]:p-1"
+      >
+        <ActionBarPrimitive.Copy asChild>
+          <TooltipIconButton tooltip="Copy">
+            <MessagePrimitive.If copied>
+              <CheckIcon />
+            </MessagePrimitive.If>
+            <MessagePrimitive.If copied={false}>
+              <CopyIcon />
+            </MessagePrimitive.If>
+          </TooltipIconButton>
+        </ActionBarPrimitive.Copy>
+        <ActionBarPrimitive.Reload asChild>
+          <TooltipIconButton tooltip="Refresh">
+            <RefreshCwIcon />
+          </TooltipIconButton>
+        </ActionBarPrimitive.Reload>
+      </ActionBarPrimitive.Root>
     </MessagePrimitive.Root>
   );
 };
 
-const BranchPicker: FC = () => {
+const BranchPicker: FC<BranchPickerPrimitiveRootProps> = ({
+  className,
+  ...rest
+}) => {
   return (
     <BranchPickerPrimitive.Root
       hideWhenSingleBranch
-      className="text-foreground/60 inline-flex items-center text-xs"
+      className={cn(
+        "text-muted-foreground inline-flex items-center text-xs",
+        className,
+      )}
+      {...rest}
     >
       <BranchPickerPrimitive.Previous asChild>
-        <IconButton tooltip="Previous">
-          <ChevronLeftIcon className="size-4" />
-        </IconButton>
+        <TooltipIconButton tooltip="Previous">
+          <ChevronLeftIcon />
+        </TooltipIconButton>
       </BranchPickerPrimitive.Previous>
-      <BranchPickerPrimitive.Number /> / <BranchPickerPrimitive.Count />
+      <span className="font-medium">
+        <BranchPickerPrimitive.Number /> / <BranchPickerPrimitive.Count />
+      </span>
       <BranchPickerPrimitive.Next asChild>
-        <IconButton tooltip="Next">
-          <ChevronRightIcon className="size-4" />
-        </IconButton>
+        <TooltipIconButton tooltip="Next">
+          <ChevronRightIcon />
+        </TooltipIconButton>
       </BranchPickerPrimitive.Next>
     </BranchPickerPrimitive.Root>
   );
 };
 
-type IconButton = ButtonProps & { tooltip: string };
+type TooltipIconButtonProps = ButtonProps & { tooltip: string };
 
-const IconButton: FC<IconButton> = ({
+const TooltipIconButton: FC<TooltipIconButtonProps> = ({
   children,
   tooltip,
   className,
@@ -229,7 +242,7 @@ const IconButton: FC<IconButton> = ({
         <Button
           variant="ghost"
           size="icon"
-          className={cn("size-auto p-1", className)}
+          className={cn("size-6 p-1", className)}
           {...rest}
         >
           {children}
@@ -240,3 +253,15 @@ const IconButton: FC<IconButton> = ({
     </Tooltip>
   );
 };
+
+const CircleStopIcon = () => {
+  return (
+    <svg
+      xmlns="http://www.w3.org/2000/svg"
+      viewBox="0 0 16 16"
+      fill="currentColor"
+    >
+      <rect width="10" height="10" x="3" y="3" rx="2" />
+    </svg>
+  );
+};
diff --git a/packages/shadcn-registry/registry/assistant-ui/markdown-text.tsx b/packages/shadcn-registry/registry/assistant-ui/markdown-text.tsx
index e09d49ada..18e26387c 100644
--- a/packages/shadcn-registry/registry/assistant-ui/markdown-text.tsx
+++ b/packages/shadcn-registry/registry/assistant-ui/markdown-text.tsx
@@ -75,81 +75,76 @@ const MarkdownTextImpl = () => {
             {...props}
           />
         ),
-        a: ({ node, ...props }) => (
+        a: ({ node, className, ...props }) => (
           <a
             target="_blank"
             className={cn(
               "text-primary font-medium underline underline-offset-4",
-              props.className,
+              className,
             )}
             {...props}
           />
         ),
-        blockquote: ({ node, ...props }) => (
+        blockquote: ({ node, className, ...props }) => (
           <blockquote
-            className={cn("border-l-2 pl-6 italic", props.className)}
+            className={cn("border-l-2 pl-6 italic", className)}
             {...props}
           />
         ),
-        ul: ({ node, ...props }) => (
+        ul: ({ node, className, ...props }) => (
           <ul
-            className={cn("my-5 ml-6 list-disc [&>li]:mt-2", props.className)}
+            className={cn("my-5 ml-6 list-disc [&>li]:mt-2", className)}
             {...props}
           />
         ),
-        ol: ({ node, ...props }) => (
+        ol: ({ node, className, ...props }) => (
           <ol
-            className={cn(
-              "my-5 ml-6 list-decimal [&>li]:mt-2",
-              props.className,
-            )}
+            className={cn("my-5 ml-6 list-decimal [&>li]:mt-2", className)}
             {...props}
           />
         ),
-        hr: ({ node, ...props }) => (
-          <hr className={cn("my-5 border-b", props.className)} {...props} />
+        hr: ({ node, className, ...props }) => (
+          <hr className={cn("my-5 border-b", className)} {...props} />
         ),
-
-        table: ({ node, ...props }) => (
+        table: ({ node, className, ...props }) => (
           <table
             className={cn(
               "my-5 w-full border-separate border-spacing-0 overflow-y-auto",
-              props.className,
+              className,
             )}
             {...props}
           />
         ),
-        th: ({ node, ...props }) => (
+        th: ({ node, className, ...props }) => (
           <th
             className={cn(
               "bg-muted px-4 py-2 text-left font-bold first:rounded-tl-lg last:rounded-tr-lg [&[align=center]]:text-center [&[align=right]]:text-right",
-              props.className,
+              className,
             )}
             {...props}
           />
         ),
-        td: ({ node, ...props }) => (
+        td: ({ node, className, ...props }) => (
           <td
             className={cn(
               "border-b border-l px-4 py-2 text-left last:border-r [&[align=center]]:text-center [&[align=right]]:text-right",
-              props.className,
+              className,
             )}
             {...props}
           />
         ),
-        tr: ({ node, ...props }) => (
+        tr: ({ node, className, ...props }) => (
           <tr
             className={cn(
               "m-0 border-b p-0 first:border-t [&:last-child>td:first-child]:rounded-bl-lg [&:last-child>td:last-child]:rounded-br-lg",
-              props.className,
+              className,
             )}
             {...props}
           />
         ),
-
-        sup: ({ node, ...props }) => (
+        sup: ({ node, className, ...props }) => (
           <sup
-            className={cn("[&>a]:text-xs [&>a]:no-underline", props.className)}
+            className={cn("[&>a]:text-xs [&>a]:no-underline", className)}
             {...props}
           />
         ),
diff --git a/packages/shadcn-registry/registry/assistant-ui/thread.tsx b/packages/shadcn-registry/registry/assistant-ui/thread.tsx
index 741d40ab4..1b6d99978 100644
--- a/packages/shadcn-registry/registry/assistant-ui/thread.tsx
+++ b/packages/shadcn-registry/registry/assistant-ui/thread.tsx
@@ -8,68 +8,87 @@ import {
 import type { FC } from "react";
 
 import { Avatar, AvatarFallback } from "@/components/ui/avatar";
+import { Button } from "@/components/ui/button";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { cn } from "@/lib/utils";
+import { TooltipProvider } from "@radix-ui/react-tooltip";
 import { SendHorizonalIcon } from "lucide-react";
 
 export const Thread: FC = () => {
   return (
-    <ThreadPrimitive.Root className="flex h-full flex-col items-center pb-3">
-      <ThreadPrimitive.Viewport className="flex w-full flex-grow flex-col items-center overflow-y-scroll scroll-smooth px-4 pt-12">
-        <ThreadPrimitive.Empty>
-          <ThreadEmpty />
-        </ThreadPrimitive.Empty>
+    <TooltipProvider>
+      <ThreadPrimitive.Root className="bg-background h-full">
+        <ThreadPrimitive.Viewport className="flex h-full flex-col items-center overflow-y-scroll scroll-smooth px-4 pt-8">
+          <ThreadWelcome />
 
-        <ThreadPrimitive.Messages
-          components={{
-            UserMessage,
-            AssistantMessage,
-          }}
-        />
-      </ThreadPrimitive.Viewport>
+          <ThreadPrimitive.Messages
+            components={{
+              UserMessage,
+              AssistantMessage,
+            }}
+          />
 
-      <Composer />
-    </ThreadPrimitive.Root>
+          <div className="sticky bottom-0 mt-4 flex w-full max-w-2xl flex-grow flex-col items-center justify-end rounded-t-lg bg-inherit pb-4">
+            <Composer />
+          </div>
+        </ThreadPrimitive.Viewport>
+      </ThreadPrimitive.Root>
+    </TooltipProvider>
   );
 };
 
-const ThreadEmpty: FC = () => {
+const ThreadWelcome: FC = () => {
   return (
-    <div className="flex flex-grow flex-col items-center justify-center">
-      <Avatar>
-        <AvatarFallback>C</AvatarFallback>
-      </Avatar>
-      <p className="mt-4 text-xl">How can I help you today?</p>
-    </div>
+    <ThreadPrimitive.Empty>
+      <div className="flex flex-grow basis-full flex-col items-center justify-center">
+        <Avatar>
+          <AvatarFallback>C</AvatarFallback>
+        </Avatar>
+        <p className="mt-4 text-lg font-medium">How can I help you today?</p>
+      </div>
+    </ThreadPrimitive.Empty>
   );
 };
 
 const Composer: FC = () => {
   return (
-    <ComposerPrimitive.Root className="flex w-[calc(100%-32px)] max-w-[42rem] items-end rounded-lg border p-0.5 transition-shadow focus-within:shadow-sm">
+    <ComposerPrimitive.Root className="relative flex w-full items-end rounded-lg border transition-shadow focus-within:shadow-sm">
       <ComposerPrimitive.Input
+        autoFocus
         placeholder="Write a message..."
-        className="placeholder:text-foreground/50 h-12 max-h-40 flex-grow resize-none bg-transparent p-3.5 text-sm outline-none"
+        rows={1}
+        className="placeholder:text-muted-foreground size-full max-h-40 resize-none bg-transparent p-4 pr-12 text-sm outline-none"
       />
-      <ThreadPrimitive.If running={false}>
-        <ComposerPrimitive.Send className="bg-foreground m-2 flex h-8 w-8 items-center justify-center rounded-md text-2xl font-bold shadow transition-opacity disabled:opacity-10">
-          <SendHorizonalIcon className="text-background size-4" />
+      <Tooltip>
+        <ComposerPrimitive.Send asChild>
+          <TooltipTrigger asChild>
+            <Button
+              variant="ghost"
+              size="icon"
+              className={cn(
+                "absolute bottom-0 right-0 m-2.5 size-8 p-2 transition-opacity",
+              )}
+            >
+              <SendHorizonalIcon />
+              <span className="sr-only">Send</span>
+            </Button>
+          </TooltipTrigger>
         </ComposerPrimitive.Send>
-      </ThreadPrimitive.If>
-      <ThreadPrimitive.If running>
-        <ComposerPrimitive.Cancel className="border-foreground m-3.5 flex size-5 items-center justify-center rounded-full border-2">
-          <div className="bg-foreground size-2 rounded-[1px]" />
-        </ComposerPrimitive.Cancel>
-      </ThreadPrimitive.If>
+        <TooltipContent side="bottom">Send</TooltipContent>
+      </Tooltip>
     </ComposerPrimitive.Root>
   );
 };
 
 const UserMessage: FC = () => {
   return (
-    <MessagePrimitive.Root className="relative mb-6 flex w-full max-w-2xl flex-col items-end gap-2 pl-24">
-      <div className="relative mr-1 flex items-start gap-3">
-        <div className="bg-foreground/5 text-foreground max-w-xl break-words rounded-3xl px-5 py-2.5">
-          <MessagePrimitive.Content />
-        </div>
+    <MessagePrimitive.Root className="my-4 grid w-full max-w-2xl auto-rows-auto grid-cols-[minmax(72px,1fr)_auto] gap-y-2">
+      <div className="bg-muted text-foreground col-start-2 row-start-1 max-w-xl break-words rounded-3xl px-5 py-2.5">
+        <MessagePrimitive.Content />
       </div>
     </MessagePrimitive.Root>
   );
@@ -77,15 +96,13 @@ const UserMessage: FC = () => {
 
 const AssistantMessage: FC = () => {
   return (
-    <MessagePrimitive.Root className="relative mb-6 flex w-full max-w-2xl gap-3">
-      <Avatar>
+    <MessagePrimitive.Root className="relative my-4 grid w-full max-w-2xl grid-cols-[auto_1fr] grid-rows-[auto_1fr]">
+      <Avatar className="col-start-1 row-span-full row-start-1 mr-4">
         <AvatarFallback>A</AvatarFallback>
       </Avatar>
 
-      <div className="mt-2 flex-grow">
-        <div className="text-foreground max-w-x break-words">
-          <MessagePrimitive.Content />
-        </div>
+      <div className="text-foreground col-start-2 row-start-1 my-1.5 max-w-xl break-words leading-7">
+        <MessagePrimitive.Content />
       </div>
     </MessagePrimitive.Root>
   );