Merge pull request #3224 from code-yeongyu/fix/oauth-wiring

fix(oauth): wire refresh mutex + post-request auth handler

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<ManagerModule> {
+  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)
+  })
+})

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<T> {
     const maxRetries = 3
     let lastError: Error | null = null
+    const refreshAttempted = new Set<string>()
 
     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
         }

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<OAuthHandlerModule> {
+  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<TValue> = {
+  promise: Promise<TValue>
+  resolve: (value: TValue) => void
+}
+
+function createDeferred<TValue>(): Deferred<TValue> {
+  let resolvePromise: ((value: TValue) => void) | null = null
+  const promise = new Promise<TValue>((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<OAuthTokenData>()
+    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<string, OAuthProviderLike>()
+    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<OAuthTokenData>()
+    const secondDeferred = createDeferred<OAuthTokenData>()
+    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<string, OAuthProviderLike>([["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)
+  })
+})

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)) {
+      if (tokenData && isTokenExpired(tokenData)) {
-      try {
-        tokenData = tokenData.refreshToken
-          ? await provider.refresh(tokenData.refreshToken)
-          : await provider.login()
-      } catch {
         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