diff --git a/src/features/skill-mcp-manager/manager-oauth-retry.test.ts b/src/features/skill-mcp-manager/manager-oauth-retry.test.ts new file mode 100644 index 000000000..f887e80e7 --- /dev/null +++ b/src/features/skill-mcp-manager/manager-oauth-retry.test.ts @@ -0,0 +1,162 @@ +import { afterAll, 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 { SkillMcpClientInfo, SkillMcpServerContext } from "./types" + +const mockGetOrCreateClient = mock(async () => { + throw new Error("not used") +}) + +const mockGetOrCreateClientWithRetryImpl = mock(async () => ({ + callTool: mock(async () => ({ content: [{ type: "text", text: "unused" }] })), + close: mock(async () => {}), +})) + +type ManagerModule = typeof import("./manager") + +async function importFreshManagerModule(): Promise { + mock.module("./connection", () => ({ + getOrCreateClient: mockGetOrCreateClient, + getOrCreateClientWithRetryImpl: mockGetOrCreateClientWithRetryImpl, + })) + + mock.module("../mcp-oauth/provider", () => ({ + McpOAuthProvider: class MockMcpOAuthProvider {}, + })) + + return await import(new URL(`./manager.ts?oauth-retry-test=${Date.now()}-${Math.random()}`, import.meta.url).href) +} + +function createInfo(): SkillMcpClientInfo { + return { + serverName: "oauth-server", + skillName: "oauth-skill", + sessionID: "session-1", + scope: "builtin", + } +} + +function createContext(): SkillMcpServerContext { + return { + skillName: "oauth-skill", + config: { + url: "https://mcp.example.com/mcp", + oauth: { clientId: "test-client" }, + } satisfies ClaudeCodeMcpServer, + } +} + +afterAll(() => { + mock.restore() +}) + +describe("SkillMcpManager post-request OAuth retry", () => { + beforeEach(() => { + mockGetOrCreateClient.mockClear() + mockGetOrCreateClientWithRetryImpl.mockClear() + }) + + it("retries the operation after a 401 refresh succeeds", async () => { + // given + const { SkillMcpManager } = await importFreshManagerModule() + const refresh = mock(async () => ({ accessToken: "refreshed-token" } satisfies OAuthTokenData)) + const manager = new SkillMcpManager({ + createOAuthProvider: () => ({ + tokens: () => ({ accessToken: "stale-token", refreshToken: "refresh-token" }), + login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)), + refresh, + }), + }) + const callTool = mock(async () => { + if (callTool.mock.calls.length === 1) { + throw new Error("401 Unauthorized") + } + + return { content: [{ type: "text", text: "success" }] } + }) + mockGetOrCreateClientWithRetryImpl.mockResolvedValue({ callTool, close: mock(async () => {}) }) + + // when + const result = await manager.callTool(createInfo(), createContext(), "test-tool", {}) + + // then + expect(result).toEqual([{ type: "text", text: "success" }]) + expect(refresh).toHaveBeenCalledTimes(1) + expect(callTool).toHaveBeenCalledTimes(2) + }) + + it("retries the operation after a 403 refresh succeeds without step-up scope", async () => { + // given + const { SkillMcpManager } = await importFreshManagerModule() + const refresh = mock(async () => ({ accessToken: "refreshed-token" } satisfies OAuthTokenData)) + const manager = new SkillMcpManager({ + createOAuthProvider: () => ({ + tokens: () => ({ accessToken: "stale-token", refreshToken: "refresh-token" }), + login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)), + refresh, + }), + }) + const callTool = mock(async () => { + if (callTool.mock.calls.length === 1) { + throw new Error("403 Forbidden") + } + + return { content: [{ type: "text", text: "success" }] } + }) + mockGetOrCreateClientWithRetryImpl.mockResolvedValue({ callTool, close: mock(async () => {}) }) + + // when + const result = await manager.callTool(createInfo(), createContext(), "test-tool", {}) + + // then + expect(result).toEqual([{ type: "text", text: "success" }]) + expect(refresh).toHaveBeenCalledTimes(1) + expect(callTool).toHaveBeenCalledTimes(2) + }) + + it("propagates the auth error without retry when refresh fails", async () => { + // given + const { SkillMcpManager } = await importFreshManagerModule() + const refresh = mock(async () => { + throw new Error("refresh failed") + }) + const manager = new SkillMcpManager({ + createOAuthProvider: () => ({ + tokens: () => ({ accessToken: "stale-token", refreshToken: "refresh-token" }), + login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)), + refresh, + }), + }) + const callTool = mock(async () => { + throw new Error("401 Unauthorized") + }) + mockGetOrCreateClientWithRetryImpl.mockResolvedValue({ callTool, close: mock(async () => {}) }) + + // when / then + await expect(manager.callTool(createInfo(), createContext(), "test-tool", {})).rejects.toThrow("401 Unauthorized") + expect(refresh).toHaveBeenCalledTimes(1) + expect(callTool).toHaveBeenCalledTimes(1) + }) + + it("only attempts one refresh when the retried operation returns 401 again", async () => { + // given + const { SkillMcpManager } = await importFreshManagerModule() + const refresh = mock(async () => ({ accessToken: "refreshed-token" } satisfies OAuthTokenData)) + const manager = new SkillMcpManager({ + createOAuthProvider: () => ({ + tokens: () => ({ accessToken: "stale-token", refreshToken: "refresh-token" }), + login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)), + refresh, + }), + }) + const callTool = mock(async () => { + throw new Error("401 Unauthorized") + }) + mockGetOrCreateClientWithRetryImpl.mockResolvedValue({ callTool, close: mock(async () => {}) }) + + // when / then + await expect(manager.callTool(createInfo(), createContext(), "test-tool", {})).rejects.toThrow("401 Unauthorized") + expect(refresh).toHaveBeenCalledTimes(1) + expect(callTool).toHaveBeenCalledTimes(2) + }) +}) diff --git a/src/features/skill-mcp-manager/manager.ts b/src/features/skill-mcp-manager/manager.ts index 473d5f390..f91524be4 100644 --- a/src/features/skill-mcp-manager/manager.ts +++ b/src/features/skill-mcp-manager/manager.ts @@ -4,7 +4,7 @@ import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" import { McpOAuthProvider } from "../mcp-oauth/provider" import { disconnectAll, disconnectSession, forceReconnect } from "./cleanup" import { getOrCreateClient, getOrCreateClientWithRetryImpl } from "./connection" -import { handleStepUpIfNeeded } from "./oauth-handler" +import { handlePostRequestAuthError, handleStepUpIfNeeded } from "./oauth-handler" import type { OAuthProviderFactory, SkillMcpClientInfo, @@ -110,6 +110,7 @@ export class SkillMcpManager { ): Promise { const maxRetries = 3 let lastError: Error | null = null + const refreshAttempted = new Set() for (let attempt = 1; attempt <= maxRetries; attempt++) { try { @@ -130,6 +131,17 @@ export class SkillMcpManager { continue } + const postRequestRefreshHandled = await handlePostRequestAuthError({ + error: lastError, + config, + authProviders: this.state.authProviders, + createOAuthProvider: this.state.createOAuthProvider, + refreshAttempted, + }) + if (postRequestRefreshHandled) { + continue + } + if (!errorMessage.includes("not connected")) { throw lastError } 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..35823c6ae --- /dev/null +++ b/src/features/skill-mcp-manager/oauth-handler.test.ts @@ -0,0 +1,141 @@ +import { 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" + +type OAuthHandlerModule = typeof import("./oauth-handler") + +async function importFreshOAuthHandlerModule(): Promise { + mock.module("../mcp-oauth/provider", () => ({ + McpOAuthProvider: class MockMcpOAuthProvider {}, + })) + + 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", () => { + 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