diff --git a/src/features/skill-mcp-manager/oauth-handler.test.ts b/src/features/skill-mcp-manager/oauth-handler.test.ts new file mode 100644 index 000000000..d6eb317bc --- /dev/null +++ b/src/features/skill-mcp-manager/oauth-handler.test.ts @@ -0,0 +1,145 @@ +import { beforeEach, describe, expect, it, mock } from "bun:test" +import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" +import type { OAuthTokenData } from "../mcp-oauth/storage" +import type { OAuthProviderFactory, OAuthProviderLike } from "./types" + +mock.module("../mcp-oauth/provider", () => ({ + McpOAuthProvider: class MockMcpOAuthProvider {}, +})) + +type OAuthHandlerModule = typeof import("./oauth-handler") + +async function importFreshOAuthHandlerModule(): Promise { + return await import(new URL(`./oauth-handler.ts?oauth-handler-test=${Date.now()}-${Math.random()}`, import.meta.url).href) +} + +type Deferred = { + promise: Promise + resolve: (value: TValue) => void +} + +function createDeferred(): Deferred { + let resolvePromise: ((value: TValue) => void) | null = null + const promise = new Promise((resolve) => { + resolvePromise = resolve + }) + + if (!resolvePromise) { + throw new Error("Failed to create deferred promise") + } + + return { promise, resolve: resolvePromise } +} + +function createConfig(serverUrl: string): ClaudeCodeMcpServer { + return { + url: serverUrl, + oauth: { + clientId: "test-client", + }, + } +} + +describe("oauth-handler refresh mutex wiring", () => { + beforeEach(() => { + mock.restore() + }) + + it("deduplicates concurrent pre-request refresh attempts for the same server", async () => { + // given + const { buildHttpRequestInit } = await importFreshOAuthHandlerModule() + const deferred = createDeferred() + const refresh = mock(() => deferred.promise) + const provider: OAuthProviderLike = { + tokens: () => ({ + accessToken: "expired-token", + refreshToken: "refresh-token", + expiresAt: Math.floor(Date.now() / 1000) - 60, + }), + login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)), + refresh, + } + const authProviders = new Map() + const createOAuthProvider: OAuthProviderFactory = () => provider + + // when + const firstRequest = buildHttpRequestInit(createConfig("https://same.example.com/mcp"), authProviders, createOAuthProvider) + const secondRequest = buildHttpRequestInit(createConfig("https://same.example.com/mcp"), authProviders, createOAuthProvider) + + // then + expect(refresh).toHaveBeenCalledTimes(1) + deferred.resolve({ accessToken: "refreshed-token" }) + await expect(firstRequest).resolves.toEqual({ headers: { Authorization: "Bearer refreshed-token" } }) + await expect(secondRequest).resolves.toEqual({ headers: { Authorization: "Bearer refreshed-token" } }) + }) + + it("allows different servers to refresh independently after request auth errors", async () => { + // given + const { handlePostRequestAuthError } = await importFreshOAuthHandlerModule() + const firstDeferred = createDeferred() + const secondDeferred = createDeferred() + const firstProvider: OAuthProviderLike = { + tokens: () => ({ accessToken: "expired-a", refreshToken: "refresh-a" }), + login: mock(async () => ({ accessToken: "login-a" } satisfies OAuthTokenData)), + refresh: mock(() => firstDeferred.promise), + } + const secondProvider: OAuthProviderLike = { + tokens: () => ({ accessToken: "expired-b", refreshToken: "refresh-b" }), + login: mock(async () => ({ accessToken: "login-b" } satisfies OAuthTokenData)), + refresh: mock(() => secondDeferred.promise), + } + const providers = new Map([ + ["https://server-a.example.com/mcp", firstProvider], + ["https://server-b.example.com/mcp", secondProvider], + ]) + + // when + const firstAttempt = handlePostRequestAuthError({ + error: new Error("401 Unauthorized"), + config: createConfig("https://server-a.example.com/mcp"), + authProviders: providers, + }) + const secondAttempt = handlePostRequestAuthError({ + error: new Error("403 Forbidden"), + config: createConfig("https://server-b.example.com/mcp"), + authProviders: providers, + }) + + // then + expect(firstProvider.refresh).toHaveBeenCalledTimes(1) + expect(secondProvider.refresh).toHaveBeenCalledTimes(1) + firstDeferred.resolve({ accessToken: "refreshed-a" }) + secondDeferred.resolve({ accessToken: "refreshed-b" }) + await expect(firstAttempt).resolves.toBe(true) + await expect(secondAttempt).resolves.toBe(true) + }) + + it("allows a new refresh after the previous same-server refresh completes", async () => { + // given + const { handlePostRequestAuthError } = await importFreshOAuthHandlerModule() + const refresh = mock(async () => ({ accessToken: `refreshed-${refresh.mock.calls.length + 1}` } satisfies OAuthTokenData)) + const provider: OAuthProviderLike = { + tokens: () => ({ accessToken: "expired-token", refreshToken: "refresh-token" }), + login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)), + refresh, + } + const authProviders = new Map([["https://same.example.com/mcp", provider]]) + + // when + const firstResult = await handlePostRequestAuthError({ + error: new Error("401 Unauthorized"), + config: createConfig("https://same.example.com/mcp"), + authProviders, + }) + const secondResult = await handlePostRequestAuthError({ + error: new Error("401 Unauthorized"), + config: createConfig("https://same.example.com/mcp"), + authProviders, + }) + + // then + expect(firstResult).toBe(true) + expect(secondResult).toBe(true) + expect(refresh).toHaveBeenCalledTimes(2) + }) +}) diff --git a/src/features/skill-mcp-manager/oauth-handler.ts b/src/features/skill-mcp-manager/oauth-handler.ts index d1b2b7513..63f3d8676 100644 --- a/src/features/skill-mcp-manager/oauth-handler.ts +++ b/src/features/skill-mcp-manager/oauth-handler.ts @@ -1,5 +1,6 @@ import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" import { McpOAuthProvider } from "../mcp-oauth/provider" +import { withRefreshMutex } from "../mcp-oauth/refresh-mutex" import type { OAuthTokenData } from "../mcp-oauth/storage" import { isStepUpRequired, mergeScopes } from "../mcp-oauth/step-up" import type { OAuthProviderFactory, OAuthProviderLike } from "./types" @@ -52,14 +53,15 @@ export async function buildHttpRequestInit( } } - if (tokenData && isTokenExpired(tokenData)) { - try { - tokenData = tokenData.refreshToken - ? await provider.refresh(tokenData.refreshToken) - : await provider.login() - } catch { + if (tokenData && isTokenExpired(tokenData)) { try { - tokenData = await provider.login() + const refreshToken = tokenData.refreshToken + tokenData = refreshToken + ? await withRefreshMutex(config.url, () => provider.refresh(refreshToken)) + : await provider.login() + } catch { + try { + tokenData = await provider.login() } catch { tokenData = null } @@ -149,7 +151,8 @@ export async function handlePostRequestAuthError(params: { refreshAttempted.add(config.url) try { - await provider.refresh(tokenData.refreshToken) + const refreshToken = tokenData.refreshToken + await withRefreshMutex(config.url, () => provider.refresh(refreshToken)) return true } catch { return false