Skip to content

Commit 8d29511

Browse files
hakonharneshakonharnes
committed
feat: improve MCP OAuth session handling and auth UX
1 parent e914a30 commit 8d29511

10 files changed

Lines changed: 171 additions & 7 deletions

File tree

‎messages/en.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -794,6 +794,10 @@
794794
"public": "Featured",
795795
"publicDescription": "Featured for all users to use",
796796
"featured": "Featured",
797-
"featuredDescription": "Featured for all users to use"
797+
"featuredDescription": "Featured for all users to use",
798+
"authorizationRequired": "Authorization Required",
799+
"authorizationRequiredDescription": "{serverName} requires authorization to access its tools.",
800+
"authorize": "Authorize",
801+
"authorizationFailed": "Authorization failed"
798802
}
799803
}

‎messages/es.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -358,7 +358,11 @@
358358
"mcpServerCustomization": "Personalización del Servidor MCP",
359359
"mcpServerCustomizationDescription": "Las instrucciones de personalización del servidor MCP se añadirán al prompt del sistema cuando el servidor MCP esté disponible.",
360360
"toolCustomizationInstructionsPlaceholder": "Las instrucciones de personalización de herramientas no están disponibles.",
361-
"mcpServerCustomizationPlaceholder": "ej. Si el valor de entrada es email, siempre ingresa el email en el formato ejemplo@ejemplo.com."
361+
"mcpServerCustomizationPlaceholder": "ej. Si el valor de entrada es email, siempre ingresa el email en el formato ejemplo@ejemplo.com.",
362+
"authorizationRequired": "Autorización Requerida",
363+
"authorizationRequiredDescription": "{serverName} requiere autorización para acceder a sus herramientas.",
364+
"authorize": "Autorizar",
365+
"authorizationFailed": "Error de autorización"
362366
},
363367
"Error": {},
364368
"Info": {

‎messages/fr.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -358,7 +358,11 @@
358358
"mcpServerCustomization": "Personnalisation du Serveur MCP",
359359
"mcpServerCustomizationDescription": "Les instructions de personnalisation du serveur MCP seront ajoutées à l'invite système lorsque le serveur MCP sera disponible.",
360360
"toolCustomizationInstructionsPlaceholder": "Les instructions de personnalisation d'outils ne sont pas disponibles.",
361-
"mcpServerCustomizationPlaceholder": "ex. Si la valeur d'entrée est un email, toujours entrer l'email au format exemple@exemple.com."
361+
"mcpServerCustomizationPlaceholder": "ex. Si la valeur d'entrée est un email, toujours entrer l'email au format exemple@exemple.com.",
362+
"authorizationRequired": "Autorisation Requise",
363+
"authorizationRequiredDescription": "{serverName} nécessite une autorisation pour accéder à ses outils.",
364+
"authorize": "Autoriser",
365+
"authorizationFailed": "Échec de l'autorisation"
362366
},
363367
"Error": {},
364368
"Info": {

‎messages/ja.json‎

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

‎messages/ko.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -372,7 +372,11 @@
372372
"mcpServerCustomization": "MCP 서버 지침 설정",
373373
"mcpServerCustomizationDescription": "MCP 서버 사용자 지정 지침은 MCP서버를 사용할 때 시스템 프롬프트에 추가됩니다.",
374374
"toolCustomizationInstructionsPlaceholder": "도구 사용자 지정 지침이 없습니다.",
375-
"mcpServerCustomizationPlaceholder": "예) 이 MCP 서버의 툴중에 입력값이 email인 경우 항상 example@example.com 형식으로 입력해주세요."
375+
"mcpServerCustomizationPlaceholder": "예) 이 MCP 서버의 툴중에 입력값이 email인 경우 항상 example@example.com 형식으로 입력해주세요.",
376+
"authorizationRequired": "인증 필요",
377+
"authorizationRequiredDescription": "{serverName}의 도구에 액세스하려면 인증이 필요합니다.",
378+
"authorize": "인증",
379+
"authorizationFailed": "인증 실패"
376380
},
377381
"Error": {},
378382
"Info": {

‎messages/no.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -794,6 +794,10 @@
794794
"public": "Offentlig",
795795
"publicDescription": "Offentlig for alle brukere",
796796
"featured": "Fremhevet",
797-
"featuredDescription": "Fremhevet for alle brukere"
797+
"featuredDescription": "Fremhevet for alle brukere",
798+
"authorizationRequired": "Autorisasjon kreves",
799+
"authorizationRequiredDescription": "{serverName} krever autorisasjon for å få tilgang til verktøyene.",
800+
"authorize": "Autoriser",
801+
"authorizationFailed": "Autorisasjon feilet"
798802
}
799803
}

‎messages/zh.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -357,7 +357,11 @@
357357
"mcpServerCustomization": "MCP 服务器自定义",
358358
"mcpServerCustomizationDescription": "MCP 服务器自定义指令将在 MCP 服务器可用时添加到系统提示中。",
359359
"toolCustomizationInstructionsPlaceholder": "工具自定义指令不可用。",
360-
"mcpServerCustomizationPlaceholder": "例如：如果输入值是电子邮件，始终以 example@example.com 格式输入电子邮件。"
360+
"mcpServerCustomizationPlaceholder": "例如：如果输入值是电子邮件，始终以 example@example.com 格式输入电子邮件。",
361+
"authorizationRequired": "需要授权",
362+
"authorizationRequiredDescription": "{serverName} 需要授权才能访问其工具。",
363+
"authorize": "授权",
364+
"authorizationFailed": "授权失败"
361365
},
362366
"Error": {},
363367
"Info": {

‎src/components/mcp-auth-dialog.tsx‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
"use client";
2+
3+
import { useCallback, useState, useEffect } from "react";
4+
import { mutate } from "swr";
5+
import { toast } from "sonner";
6+
import { Loader } from "lucide-react";
7+
import { useTranslations } from "next-intl";
8+
import {
9+
Dialog,
10+
DialogClose,
11+
DialogContent,
12+
DialogDescription,
13+
DialogFooter,
14+
DialogHeader,
15+
DialogTitle,
16+
} from "ui/dialog";
17+
import { Button } from "ui/button";
18+
import { redriectMcpOauth } from "lib/ai/mcp/oauth-redirect";
19+
import type { MCPServerInfo } from "app-types/mcp";
20+
21+
interface MCPAuthDialogProps {
22+
server: MCPServerInfo | undefined;
23+
serverName: string;
24+
}
25+
26+
export function MCPAuthDialog({ server, serverName }: MCPAuthDialogProps) {
27+
const t = useTranslations();
28+
const [loading, setLoading] = useState(false);
29+
const [dismissed, setDismissed] = useState(false);
30+
31+
const needsAuth = server?.status === "authorizing";
32+
33+
useEffect(() => {
34+
if (!needsAuth) {
35+
setDismissed(false);
36+
}
37+
}, [needsAuth]);
38+
39+
const handleAuthorize = useCallback(() => {
40+
if (!server) return;
41+
setLoading(true);
42+
redriectMcpOauth(server.id)
43+
.then(() => mutate("/api/mcp/list"))
44+
.catch((err) => toast.error(err?.message || t("MCP.authorizationFailed")))
45+
.finally(() => setLoading(false));
46+
}, [server, t]);
47+
48+
if (!server || !needsAuth || dismissed) return null;
49+
50+
return (
51+
<Dialog open onOpenChange={(open) => !open && setDismissed(true)}>
52+
<DialogContent className="sm:max-w-md">
53+
<DialogHeader>
54+
<DialogTitle>{t("MCP.authorizationRequired")}</DialogTitle>
55+
<DialogDescription>
56+
{t("MCP.authorizationRequiredDescription", { serverName })}
57+
</DialogDescription>
58+
</DialogHeader>
59+
<DialogFooter>
60+
<DialogClose asChild>
61+
<Button variant="outline" disabled={loading}>
62+
{t("Common.cancel")}
63+
</Button>
64+
</DialogClose>
65+
<Button onClick={handleAuthorize} disabled={loading} autoFocus>
66+
{t("MCP.authorize")}
67+
{loading && <Loader className="size-4 animate-spin" />}
68+
</Button>
69+
</DialogFooter>
70+
</DialogContent>
71+
</Dialog>
72+
);
73+
}

‎src/components/message-parts.tsx‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"use client";
22

3+
import { mutate } from "swr";
34
import { FileUIPart, getToolName, ToolUIPart, UIMessage } from "ai";
45
import {
56
Check,
@@ -43,6 +44,8 @@ import { ChatMetadata, ChatModel, ManualToolConfirmTag } from "app-types/chat";
4344
import { useTranslations } from "next-intl";
4445
import { extractMCPToolId } from "lib/ai/mcp/mcp-tool-id";
4546
import { Separator } from "ui/separator";
47+
import { useMcpList } from "@/hooks/queries/use-mcp-list";
48+
import { MCPAuthDialog } from "./mcp-auth-dialog";
4649

4750
import { TextShimmer } from "ui/text-shimmer";
4851
import equal from "lib/equal";
@@ -755,6 +758,7 @@ export const ToolMessagePart = memo(
755758
isManualToolInvocation,
756759
}: ToolMessagePartProps) => {
757760
const t = useTranslations("");
761+
const { data: mcpList } = useMcpList();
758762

759763
const { output, toolCallId, state, input, errorText } = part;
760764

@@ -935,6 +939,16 @@ export const ToolMessagePart = memo(
935939
return extractMCPToolId(toolName);
936940
}, [toolName]);
937941

942+
const mcpServer = useMemo(() => {
943+
return mcpList?.find((s) => s.name === mcpServerName);
944+
}, [mcpList, mcpServerName]);
945+
946+
useEffect(() => {
947+
if (mcpServerName && (output as any)?.isError) {
948+
mutate("/api/mcp/list");
949+
}
950+
}, [mcpServerName, output]);
951+
938952
const isExpanded = useMemo(() => {
939953
return expanded || result === null || isWorkflowTool;
940954
}, [expanded, result, isWorkflowTool]);
@@ -949,6 +963,8 @@ export const ToolMessagePart = memo(
949963

950964
return (
951965
<div className="group w-full">
966+
<MCPAuthDialog server={mcpServer} serverName={mcpServerName} />
967+
952968
{CustomToolComponent ? (
953969
CustomToolComponent
954970
) : (

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

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import {
2626
import { safe } from "ts-safe";
2727
import { BASE_URL, IS_MCP_SERVER_REMOTE_ONLY, IS_VERCEL_ENV } from "lib/const";
2828
import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
29+
import { InvalidGrantError } from "@modelcontextprotocol/sdk/server/auth/errors.js";
2930
import { PgOAuthClientProvider } from "./pg-oauth-provider";
3031
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
3132

@@ -54,6 +55,7 @@ export class MCPClient {
5455
toolInfo: MCPToolInfo[] = [];
5556
private disconnectDebounce = createDebounce();
5657
private needOauthProvider = false;
58+
private hasTriedRefreshTokenRecovery = false;
5759
private inProgressToolCallIds: string[] = [];
5860
constructor(
5961
private id: string,
@@ -258,6 +260,20 @@ export class MCPClient {
258260
return this.connect(oauthState); // Recursive call with OAuth
259261
}
260262

263+
if (
264+
!this.hasTriedRefreshTokenRecovery &&
265+
isRefreshTokenError(streamableHttpError)
266+
) {
267+
this.logger.info(
268+
"Refresh token expired during connect, clearing credentials and retrying",
269+
);
270+
this.hasTriedRefreshTokenRecovery = true;
271+
await this.oauthProvider?.invalidateCredentials("all");
272+
this.locker.unlock();
273+
await this.disconnect();
274+
return this.connect(oauthState);
275+
}
276+
261277
if (!isOAuthAuthorizationRequired(streamableHttpError)) {
262278
this.logger.warn(
263279
`Streamable HTTP connection failed, Because ${streamableHttpError.message}, falling back to SSE transport`,
@@ -302,6 +318,7 @@ export class MCPClient {
302318
);
303319
this.client = client;
304320
this.isConnected = true;
321+
this.hasTriedRefreshTokenRecovery = false;
305322

306323
this.scheduleAutoDisconnect();
307324
} catch (error) {
@@ -374,6 +391,21 @@ export class MCPClient {
374391
await this.disconnect();
375392
return execute();
376393
}
394+
if (isSessionNotFoundError(err)) {
395+
this.logger.info("Session not found, reconnecting...");
396+
await this.disconnect();
397+
return execute();
398+
}
399+
if (isRefreshTokenError(err)) {
400+
this.logger.info("Refresh token expired, requires re-authorization");
401+
await this.oauthProvider?.invalidateCredentials("all");
402+
await this.disconnect();
403+
try {
404+
await this.connect();
405+
} catch {
406+
// Expected - triggers OAuth flow
407+
}
408+
}
377409
throw err;
378410
})
379411
.ifOk((v) => {
@@ -446,3 +478,18 @@ function isUnauthorized(error: any): boolean {
446478
function isOAuthAuthorizationRequired(error: any): boolean {
447479
return error instanceof OAuthAuthorizationRequiredError;
448480
}
481+
482+
function isRefreshTokenError(error: any): boolean {
483+
return (
484+
error instanceof InvalidGrantError ||
485+
error?.message?.includes("Refresh token") ||
486+
error?.message?.includes("invalid_grant")
487+
);
488+
}
489+
490+
function isSessionNotFoundError(error: any): boolean {
491+
return (
492+
(error as any)?.code === 404 ||
493+
error?.message?.includes("Session not found")
494+
);
495+
}

0 commit comments

Comments
 (0)