🎄 merry christmas: ai sdk v6 beta + tool approval (#1361)

2025-12-19 23:24:24 +00:00 · 2025-12-19 23:24:24 +00:00 · 4d3ba8d9fe
commit 4d3ba8d9fe
parent 6e5b883cf2
21 changed files with 429 additions and 202 deletions
--- a/components/ai-elements/confirmation.tsx
+++ b/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;
  }
--- a/components/ai-elements/tool.tsx
+++ b/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" />,
--- a/components/artifact-messages.tsx
+++ b/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
--- a/components/artifact.tsx
+++ b/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}
--- a/components/chat.tsx
+++ b/components/chat.tsx
@ -85,19 +85,54 @@ export function Chat({
    stop,
    regenerate,
    resumeStream,
+    addToolApprovalResponse,
  } = useChat<ChatMessage>({
    id,
    messages: initialMessages,
    experimental_throttle: 100,
    generateId: generateUUID,
+    // Auto-continue after tool approval (only for APPROVED tools)
+    // Denied tools don't need server continuation - state is saved on next user message
+    sendAutomaticallyWhen: ({ messages: currentMessages }) => {
+      const lastMessage = currentMessages.at(-1);
+      // Only continue if a tool was APPROVED (not denied)
+      const shouldContinue =
+        lastMessage?.parts?.some(
+          (part) =>
+            "state" in part &&
+            part.state === "approval-responded" &&
+            "approval" in part &&
+            (part.approval as { approved?: boolean })?.approved === true
+        ) ?? false;
+      return shouldContinue;
+    },
    transport: new DefaultChatTransport({
      api: "/api/chat",
      fetch: fetchWithErrorHandlers,
      prepareSendMessagesRequest(request) {
+        const lastMessage = request.messages.at(-1);
+
+        // Check if this is a tool approval continuation:
+        // - Last message is NOT a user message (meaning no new user input)
+        // - OR any message has tool parts that were responded to (approved or denied)
+        const isToolApprovalContinuation =
+          lastMessage?.role !== "user" ||
+          request.messages.some((msg) =>
+            msg.parts?.some((part) => {
+              const state = (part as { state?: string }).state;
+              return (
+                state === "approval-responded" || state === "output-denied"
+              );
+            })
+          );
+
        return {
          body: {
            id: request.id,
-            message: request.messages.at(-1),
+            // Send all messages for tool approval continuation, otherwise just the last user message
+            ...(isToolApprovalContinuation
+              ? { messages: request.messages }
+              : { message: lastMessage }),
            selectedChatModel: currentModelIdRef.current,
            selectedVisibilityType: visibilityType,
            ...request.body,
@ -170,6 +205,7 @@ export function Chat({
        />

        <Messages
+          addToolApprovalResponse={addToolApprovalResponse}
          chatId={id}
          isArtifactVisible={isArtifactVisible}
          isReadonly={isReadonly}
@ -203,6 +239,7 @@ export function Chat({
      </div>

      <Artifact
+        addToolApprovalResponse={addToolApprovalResponse}
        attachments={attachments}
        chatId={id}
        input={input}
--- a/components/document-preview.tsx
+++ b/components/document-preview.tsx
@ -102,7 +102,7 @@ export function DocumentPreview({
  }

  return (
-    <div className="relative w-full cursor-pointer">
+    <div className="relative w-full max-w-[450px] cursor-pointer">
      <HitboxLayer
        hitboxRef={hitboxRef}
        result={result}
@ -119,7 +119,7 @@ export function DocumentPreview({
 }

 const LoadingSkeleton = ({ artifactKind }: { artifactKind: ArtifactKind }) => (
-  <div className="w-full">
+  <div className="w-full max-w-[450px]">
    <div className="flex h-[57px] flex-row items-center justify-between gap-2 rounded-t-2xl border border-b-0 p-4 dark:border-zinc-700 dark:bg-muted">
      <div className="flex flex-row items-center gap-3">
        <div className="text-muted-foreground">
--- a/components/elements/tool.tsx
+++ b/components/elements/tool.tsx
@ -35,19 +35,25 @@ export type ToolHeaderProps = {
 };

 const getStatusBadge = (status: ToolUIPart["state"]) => {
-  const labels = {
+  const labels: Record<ToolUIPart["state"], string> = {
    "input-streaming": "Pending",
    "input-available": "Running",
+    "approval-requested": "Pending",
+    "approval-responded": "Approved",
    "output-available": "Completed",
    "output-error": "Error",
-  } as const;
+    "output-denied": "Denied",
+  };

-  const icons = {
+  const icons: Record<ToolUIPart["state"], ReactNode> = {
    "input-streaming": <CircleIcon className="size-4" />,
    "input-available": <ClockIcon className="size-4 animate-pulse" />,
+    "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" />,
    "output-error": <XCircleIcon className="size-4 text-red-600" />,
-  } as const;
+    "output-denied": <XCircleIcon className="size-4 text-orange-600" />,
+  };

  return (
    <Badge
--- a/components/message.tsx
+++ b/components/message.tsx
@ -25,6 +25,7 @@ import { PreviewAttachment } from "./preview-attachment";
 import { Weather } from "./weather";

 const PurePreviewMessage = ({
+  addToolApprovalResponse,
  chatId,
  message,
  vote,
@ -34,6 +35,7 @@ const PurePreviewMessage = ({
  isReadonly,
  requiresScrollPadding: _requiresScrollPadding,
 }: {
+  addToolApprovalResponse: UseChatHelpers<ChatMessage>["addToolApprovalResponse"];
  chatId: string;
  message: ChatMessage;
  vote: Vote | undefined;
@ -76,9 +78,10 @@ const PurePreviewMessage = ({
            ),
            "w-full":
              (message.role === "assistant" &&
-                message.parts?.some(
+                (message.parts?.some(
                  (p) => p.type === "text" && p.text?.trim()
-                )) ||
+                ) ||
+                  message.parts?.some((p) => p.type.startsWith("tool-")))) ||
              mode === "edit",
            "max-w-[calc(100%-2.5rem)] sm:max-w-[min(fit-content,80%)]":
              message.role === "user" && mode !== "edit",
@ -122,7 +125,7 @@ const PurePreviewMessage = ({
                  <div key={key}>
                    <MessageContent
                      className={cn({
-                        "w-fit break-words rounded-2xl px-3 py-2 text-right text-white":
+                        "wrap-break-word w-fit rounded-2xl px-3 py-2 text-right text-white":
                          message.role === "user",
                        "bg-transparent px-0 py-0 text-left":
                          message.role === "assistant",
@ -163,22 +166,95 @@ const PurePreviewMessage = ({

            if (type === "tool-getWeather") {
              const { toolCallId, state } = part;
+              const approvalId = (part as { approval?: { id: string } })
+                .approval?.id;
+              const isDenied =
+                state === "output-denied" ||
+                (state === "approval-responded" &&
+                  (part as { approval?: { approved?: boolean } }).approval
+                    ?.approved === false);
+              const widthClass = "w-[min(100%,450px)]";
+
+              if (state === "output-available") {
+                return (
+                  <div className={widthClass} key={toolCallId}>
+                    <Weather weatherAtLocation={part.output} />
+                  </div>
+                );
+              }
+
+              if (isDenied) {
+                return (
+                  <div className={widthClass} key={toolCallId}>
+                    <Tool className="w-full" defaultOpen={true}>
+                      <ToolHeader
+                        state="output-denied"
+                        type="tool-getWeather"
+                      />
+                      <ToolContent>
+                        <div className="px-4 py-3 text-muted-foreground text-sm">
+                          Weather lookup was denied.
+                        </div>
+                      </ToolContent>
+                    </Tool>
+                  </div>
+                );
+              }
+
+              if (state === "approval-responded") {
+                return (
+                  <div className={widthClass} key={toolCallId}>
+                    <Tool className="w-full" defaultOpen={true}>
+                      <ToolHeader state={state} type="tool-getWeather" />
+                      <ToolContent>
+                        <ToolInput input={part.input} />
+                      </ToolContent>
+                    </Tool>
+                  </div>
+                );
+              }

              return (
-                <Tool defaultOpen={true} key={toolCallId}>
-                  <ToolHeader state={state} type="tool-getWeather" />
-                  <ToolContent>
-                    {state === "input-available" && (
-                      <ToolInput input={part.input} />
-                    )}
-                    {state === "output-available" && (
-                      <ToolOutput
-                        errorText={undefined}
-                        output={<Weather weatherAtLocation={part.output} />}
-                      />
-                    )}
-                  </ToolContent>
-                </Tool>
+                <div className={widthClass} key={toolCallId}>
+                  <Tool className="w-full" defaultOpen={true}>
+                    <ToolHeader state={state} type="tool-getWeather" />
+                    <ToolContent>
+                      {(state === "input-available" ||
+                        state === "approval-requested") && (
+                        <ToolInput input={part.input} />
+                      )}
+                      {state === "approval-requested" && approvalId && (
+                        <div className="flex items-center justify-end gap-2 border-t px-4 py-3">
+                          <button
+                            className="rounded-md px-3 py-1.5 text-muted-foreground text-sm transition-colors hover:bg-muted hover:text-foreground"
+                            onClick={() => {
+                              addToolApprovalResponse({
+                                id: approvalId,
+                                approved: false,
+                                reason: "User denied weather lookup",
+                              });
+                            }}
+                            type="button"
+                          >
+                            Deny
+                          </button>
+                          <button
+                            className="rounded-md bg-primary px-3 py-1.5 text-primary-foreground text-sm transition-colors hover:bg-primary/90"
+                            onClick={() => {
+                              addToolApprovalResponse({
+                                id: approvalId,
+                                approved: true,
+                              });
+                            }}
+                            type="button"
+                          >
+                            Allow
+                          </button>
+                        </div>
+                      )}
+                    </ToolContent>
+                  </Tool>
+                </div>
              );
            }

@ -285,22 +361,15 @@ const PurePreviewMessage = ({
 export const PreviewMessage = memo(
  PurePreviewMessage,
  (prevProps, nextProps) => {
-    if (prevProps.isLoading !== nextProps.isLoading) {
-      return false;
+    if (
+      prevProps.isLoading === nextProps.isLoading &&
+      prevProps.message.id === nextProps.message.id &&
+      prevProps.requiresScrollPadding === nextProps.requiresScrollPadding &&
+      equal(prevProps.message.parts, nextProps.message.parts) &&
+      equal(prevProps.vote, nextProps.vote)
+    ) {
+      return true;
    }
-    if (prevProps.message.id !== nextProps.message.id) {
-      return false;
-    }
-    if (prevProps.requiresScrollPadding !== nextProps.requiresScrollPadding) {
-      return false;
-    }
-    if (!equal(prevProps.message.parts, nextProps.message.parts)) {
-      return false;
-    }
-    if (!equal(prevProps.vote, nextProps.vote)) {
-      return false;
-    }
-
    return false;
  }
 );
--- a/components/messages.tsx
+++ b/components/messages.tsx
@ -10,6 +10,7 @@ import { Greeting } from "./greeting";
 import { PreviewMessage, ThinkingMessage } from "./message";

 type MessagesProps = {
+  addToolApprovalResponse: UseChatHelpers<ChatMessage>["addToolApprovalResponse"];
  chatId: string;
  status: UseChatHelpers<ChatMessage>["status"];
  votes: Vote[] | undefined;
@ -22,6 +23,7 @@ type MessagesProps = {
 };

 function PureMessages({
+  addToolApprovalResponse,
  chatId,
  status,
  votes,
@ -54,6 +56,7 @@ function PureMessages({

          {messages.map((message, index) => (
            <PreviewMessage
+              addToolApprovalResponse={addToolApprovalResponse}
              chatId={chatId}
              isLoading={
                status === "streaming" && messages.length - 1 === index
@ -74,7 +77,12 @@ function PureMessages({
            />
          ))}

-          {status === "submitted" && <ThinkingMessage />}
+          {status === "submitted" &&
+            !messages.some((msg) =>
+              msg.parts?.some(
+                (part) => "state" in part && part.state === "approval-responded"
+              )
+            ) && <ThinkingMessage />}

          <div
            className="min-h-[24px] min-w-[24px] shrink-0"