fix: init mcp list on chat load; stop on auth requested

Author: Cloud User

src/app/api/chat/route.ts

@@ -27,6 +27,7 @@ import {
   ChatMention,
   ChatMetadata,
 } from "app-types/chat";
+import { isMcpAuthRequiredToolResult } from "app-types/mcp";
 
 import { errorIf, safe } from "ts-safe";
 
@@ -89,6 +90,31 @@ function isExistingAttachmentPart(part: unknown): part is ChatAttachmentPart {
   );
 }
 
+function wrapToolsWithAuthAbort(
+  tools: Record<string, Tool>,
+  abortController: AbortController,
+): Record<string, Tool> {
+  const wrapped: Record<string, Tool> = {};
+  for (const [name, tool] of Object.entries(tools)) {
+    if (!tool.execute) {
+      wrapped[name] = tool;
+      continue;
+    }
+    const originalExecute = tool.execute.bind(tool);
+    wrapped[name] = {
+      ...tool,
+      execute: async (input, options) => {
+        const result = await originalExecute(input, options);
+        if (isMcpAuthRequiredToolResult(result)) {
+          abortController.abort();
+        }
+        return result;
+      },
+    };
+  }
+  return wrapped;
+}
+
 export async function POST(request: Request) {
   try {
     const json = await request.json();
@@ -281,7 +307,7 @@ export async function POST(request: Request) {
           .orElse({});
         const inProgressToolParts = extractInProgressToolPart(message);
         if (inProgressToolParts.length) {
-          await Promise.all(
+          const manualToolOutputs = await Promise.all(
             inProgressToolParts.map(async (part) => {
               const output = await manualToolExecuteByLastMessage(
                 part,
@@ -299,6 +325,9 @@ export async function POST(request: Request) {
               return output;
             }),
           );
+          if (manualToolOutputs.some(isMcpAuthRequiredToolResult)) {
+            return;
+          }
         }
 
         const userPreferences = thread?.userPreferences || undefined;
@@ -365,16 +394,19 @@ export async function POST(request: Request) {
         }
         logger.info(`model: ${chatModel?.provider}/${chatModel?.model}`);
 
+        const mcpAuthAbort = new AbortController();
+        request.signal.addEventListener("abort", () => mcpAuthAbort.abort());
+
         const result = streamText({
           model,
           system: systemPrompt,
           messages: convertToModelMessages(messages),
           experimental_transform: smoothStream({ chunking: "word" }),
           maxRetries: 2,
-          tools: vercelAITooles,
+          tools: wrapToolsWithAuthAbort(vercelAITooles, mcpAuthAbort),
           stopWhen: stepCountIs(10),
           toolChoice: "auto",
-          abortSignal: request.signal,
+          abortSignal: mcpAuthAbort.signal,
         });
         result.consumeStream();
         dataStream.merge(

src/components/chat-bot.tsx

@@ -49,6 +49,7 @@ import { getStorageManager } from "lib/browser-stroage";
 import { AnimatePresence, motion } from "framer-motion";
 import { useThreadFileUploader } from "@/hooks/use-thread-file-uploader";
 import { useFileDragOverlay } from "@/hooks/use-file-drag-overlay";
+import { useMcpList } from "@/hooks/queries/use-mcp-list";
 
 type Props = {
   threadId: string;
@@ -71,6 +72,7 @@ const isFirstTime = firstTimeStorage.get() ?? true;
 firstTimeStorage.set(false);
 
 export default function ChatBot({ threadId, initialMessages }: Props) {
+  useMcpList();
   const containerRef = useRef<HTMLDivElement>(null);
   const [isAtBottom, setIsAtBottom] = useState(true);
   const { uploadFiles } = useThreadFileUploader(threadId);

src/components/mcp-auth-dialog.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { useCallback, useState, useEffect } from "react";
+import { useCallback, useState, useEffect, useRef } from "react";
 import { mutate } from "swr";
 import { toast } from "sonner";
 import { Loader } from "lucide-react";
@@ -18,17 +18,40 @@ import { Button } from "ui/button";
 import { redriectMcpOauth } from "lib/ai/mcp/oauth-redirect";
 import type { MCPServerInfo } from "app-types/mcp";
 
+const activeAuthDialogs = new Set<string>();
+
 interface MCPAuthDialogProps {
   server: MCPServerInfo | undefined;
   serverName: string;
+  onAuthorized?: () => Promise<void> | void;
 }
 
-export function MCPAuthDialog({ server, serverName }: MCPAuthDialogProps) {
+export function MCPAuthDialog({
+  server,
+  serverName,
+  onAuthorized,
+}: MCPAuthDialogProps) {
   const t = useTranslations();
   const [loading, setLoading] = useState(false);
   const [dismissed, setDismissed] = useState(false);
+  const claimedRef = useRef(false);
 
   const needsAuth = server?.status === "authorizing";
+  const serverId = server?.id;
+
+  const isDuplicate = needsAuth && serverId && activeAuthDialogs.has(serverId);
+
+  useEffect(() => {
+    if (!needsAuth || !serverId || isDuplicate) {
+      return;
+    }
+    activeAuthDialogs.add(serverId);
+    claimedRef.current = true;
+    return () => {
+      activeAuthDialogs.delete(serverId);
+      claimedRef.current = false;
+    };
+  }, [needsAuth, serverId, isDuplicate]);
 
   useEffect(() => {
     if (!needsAuth) {
@@ -41,11 +64,12 @@ export function MCPAuthDialog({ server, serverName }: MCPAuthDialogProps) {
     setLoading(true);
     redriectMcpOauth(server.id)
       .then(() => mutate("/api/mcp/list"))
+      .then(() => onAuthorized?.())
       .catch((err) => toast.error(err?.message || t("MCP.authorizationFailed")))
       .finally(() => setLoading(false));
-  }, [server, t]);
+  }, [server, t, onAuthorized]);
 
-  if (!server || !needsAuth || dismissed) return null;
+  if (!server || !needsAuth || dismissed || isDuplicate) return null;
 
   return (
     <Dialog open onOpenChange={(open) => !open && setDismissed(true)}>

src/components/message-parts.tsx

@@ -118,6 +118,40 @@ interface ToolMessagePartProps {
 
 const MAX_TEXT_LENGTH = 600;
 
+async function replayMessageAfterMcpAuth({
+  replayMessage,
+  blockedMessageId,
+  threadId,
+  setMessages,
+  sendMessage,
+}: {
+  replayMessage?: UIMessage;
+  blockedMessageId: string;
+  threadId?: string;
+  setMessages?: UseChatHelpers<UIMessage>["setMessages"];
+  sendMessage?: UseChatHelpers<UIMessage>["sendMessage"];
+}) {
+  if (!replayMessage || !setMessages || !sendMessage) {
+    return;
+  }
+
+  if (threadId) {
+    await deleteMessagesByChatIdAfterTimestampAction(blockedMessageId);
+  }
+
+  setMessages((messages) => {
+    const replayIndex = messages.findIndex((m) => m.id === replayMessage.id);
+
+    if (replayIndex === -1) {
+      return messages;
+    }
+
+    return messages.slice(0, replayIndex);
+  });
+
+  sendMessage(replayMessage);
+}
+
 export const UserMessagePart = memo(
   function UserMessagePart({
     part,
@@ -762,7 +796,10 @@ export const ToolMessagePart = memo(
     addToolResult,
     isError,
     messageId,
+    prevMessage,
+    threadId,
     setMessages,
+    sendMessage,
     isManualToolInvocation,
   }: ToolMessagePartProps) => {
     const t = useTranslations("");
@@ -977,7 +1014,21 @@ export const ToolMessagePart = memo(
 
     return (
       <div className="group w-full">
-        <MCPAuthDialog server={mcpServer} serverName={mcpServerName} />
+        {isLast && (
+          <MCPAuthDialog
+            server={mcpServer}
+            serverName={mcpServerName}
+            onAuthorized={() =>
+              replayMessageAfterMcpAuth({
+                replayMessage: prevMessage,
+                blockedMessageId: messageId,
+                threadId,
+                setMessages,
+                sendMessage,
+              })
+            }
+          />
+        )}
 
         {CustomToolComponent ? (
           CustomToolComponent

src/lib/ai/mcp/create-mcp-client.ts

@@ -388,12 +388,30 @@ export class MCPClient {
   async callTool(toolName: string, input?: unknown) {
     const id = generateUUID();
     this.inProgressToolCallIds.push(id);
+    const authRequiredResult = () => ({
+      isError: true,
+      content: [
+        {
+          type: "text",
+          text: "OAuth authorization required",
+        },
+      ],
+      _mcpAuthRequired: true,
+      _mcpServerId: this.id,
+    });
     const execute = async () => {
-      const client = await this.connect();
-      return await client?.callTool({
-        name: toolName,
-        arguments: input as Record<string, unknown>,
-      });
+      try {
+        const client = await this.connect();
+        return await client?.callTool({
+          name: toolName,
+          arguments: input as Record<string, unknown>,
+        });
+      } catch (err) {
+        if (this.status === "authorizing") {
+          return authRequiredResult();
+        }
+        throw err;
+      }
     };
     return safe(() => this.logger.info("tool call", toolName))
       .ifOk(() => this.scheduleAutoDisconnect()) // disconnect if autoDisconnectSeconds is set
@@ -419,6 +437,9 @@ export class MCPClient {
             // Expected - triggers OAuth flow
           }
         }
+        if (this.status === "authorizing") {
+          return authRequiredResult();
+        }
         throw err;
       })
       .ifOk((v) => {

src/types/mcp.ts

@@ -248,10 +248,30 @@ export const CallToolResultSchema = z.object({
   content: z.array(ContentUnion).default([]),
   structuredContent: z.object({}).passthrough().optional(),
   isError: z.boolean().optional(),
+  _mcpAuthRequired: z.boolean().optional(),
+  _mcpServerId: z.string().optional(),
 });
 
 export type CallToolResult = z.infer<typeof CallToolResultSchema>;
 
+export type McpAuthRequiredToolResult = CallToolResult & {
+  _mcpAuthRequired: true;
+  _mcpServerId: string;
+};
+
+export function isMcpAuthRequiredToolResult(
+  value: unknown,
+): value is McpAuthRequiredToolResult {
+  return (
+    typeof value === "object" &&
+    value !== null &&
+    "_mcpAuthRequired" in value &&
+    "_mcpServerId" in value &&
+    (value as { _mcpAuthRequired?: unknown })._mcpAuthRequired === true &&
+    typeof (value as { _mcpServerId?: unknown })._mcpServerId === "string"
+  );
+}
+
 export type McpOAuthSession = {
   id: string;
   mcpServerId: string;