chore: upgrade ai sdk v6 + tool approval (vercel#1361)

Author: DavideGCosta

app/(chat)/api/chat/route.ts

@@ -14,7 +14,6 @@ import {
 } from "resumable-stream";
 import type { VisibilityType } from "@/components/visibility-selector";
 import { entitlementsByUserType } from "@/lib/ai/entitlements";
-import type { ChatModel } from "@/lib/ai/models";
 import { type RequestHints, systemPrompt } from "@/lib/ai/prompts";
 import { getLanguageModel } from "@/lib/ai/providers";
 import { createDocument } from "@/lib/ai/tools/create-document";
@@ -32,6 +31,7 @@ import {
   saveChat,
   saveMessages,
   updateChatTitleById,
+  updateMessage,
 } from "@/lib/db/queries";
 import type { DBMessage } from "@/lib/db/schema";
 import { ChatSDKError } from "@/lib/errors";
@@ -78,13 +78,9 @@ export async function POST(request: Request) {
     const {
       id,
       message,
+      messages,
       selectedChatModel,
       selectedVisibilityType,
-    }: {
-      id: string;
-      message: ChatMessage;
-      selectedChatModel: ChatModel["id"];
-      selectedVisibilityType: VisibilityType;
     } = requestBody;
 
     const session = await getSession();
@@ -104,6 +100,10 @@ export async function POST(request: Request) {
       return new ChatSDKError("rate_limit:chat").toResponse();
     }
 
+    const isToolApprovalFlow = Boolean(messages);
+    if (!isToolApprovalFlow && !message) {
+      return new ChatSDKError("bad_request:api").toResponse();
+    }
     const chat = await getChatById({ id });
     let messagesFromDb: DBMessage[] = [];
     let titlePromise: Promise<string> | null = null;
@@ -112,23 +112,25 @@ export async function POST(request: Request) {
       if (chat.userId !== session.user.id) {
         return new ChatSDKError("forbidden:chat").toResponse();
       }
-      // Only fetch messages if chat already exists
-      messagesFromDb = await getMessagesByChatId({ id });
-    } else {
-      // Save chat immediately with a placeholder title
+      if (!isToolApprovalFlow) {
+        messagesFromDb = await getMessagesByChatId({ id });
+      }
+    } else if (message?.role === "user") {
       await saveChat({
         id,
         userId: session.user.id,
         title: "New chat",
         visibility: selectedVisibilityType,
       });
 
-      // Kick off title generation in the background
       titlePromise = generateTitleFromUserMessage({ message });
-      // New chat - no need to fetch messages, it's empty
+    } else {
+      return new ChatSDKError("bad_request:chat").toResponse();
     }
 
-    const uiMessages = [...convertToUIMessages(messagesFromDb), message];
+    const uiMessages: ChatMessage[] = isToolApprovalFlow
+      ? (messages as ChatMessage[])
+      : [...convertToUIMessages(messagesFromDb), message as ChatMessage];
 
     const { longitude, latitude, city, country } = geolocation(request);
 
@@ -139,24 +141,27 @@ export async function POST(request: Request) {
       country,
     };
 
-    await saveMessages({
-      messages: [
-        {
-          chatId: id,
-          id: message.id,
-          role: "user",
-          parts: message.parts,
-          attachments: [],
-          createdAt: new Date(),
-        },
-      ],
-    });
+    if (message?.role === "user") {
+      await saveMessages({
+        messages: [
+          {
+            chatId: id,
+            id: message.id,
+            role: "user",
+            parts: message.parts,
+            attachments: [],
+            createdAt: new Date(),
+          },
+        ],
+      });
+    }
 
     const streamId = generateUUID();
     await createStreamId({ streamId, chatId: id });
 
     const stream = createUIMessageStream({
-      execute: ({ writer: dataStream }) => {
+      originalMessages: isToolApprovalFlow ? uiMessages : undefined,
+      execute: async ({ writer: dataStream }) => {
         // Handle async title generation in parallel for new chats
         if (titlePromise) {
           titlePromise.then((title) => {
@@ -172,7 +177,7 @@ export async function POST(request: Request) {
         const result = streamText({
           model: getLanguageModel(selectedChatModel),
           system: systemPrompt({ selectedChatModel, requestHints }),
-          messages: convertToModelMessages(uiMessages),
+          messages: await convertToModelMessages(uiMessages),
           stopWhen: stepCountIs(5),
           experimental_activeTools: isReasoningModel
             ? []
@@ -216,32 +221,63 @@ export async function POST(request: Request) {
         );
       },
       generateId: generateUUID,
-      onFinish: async ({ messages }) => {
-        await saveMessages({
-          messages: messages.map((currentMessage) => ({
-            id: currentMessage.id,
-            role: currentMessage.role,
-            parts: currentMessage.parts,
-            createdAt: new Date(),
-            attachments: [],
-            chatId: id,
-          })),
-        });
+      onFinish: async ({ messages: finishedMessages }) => {
+        if (isToolApprovalFlow) {
+          for (const finishedMsg of finishedMessages) {
+            const existingMsg = uiMessages.find((m) => m.id === finishedMsg.id);
+            if (existingMsg) {
+              await updateMessage({
+                id: finishedMsg.id,
+                parts: finishedMsg.parts,
+              });
+            } else {
+              await saveMessages({
+                messages: [
+                  {
+                    id: finishedMsg.id,
+                    role: finishedMsg.role,
+                    parts: finishedMsg.parts,
+                    createdAt: new Date(),
+                    attachments: [],
+                    chatId: id,
+                  },
+                ],
+              });
+            }
+          }
+        } else if (finishedMessages.length > 0) {
+          await saveMessages({
+            messages: finishedMessages.map((currentMessage) => ({
+              id: currentMessage.id,
+              role: currentMessage.role,
+              parts: currentMessage.parts,
+              createdAt: new Date(),
+              attachments: [],
+              chatId: id,
+            })),
+          });
+        }
       },
       onError: () => {
         return "Oops, an error occurred!";
       },
     });
 
-    // const streamContext = getStreamContext();
+    const streamContext = getStreamContext();
 
-    // if (streamContext) {
-    //   return new Response(
-    //     await streamContext.resumableStream(streamId, () =>
-    //       stream.pipeThrough(new JsonToSseTransformStream())
-    //     )
-    //   );
-    // }
+    if (streamContext) {
+      try {
+        const resumableStream = await streamContext.resumableStream(
+          streamId,
+          () => stream.pipeThrough(new JsonToSseTransformStream())
+        );
+        if (resumableStream) {
+          return new Response(resumableStream);
+        }
+      } catch (error) {
+        console.error("Failed to create resumable stream:", error);
+      }
+    }
 
     return new Response(stream.pipeThrough(new JsonToSseTransformStream()));
   } catch (error) {

app/(chat)/api/chat/schema.ts

@@ -14,13 +14,24 @@ const filePartSchema = z.object({
 
 const partSchema = z.union([textPartSchema, filePartSchema]);
 
+const userMessageSchema = z.object({
+  id: z.string().uuid(),
+  role: z.enum(["user"]),
+  parts: z.array(partSchema),
+});
+
+// For tool approval flows, we accept all messages (more permissive schema)
+const messageSchema = z.object({
+  id: z.string(),
+  role: z.string(),
+  parts: z.array(z.any()),
+});
+
 export const postRequestBodySchema = z.object({
   id: z.string().uuid(),
-  message: z.object({
-    id: z.string().uuid(),
-    role: z.enum(["user"]),
-    parts: z.array(partSchema),
-  }),
+  // Either a single new message or all messages (for tool approvals)
+  message: userMessageSchema.optional(),
+  messages: z.array(messageSchema).optional(),
   selectedChatModel: z.string(),
   selectedVisibilityType: z.enum(["public", "private"]),
 });

components/ai-elements/confirmation.tsx

@@ -97,7 +97,6 @@ export const ConfirmationRequest = ({ children }: ConfirmationRequestProps) => {
   const { state } = useConfirmation();
 
   // Only show when approval is requested
-  // @ts-expect-error state only available in AI SDK v6
   if (state !== "approval-requested") {
     return null;
   }
@@ -117,9 +116,7 @@ export const ConfirmationAccepted = ({
   // Only show when approved and in response states
   if (
     !approval?.approved ||
-    // @ts-expect-error state only available in AI SDK v6
     (state !== "approval-responded" &&
-      // @ts-expect-error state only available in AI SDK v6
       state !== "output-denied" &&
       state !== "output-available")
   ) {
@@ -141,9 +138,7 @@ export const ConfirmationRejected = ({
   // Only show when rejected and in response states
   if (
     approval?.approved !== false ||
-    // @ts-expect-error state only available in AI SDK v6
     (state !== "approval-responded" &&
-      // @ts-expect-error state only available in AI SDK v6
       state !== "output-denied" &&
       state !== "output-available")
   ) {
@@ -162,7 +157,6 @@ export const ConfirmationActions = ({
   const { state } = useConfirmation();
 
   // Only show when approval is requested
-  // @ts-expect-error state only available in AI SDK v6
   if (state !== "approval-requested") {
     return null;
   }

components/ai-elements/tool.tsx

@@ -40,7 +40,6 @@ const getStatusBadge = (status: ToolUIPart["state"]) => {
   const labels: Record<ToolUIPart["state"], string> = {
     "input-streaming": "Pending",
     "input-available": "Running",
-    // @ts-expect-error state only available in AI SDK v6
     "approval-requested": "Awaiting Approval",
     "approval-responded": "Responded",
     "output-available": "Completed",
@@ -51,7 +50,6 @@ const getStatusBadge = (status: ToolUIPart["state"]) => {
   const icons: Record<ToolUIPart["state"], ReactNode> = {
     "input-streaming": <CircleIcon className="size-4" />,
     "input-available": <ClockIcon className="size-4 animate-pulse" />,
-    // @ts-expect-error state only available in AI SDK v6
     "approval-requested": <ClockIcon className="size-4 text-yellow-600" />,
     "approval-responded": <CheckCircleIcon className="size-4 text-blue-600" />,
     "output-available": <CheckCircleIcon className="size-4 text-green-600" />,

components/artifact-messages.tsx

@@ -9,6 +9,7 @@ import type { UIArtifact } from "./artifact";
 import { PreviewMessage, ThinkingMessage } from "./message";
 
 type ArtifactMessagesProps = {
+  addToolApprovalResponse: UseChatHelpers<ChatMessage>["addToolApprovalResponse"];
   chatId: string;
   status: UseChatHelpers<ChatMessage>["status"];
   votes: Vote[] | undefined;
@@ -20,6 +21,7 @@ type ArtifactMessagesProps = {
 };
 
 function PureArtifactMessages({
+  addToolApprovalResponse,
   chatId,
   status,
   votes,
@@ -45,6 +47,7 @@ function PureArtifactMessages({
     >
       {messages.map((message, index) => (
         <PreviewMessage
+          addToolApprovalResponse={addToolApprovalResponse}
           chatId={chatId}
           isLoading={status === "streaming" && index === messages.length - 1}
           isReadonly={isReadonly}
@@ -64,7 +67,12 @@ function PureArtifactMessages({
       ))}
 
       <AnimatePresence mode="wait">
-        {status === "submitted" && <ThinkingMessage key="thinking" />}
+        {status === "submitted" &&
+          !messages.some((msg) =>
+            msg.parts?.some(
+              (part) => "state" in part && part.state === "approval-responded"
+            )
+          ) && <ThinkingMessage key="thinking" />}
       </AnimatePresence>
 
       <motion.div

components/artifact.tsx

@@ -53,6 +53,7 @@ export type UIArtifact = {
 };
 
 function PureArtifact({
+  addToolApprovalResponse,
   chatId,
   input,
   setInput,
@@ -69,6 +70,7 @@ function PureArtifact({
   selectedVisibilityType,
   selectedModelId,
 }: {
+  addToolApprovalResponse: UseChatHelpers<ChatMessage>["addToolApprovalResponse"];
   chatId: string;
   input: string;
   setInput: Dispatch<SetStateAction<string>>;
@@ -320,6 +322,7 @@ function PureArtifact({
 
               <div className="flex h-full flex-col items-center justify-between">
                 <ArtifactMessages
+                  addToolApprovalResponse={addToolApprovalResponse}
                   artifactStatus={artifact.status}
                   chatId={chatId}
                   isReadonly={isReadonly}