From dff219161182f8fca99e8df6dbd7c27ee9dd6118 Mon Sep 17 00:00:00 2001
From: josh <144584931+dancer@users.noreply.github.com>
Date: Mon, 1 Sep 2025 11:07:07 +0100
Subject: [PATCH] integrate ai gateway and improve ui components (#1145)

---
 .env.example                       |   6 +-
 .github/workflows/lint.yml         |  29 +++---
 README.md                          |  10 ++-
 app/(chat)/api/chat/route.ts       |   3 +
 artifacts/image/server.ts          |  45 ----------
 components/artifact-actions.tsx    |   6 +-
 components/artifact-messages.tsx   |   3 +-
 components/artifact.tsx            |  15 ++--
 components/chat-header.tsx         |  25 ++----
 components/chat.tsx                |   5 +-
 components/code-editor.tsx         |   2 +-
 components/document-skeleton.tsx   |   2 +-
 components/elements/code-block.tsx |   6 ++
 components/elements/tool.tsx       |  10 +--
 components/icons.tsx               |  34 +++----
 components/message.tsx             | 139 +++++++++++++----------------
 components/messages.tsx            |   5 +-
 components/multimodal-input.tsx    |  58 ++++++++++++
 components/preview-attachment.tsx  |   2 +-
 components/sidebar-history.tsx     |  10 +--
 components/ui/sidebar.tsx          |   2 +-
 components/visibility-selector.tsx |   2 +-
 hooks/use-artifact.ts              |   2 +-
 lib/ai/models.ts                   |   8 +-
 lib/ai/providers.ts                |  13 ++-
 lib/artifacts/server.ts            |   4 +-
 package.json                       |   1 +
 pnpm-lock.yaml                     |   3 +
 tests/pages/artifact.ts            |   2 +-
 29 files changed, 235 insertions(+), 217 deletions(-)
 delete mode 100644 artifacts/image/server.ts

diff --git a/.env.example b/.env.example
index 6a40fdd..b25ecc4 100644
--- a/.env.example
+++ b/.env.example
@@ -4,8 +4,10 @@ AUTH_SECRET=****
 # The following keys below are automatically created and
 # added to your environment when you deploy on vercel
 
-# Get your xAI API Key here for chat and image models: https://console.x.ai/
-XAI_API_KEY=****
+# AI Gateway API Key (required for non-Vercel deployments)
+# For Vercel deployments, OIDC tokens are used automatically
+AI_GATEWAY_API_KEY=****
+
 
 # Instructions to create a Vercel Blob Store here: https://vercel.com/docs/storage/vercel-blob
 BLOB_READ_WRITE_TOKEN=****
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index deda7f5..e24fc77 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -9,17 +9,18 @@ jobs:
       matrix:
         node-version: [20]
     steps:
-    - uses: actions/checkout@v4
-    - name: Install pnpm
-      uses: pnpm/action-setup@v4
-      with:
-        version: 9.12.3
-    - name: Use Node.js ${{ matrix.node-version }}
-      uses: actions/setup-node@v4
-      with:
-        node-version: ${{ matrix.node-version }}
-        cache: 'pnpm'
-    - name: Install dependencies
-      run: pnpm install
-    - name: Run lint
-      run: pnpm lint
\ No newline at end of file
+      - uses: actions/checkout@v4
+      - name: Install pnpm
+        uses: pnpm/action-setup@v4
+        with:
+          version: 9.12.3
+      - name: Use Node.js ${{ matrix.node-version }}
+        uses: actions/setup-node@v4
+        with:
+          node-version: ${{ matrix.node-version }}
+          cache: "pnpm"
+      - name: Install dependencies
+        run: pnpm install
+      - name: Run lint
+        run: pnpm lint
+1
\ No newline at end of file
diff --git a/README.md b/README.md
index 44a8cb2..8202559 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,15 @@
 
 ## Model Providers
 
-This template ships with [xAI](https://x.ai) `grok-2-1212` as the default chat model. However, with the [AI SDK](https://sdk.vercel.ai/docs), you can switch LLM providers to [OpenAI](https://openai.com), [Anthropic](https://anthropic.com), [Cohere](https://cohere.com/), and [many more](https://sdk.vercel.ai/providers/ai-sdk-providers) with just a few lines of code.
+This template uses the [Vercel AI Gateway](https://vercel.com/docs/ai-gateway) to access multiple AI models through a unified interface. The default configuration includes [xAI](https://x.ai) models (`grok-2-vision-1212`, `grok-3-mini-beta`) routed through the gateway.
+
+### AI Gateway Authentication
+
+**For Vercel deployments**: Authentication is handled automatically via OIDC tokens.
+
+**For non-Vercel deployments**: You need to provide an AI Gateway API key by setting the `AI_GATEWAY_API_KEY` environment variable in your `.env.local` file.
+
+With the [AI SDK](https://ai-sdk.dev/docs/introduction), you can also switch to direct LLM providers like [OpenAI](https://openai.com), [Anthropic](https://anthropic.com), [Cohere](https://cohere.com/), and [many more](https://ai-sdk.dev/providers/ai-sdk-providers) with just a few lines of code.
 
 ## Deploy Your Own
 
diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts
index b7ea099..8e16929 100644
--- a/app/(chat)/api/chat/route.ts
+++ b/app/(chat)/api/chat/route.ts
@@ -222,6 +222,9 @@ export async function POST(request: Request) {
     if (error instanceof ChatSDKError) {
       return error.toResponse();
     }
+
+    console.error('Unhandled error in chat API:', error);
+    return new ChatSDKError('offline:chat').toResponse();
   }
 }
 
diff --git a/artifacts/image/server.ts b/artifacts/image/server.ts
deleted file mode 100644
index 14a994c..0000000
--- a/artifacts/image/server.ts
+++ /dev/null
@@ -1,45 +0,0 @@
-import { myProvider } from '@/lib/ai/providers';
-import { createDocumentHandler } from '@/lib/artifacts/server';
-import { experimental_generateImage } from 'ai';
-
-export const imageDocumentHandler = createDocumentHandler<'image'>({
-  kind: 'image',
-  onCreateDocument: async ({ title, dataStream }) => {
-    let draftContent = '';
-
-    const { image } = await experimental_generateImage({
-      model: myProvider.imageModel('small-model'),
-      prompt: title,
-      n: 1,
-    });
-
-    draftContent = image.base64;
-
-    dataStream.write({
-      type: 'data-imageDelta',
-      data: image.base64,
-      transient: true,
-    });
-
-    return draftContent;
-  },
-  onUpdateDocument: async ({ description, dataStream }) => {
-    let draftContent = '';
-
-    const { image } = await experimental_generateImage({
-      model: myProvider.imageModel('small-model'),
-      prompt: description,
-      n: 1,
-    });
-
-    draftContent = image.base64;
-
-    dataStream.write({
-      type: 'data-imageDelta',
-      data: image.base64,
-      transient: true,
-    });
-
-    return draftContent;
-  },
-});
diff --git a/components/artifact-actions.tsx b/components/artifact-actions.tsx
index 009cbed..a5689d3 100644
--- a/components/artifact-actions.tsx
+++ b/components/artifact-actions.tsx
@@ -1,8 +1,8 @@
 import { Button } from './ui/button';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
-import { artifactDefinitions, UIArtifact } from './artifact';
-import { Dispatch, memo, SetStateAction, useState } from 'react';
-import { ArtifactActionContext } from './create-artifact';
+import { artifactDefinitions, type UIArtifact } from './artifact';
+import { type Dispatch, memo, type SetStateAction, useState } from 'react';
+import type { ArtifactActionContext } from './create-artifact';
 import { cn } from '@/lib/utils';
 import { toast } from 'sonner';
 
diff --git a/components/artifact-messages.tsx b/components/artifact-messages.tsx
index 4d71435..5cb6f93 100644
--- a/components/artifact-messages.tsx
+++ b/components/artifact-messages.tsx
@@ -42,7 +42,7 @@ function PureArtifactMessages({
   return (
     <div
       ref={messagesContainerRef}
-      className="flex flex-col gap-4 h-full items-center overflow-y-scroll px-4 pt-20"
+      className="flex overflow-y-scroll flex-col gap-4 items-center px-4 pt-20 h-full"
     >
       {messages.map((message, index) => (
         <PreviewMessage
@@ -61,6 +61,7 @@ function PureArtifactMessages({
           requiresScrollPadding={
             hasSentMessage && index === messages.length - 1
           }
+          isArtifactVisible={true}
         />
       ))}
 
diff --git a/components/artifact.tsx b/components/artifact.tsx
index 6b266d9..9f5d6d7 100644
--- a/components/artifact.tsx
+++ b/components/artifact.tsx
@@ -67,6 +67,7 @@ function PureArtifact({
   votes,
   isReadonly,
   selectedVisibilityType,
+  selectedModelId,
 }: {
   chatId: string;
   input: string;
@@ -82,6 +83,7 @@ function PureArtifact({
   regenerate: UseChatHelpers<ChatMessage>['regenerate'];
   isReadonly: boolean;
   selectedVisibilityType: VisibilityType;
+  selectedModelId: string;
 }) {
   const { artifact, setArtifact, metadata, setMetadata } = useArtifact();
 
@@ -286,9 +288,9 @@ function PureArtifact({
                 x: 0,
                 scale: 1,
                 transition: {
-                  delay: 0.2,
+                  delay: 0.1,
                   type: 'spring',
-                  stiffness: 200,
+                  stiffness: 300,
                   damping: 30,
                 },
               }}
@@ -336,6 +338,7 @@ function PureArtifact({
                     className="bg-background dark:bg-muted"
                     setMessages={setMessages}
                     selectedVisibilityType={selectedVisibilityType}
+                    selectedModelId={selectedModelId}
                   />
                 </div>
               </div>
@@ -375,9 +378,9 @@ function PureArtifact({
                     transition: {
                       delay: 0,
                       type: 'spring',
-                      stiffness: 200,
+                      stiffness: 300,
                       damping: 30,
-                      duration: 5000,
+                      duration: 0.8,
                     },
                   }
                 : {
@@ -392,9 +395,9 @@ function PureArtifact({
                     transition: {
                       delay: 0,
                       type: 'spring',
-                      stiffness: 200,
+                      stiffness: 300,
                       damping: 30,
-                      duration: 5000,
+                      duration: 0.8,
                     },
                   }
             }
diff --git a/components/chat-header.tsx b/components/chat-header.tsx
index f21f151..8370906 100644
--- a/components/chat-header.tsx
+++ b/components/chat-header.tsx
@@ -4,7 +4,6 @@ import Link from 'next/link';
 import { useRouter } from 'next/navigation';
 import { useWindowSize } from 'usehooks-ts';
 
-import { ModelSelector } from '@/components/model-selector';
 import { SidebarToggle } from '@/components/sidebar-toggle';
 import { Button } from '@/components/ui/button';
 import { PlusIcon, VercelIcon } from './icons';
@@ -16,13 +15,11 @@ import type { Session } from 'next-auth';
 
 function PureChatHeader({
   chatId,
-  selectedModelId,
   selectedVisibilityType,
   isReadonly,
   session,
 }: {
   chatId: string;
-  selectedModelId: string;
   selectedVisibilityType: VisibilityType;
   isReadonly: boolean;
   session: Session;
@@ -56,23 +53,15 @@ function PureChatHeader({
       )}
 
       {!isReadonly && (
-        <ModelSelector
-          session={session}
-          selectedModelId={selectedModelId}
+        <VisibilitySelector
+          chatId={chatId}
+          selectedVisibilityType={selectedVisibilityType}
           className="order-1 md:order-2"
         />
       )}
 
-      {!isReadonly && (
-        <VisibilitySelector
-          chatId={chatId}
-          selectedVisibilityType={selectedVisibilityType}
-          className="order-1 md:order-3"
-        />
-      )}
-
       <Button
-        className="bg-zinc-900 dark:bg-zinc-100 hover:bg-zinc-800 dark:hover:bg-zinc-200 text-zinc-50 dark:text-zinc-900 hidden md:flex py-1.5 px-2 h-fit md:h-[34px] order-4 md:ml-auto"
+        className="bg-zinc-900 dark:bg-zinc-100 hover:bg-zinc-800 dark:hover:bg-zinc-200 text-zinc-50 dark:text-zinc-900 hidden md:flex py-1.5 px-2 h-fit md:h-[34px] order-3 md:ml-auto"
         asChild
       >
         <Link
@@ -88,5 +77,9 @@ function PureChatHeader({
 }
 
 export const ChatHeader = memo(PureChatHeader, (prevProps, nextProps) => {
-  return prevProps.selectedModelId === nextProps.selectedModelId;
+  return (
+    prevProps.chatId === nextProps.chatId &&
+    prevProps.selectedVisibilityType === nextProps.selectedVisibilityType &&
+    prevProps.isReadonly === nextProps.isReadonly
+  );
 });
diff --git a/components/chat.tsx b/components/chat.tsx
index f187ad3..959373b 100644
--- a/components/chat.tsx
+++ b/components/chat.tsx
@@ -6,7 +6,7 @@ import { useEffect, useState } from 'react';
 import useSWR, { useSWRConfig } from 'swr';
 import { ChatHeader } from '@/components/chat-header';
 import type { Vote } from '@/lib/db/schema';
-import { fetcher, fetchWithErrorHandlers, generateUUID } from '@/lib/utils';
+import { fetcher, fetchWithErrorHandlers, generateUUID, cn } from '@/lib/utils';
 import { Artifact } from './artifact';
 import { MultimodalInput } from './multimodal-input';
 import { Messages } from './messages';
@@ -131,7 +131,6 @@ export function Chat({
       <div className="flex flex-col min-w-0 h-dvh bg-background">
         <ChatHeader
           chatId={id}
-          selectedModelId={initialChatModel}
           selectedVisibilityType={initialVisibilityType}
           isReadonly={isReadonly}
           session={session}
@@ -162,6 +161,7 @@ export function Chat({
               setMessages={setMessages}
               sendMessage={sendMessage}
               selectedVisibilityType={visibilityType}
+              selectedModelId={initialChatModel}
             />
           )}
         </div>
@@ -182,6 +182,7 @@ export function Chat({
         votes={votes}
         isReadonly={isReadonly}
         selectedVisibilityType={visibilityType}
+        selectedModelId={initialChatModel}
       />
     </>
   );
diff --git a/components/code-editor.tsx b/components/code-editor.tsx
index ab81e13..325e5db 100644
--- a/components/code-editor.tsx
+++ b/components/code-editor.tsx
@@ -6,7 +6,7 @@ import { python } from '@codemirror/lang-python';
 import { oneDark } from '@codemirror/theme-one-dark';
 import { basicSetup } from 'codemirror';
 import React, { memo, useEffect, useRef } from 'react';
-import { Suggestion } from '@/lib/db/schema';
+import type { Suggestion } from '@/lib/db/schema';
 
 type EditorProps = {
   content: string;
diff --git a/components/document-skeleton.tsx b/components/document-skeleton.tsx
index 5a9a9f7..900c3c1 100644
--- a/components/document-skeleton.tsx
+++ b/components/document-skeleton.tsx
@@ -1,6 +1,6 @@
 'use client';
 
-import { ArtifactKind } from './artifact';
+import type { ArtifactKind } from './artifact';
 
 export const DocumentSkeleton = ({
   artifactKind,
diff --git a/components/elements/code-block.tsx b/components/elements/code-block.tsx
index c6df4b2..86652e7 100644
--- a/components/elements/code-block.tsx
+++ b/components/elements/code-block.tsx
@@ -54,6 +54,9 @@ export const CodeBlock = ({
             fontSize: '0.875rem',
             background: 'hsl(var(--background))',
             color: 'hsl(var(--foreground))',
+            overflowX: 'auto',
+            overflowWrap: 'break-word',
+            wordBreak: 'break-all',
           }}
           language={language}
           lineNumberStyle={{
@@ -77,6 +80,9 @@ export const CodeBlock = ({
             fontSize: '0.875rem',
             background: 'hsl(var(--background))',
             color: 'hsl(var(--foreground))',
+            overflowX: 'auto',
+            overflowWrap: 'break-word',
+            wordBreak: 'break-all',
           }}
           language={language}
           lineNumberStyle={{
diff --git a/components/elements/tool.tsx b/components/elements/tool.tsx
index 12bd962..4570911 100644
--- a/components/elements/tool.tsx
+++ b/components/elements/tool.tsx
@@ -65,16 +65,16 @@ export const ToolHeader = ({
 }: ToolHeaderProps) => (
   <CollapsibleTrigger
     className={cn(
-      'flex w-full items-center justify-between gap-4 p-3',
+      'flex w-full items-center justify-between gap-2 p-3 min-w-0',
       className,
     )}
     {...props}
   >
-    <div className="flex items-center gap-2">
-      <WrenchIcon className="size-4 text-muted-foreground" />
-      <span className="font-medium text-sm">{type}</span>
+    <div className="flex items-center gap-2 min-w-0 flex-1">
+      <WrenchIcon className="size-4 text-muted-foreground shrink-0" />
+      <span className="font-medium text-sm truncate">{type}</span>
     </div>
-    <div className="flex items-center gap-2">
+    <div className="flex items-center gap-2 shrink-0">
       {getStatusBadge(state)}
       <ChevronDownIcon className="size-4 text-muted-foreground transition-transform group-data-[state=open]:rotate-180" />
     </div>
diff --git a/components/icons.tsx b/components/icons.tsx
index a6edb5a..da76998 100644
--- a/components/icons.tsx
+++ b/components/icons.tsx
@@ -640,7 +640,7 @@ export const CrossSmallIcon = ({ size = 16 }: { size?: number }) => (
       clipRule="evenodd"
       d="M9.96966 11.0303L10.5 11.5607L11.5607 10.5L11.0303 9.96966L9.06065 7.99999L11.0303 6.03032L11.5607 5.49999L10.5 4.43933L9.96966 4.96966L7.99999 6.93933L6.03032 4.96966L5.49999 4.43933L4.43933 5.49999L4.96966 6.03032L6.93933 7.99999L4.96966 9.96966L4.43933 10.5L5.49999 11.5607L6.03032 11.0303L7.99999 9.06065L9.96966 11.0303Z"
       fill="currentColor"
-    ></path>
+    />
   </svg>
 );
 
@@ -887,7 +887,7 @@ export const GlobeIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M10.268 14.0934C11.9051 13.4838 13.2303 12.2333 13.9384 10.6469C13.1192 10.7941 12.2138 10.9111 11.2469 10.9925C11.0336 12.2005 10.695 13.2621 10.268 14.0934ZM8 16C12.4183 16 16 12.4183 16 8C16 3.58172 12.4183 0 8 0C3.58172 0 0 3.58172 0 8C0 12.4183 3.58172 16 8 16ZM8.48347 14.4823C8.32384 14.494 8.16262 14.5 8 14.5C7.83738 14.5 7.67616 14.494 7.51654 14.4823C7.5132 14.4791 7.50984 14.4759 7.50647 14.4726C7.2415 14.2165 6.94578 13.7854 6.67032 13.1558C6.41594 12.5744 6.19979 11.8714 6.04101 11.0778C6.67605 11.1088 7.33104 11.125 8 11.125C8.66896 11.125 9.32395 11.1088 9.95899 11.0778C9.80021 11.8714 9.58406 12.5744 9.32968 13.1558C9.05422 13.7854 8.7585 14.2165 8.49353 14.4726C8.49016 14.4759 8.4868 14.4791 8.48347 14.4823ZM11.4187 9.72246C12.5137 9.62096 13.5116 9.47245 14.3724 9.28806C14.4561 8.87172 14.5 8.44099 14.5 8C14.5 7.55901 14.4561 7.12828 14.3724 6.71194C13.5116 6.52755 12.5137 6.37904 11.4187 6.27753C11.4719 6.83232 11.5 7.40867 11.5 8C11.5 8.59133 11.4719 9.16768 11.4187 9.72246ZM10.1525 6.18401C10.2157 6.75982 10.25 7.36805 10.25 8C10.25 8.63195 10.2157 9.24018 10.1525 9.81598C9.46123 9.85455 8.7409 9.875 8 9.875C7.25909 9.875 6.53877 9.85455 5.84749 9.81598C5.7843 9.24018 5.75 8.63195 5.75 8C5.75 7.36805 5.7843 6.75982 5.84749 6.18401C6.53877 6.14545 7.25909 6.125 8 6.125C8.74091 6.125 9.46123 6.14545 10.1525 6.18401ZM11.2469 5.00748C12.2138 5.08891 13.1191 5.20593 13.9384 5.35306C13.2303 3.7667 11.9051 2.51622 10.268 1.90662C10.695 2.73788 11.0336 3.79953 11.2469 5.00748ZM8.48347 1.51771C8.4868 1.52089 8.49016 1.52411 8.49353 1.52737C8.7585 1.78353 9.05422 2.21456 9.32968 2.84417C9.58406 3.42562 9.80021 4.12856 9.95899 4.92219C9.32395 4.89118 8.66896 4.875 8 4.875C7.33104 4.875 6.67605 4.89118 6.04101 4.92219C6.19978 4.12856 6.41594 3.42562 6.67032 2.84417C6.94578 2.21456 7.2415 1.78353 7.50647 1.52737C7.50984 1.52411 7.51319 1.52089 7.51653 1.51771C7.67615 1.50597 7.83738 1.5 8 1.5C8.16262 1.5 8.32384 1.50597 8.48347 1.51771ZM5.73202 1.90663C4.0949 2.51622 2.76975 3.7667 2.06159 5.35306C2.88085 5.20593 3.78617 5.08891 4.75309 5.00748C4.96639 3.79953 5.30497 2.73788 5.73202 1.90663ZM4.58133 6.27753C3.48633 6.37904 2.48837 6.52755 1.62761 6.71194C1.54392 7.12828 1.5 7.55901 1.5 8C1.5 8.44099 1.54392 8.87172 1.62761 9.28806C2.48837 9.47245 3.48633 9.62096 4.58133 9.72246C4.52807 9.16768 4.5 8.59133 4.5 8C4.5 7.40867 4.52807 6.83232 4.58133 6.27753ZM4.75309 10.9925C3.78617 10.9111 2.88085 10.7941 2.06159 10.6469C2.76975 12.2333 4.0949 13.4838 5.73202 14.0934C5.30497 13.2621 4.96639 12.2005 4.75309 10.9925Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -906,7 +906,7 @@ export const LockIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M10 4.5V6H6V4.5C6 3.39543 6.89543 2.5 8 2.5C9.10457 2.5 10 3.39543 10 4.5ZM4.5 6V4.5C4.5 2.567 6.067 1 8 1C9.933 1 11.5 2.567 11.5 4.5V6H12.5H14V7.5V12.5C14 13.8807 12.8807 15 11.5 15H4.5C3.11929 15 2 13.8807 2 12.5V7.5V6H3.5H4.5ZM11.5 7.5H10H6H4.5H3.5V12.5C3.5 13.0523 3.94772 13.5 4.5 13.5H11.5C12.0523 13.5 12.5 13.0523 12.5 12.5V7.5H11.5Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -925,7 +925,7 @@ export const EyeIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M4.02168 4.76932C6.11619 2.33698 9.88374 2.33698 11.9783 4.76932L14.7602 7.99999L11.9783 11.2307C9.88374 13.663 6.1162 13.663 4.02168 11.2307L1.23971 7.99999L4.02168 4.76932ZM13.1149 3.79054C10.422 0.663244 5.57797 0.663247 2.88503 3.79054L-0.318359 7.5106V8.48938L2.88503 12.2094C5.57797 15.3367 10.422 15.3367 13.1149 12.2094L16.3183 8.48938V7.5106L13.1149 3.79054ZM6.49997 7.99999C6.49997 7.17157 7.17154 6.49999 7.99997 6.49999C8.82839 6.49999 9.49997 7.17157 9.49997 7.99999C9.49997 8.82842 8.82839 9.49999 7.99997 9.49999C7.17154 9.49999 6.49997 8.82842 6.49997 7.99999ZM7.99997 4.99999C6.34311 4.99999 4.99997 6.34314 4.99997 7.99999C4.99997 9.65685 6.34311 11 7.99997 11C9.65682 11 11 9.65685 11 7.99999C11 6.34314 9.65682 4.99999 7.99997 4.99999Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -944,7 +944,7 @@ export const ShareIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M15 11.25V10.5H13.5V11.25V12.75C13.5 13.1642 13.1642 13.5 12.75 13.5H3.25C2.83579 13.5 2.5 13.1642 2.5 12.75L2.5 3.25C2.5 2.83579 2.83579 2.5 3.25 2.5H5.75H6.5V1H5.75H3.25C2.00736 1 1 2.00736 1 3.25V12.75C1 13.9926 2.00736 15 3.25 15H12.75C13.9926 15 15 13.9926 15 12.75V11.25ZM15 5.5L10.5 1V4C7.46243 4 5 6.46243 5 9.5V10L5.05855 9.91218C6.27146 8.09281 8.31339 7 10.5 7V10L15 5.5Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -963,7 +963,7 @@ export const CodeIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M4.21969 12.5303L4.75002 13.0607L5.81068 12L5.28035 11.4697L1.81068 7.99999L5.28035 4.53032L5.81068 3.99999L4.75002 2.93933L4.21969 3.46966L0.39647 7.29289C0.00594562 7.68341 0.00594562 8.31658 0.39647 8.7071L4.21969 12.5303ZM11.7804 12.5303L11.25 13.0607L10.1894 12L10.7197 11.4697L14.1894 7.99999L10.7197 4.53032L10.1894 3.99999L11.25 2.93933L11.7804 3.46966L15.6036 7.29289C15.9941 7.68341 15.9941 8.31658 15.6036 8.7071L11.7804 12.5303Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -982,7 +982,7 @@ export const PlayIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M13.4549 7.22745L13.3229 7.16146L2.5 1.74999L2.4583 1.72914L1.80902 1.4045L1.3618 1.18089C1.19558 1.09778 1 1.21865 1 1.4045L1 1.9045L1 2.63041L1 2.67704L1 13.3229L1 13.3696L1 14.0955L1 14.5955C1 14.7813 1.19558 14.9022 1.3618 14.8191L1.80902 14.5955L2.4583 14.2708L2.5 14.25L13.3229 8.83852L13.4549 8.77253L14.2546 8.37267L14.5528 8.2236C14.737 8.13147 14.737 7.86851 14.5528 7.77638L14.2546 7.62731L13.4549 7.22745ZM11.6459 7.99999L2.5 3.42704L2.5 12.5729L11.6459 7.99999Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -999,11 +999,11 @@ export const PythonIcon = ({ size = 16 }: { size?: number }) => {
       <path
         d="M7.90474 0.00013087C7.24499 0.00316291 6.61494 0.0588153 6.06057 0.15584C4.42745 0.441207 4.13094 1.0385 4.13094 2.14002V3.59479H7.9902V4.07971H4.13094H2.68259C1.56099 4.07971 0.578874 4.7465 0.271682 6.01496C-0.0826597 7.4689 -0.0983767 8.37619 0.271682 9.89434C0.546012 11.0244 1.20115 11.8296 2.32276 11.8296H3.64966V10.0856C3.64966 8.82574 4.75179 7.71441 6.06057 7.71441H9.91533C10.9884 7.71441 11.845 6.84056 11.845 5.77472V2.14002C11.845 1.10556 10.9626 0.328487 9.91533 0.15584C9.25237 0.046687 8.56448 -0.00290121 7.90474 0.00013087ZM5.81768 1.17017C6.21631 1.17017 6.54185 1.49742 6.54185 1.89978C6.54185 2.30072 6.21631 2.62494 5.81768 2.62494C5.41761 2.62494 5.09351 2.30072 5.09351 1.89978C5.09351 1.49742 5.41761 1.17017 5.81768 1.17017Z"
         fill="currentColor"
-      ></path>
+      />
       <path
         d="M12.3262 4.07971V5.77472C12.3262 7.08883 11.1997 8.19488 9.91525 8.19488H6.06049C5.0046 8.19488 4.13086 9.0887 4.13086 10.1346V13.7693C4.13086 14.8037 5.04033 15.4122 6.06049 15.709C7.28211 16.0642 8.45359 16.1285 9.91525 15.709C10.8868 15.4307 11.8449 14.8708 11.8449 13.7693V12.3145H7.99012V11.8296H11.8449H13.7745C14.8961 11.8296 15.3141 11.0558 15.7041 9.89434C16.1071 8.69865 16.0899 7.5488 15.7041 6.01495C15.4269 4.91058 14.8975 4.07971 13.7745 4.07971H12.3262ZM10.1581 13.2843C10.5582 13.2843 10.8823 13.6086 10.8823 14.0095C10.8823 14.4119 10.5582 14.7391 10.1581 14.7391C9.7595 14.7391 9.43397 14.4119 9.43397 14.0095C9.43397 13.6086 9.7595 13.2843 10.1581 13.2843Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -1022,7 +1022,7 @@ export const TerminalWindowIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M1.5 2.5H14.5V12.5C14.5 13.0523 14.0523 13.5 13.5 13.5H2.5C1.94772 13.5 1.5 13.0523 1.5 12.5V2.5ZM0 1H1.5H14.5H16V2.5V12.5C16 13.8807 14.8807 15 13.5 15H2.5C1.11929 15 0 13.8807 0 12.5V2.5V1ZM4 11.1339L4.44194 10.6919L6.51516 8.61872C6.85687 8.27701 6.85687 7.72299 6.51517 7.38128L4.44194 5.30806L4 4.86612L3.11612 5.75L3.55806 6.19194L5.36612 8L3.55806 9.80806L3.11612 10.25L4 11.1339ZM8 9.75494H8.6225H11.75H12.3725V10.9999H11.75H8.6225H8V9.75494Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -1041,7 +1041,7 @@ export const TerminalIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M1.53035 12.7804L1.00002 13.3108L-0.0606384 12.2501L0.469692 11.7198L4.18936 8.00011L0.469692 4.28044L-0.0606384 3.75011L1.00002 2.68945L1.53035 3.21978L5.60358 7.29301C5.9941 7.68353 5.9941 8.3167 5.60357 8.70722L1.53035 12.7804ZM8.75002 12.5001H8.00002V14.0001H8.75002H15.25H16V12.5001H15.25H8.75002Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -1060,7 +1060,7 @@ export const ClockRewind = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M7.96452 2.5C11.0257 2.5 13.5 4.96643 13.5 8C13.5 11.0336 11.0257 13.5 7.96452 13.5C6.12055 13.5 4.48831 12.6051 3.48161 11.2273L3.03915 10.6217L1.828 11.5066L2.27046 12.1122C3.54872 13.8617 5.62368 15 7.96452 15C11.8461 15 15 11.87 15 8C15 4.13001 11.8461 1 7.96452 1C5.06835 1 2.57851 2.74164 1.5 5.23347V3.75V3H0V3.75V7.25C0 7.66421 0.335786 8 0.75 8H3.75H4.5V6.5H3.75H2.63724C3.29365 4.19393 5.42843 2.5 7.96452 2.5ZM8.75 5.25V4.5H7.25V5.25V7.8662C7.25 8.20056 7.4171 8.51279 7.6953 8.69825L9.08397 9.62404L9.70801 10.0401L10.5401 8.79199L9.91603 8.37596L8.75 7.59861V5.25Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -1079,7 +1079,7 @@ export const LogsIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M9 2H9.75H14.25H15V3.5H14.25H9.75H9V2ZM9 12.5H9.75H14.25H15V14H14.25H9.75H9V12.5ZM9.75 7.25H9V8.75H9.75H14.25H15V7.25H14.25H9.75ZM1 12.5H1.75H2.25H3V14H2.25H1.75H1V12.5ZM1.75 2H1V3.5H1.75H2.25H3V2H2.25H1.75ZM1 7.25H1.75H2.25H3V8.75H2.25H1.75H1V7.25ZM5.75 12.5H5V14H5.75H6.25H7V12.5H6.25H5.75ZM5 2H5.75H6.25H7V3.5H6.25H5.75H5V2ZM5.75 7.25H5V8.75H5.75H6.25H7V7.25H6.25H5.75Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -1098,7 +1098,7 @@ export const ImageIcon = ({ size = 16 }: { size?: number }) => {
         clipRule="evenodd"
         d="M14.5 2.5H1.5V9.18933L2.96966 7.71967L3.18933 7.5H3.49999H6.63001H6.93933L6.96966 7.46967L10.4697 3.96967L11.5303 3.96967L14.5 6.93934V2.5ZM8.00066 8.55999L9.53034 10.0897L10.0607 10.62L9.00001 11.6807L8.46968 11.1503L6.31935 9H3.81065L1.53032 11.2803L1.5 11.3106V12.5C1.5 13.0523 1.94772 13.5 2.5 13.5H13.5C14.0523 13.5 14.5 13.0523 14.5 12.5V9.06066L11 5.56066L8.03032 8.53033L8.00066 8.55999ZM4.05312e-06 10.8107V12.5C4.05312e-06 13.8807 1.11929 15 2.5 15H13.5C14.8807 15 16 13.8807 16 12.5V9.56066L16.5607 9L16.0303 8.46967L16 8.43934V2.5V1H14.5H1.5H4.05312e-06V2.5V10.6893L-0.0606689 10.75L4.05312e-06 10.8107Z"
         fill="currentColor"
-      ></path>
+      />
     </svg>
   );
 };
@@ -1116,7 +1116,7 @@ export const FullscreenIcon = ({ size = 16 }: { size?: number }) => (
       clipRule="evenodd"
       d="M1 5.25V6H2.5V5.25V2.5H5.25H6V1H5.25H2C1.44772 1 1 1.44772 1 2V5.25ZM5.25 14.9994H6V13.4994H5.25H2.5V10.7494V9.99939H1V10.7494V13.9994C1 14.5517 1.44772 14.9994 2 14.9994H5.25ZM15 10V10.75V14C15 14.5523 14.5523 15 14 15H10.75H10V13.5H10.75H13.5V10.75V10H15ZM10.75 1H10V2.5H10.75H13.5V5.25V6H15V5.25V2C15 1.44772 14.5523 1 14 1H10.75Z"
       fill="currentColor"
-    ></path>
+    />
   </svg>
 );
 
@@ -1133,7 +1133,7 @@ export const DownloadIcon = ({ size = 16 }: { size?: number }) => (
       clipRule="evenodd"
       d="M8.75 1V1.75V8.68934L10.7197 6.71967L11.25 6.18934L12.3107 7.25L11.7803 7.78033L8.70711 10.8536C8.31658 11.2441 7.68342 11.2441 7.29289 10.8536L4.21967 7.78033L3.68934 7.25L4.75 6.18934L5.28033 6.71967L7.25 8.68934V1.75V1H8.75ZM13.5 9.25V13.5H2.5V9.25V8.5H1V9.25V14C1 14.5523 1.44771 15 2 15H14C14.5523 15 15 14.5523 15 14V9.25V8.5H13.5V9.25Z"
       fill="currentColor"
-    ></path>
+    />
   </svg>
 );
 
@@ -1150,7 +1150,7 @@ export const LineChartIcon = ({ size = 16 }: { size?: number }) => (
       fillRule="evenodd"
       d="M1 1v11.75A2.25 2.25 0 0 0 3.25 15H15v-1.5H3.25a.75.75 0 0 1-.75-.75V1H1Zm13.297 5.013.513-.547-1.094-1.026-.513.547-3.22 3.434-2.276-2.275a1 1 0 0 0-1.414 0L4.22 8.22l-.53.53 1.06 1.06.53-.53L7 7.56l2.287 2.287a1 1 0 0 0 1.437-.023l3.573-3.811Z"
       clipRule="evenodd"
-    ></path>
+    />
   </svg>
 );
 
diff --git a/components/message.tsx b/components/message.tsx
index 085bfb0..f46bc89 100644
--- a/components/message.tsx
+++ b/components/message.tsx
@@ -4,7 +4,7 @@ import { AnimatePresence, motion } from 'framer-motion';
 import { memo, useState } from 'react';
 import type { Vote } from '@/lib/db/schema';
 import { DocumentToolResult } from './document';
-import { PencilEditIcon, SparklesIcon } from './icons';
+import { PencilEditIcon, SparklesIcon, LoaderIcon } from './icons';
 import { Response } from './elements/response';
 import { MessageContent } from './elements/message';
 import {
@@ -40,6 +40,7 @@ const PurePreviewMessage = ({
   regenerate,
   isReadonly,
   requiresScrollPadding,
+  isArtifactVisible,
 }: {
   chatId: string;
   message: ChatMessage;
@@ -49,6 +50,7 @@ const PurePreviewMessage = ({
   regenerate: UseChatHelpers<ChatMessage>['regenerate'];
   isReadonly: boolean;
   requiresScrollPadding: boolean;
+  isArtifactVisible: boolean;
 }) => {
   const [mode, setMode] = useState<'view' | 'edit'>('view');
 
@@ -62,31 +64,30 @@ const PurePreviewMessage = ({
     <AnimatePresence>
       <motion.div
         data-testid={`message-${message.role}`}
-        className="px-4 mx-auto w-full max-w-3xl group/message"
+        className="w-full group/message"
         initial={{ y: 5, opacity: 0 }}
         animate={{ y: 0, opacity: 1 }}
         data-role={message.role}
       >
         <div
-          className={cn(
-            'flex gap-4 w-full group-data-[role=user]/message:ml-auto group-data-[role=user]/message:max-w-2xl',
-            {
-              'w-full': mode === 'edit',
-              'group-data-[role=user]/message:w-fit': mode !== 'edit',
-            },
-          )}
+          className={cn('flex items-start gap-3', {
+            'w-full': mode === 'edit',
+            'max-w-xl ml-auto justify-end mr-10':
+              message.role === 'user' && mode !== 'edit',
+            'justify-start': message.role === 'assistant',
+          })}
         >
           {message.role === 'assistant' && (
-            <div className="flex justify-center items-center rounded-full ring-1 size-8 shrink-0 ring-border bg-background">
-              <div className="translate-y-px">
-                <SparklesIcon size={14} />
-              </div>
+            <div className="flex justify-center items-center mt-1 rounded-full ring-1 size-8 shrink-0 ring-border bg-background">
+              <SparklesIcon size={14} />
             </div>
           )}
 
           <div
-            className={cn('flex flex-col gap-4 w-full', {
+            className={cn('flex flex-col gap-4', {
               'min-h-96': message.role === 'assistant' && requiresScrollPadding,
+              'w-full': message.role === 'assistant',
+              'w-fit': message.role === 'user',
             })}
           >
             {attachmentsFromMessage.length > 0 && (
@@ -159,16 +160,20 @@ const PurePreviewMessage = ({
 
                 if (mode === 'edit') {
                   return (
-                    <div key={key} className="flex flex-row gap-2 items-start">
+                    <div
+                      key={key}
+                      className="flex flex-row gap-3 items-start w-full"
+                    >
                       <div className="size-8" />
-
-                      <MessageEditor
-                        key={message.id}
-                        message={message}
-                        setMode={setMode}
-                        setMessages={setMessages}
-                        regenerate={regenerate}
-                      />
+                      <div className="flex-1 min-w-0">
+                        <MessageEditor
+                          key={message.id}
+                          message={message}
+                          setMode={setMode}
+                          setMessages={setMessages}
+                          regenerate={regenerate}
+                        />
+                      </div>
                     </div>
                   );
                 }
@@ -196,67 +201,47 @@ const PurePreviewMessage = ({
               }
 
               if (type === 'tool-createDocument') {
-                const { toolCallId, state } = part;
+                const { toolCallId } = part;
+
+                if (part.output && 'error' in part.output) {
+                  return (
+                    <div
+                      key={toolCallId}
+                      className="p-4 text-red-500 bg-red-50 rounded-lg border border-red-200 dark:bg-red-950/50"
+                    >
+                      Error creating document: {String(part.output.error)}
+                    </div>
+                  );
+                }
 
                 return (
-                  <Tool key={toolCallId} defaultOpen={true}>
-                    <ToolHeader type="tool-createDocument" state={state} />
-                    <ToolContent>
-                      {state === 'input-available' && (
-                        <ToolInput input={part.input} />
-                      )}
-                      {state === 'output-available' && (
-                        <ToolOutput
-                          output={
-                            'error' in part.output ? (
-                              <div className="p-2 text-red-500 rounded border">
-                                Error: {String(part.output.error)}
-                              </div>
-                            ) : (
-                              <DocumentPreview
-                                isReadonly={isReadonly}
-                                result={part.output}
-                              />
-                            )
-                          }
-                          errorText={undefined}
-                        />
-                      )}
-                    </ToolContent>
-                  </Tool>
+                  <DocumentPreview
+                    key={toolCallId}
+                    isReadonly={isReadonly}
+                    result={part.output}
+                  />
                 );
               }
 
               if (type === 'tool-updateDocument') {
-                const { toolCallId, state } = part;
+                const { toolCallId } = part;
 
+                if (part.output && 'error' in part.output) {
+                  return (
+                    <div key={toolCallId} className="p-4 text-red-500 rounded-lg border border-red-200 bg-red-50 dark:bg-red-950/50">
+                      Error updating document: {String(part.output.error)}
+                    </div>
+                  );
+                }
+                
                 return (
-                  <Tool key={toolCallId} defaultOpen={true}>
-                    <ToolHeader type="tool-updateDocument" state={state} />
-                    <ToolContent>
-                      {state === 'input-available' && (
-                        <ToolInput input={part.input} />
-                      )}
-                      {state === 'output-available' && (
-                        <ToolOutput
-                          output={
-                            'error' in part.output ? (
-                              <div className="p-2 text-red-500 rounded border">
-                                Error: {String(part.output.error)}
-                              </div>
-                            ) : (
-                              <DocumentToolResult
-                                type="update"
-                                result={part.output}
-                                isReadonly={isReadonly}
-                              />
-                            )
-                          }
-                          errorText={undefined}
-                        />
-                      )}
-                    </ToolContent>
-                  </Tool>
+                  <div key={toolCallId} className="relative">
+                    <DocumentPreview
+                      isReadonly={isReadonly}
+                      result={part.output}
+                      args={{ ...part.output, isUpdate: true }}
+                    />
+                  </div>
                 );
               }
 
@@ -343,7 +328,7 @@ export const ThinkingMessage = () => {
           },
         )}
       >
-        <div className="flex justify-center items-center rounded-full ring-1 size-8 shrink-0 ring-border">
+        <div className="flex justify-center items-center mt-1 rounded-full ring-1 size-8 shrink-0 ring-border bg-background">
           <SparklesIcon size={14} />
         </div>
 
diff --git a/components/messages.tsx b/components/messages.tsx
index 83db795..c3239dc 100644
--- a/components/messages.tsx
+++ b/components/messages.tsx
@@ -9,6 +9,7 @@ import { useMessages } from '@/hooks/use-messages';
 import type { ChatMessage } from '@/lib/types';
 import { useDataStream } from './data-stream-provider';
 import { Conversation, ConversationContent, ConversationScrollButton } from './elements/conversation';
+import { cn } from '@/lib/utils';
 
 interface MessagesProps {
   chatId: string;
@@ -29,6 +30,7 @@ function PureMessages({
   setMessages,
   regenerate,
   isReadonly,
+  isArtifactVisible,
 }: MessagesProps) {
   const {
     containerRef: messagesContainerRef,
@@ -45,7 +47,7 @@ function PureMessages({
 
   return (
     <div ref={messagesContainerRef} className="flex-1 overflow-y-auto">
-      <Conversation className="flex flex-col min-w-0 gap-6 pt-4 pb-32">
+      <Conversation className="flex flex-col min-w-0 gap-6 pt-4 pb-32 px-4 max-w-4xl mx-auto">
         <ConversationContent className="flex flex-col gap-6">
           {messages.length === 0 && <Greeting />}
 
@@ -66,6 +68,7 @@ function PureMessages({
               requiresScrollPadding={
                 hasSentMessage && index === messages.length - 1
               }
+              isArtifactVisible={isArtifactVisible}
             />
           ))}
 
diff --git a/components/multimodal-input.tsx b/components/multimodal-input.tsx
index f0a8e4d..3d5c94d 100644
--- a/components/multimodal-input.tsx
+++ b/components/multimodal-input.tsx
@@ -24,7 +24,14 @@ import {
   PromptInputToolbar,
   PromptInputTools,
   PromptInputSubmit,
+  PromptInputModelSelect,
+  PromptInputModelSelectTrigger,
+  PromptInputModelSelectContent,
 } from './elements/prompt-input';
+import {
+  SelectItem,
+  SelectValue,
+} from '@/components/ui/select';
 import equal from 'fast-deep-equal';
 import type { UseChatHelpers } from '@ai-sdk/react';
 import { AnimatePresence, motion } from 'framer-motion';
@@ -32,6 +39,9 @@ import { ArrowDown } from 'lucide-react';
 import { useScrollToBottom } from '@/hooks/use-scroll-to-bottom';
 import type { VisibilityType } from './visibility-selector';
 import type { Attachment, ChatMessage } from '@/lib/types';
+import { chatModels } from '@/lib/ai/models';
+import { saveChatModelAsCookie } from '@/app/(chat)/actions';
+import { startTransition } from 'react';
 
 function PureMultimodalInput({
   chatId,
@@ -46,6 +56,7 @@ function PureMultimodalInput({
   sendMessage,
   className,
   selectedVisibilityType,
+  selectedModelId,
 }: {
   chatId: string;
   input: string;
@@ -59,6 +70,7 @@ function PureMultimodalInput({
   sendMessage: UseChatHelpers<ChatMessage>['sendMessage'];
   className?: string;
   selectedVisibilityType: VisibilityType;
+  selectedModelId: string;
 }) {
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const { width } = useWindowSize();
@@ -318,6 +330,7 @@ function PureMultimodalInput({
         <PromptInputToolbar className="px-2 py-1">
           <PromptInputTools className="gap-2">
             <AttachmentsButton fileInputRef={fileInputRef} status={status} />
+            <ModelSelectorCompact selectedModelId={selectedModelId} />
           </PromptInputTools>
           {status === 'submitted' ? (
             <StopButton stop={stop} setMessages={setMessages} />
@@ -343,6 +356,7 @@ export const MultimodalInput = memo(
     if (!equal(prevProps.attachments, nextProps.attachments)) return false;
     if (prevProps.selectedVisibilityType !== nextProps.selectedVisibilityType)
       return false;
+    if (prevProps.selectedModelId !== nextProps.selectedModelId) return false;
 
     return true;
   },
@@ -373,6 +387,50 @@ function PureAttachmentsButton({
 
 const AttachmentsButton = memo(PureAttachmentsButton);
 
+function PureModelSelectorCompact({
+  selectedModelId,
+}: {
+  selectedModelId: string;
+}) {
+  const [optimisticModelId, setOptimisticModelId] = useState(selectedModelId);
+
+  const selectedModel = chatModels.find(model => model.id === optimisticModelId);
+
+  return (
+    <PromptInputModelSelect
+      value={selectedModel?.name}
+      onValueChange={(modelName) => {
+        const model = chatModels.find(m => m.name === modelName);
+        if (model) {
+          setOptimisticModelId(model.id);
+          startTransition(() => {
+            saveChatModelAsCookie(model.id);
+          });
+        }
+      }}
+    >
+      <PromptInputModelSelectTrigger 
+        type="button"
+        className="text-xs focus:outline-none focus:ring-0 focus:ring-offset-0 focus-visible:ring-0 focus-visible:ring-offset-0 data-[state=open]:ring-0 data-[state=closed]:ring-0"
+      >
+        {selectedModel?.name || "Select model"}
+      </PromptInputModelSelectTrigger>
+      <PromptInputModelSelectContent>
+        {chatModels.map((model) => (
+          <SelectItem key={model.id} value={model.name}>
+            <div className="flex flex-col items-start gap-1 py-1">
+              <div className="font-medium">{model.name}</div>
+              <div className="text-xs text-muted-foreground">{model.description}</div>
+            </div>
+          </SelectItem>
+        ))}
+      </PromptInputModelSelectContent>
+    </PromptInputModelSelect>
+  );
+}
+
+const ModelSelectorCompact = memo(PureModelSelectorCompact);
+
 function PureStopButton({
   stop,
   setMessages,
diff --git a/components/preview-attachment.tsx b/components/preview-attachment.tsx
index ccd41be..caf150b 100644
--- a/components/preview-attachment.tsx
+++ b/components/preview-attachment.tsx
@@ -1,6 +1,6 @@
 import type { Attachment } from '@/lib/types';
 import { Loader } from './elements/loader';
-import { CrossSmallIcon, PencilEditIcon } from './icons';
+import { CrossSmallIcon, } from './icons';
 import { Button } from './ui/button';
 
 export const PreviewAttachment = ({
diff --git a/components/sidebar-history.tsx b/components/sidebar-history.tsx
index 04b441f..d1441da 100644
--- a/components/sidebar-history.tsx
+++ b/components/sidebar-history.tsx
@@ -152,7 +152,7 @@ export function SidebarHistory({ user }: { user: User | undefined }) {
     return (
       <SidebarGroup>
         <SidebarGroupContent>
-          <div className="px-2 text-zinc-500 w-full flex flex-row justify-center items-center text-sm gap-2">
+          <div className="flex flex-row gap-2 justify-center items-center px-2 w-full text-sm text-zinc-500">
             Login to save and revisit previous chats!
           </div>
         </SidebarGroupContent>
@@ -171,7 +171,7 @@ export function SidebarHistory({ user }: { user: User | undefined }) {
             {[44, 32, 28, 64, 52].map((item) => (
               <div
                 key={item}
-                className="rounded-md h-8 flex gap-2 px-2 items-center"
+                className="flex gap-2 items-center px-2 h-8 rounded-md"
               >
                 <div
                   className="h-4 rounded-md flex-1 max-w-[--skeleton-width] bg-sidebar-accent-foreground/10"
@@ -193,7 +193,7 @@ export function SidebarHistory({ user }: { user: User | undefined }) {
     return (
       <SidebarGroup>
         <SidebarGroupContent>
-          <div className="px-2 text-zinc-500 w-full flex flex-row justify-center items-center text-sm gap-2">
+          <div className="flex flex-row gap-2 justify-center items-center px-2 w-full text-sm text-zinc-500">
             Your conversations will appear here once you start chatting!
           </div>
         </SidebarGroupContent>
@@ -329,11 +329,11 @@ export function SidebarHistory({ user }: { user: User | undefined }) {
           />
 
           {hasReachedEnd ? (
-            <div className="px-2 text-zinc-500 w-full flex flex-row justify-center items-center text-sm gap-2 mt-8">
+            <div className="flex flex-row gap-2 justify-center items-center px-2 mt-8 w-full text-sm text-zinc-500">
               You have reached the end of your chat history.
             </div>
           ) : (
-            <div className="p-2 text-zinc-500 dark:text-zinc-400 flex flex-row gap-2 items-center mt-8">
+            <div className="flex flex-row gap-2 items-center p-2 mt-8 text-zinc-500 dark:text-zinc-400">
               <div className="animate-spin">
                 <LoaderIcon />
               </div>
diff --git a/components/ui/sidebar.tsx b/components/ui/sidebar.tsx
index 16e1437..dd0d59c 100644
--- a/components/ui/sidebar.tsx
+++ b/components/ui/sidebar.tsx
@@ -2,7 +2,7 @@
 
 import * as React from 'react';
 import { Slot } from '@radix-ui/react-slot';
-import { VariantProps, cva } from 'class-variance-authority';
+import { type VariantProps, cva } from 'class-variance-authority';
 import { PanelLeft } from 'lucide-react';
 
 import { useIsMobile } from '@/hooks/use-mobile';
diff --git a/components/visibility-selector.tsx b/components/visibility-selector.tsx
index b2d6dec..dd20e88 100644
--- a/components/visibility-selector.tsx
+++ b/components/visibility-selector.tsx
@@ -71,7 +71,7 @@ export function VisibilitySelector({
         <Button
           data-testid="visibility-selector"
           variant="outline"
-          className="hidden md:flex md:px-2 md:h-[34px]"
+          className="hidden md:flex md:px-2 md:h-[34px] focus:outline-none focus:ring-0"
         >
           {selectedVisibility?.icon}
           {selectedVisibility?.label}
diff --git a/hooks/use-artifact.ts b/hooks/use-artifact.ts
index 9133fe8..cf1f333 100644
--- a/hooks/use-artifact.ts
+++ b/hooks/use-artifact.ts
@@ -1,7 +1,7 @@
 'use client';
 
 import useSWR from 'swr';
-import { UIArtifact } from '@/components/artifact';
+import type { UIArtifact } from '@/components/artifact';
 import { useCallback, useMemo } from 'react';
 
 export const initialArtifactData: UIArtifact = {
diff --git a/lib/ai/models.ts b/lib/ai/models.ts
index 91fb137..a2e5edb 100644
--- a/lib/ai/models.ts
+++ b/lib/ai/models.ts
@@ -9,12 +9,12 @@ export interface ChatModel {
 export const chatModels: Array<ChatModel> = [
   {
     id: 'chat-model',
-    name: 'Chat model',
-    description: 'Primary model for all-purpose chat',
+    name: 'Grok Vision',
+    description: 'Advanced multimodal model with vision and text capabilities',
   },
   {
     id: 'chat-model-reasoning',
-    name: 'Reasoning model',
-    description: 'Uses advanced reasoning',
+    name: 'Grok Reasoning',
+    description: 'Uses advanced chain-of-thought reasoning for complex problems',
   },
 ];
diff --git a/lib/ai/providers.ts b/lib/ai/providers.ts
index 77fc284..e65234d 100644
--- a/lib/ai/providers.ts
+++ b/lib/ai/providers.ts
@@ -3,7 +3,7 @@ import {
   extractReasoningMiddleware,
   wrapLanguageModel,
 } from 'ai';
-import { xai } from '@ai-sdk/xai';
+import { gateway } from '@ai-sdk/gateway';
 import {
   artifactModel,
   chatModel,
@@ -23,15 +23,12 @@ export const myProvider = isTestEnvironment
     })
   : customProvider({
       languageModels: {
-        'chat-model': xai('grok-2-vision-1212'),
+        'chat-model': gateway.languageModel('xai/grok-2-vision-1212'),
         'chat-model-reasoning': wrapLanguageModel({
-          model: xai('grok-3-mini-beta'),
+          model: gateway.languageModel('xai/grok-3-mini-beta'),
           middleware: extractReasoningMiddleware({ tagName: 'think' }),
         }),
-        'title-model': xai('grok-2-1212'),
-        'artifact-model': xai('grok-2-1212'),
-      },
-      imageModels: {
-        'small-model': xai.imageModel('grok-2-image'),
+        'title-model': gateway.languageModel('xai/grok-2-1212'),
+        'artifact-model': gateway.languageModel('xai/grok-2-1212'),
       },
     });
diff --git a/lib/artifacts/server.ts b/lib/artifacts/server.ts
index dcb0312..575cb1e 100644
--- a/lib/artifacts/server.ts
+++ b/lib/artifacts/server.ts
@@ -1,5 +1,4 @@
 import { codeDocumentHandler } from '@/artifacts/code/server';
-import { imageDocumentHandler } from '@/artifacts/image/server';
 import { sheetDocumentHandler } from '@/artifacts/sheet/server';
 import { textDocumentHandler } from '@/artifacts/text/server';
 import type { ArtifactKind } from '@/components/artifact';
@@ -93,8 +92,7 @@ export function createDocumentHandler<T extends ArtifactKind>(config: {
 export const documentHandlersByArtifactKind: Array<DocumentHandler> = [
   textDocumentHandler,
   codeDocumentHandler,
-  imageDocumentHandler,
   sheetDocumentHandler,
 ];
 
-export const artifactKinds = ['text', 'code', 'image', 'sheet'] as const;
+export const artifactKinds = ['text', 'code', 'sheet'] as const;
diff --git a/package.json b/package.json
index 293a639..576765d 100644
--- a/package.json
+++ b/package.json
@@ -19,6 +19,7 @@
     "test": "export PLAYWRIGHT=True && pnpm exec playwright test"
   },
   "dependencies": {
+    "@ai-sdk/gateway": "^1.0.15",
     "@ai-sdk/provider": "2.0.0",
     "@ai-sdk/react": "2.0.26",
     "@ai-sdk/xai": "2.0.13",
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 0bd2f2a..097f339 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -8,6 +8,9 @@ importers:
 
   .:
     dependencies:
+      '@ai-sdk/gateway':
+        specifier: ^1.0.15
+        version: 1.0.15(zod@3.25.76)
       '@ai-sdk/provider':
         specifier: 2.0.0
         version: 2.0.0
diff --git a/tests/pages/artifact.ts b/tests/pages/artifact.ts
index 103603a..14ef5a3 100644
--- a/tests/pages/artifact.ts
+++ b/tests/pages/artifact.ts
@@ -1,4 +1,4 @@
-import { expect, Page } from '@playwright/test';
+import { expect, type Page } from '@playwright/test';
 
 export class ArtifactPage {
   constructor(private page: Page) {}