feat: improve MCP OAuth session handling and auth UX

messages/en.json

@@ -794,6 +794,10 @@
     "public": "Featured",
     "publicDescription": "Featured for all users to use",
     "featured": "Featured",
-    "featuredDescription": "Featured for all users to use"
+    "featuredDescription": "Featured for all users to use",
+    "authorizationRequired": "Authorization Required",
+    "authorizationRequiredDescription": "{serverName} requires authorization to access its tools.",
+    "authorize": "Authorize",
+    "authorizationFailed": "Authorization failed"
   }
 }

messages/es.json

@@ -358,7 +358,11 @@
     "mcpServerCustomization": "Personalización del Servidor MCP",
     "mcpServerCustomizationDescription": "Las instrucciones de personalización del servidor MCP se añadirán al prompt del sistema cuando el servidor MCP esté disponible.",
     "toolCustomizationInstructionsPlaceholder": "Las instrucciones de personalización de herramientas no están disponibles.",
-    "mcpServerCustomizationPlaceholder": "ej. Si el valor de entrada es email, siempre ingresa el email en el formato ejemplo@ejemplo.com."
+    "mcpServerCustomizationPlaceholder": "ej. Si el valor de entrada es email, siempre ingresa el email en el formato ejemplo@ejemplo.com.",
+    "authorizationRequired": "Autorización Requerida",
+    "authorizationRequiredDescription": "{serverName} requiere autorización para acceder a sus herramientas.",
+    "authorize": "Autorizar",
+    "authorizationFailed": "Error de autorización"
   },
   "Error": {},
   "Info": {

messages/fr.json

@@ -358,7 +358,11 @@
     "mcpServerCustomization": "Personnalisation du Serveur MCP",
     "mcpServerCustomizationDescription": "Les instructions de personnalisation du serveur MCP seront ajoutées à l'invite système lorsque le serveur MCP sera disponible.",
     "toolCustomizationInstructionsPlaceholder": "Les instructions de personnalisation d'outils ne sont pas disponibles.",
-    "mcpServerCustomizationPlaceholder": "ex. Si la valeur d'entrée est un email, toujours entrer l'email au format exemple@exemple.com."
+    "mcpServerCustomizationPlaceholder": "ex. Si la valeur d'entrée est un email, toujours entrer l'email au format exemple@exemple.com.",
+    "authorizationRequired": "Autorisation Requise",
+    "authorizationRequiredDescription": "{serverName} nécessite une autorisation pour accéder à ses outils.",
+    "authorize": "Autoriser",
+    "authorizationFailed": "Échec de l'autorisation"
   },
   "Error": {},
   "Info": {

messages/ja.json

@@ -357,7 +357,11 @@
     "mcpServerCustomization": "MCPサーバーカスタマイゼーション",
     "mcpServerCustomizationDescription": "MCPサーバーカスタマイゼーション指示は、MCPサーバーが利用可能な場合にシステムプロンプトに追加されます。",
     "toolCustomizationInstructionsPlaceholder": "ツールカスタマイゼーション指示は利用できません。",
-    "mcpServerCustomizationPlaceholder": "例：入力値がメールの場合、常にexample@example.comの形式でメールを入力してください。"
+    "mcpServerCustomizationPlaceholder": "例：入力値がメールの場合、常にexample@example.comの形式でメールを入力してください。",
+    "authorizationRequired": "認証が必要です",
+    "authorizationRequiredDescription": "{serverName} のツールにアクセスするには認証が必要です。",
+    "authorize": "認証",
+    "authorizationFailed": "認証に失敗しました"
   },
   "Error": {},
   "Info": {

messages/ko.json

@@ -372,7 +372,11 @@
     "mcpServerCustomization": "MCP 서버 지침 설정",
     "mcpServerCustomizationDescription": "MCP 서버 사용자 지정 지침은 MCP서버를 사용할 때 시스템 프롬프트에 추가됩니다.",
     "toolCustomizationInstructionsPlaceholder": "도구 사용자 지정 지침이 없습니다.",
-    "mcpServerCustomizationPlaceholder": "예) 이 MCP 서버의 툴중에 입력값이 email인 경우 항상 example@example.com 형식으로 입력해주세요."
+    "mcpServerCustomizationPlaceholder": "예) 이 MCP 서버의 툴중에 입력값이 email인 경우 항상 example@example.com 형식으로 입력해주세요.",
+    "authorizationRequired": "인증 필요",
+    "authorizationRequiredDescription": "{serverName}의 도구에 액세스하려면 인증이 필요합니다.",
+    "authorize": "인증",
+    "authorizationFailed": "인증 실패"
   },
   "Error": {},
   "Info": {

messages/no.json

@@ -794,6 +794,10 @@
     "public": "Offentlig",
     "publicDescription": "Offentlig for alle brukere",
     "featured": "Fremhevet",
-    "featuredDescription": "Fremhevet for alle brukere"
+    "featuredDescription": "Fremhevet for alle brukere",
+    "authorizationRequired": "Autorisasjon kreves",
+    "authorizationRequiredDescription": "{serverName} krever autorisasjon for å få tilgang til verktøyene.",
+    "authorize": "Autoriser",
+    "authorizationFailed": "Autorisasjon feilet"
   }
 }

messages/zh.json

@@ -357,7 +357,11 @@
     "mcpServerCustomization": "MCP 服务器自定义",
     "mcpServerCustomizationDescription": "MCP 服务器自定义指令将在 MCP 服务器可用时添加到系统提示中。",
     "toolCustomizationInstructionsPlaceholder": "工具自定义指令不可用。",
-    "mcpServerCustomizationPlaceholder": "例如：如果输入值是电子邮件，始终以 example@example.com 格式输入电子邮件。"
+    "mcpServerCustomizationPlaceholder": "例如：如果输入值是电子邮件，始终以 example@example.com 格式输入电子邮件。",
+    "authorizationRequired": "需要授权",
+    "authorizationRequiredDescription": "{serverName} 需要授权才能访问其工具。",
+    "authorize": "授权",
+    "authorizationFailed": "授权失败"
   },
   "Error": {},
   "Info": {

src/components/mcp-auth-dialog.tsx

@@ -0,0 +1,73 @@
+"use client";
+
+import { useCallback, useState, useEffect } from "react";
+import { mutate } from "swr";
+import { toast } from "sonner";
+import { Loader } from "lucide-react";
+import { useTranslations } from "next-intl";
+import {
+  Dialog,
+  DialogClose,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+} from "ui/dialog";
+import { Button } from "ui/button";
+import { redriectMcpOauth } from "lib/ai/mcp/oauth-redirect";
+import type { MCPServerInfo } from "app-types/mcp";
+
+interface MCPAuthDialogProps {
+  server: MCPServerInfo | undefined;
+  serverName: string;
+}
+
+export function MCPAuthDialog({ server, serverName }: MCPAuthDialogProps) {
+  const t = useTranslations();
+  const [loading, setLoading] = useState(false);
+  const [dismissed, setDismissed] = useState(false);
+
+  const needsAuth = server?.status === "authorizing";
+
+  useEffect(() => {
+    if (!needsAuth) {
+      setDismissed(false);
+    }
+  }, [needsAuth]);
+
+  const handleAuthorize = useCallback(() => {
+    if (!server) return;
+    setLoading(true);
+    redriectMcpOauth(server.id)
+      .then(() => mutate("/api/mcp/list"))
+      .catch((err) => toast.error(err?.message || t("MCP.authorizationFailed")))
+      .finally(() => setLoading(false));
+  }, [server, t]);
+
+  if (!server || !needsAuth || dismissed) return null;
+
+  return (
+    <Dialog open onOpenChange={(open) => !open && setDismissed(true)}>
+      <DialogContent className="sm:max-w-md">
+        <DialogHeader>
+          <DialogTitle>{t("MCP.authorizationRequired")}</DialogTitle>
+          <DialogDescription>
+            {t("MCP.authorizationRequiredDescription", { serverName })}
+          </DialogDescription>
+        </DialogHeader>
+        <DialogFooter>
+          <DialogClose asChild>
+            <Button variant="outline" disabled={loading}>
+              {t("Common.cancel")}
+            </Button>
+          </DialogClose>
+          <Button onClick={handleAuthorize} disabled={loading} autoFocus>
+            {t("MCP.authorize")}
+            {loading && <Loader className="size-4 animate-spin" />}
+          </Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
+  );
+}

src/components/message-parts.tsx

@@ -1,5 +1,6 @@
 "use client";
 
+import { mutate } from "swr";
 import { FileUIPart, getToolName, ToolUIPart, UIMessage } from "ai";
 import {
   Check,
@@ -43,6 +44,8 @@ import { ChatMetadata, ChatModel, ManualToolConfirmTag } from "app-types/chat";
 import { useTranslations } from "next-intl";
 import { extractMCPToolId } from "lib/ai/mcp/mcp-tool-id";
 import { Separator } from "ui/separator";
+import { useMcpList } from "@/hooks/queries/use-mcp-list";
+import { MCPAuthDialog } from "./mcp-auth-dialog";
 
 import { TextShimmer } from "ui/text-shimmer";
 import equal from "lib/equal";
@@ -755,6 +758,7 @@ export const ToolMessagePart = memo(
     isManualToolInvocation,
   }: ToolMessagePartProps) => {
     const t = useTranslations("");
+    const { data: mcpList } = useMcpList();
 
     const { output, toolCallId, state, input, errorText } = part;
 
@@ -935,6 +939,16 @@ export const ToolMessagePart = memo(
       return extractMCPToolId(toolName);
     }, [toolName]);
 
+    const mcpServer = useMemo(() => {
+      return mcpList?.find((s) => s.name === mcpServerName);
+    }, [mcpList, mcpServerName]);
+
+    useEffect(() => {
+      if (mcpServerName && (output as any)?.isError) {
+        mutate("/api/mcp/list");
+      }
+    }, [mcpServerName, output]);
+
     const isExpanded = useMemo(() => {
       return expanded || result === null || isWorkflowTool;
     }, [expanded, result, isWorkflowTool]);
@@ -949,6 +963,8 @@ export const ToolMessagePart = memo(
 
     return (
       <div className="group w-full">
+        <MCPAuthDialog server={mcpServer} serverName={mcpServerName} />
+
         {CustomToolComponent ? (
           CustomToolComponent
         ) : (

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

@@ -26,6 +26,7 @@ import {
 import { safe } from "ts-safe";
 import { BASE_URL, IS_MCP_SERVER_REMOTE_ONLY, IS_VERCEL_ENV } from "lib/const";
 import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
+import { InvalidGrantError } from "@modelcontextprotocol/sdk/server/auth/errors.js";
 import { PgOAuthClientProvider } from "./pg-oauth-provider";
 import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 
@@ -54,6 +55,7 @@ export class MCPClient {
   toolInfo: MCPToolInfo[] = [];
   private disconnectDebounce = createDebounce();
   private needOauthProvider = false;
+  private hasTriedRefreshTokenRecovery = false;
   private inProgressToolCallIds: string[] = [];
   constructor(
     private id: string,
@@ -258,6 +260,20 @@ export class MCPClient {
             return this.connect(oauthState); // Recursive call with OAuth
           }
 
+          if (
+            !this.hasTriedRefreshTokenRecovery &&
+            isRefreshTokenError(streamableHttpError)
+          ) {
+            this.logger.info(
+              "Refresh token expired during connect, clearing credentials and retrying",
+            );
+            this.hasTriedRefreshTokenRecovery = true;
+            await this.oauthProvider?.invalidateCredentials("all");
+            this.locker.unlock();
+            await this.disconnect();
+            return this.connect(oauthState);
+          }
+
           if (!isOAuthAuthorizationRequired(streamableHttpError)) {
             this.logger.warn(
               `Streamable HTTP connection failed, Because ${streamableHttpError.message}, falling back to SSE transport`,
@@ -302,6 +318,7 @@ export class MCPClient {
       );
       this.client = client;
       this.isConnected = true;
+      this.hasTriedRefreshTokenRecovery = false;
 
       this.scheduleAutoDisconnect();
     } catch (error) {
@@ -374,6 +391,21 @@ export class MCPClient {
           await this.disconnect();
           return execute();
         }
+        if (isSessionNotFoundError(err)) {
+          this.logger.info("Session not found, reconnecting...");
+          await this.disconnect();
+          return execute();
+        }
+        if (isRefreshTokenError(err)) {
+          this.logger.info("Refresh token expired, requires re-authorization");
+          await this.oauthProvider?.invalidateCredentials("all");
+          await this.disconnect();
+          try {
+            await this.connect();
+          } catch {
+            // Expected - triggers OAuth flow
+          }
+        }
         throw err;
       })
       .ifOk((v) => {
@@ -446,3 +478,18 @@ function isUnauthorized(error: any): boolean {
 function isOAuthAuthorizationRequired(error: any): boolean {
   return error instanceof OAuthAuthorizationRequiredError;
 }
+
+function isRefreshTokenError(error: any): boolean {
+  return (
+    error instanceof InvalidGrantError ||
+    error?.message?.includes("Refresh token") ||
+    error?.message?.includes("invalid_grant")
+  );
+}
+
+function isSessionNotFoundError(error: any): boolean {
+  return (
+    (error as any)?.code === 404 ||
+    error?.message?.includes("Session not found")
+  );
+}