diff --git a/src/mcp/client-auth-policy.test.ts b/src/mcp/client-auth-policy.test.ts index 54bc14f0b..8c2e93ccc 100644 --- a/src/mcp/client-auth-policy.test.ts +++ b/src/mcp/client-auth-policy.test.ts @@ -17,6 +17,7 @@ let blockTokenExchange = false; let tokenExchangeSignals: (AbortSignal | null | undefined)[] = []; let tokenExchangeAborts = 0; let lastTransportAuth: (() => Promise) | undefined; +let lastTransportRedirect: (() => Promise) | undefined; let tokenRefreshSignals: (AbortSignal | null | undefined)[] = []; let tokenRefreshAborts = 0; @@ -35,7 +36,11 @@ function hangUntilAbort( }); } -const authProvider = { resetAuthorization: async () => undefined }; +const redirectToAuthorization = () => undefined; +const authProvider = { + resetAuthorization: async () => undefined, + redirectToAuthorization, +}; await withMockedModule( import.meta.resolve("@modelcontextprotocol/sdk/client/index.js"), @@ -43,7 +48,10 @@ await withMockedModule( ...real, Client: class { async connect(): Promise { - if (clientConnectError !== undefined) throw clientConnectError; + if (clientConnectError !== undefined) { + if (clientConnectError instanceof UnauthorizedError) await lastTransportRedirect?.(); + throw clientConnectError; + } } async listTools( _params?: unknown, @@ -83,12 +91,18 @@ await withMockedModule( constructor( _url: URL, private readonly options?: { + authProvider?: { redirectToAuthorization?: (url: URL) => void | Promise }; requestInit?: RequestInit; fetch?: (url: string | URL, init?: RequestInit) => Promise; }, ) { transportOptions.push(options); lastTransportAuth = () => this.auth(); + lastTransportRedirect = async () => { + await this.options?.authProvider?.redirectToAuthorization?.( + new URL("https://auth.test/authorize"), + ); + }; } async finishAuth(): Promise { const signal = this.options?.requestInit?.signal; @@ -174,8 +188,10 @@ describe("HTTP MCP auth policy", () => { tokenExchangeSignals = []; tokenExchangeAborts = 0; lastTransportAuth = undefined; + lastTransportRedirect = undefined; tokenRefreshSignals = []; tokenRefreshAborts = 0; + authProvider.redirectToAuthorization = redirectToAuthorization; }); test("built-in anonymous Exa treats 401 as a normal failure without OAuth machinery", async () => { @@ -243,7 +259,8 @@ describe("HTTP MCP auth policy", () => { ); while (tokenExchangeSignals.length === 0) await Promise.resolve(); - expect(tokenExchangeSignals[0]).toBe(abort.signal); + expect(tokenExchangeSignals[0]).toBeDefined(); + expect(tokenExchangeSignals[0]?.aborted).toBe(false); abort.abort(new Error("toolset disposed")); const result = await connection; @@ -264,12 +281,17 @@ describe("HTTP MCP auth policy", () => { expect(result.ok).toBe(true); if (!result.ok) return; expect(transportOptions).toEqual([ - { authProvider, requestInit: { signal: abort.signal }, fetch: expect.any(Function) }, + { + authProvider, + requestInit: { signal: expect.any(AbortSignal) }, + fetch: expect.any(Function), + }, ]); const call = result.client.call("ping", {}, abort.signal); while (tokenRefreshSignals.length === 0) await Promise.resolve(); - expect(tokenRefreshSignals[0]).toBe(abort.signal); + expect(tokenRefreshSignals[0]).toBeDefined(); + expect(tokenRefreshSignals[0]?.aborted).toBe(false); abort.abort(new Error("toolset disposed")); await expect(call).rejects.toThrow("toolset disposed"); expect(tokenRefreshAborts).toBe(1); @@ -294,7 +316,49 @@ describe("HTTP MCP auth policy", () => { expect(callbackStarts).toBe(1); expect(providerCreates).toBe(1); expect(providerServerURL).toBe("https://custom.example/mcp?mode=full"); - expect(transportOptions).toEqual([{ authProvider }]); + expect(transportOptions).toEqual([ + { + authProvider, + requestInit: { signal: expect.any(AbortSignal) }, + fetch: expect.any(Function), + }, + ]); + }); + + test("client close aborts in-flight OAuth without aborting the connect signal", async () => { + const connectAbort = new AbortController(); + const result = await connectMCPServer( + { name: "linear", type: "http", url: "https://mcp.linear.app/mcp" }, + { signal: connectAbort.signal }, + ); + expect(result.ok).toBe(true); + if (!result.ok) return; + + const transport = transportOptions[0] as { + requestInit?: { signal?: AbortSignal }; + fetch?: (url: string | URL, init?: RequestInit) => Promise; + }; + expect(transport.requestInit?.signal).toBeDefined(); + expect(transport.requestInit?.signal).not.toBe(connectAbort.signal); + expect(transport.fetch).toBeTypeOf("function"); + + const call = result.client.call("ping", {}, new AbortController().signal); + while (tokenRefreshSignals.length === 0) await Promise.resolve(); + expect(tokenRefreshSignals[0]).not.toBe(connectAbort.signal); + expect(tokenRefreshSignals[0]?.aborted).toBe(false); + expect(connectAbort.signal.aborted).toBe(false); + + await result.client.close(); + + await expect(call).rejects.toThrow(); + expect(tokenRefreshAborts).toBe(1); + expect(connectAbort.signal.aborted).toBe(false); + expect(transport.requestInit?.signal?.aborted).toBe(true); + const fetchFn = transport.fetch; + expect(fetchFn).toBeTypeOf("function"); + if (fetchFn === undefined) return; + await expect(fetchFn("https://auth.test/token")).rejects.toThrow(); + expect(connectAbort.signal.aborted).toBe(false); }); }); diff --git a/src/mcp/client-auth-reauth-cap.test.ts b/src/mcp/client-auth-reauth-cap.test.ts new file mode 100644 index 000000000..611f38234 --- /dev/null +++ b/src/mcp/client-auth-reauth-cap.test.ts @@ -0,0 +1,924 @@ +import { afterEach, beforeEach, describe, expect, setSystemTime, test } from "bun:test"; +import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"; +import { withMockedModule } from "../../tests/helpers/mock-module.js"; + +let finishAuthCalls = 0; +let finishAuthError: Error | undefined = new Error("finishAuth exploded"); +let connectFailuresLeft = 0; +let listFailuresLeft = 0; +let callFailuresLeft = 0; +let callToolCalls = 0; +let callRedirectsLeft = Number.POSITIVE_INFINITY; +let redirectsPerFailure = 1; +let redirectOnListFailure = false; +let redirectConcurrently = false; +let saveThenRedirectPair = false; +let overlappingSDKSaves = false; +let redirectVerifier: string | undefined; +let providerCreates = 0; +let refreshCalls = 0; +let refreshSucceeds = false; +let authEvents: string[] = []; +let authURLCount = 0; +let authorizedCount = 0; +let waitForCodeCalls = 0; +let storedCodeVerifier: string | undefined; +let exchangedCodeVerifier: string | undefined; +let emittedAuthURL: string | undefined; +let saveStarted = 0; +let refreshGate: Promise | undefined; +let releaseRefresh: (() => void) | undefined; +let callbackGate: Promise | undefined; +let releaseCallback: (() => void) | undefined; +let lastRequestSignal: AbortSignal | undefined; +let retryGate: Promise | undefined; +let releaseRetry: (() => void) | undefined; +let saveGate: Promise | undefined; +let releaseSave: (() => void) | undefined; +interface MockAuthProvider { + redirectToAuthorization?: (url: URL) => void | Promise; + saveCodeVerifier?: (codeVerifier: string) => void | Promise; + codeVerifier?: () => string | undefined; +} +let liveProvider: MockAuthProvider | undefined; + +const fakeProvider = { + resetAuthorization: async () => undefined, + tokens: () => ({ access_token: "stale", refresh_token: "refresh-me" }), + refreshToken: async () => { + refreshCalls += 1; + authEvents.push("refresh"); + await waitForOptionalGate(refreshGate, lastRequestSignal); + if (!refreshSucceeds) throw new UnauthorizedError("refresh rejected"); + return { access_token: "fresh", refresh_token: "refresh-me" }; + }, + saveCodeVerifier: async (codeVerifier: string) => { + storedCodeVerifier = codeVerifier; + saveStarted += 1; + await saveGate; + }, + codeVerifier: () => storedCodeVerifier, +}; + +async function saveThenRedirect( + provider: MockAuthProvider | undefined, + verifier: string, +): Promise { + await provider?.saveCodeVerifier?.(verifier); + await provider?.redirectToAuthorization?.( + new URL(`https://auth.test/authorize?v=${encodeURIComponent(verifier)}`), + ); +} + +async function emitRedirects(provider: MockAuthProvider | undefined): Promise { + if (overlappingSDKSaves) { + const first = saveThenRedirect(provider, "v1"); + while (saveStarted === 0) await Promise.resolve(); + const second = saveThenRedirect(provider, "v2"); + await Promise.resolve(); + releaseSave?.(); + await Promise.all([first, second]); + return; + } + if (saveThenRedirectPair) { + const first = saveThenRedirect(provider, "v1"); + await Promise.resolve(); + await Promise.all([first, saveThenRedirect(provider, "v2")]); + return; + } + if (redirectVerifier !== undefined) { + const verifier = redirectVerifier; + redirectVerifier = undefined; + await saveThenRedirect(provider, verifier); + return; + } + const redirect = () => + provider?.redirectToAuthorization?.(new URL("https://auth.test/authorize")); + if (redirectConcurrently) { + await Promise.all(Array.from({ length: redirectsPerFailure }, redirect)); + } else { + for (let call = 0; call < redirectsPerFailure; call += 1) await redirect(); + } +} + +function waitForOptionalGate( + gate: Promise | undefined, + signal: AbortSignal | undefined, +): Promise { + if (gate === undefined) return Promise.resolve(); + if (signal === undefined) return gate; + if (signal.aborted) return Promise.reject(signal.reason); + return new Promise((resolve, reject) => { + const onAbort = () => { + signal.removeEventListener("abort", onAbort); + reject(signal.reason); + }; + signal.addEventListener("abort", onAbort, { once: true }); + void gate.then(() => { + signal.removeEventListener("abort", onAbort); + resolve(); + }); + }); +} + +function waitForGate(signal: AbortSignal): Promise { + return waitForOptionalGate(callbackGate, signal); +} + +await withMockedModule( + import.meta.resolve("@modelcontextprotocol/sdk/client/index.js"), + (real: typeof import("@modelcontextprotocol/sdk/client/index.js")) => ({ + ...real, + Client: class { + async connect(transport?: { provider?: MockAuthProvider }): Promise { + liveProvider = transport?.provider; + if (connectFailuresLeft > 0) { + connectFailuresLeft -= 1; + await emitRedirects(transport?.provider); + throw new UnauthorizedError("authorization required"); + } + } + async listTools(): Promise<{ tools: [] }> { + if (listFailuresLeft > 0) { + listFailuresLeft -= 1; + if (redirectOnListFailure) { + await liveProvider?.redirectToAuthorization?.(new URL("https://auth.test/authorize")); + } + throw new UnauthorizedError("authorization required"); + } + return { tools: [] }; + } + async callTool(): Promise<{ content: [] }> { + callToolCalls += 1; + if (callFailuresLeft > 0) { + callFailuresLeft -= 1; + if (callRedirectsLeft > 0) { + callRedirectsLeft -= 1; + await emitRedirects(liveProvider); + } + throw new UnauthorizedError("authorization required"); + } + await waitForOptionalGate(retryGate, lastRequestSignal); + return { content: [] }; + } + async close(): Promise {} + }, + }), +); + +await withMockedModule( + import.meta.resolve("@modelcontextprotocol/sdk/client/streamableHttp.js"), + (real: typeof import("@modelcontextprotocol/sdk/client/streamableHttp.js")) => ({ + ...real, + StreamableHTTPClientTransport: class { + provider?: MockAuthProvider; + constructor( + _url: URL, + options?: { authProvider?: MockAuthProvider; requestInit?: RequestInit }, + ) { + if (options?.authProvider !== undefined) this.provider = options.authProvider; + const signal = options?.requestInit?.signal; + if (signal !== undefined && signal !== null) lastRequestSignal = signal; + } + async finishAuth(): Promise { + finishAuthCalls += 1; + exchangedCodeVerifier = this.provider?.codeVerifier?.(); + if (finishAuthError !== undefined) throw finishAuthError; + } + get sessionId(): string | undefined { + return undefined; + } + }, + }), +); + +await withMockedModule( + import.meta.resolve("./callback-server.js"), + (real: typeof import("./callback-server.js")) => ({ + ...real, + startCallbackServer: async () => ({ + redirectUrl: "http://127.0.0.1:12345/callback", + expectState: () => undefined, + waitForCode: async (signal: AbortSignal) => { + waitForCodeCalls += 1; + await waitForGate(signal); + return "code"; + }, + close: () => undefined, + }), + }), +); + +await withMockedModule( + import.meta.resolve("./oauth-provider.js"), + (real: typeof import("./oauth-provider.js")) => ({ + ...real, + createOAuthProvider: async (options: { + serverName: string; + onAuthURL: (serverName: string, authorizationUrl: string) => void; + }) => { + providerCreates += 1; + return { + ...fakeProvider, + redirectToAuthorization: (url: URL) => { + options.onAuthURL(options.serverName, url.toString()); + }, + }; + }, + }), +); + +const { + connectMCPServer, + resetBrowserAuthState, + setBrowserAuthWaitMs, + MAX_BROWSER_AUTH_ATTEMPTS, + BROWSER_AUTH_COOLDOWN_MS, +} = await import("./client.js"); + +const config = { name: "linear", type: "http" as const, url: "https://mcp.linear.app/mcp" }; + +async function connectWithAuthPrompt(): Promise<{ ok: boolean; error?: string }> { + const result = await connectMCPServer(config, { + onAuthURL: () => { + authURLCount += 1; + authEvents.push("authURL"); + }, + }); + return result.ok ? { ok: true } : { ok: false, error: result.error }; +} + +describe("HTTP MCP re-auth loop prevention", () => { + beforeEach(() => { + connectFailuresLeft = 0; + listFailuresLeft = 0; + callFailuresLeft = 0; + callToolCalls = 0; + callRedirectsLeft = Number.POSITIVE_INFINITY; + redirectsPerFailure = 1; + redirectOnListFailure = false; + redirectConcurrently = false; + saveThenRedirectPair = false; + overlappingSDKSaves = false; + redirectVerifier = undefined; + finishAuthCalls = 0; + finishAuthError = new Error("finishAuth exploded"); + providerCreates = 0; + liveProvider = undefined; + refreshCalls = 0; + refreshSucceeds = false; + authEvents = []; + authURLCount = 0; + authorizedCount = 0; + waitForCodeCalls = 0; + storedCodeVerifier = undefined; + exchangedCodeVerifier = undefined; + emittedAuthURL = undefined; + saveStarted = 0; + refreshGate = undefined; + releaseRefresh = undefined; + callbackGate = undefined; + releaseCallback = undefined; + saveGate = undefined; + releaseSave = undefined; + lastRequestSignal = undefined; + retryGate = undefined; + releaseRetry = undefined; + resetBrowserAuthState(); + setSystemTime(); + }); + + afterEach(() => { + resetBrowserAuthState(); + setSystemTime(); + }); + + test("uses one custom refresh before prompting when recovery starts without a redirect", async () => { + refreshSucceeds = true; + redirectsPerFailure = 0; + connectFailuresLeft = 1; + + const result = await connectWithAuthPrompt(); + + expect(result.ok).toBe(true); + expect(authEvents).toEqual(["refresh"]); + expect(refreshCalls).toBe(1); + expect(finishAuthCalls).toBe(0); + expect(authURLCount).toBe(0); + }); + + test("rejects promptly when refresh and an auth probe fail without emitting a URL", async () => { + const connected = await connectMCPServer(config, { + onAuthURL: () => (authURLCount += 1), + onAuthorized: () => (authorizedCount += 1), + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + redirectsPerFailure = 0; + callFailuresLeft = 1; + listFailuresLeft = 1; + + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "authorization required", + ); + + expect(refreshCalls).toBe(1); + expect(authURLCount).toBe(0); + expect(waitForCodeCalls).toBe(0); + expect(authorizedCount).toBe(0); + }); + + test("shares live-call recovery across concurrent unauthorized calls", async () => { + finishAuthError = undefined; + const connected = await connectMCPServer(config, { + onAuthURL: () => (authURLCount += 1), + onAuthorized: () => (authorizedCount += 1), + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + redirectsPerFailure = 0; + callFailuresLeft = 2; + listFailuresLeft = 1; + redirectOnListFailure = true; + + const calls = [ + connected.client.call("first", { value: 1 }, new AbortController().signal), + connected.client.call("second", { value: 2 }, new AbortController().signal), + ]; + + await expect(Promise.all(calls)).resolves.toEqual(["", ""]); + expect(refreshCalls).toBe(1); + expect(authURLCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + expect(finishAuthCalls).toBe(1); + expect(callToolCalls).toBe(4); + expect(authorizedCount).toBe(1); + }); + + test("does not start browser fallback while shared refresh is pending", async () => { + const connected = await connectMCPServer(config, { onAuthURL: () => (authURLCount += 1) }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + refreshGate = new Promise((resolve) => { + releaseRefresh = resolve; + }); + redirectsPerFailure = 0; + callFailuresLeft = 1; + const first = connected.client.call("first", {}, new AbortController().signal); + while (refreshCalls === 0) await Promise.resolve(); + + redirectsPerFailure = 1; + callFailuresLeft = 1; + const second = connected.client.call("second", {}, new AbortController().signal); + await Promise.resolve(); + expect(authURLCount).toBe(0); + + releaseRefresh?.(); + await expect(Promise.all([first, second])).rejects.toThrow("finishAuth exploded"); + expect(refreshCalls).toBe(1); + expect(authURLCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + }); + + test("caller abort does not cancel shared recovery for another call", async () => { + finishAuthError = undefined; + callbackGate = new Promise((resolve) => { + releaseCallback = resolve; + }); + const connected = await connectMCPServer(config, { onAuthURL: () => (authURLCount += 1) }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + callFailuresLeft = 2; + const firstAbort = new AbortController(); + const first = connected.client.call("first", {}, firstAbort.signal); + const second = connected.client.call("second", {}, new AbortController().signal); + while (waitForCodeCalls === 0) await Promise.resolve(); + + firstAbort.abort(new Error("caller stopped")); + await expect(first).rejects.toThrow("caller stopped"); + releaseCallback?.(); + + await expect(second).resolves.toBe(""); + expect(waitForCodeCalls).toBe(1); + expect(finishAuthCalls).toBe(1); + expect(authURLCount).toBe(1); + }); + + test("aborted waiter still fires onAuthorized when background finishAuth succeeds", async () => { + finishAuthError = undefined; + callbackGate = new Promise((resolve) => { + releaseCallback = resolve; + }); + const connected = await connectMCPServer(config, { + onAuthURL: () => (authURLCount += 1), + onAuthorized: () => (authorizedCount += 1), + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + callFailuresLeft = 1; + const abort = new AbortController(); + const call = connected.client.call("ping", {}, abort.signal); + while (authURLCount === 0 || waitForCodeCalls === 0) await Promise.resolve(); + + abort.abort(new Error("caller stopped")); + await expect(call).rejects.toThrow("caller stopped"); + expect(authorizedCount).toBe(0); + + releaseCallback?.(); + while (finishAuthCalls === 0) await Promise.resolve(); + for (let tick = 0; tick < 20 && authorizedCount === 0; tick += 1) await Promise.resolve(); + expect(authorizedCount).toBe(1); + expect(finishAuthCalls).toBe(1); + + finishAuthError = new Error("finishAuth exploded"); + for (let episode = 0; episode < MAX_BROWSER_AUTH_ATTEMPTS; episode += 1) { + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "finishAuth exploded", + ); + } + expect(authURLCount).toBe(1 + MAX_BROWSER_AUTH_ATTEMPTS); + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "retrying paused", + ); + expect(authURLCount).toBe(1 + MAX_BROWSER_AUTH_ATTEMPTS); + }); + + test("refresh-only recovery clears prior browser-cap counts", async () => { + const connected = await connectMCPServer(config, { + onAuthURL: () => (authURLCount += 1), + onAuthorized: () => (authorizedCount += 1), + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "finishAuth exploded", + ); + expect(authURLCount).toBe(1); + expect(authorizedCount).toBe(0); + + refreshSucceeds = true; + finishAuthError = undefined; + redirectsPerFailure = 0; + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).resolves.toBe(""); + expect(authorizedCount).toBe(1); + expect(authURLCount).toBe(1); + + refreshSucceeds = false; + finishAuthError = new Error("finishAuth exploded"); + redirectsPerFailure = 1; + for (let episode = 0; episode < MAX_BROWSER_AUTH_ATTEMPTS; episode += 1) { + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "finishAuth exploded", + ); + } + expect(authURLCount).toBe(1 + MAX_BROWSER_AUTH_ATTEMPTS); + + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "retrying paused", + ); + expect(authURLCount).toBe(1 + MAX_BROWSER_AUTH_ATTEMPTS); + }); + + test("client close aborts the shared callback waiter", async () => { + finishAuthError = undefined; + callbackGate = new Promise((resolve) => { + releaseCallback = resolve; + }); + const connected = await connectMCPServer(config, { onAuthURL: () => (authURLCount += 1) }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + callFailuresLeft = 1; + const call = connected.client.call("ping", {}, new AbortController().signal); + while (waitForCodeCalls === 0) await Promise.resolve(); + + await connected.client.close(); + + await expect(call).rejects.toHaveProperty("name", "AbortError"); + expect(finishAuthCalls).toBe(0); + }); + + test("client close during hung refresh does not emit a browser prompt", async () => { + const connected = await connectMCPServer(config, { onAuthURL: () => (authURLCount += 1) }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + refreshGate = new Promise(() => undefined); + redirectsPerFailure = 0; + callFailuresLeft = 1; + const call = connected.client.call("ping", {}, new AbortController().signal); + while (refreshCalls === 0) await Promise.resolve(); + + await connected.client.close(); + + await expect(call).rejects.toHaveProperty("name", "AbortError"); + expect(authURLCount).toBe(0); + expect(waitForCodeCalls).toBe(0); + }); + + test("client close after recovery during retry still fires onAuthorized", async () => { + finishAuthError = undefined; + retryGate = new Promise(() => undefined); + const connected = await connectMCPServer(config, { + onAuthURL: () => (authURLCount += 1), + onAuthorized: () => (authorizedCount += 1), + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + callFailuresLeft = 1; + const call = connected.client.call("ping", {}, new AbortController().signal); + while (finishAuthCalls === 0 || callToolCalls < 2) await Promise.resolve(); + expect(authorizedCount).toBe(0); + + await connected.client.close(); + + await expect(call).rejects.toHaveProperty("name", "AbortError"); + expect(authorizedCount).toBe(1); + expect(authURLCount).toBe(1); + }); + + test("late retry success from a prior recovery does not fire onAuthorized after a new recovery starts", async () => { + finishAuthError = undefined; + const connected = await connectMCPServer(config, { + onAuthURL: () => (authURLCount += 1), + onAuthorized: () => (authorizedCount += 1), + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + retryGate = new Promise((resolve) => { + releaseRetry = resolve; + }); + callFailuresLeft = 2; + const first = connected.client.call("first", {}, new AbortController().signal); + const second = connected.client.call("second", {}, new AbortController().signal); + while (callToolCalls < 4) await Promise.resolve(); + + retryGate = undefined; + refreshSucceeds = true; + redirectsPerFailure = 0; + callFailuresLeft = 1; + const third = connected.client.call("third", {}, new AbortController().signal); + await expect(third).resolves.toBe(""); + expect(authorizedCount).toBe(1); + expect(authURLCount).toBe(1); + + releaseRetry?.(); + await expect(first).resolves.toBe(""); + await expect(second).resolves.toBe(""); + expect(authorizedCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + }); + + test("does not repeat the SDK refresh after redirecting to authorization", async () => { + finishAuthError = undefined; + connectFailuresLeft = 1; + + const result = await connectWithAuthPrompt(); + + expect(result.ok).toBe(true); + expect(authEvents).toEqual(["authURL"]); + expect(refreshCalls).toBe(0); + expect(finishAuthCalls).toBe(1); + expect(authURLCount).toBe(1); + }); + + test("reconnect during in-flight waitForCode does not emit a second prompt", async () => { + connectFailuresLeft = Number.POSITIVE_INFINITY; + callbackGate = new Promise(() => undefined); + const firstAbort = new AbortController(); + const first = connectMCPServer(config, { + onAuthURL: () => (authURLCount += 1), + signal: firstAbort.signal, + }); + while (authURLCount === 0 || waitForCodeCalls === 0) await Promise.resolve(); + expect(authURLCount).toBe(1); + + const second = await connectWithAuthPrompt(); + + expect(second.ok).toBe(false); + expect(second.error).toContain("retrying paused"); + expect(authURLCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + + firstAbort.abort(); + await expect(first).resolves.toMatchObject({ ok: false }); + }); + + test("live-call auth episodes share redirect state and stop after one prompt", async () => { + const connected = await connectMCPServer(config, { + onAuthURL: () => { + authURLCount += 1; + }, + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + + for (let episode = 0; episode < MAX_BROWSER_AUTH_ATTEMPTS; episode += 1) { + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "finishAuth exploded", + ); + } + expect(authURLCount).toBe(MAX_BROWSER_AUTH_ATTEMPTS); + + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "retrying paused", + ); + expect(authURLCount).toBe(MAX_BROWSER_AUTH_ATTEMPTS); + }); + + test("repeated concurrent redirects emit one counted prompt", async () => { + const connected = await connectMCPServer(config, { + onAuthURL: () => { + authURLCount += 1; + }, + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + redirectsPerFailure = 3; + redirectConcurrently = true; + callFailuresLeft = 1; + + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "finishAuth exploded", + ); + + expect(authURLCount).toBe(1); + expect(finishAuthCalls).toBe(1); + expect(refreshCalls).toBe(0); + }); + + test("failed browser re-auth is capped and pauses re-prompting across episodes", async () => { + connectFailuresLeft = Number.POSITIVE_INFINITY; + + for (let episode = 0; episode < MAX_BROWSER_AUTH_ATTEMPTS; episode += 1) { + const result = await connectWithAuthPrompt(); + expect(result.ok).toBe(false); + expect(result.error).toContain("finishAuth exploded"); + } + expect(authURLCount).toBe(MAX_BROWSER_AUTH_ATTEMPTS); + expect(finishAuthCalls).toBe(MAX_BROWSER_AUTH_ATTEMPTS); + + for (let episode = 0; episode < 2; episode += 1) { + const result = await connectWithAuthPrompt(); + expect(result.ok).toBe(false); + expect(result.error).toContain( + `MCP authorization for linear failed after ${MAX_BROWSER_AUTH_ATTEMPTS} ${MAX_BROWSER_AUTH_ATTEMPTS === 1 ? "attempt" : "attempts"}`, + ); + expect(result.error).toContain("retrying paused for 5 minutes"); + expect(result.error).toContain("Retry later after the cooldown"); + expect(result.error).not.toContain("Reconnect the server"); + } + + expect(authURLCount).toBe(MAX_BROWSER_AUTH_ATTEMPTS); + expect(finishAuthCalls).toBe(MAX_BROWSER_AUTH_ATTEMPTS); + expect(providerCreates).toBeGreaterThan(MAX_BROWSER_AUTH_ATTEMPTS); + }); + + test("successful interactive auth clears the cap so a later failure can prompt", async () => { + finishAuthError = undefined; + connectFailuresLeft = 1; + const recovered = await connectWithAuthPrompt(); + expect(recovered.ok).toBe(true); + expect(authURLCount).toBe(1); + + finishAuthError = new Error("finishAuth exploded"); + connectFailuresLeft = Number.POSITIVE_INFINITY; + for (let episode = 0; episode < MAX_BROWSER_AUTH_ATTEMPTS; episode += 1) { + const result = await connectWithAuthPrompt(); + expect(result.ok).toBe(false); + expect(result.error).toContain("finishAuth exploded"); + } + expect(authURLCount).toBe(1 + MAX_BROWSER_AUTH_ATTEMPTS); + + const capped = await connectWithAuthPrompt(); + expect(capped.ok).toBe(false); + expect(capped.error).toContain("retrying paused"); + expect(authURLCount).toBe(1 + MAX_BROWSER_AUTH_ATTEMPTS); + }); + + test("prompts resume five minutes after the capped failed episode", async () => { + const thirdEpisodeAt = new Date("2026-01-01T00:00:00Z").getTime(); + setSystemTime(thirdEpisodeAt); + connectFailuresLeft = Number.POSITIVE_INFINITY; + for (let episode = 0; episode < MAX_BROWSER_AUTH_ATTEMPTS; episode += 1) { + const result = await connectWithAuthPrompt(); + expect(result.ok).toBe(false); + } + expect(authURLCount).toBe(MAX_BROWSER_AUTH_ATTEMPTS); + + setSystemTime(thirdEpisodeAt + BROWSER_AUTH_COOLDOWN_MS + 1); + const afterCooldown = await connectWithAuthPrompt(); + expect(afterCooldown.ok).toBe(false); + expect(afterCooldown.error).toContain("finishAuth exploded"); + expect(authURLCount).toBe(MAX_BROWSER_AUTH_ATTEMPTS + 1); + }); + + test("resetBrowserAuthState clears the cap within the process", async () => { + connectFailuresLeft = Number.POSITIVE_INFINITY; + for (let episode = 0; episode < MAX_BROWSER_AUTH_ATTEMPTS; episode += 1) { + const result = await connectWithAuthPrompt(); + expect(result.ok).toBe(false); + expect(result.error).toContain("finishAuth exploded"); + } + expect(await connectWithAuthPrompt()).toEqual({ + ok: false, + error: expect.stringContaining("retrying paused"), + }); + expect(authURLCount).toBe(MAX_BROWSER_AUTH_ATTEMPTS); + + resetBrowserAuthState(); + const promptsBefore = finishAuthCalls; + const result = await connectWithAuthPrompt(); + expect(result.ok).toBe(false); + expect(result.error).toContain("finishAuth exploded"); + expect(finishAuthCalls).toBe(promptsBefore + 1); + expect(authURLCount).toBe(MAX_BROWSER_AUTH_ATTEMPTS + 1); + }); + + test("first in-episode saveCodeVerifier wins until the episode ends", async () => { + finishAuthError = undefined; + saveThenRedirectPair = true; + const connected = await connectMCPServer(config, { + onAuthURL: () => { + authURLCount += 1; + }, + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + callFailuresLeft = 1; + + await expect(connected.client.call("ping", {}, new AbortController().signal)).resolves.toBe(""); + + expect(authURLCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + expect(finishAuthCalls).toBe(1); + expect(exchangedCodeVerifier).toBe("v1"); + expect(storedCodeVerifier).toBe("v1"); + }); + + test("overlapping SDK-order saves emit the first verifier's authorize URL", async () => { + finishAuthError = undefined; + overlappingSDKSaves = true; + saveGate = new Promise((resolve) => { + releaseSave = resolve; + }); + const connected = await connectMCPServer(config, { + onAuthURL: (_name, url) => { + authURLCount += 1; + emittedAuthURL = url; + }, + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + callFailuresLeft = 1; + + await expect(connected.client.call("ping", {}, new AbortController().signal)).resolves.toBe(""); + + expect(authURLCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + expect(finishAuthCalls).toBe(1); + expect(exchangedCodeVerifier).toBe("v1"); + expect(storedCodeVerifier).toBe("v1"); + expect(emittedAuthURL).toBe("https://auth.test/authorize?v=v1"); + }); + + test("refresh-skip unfreezes so a later browser episode can save a new verifier", async () => { + refreshSucceeds = true; + finishAuthError = undefined; + const connected = await connectMCPServer(config, { + onAuthURL: () => { + authURLCount += 1; + }, + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + + refreshGate = new Promise((resolve) => { + releaseRefresh = resolve; + }); + redirectsPerFailure = 0; + callFailuresLeft = 1; + const first = connected.client.call("first", {}, new AbortController().signal); + while (refreshCalls === 0) await Promise.resolve(); + + const skipped = saveThenRedirect(liveProvider, "v-refresh"); + await Promise.resolve(); + expect(authURLCount).toBe(0); + + releaseRefresh?.(); + await skipped; + await expect(first).resolves.toBe(""); + expect(authURLCount).toBe(0); + expect(waitForCodeCalls).toBe(0); + expect(storedCodeVerifier).toBe("v-refresh"); + + refreshSucceeds = false; + redirectVerifier = "v-later"; + redirectsPerFailure = 1; + callFailuresLeft = 1; + await expect(connected.client.call("second", {}, new AbortController().signal)).resolves.toBe( + "", + ); + + expect(authURLCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + expect(exchangedCodeVerifier).toBe("v-later"); + expect(storedCodeVerifier).toBe("v-later"); + }); + + test("does not clear the cap or fire onAuthorized until a retried tool call succeeds", async () => { + finishAuthError = undefined; + const connected = await connectMCPServer(config, { + onAuthURL: () => { + authURLCount += 1; + }, + onAuthorized: () => { + authorizedCount += 1; + }, + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + + callFailuresLeft = 2; + callRedirectsLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "authorization required", + ); + expect(authorizedCount).toBe(0); + expect(authURLCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + expect(finishAuthCalls).toBe(1); + + finishAuthError = new Error("finishAuth exploded"); + callRedirectsLeft = Number.POSITIVE_INFINITY; + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "retrying paused", + ); + expect(authURLCount).toBe(1); + expect(authorizedCount).toBe(0); + }); + + test("clears the cap and fires onAuthorized after a retried tool call succeeds", async () => { + finishAuthError = undefined; + const connected = await connectMCPServer(config, { + onAuthURL: () => { + authURLCount += 1; + }, + onAuthorized: () => { + authorizedCount += 1; + }, + }); + expect(connected.ok).toBe(true); + if (!connected.ok) return; + + callFailuresLeft = 1; + callRedirectsLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).resolves.toBe(""); + expect(authorizedCount).toBe(1); + expect(authURLCount).toBe(1); + + finishAuthError = new Error("finishAuth exploded"); + callRedirectsLeft = Number.POSITIVE_INFINITY; + for (let episode = 0; episode < MAX_BROWSER_AUTH_ATTEMPTS; episode += 1) { + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "finishAuth exploded", + ); + } + expect(authURLCount).toBe(1 + MAX_BROWSER_AUTH_ATTEMPTS); + + callFailuresLeft = 1; + await expect(connected.client.call("ping", {}, new AbortController().signal)).rejects.toThrow( + "retrying paused", + ); + expect(authURLCount).toBe(1 + MAX_BROWSER_AUTH_ATTEMPTS); + expect(authorizedCount).toBe(1); + }); + + test("ignored browser wait times out as disconnected and still counts the prompt", async () => { + setBrowserAuthWaitMs(50); + connectFailuresLeft = Number.POSITIVE_INFINITY; + callbackGate = new Promise(() => undefined); + + const result = await connectWithAuthPrompt(); + + expect(result.ok).toBe(false); + expect(result.error).toContain("timed out waiting for the browser"); + expect(result.error).toContain("disconnected"); + expect(authURLCount).toBe(1); + expect(waitForCodeCalls).toBe(1); + + const capped = await connectWithAuthPrompt(); + expect(capped.ok).toBe(false); + expect(capped.error).toContain("retrying paused"); + expect(authURLCount).toBe(1); + }); +}); diff --git a/src/mcp/client-auth-retry.test.ts b/src/mcp/client-auth-retry.test.ts index 98b63152f..ddfde5b9d 100644 --- a/src/mcp/client-auth-retry.test.ts +++ b/src/mcp/client-auth-retry.test.ts @@ -48,4 +48,20 @@ describe("retryAfterInteractiveAuth", () => { expect(operationRan).toBe(false); expect(notified).toBe(false); }); + + test("notifies when the retried operation is aborted after auth completes", async () => { + let notified = false; + await expect( + retryAfterInteractiveAuth( + async () => undefined, + async () => { + throw new DOMException("toolset disposed", "AbortError"); + }, + () => { + notified = true; + }, + ), + ).rejects.toHaveProperty("name", "AbortError"); + expect(notified).toBe(true); + }); }); diff --git a/src/mcp/client.ts b/src/mcp/client.ts index 9da4b4e16..1f25a126d 100644 --- a/src/mcp/client.ts +++ b/src/mcp/client.ts @@ -62,16 +62,192 @@ interface HTTPAuthContext { url: URL; authProvider: CorbitsOAuthProvider; callback: CallbackServer; - signal?: AbortSignal; + coordinator: HTTPAuthCoordinator; interactive: boolean; serverName: string; onAuthorized?: (serverName: string) => void; } -function isRecoverableAuthError(err: unknown): boolean { +interface BrowserAuthFlow { + attempt: { clear(): void }; + promptEmitted: Promise; +} + +interface HTTPAuthCoordinator { + lifecycle: AbortController; + inFlight?: Promise; + refreshInFlight?: Promise; + browserFlow?: BrowserAuthFlow; + pkceFrozen?: boolean; + pkceSavePromise?: Promise; + pendingAttempt?: { clear(): void }; + notified?: boolean; + recoveryGeneration?: number; + waiters?: number; + probe(): Promise; +} + +function isRecoverableAuthError(err: unknown): err is UnauthorizedError | OAuthError { return err instanceof UnauthorizedError || err instanceof OAuthError; } +function isAbortError(err: unknown): boolean { + return typeof err === "object" && err !== null && "name" in err && err.name === "AbortError"; +} + +export const MAX_BROWSER_AUTH_ATTEMPTS = 1; +export const BROWSER_AUTH_COOLDOWN_MS = 5 * 60_000; +export const BROWSER_AUTH_WAIT_MS = 2 * 60_000; + +interface BrowserAuthAttempts { + count: number; + cooldownUntil?: number | undefined; +} +// Keyed by server identity, not provider instance, so the cap survives the +// provider re-creation that every reconnect performs. +const browserAuthAttempts = new Map(); +let browserAuthWaitMs = BROWSER_AUTH_WAIT_MS; + +export function resetBrowserAuthState(): void { + browserAuthAttempts.clear(); + browserAuthWaitMs = BROWSER_AUTH_WAIT_MS; +} + +export function setBrowserAuthWaitMs(ms: number): void { + browserAuthWaitMs = ms; +} + +function browserAuthCapError(serverName: string): Error { + const minutes = Math.round(BROWSER_AUTH_COOLDOWN_MS / 60_000); + const attempts = + MAX_BROWSER_AUTH_ATTEMPTS === 1 ? "1 attempt" : `${String(MAX_BROWSER_AUTH_ATTEMPTS)} attempts`; + return new Error( + `MCP authorization for ${serverName} failed after ${attempts}; ` + + `retrying paused for ${minutes} minutes. Retry later after the cooldown.`, + ); +} + +function browserAuthWaitError(serverName: string): Error { + return new Error( + `MCP authorization for ${serverName} timed out waiting for the browser; ` + + `the server is disconnected. Retry later after the cooldown.`, + ); +} + +async function waitForBrowserAuthCode(context: HTTPAuthContext): Promise { + const lifecycle = context.coordinator.lifecycle.signal; + const deadline = AbortSignal.any([lifecycle, AbortSignal.timeout(browserAuthWaitMs)]); + try { + return await context.callback.waitForCode(deadline); + } catch (err) { + if (lifecycle.aborted) throw err; + if (deadline.aborted) throw browserAuthWaitError(context.serverName); + throw err; + } +} + +function browserAuthKey(context: HTTPAuthContext): string { + return `${context.serverName}|${context.url.toString()}`; +} + +function beginBrowserAuth(context: HTTPAuthContext): { clear(): void } { + const key = browserAuthKey(context); + const entry = browserAuthAttempts.get(key) ?? { count: 0 }; + const now = Date.now(); + if (entry.cooldownUntil !== undefined) { + if (now < entry.cooldownUntil) throw browserAuthCapError(context.serverName); + entry.cooldownUntil = undefined; + entry.count = 0; + } + if (entry.count >= MAX_BROWSER_AUTH_ATTEMPTS) { + throw browserAuthCapError(context.serverName); + } + entry.count += 1; + if (entry.count === MAX_BROWSER_AUTH_ATTEMPTS) { + entry.cooldownUntil = now + BROWSER_AUTH_COOLDOWN_MS; + } + browserAuthAttempts.set(key, entry); + return { + clear: () => browserAuthAttempts.delete(key), + }; +} + +function unfreezePkce(coordinator: HTTPAuthCoordinator): void { + coordinator.pkceFrozen = false; + delete coordinator.pkceSavePromise; +} + +async function tryTokenRefresh(context: HTTPAuthContext): Promise { + const refreshToken = (await context.authProvider.tokens?.())?.refresh_token; + if (refreshToken === undefined) return false; + try { + await context.authProvider.refreshToken(refreshToken); + return true; + } catch (err) { + if (context.coordinator.lifecycle.signal.aborted) throw err; + return false; + } +} + +function getOrStartRefresh(context: HTTPAuthContext): Promise { + const coordinator = context.coordinator; + if (coordinator.refreshInFlight !== undefined) return coordinator.refreshInFlight; + const shared = Promise.resolve().then(() => tryTokenRefresh(context)); + coordinator.refreshInFlight = shared; + const clear = () => { + if (coordinator.refreshInFlight === shared) delete coordinator.refreshInFlight; + }; + void shared.then(clear, clear); + return shared; +} + +function gateRedirectToAuthorization(context: HTTPAuthContext): void { + const inner = context.authProvider.redirectToAuthorization.bind(context.authProvider); + const innerSave = context.authProvider.saveCodeVerifier?.bind(context.authProvider); + if (innerSave !== undefined) { + context.authProvider.saveCodeVerifier = (codeVerifier: string) => { + const coordinator = context.coordinator; + if (coordinator.pkceFrozen) { + const pending = coordinator.pkceSavePromise; + if (pending === undefined) + throw new Error("PKCE verifier save is frozen without a pending write"); + return pending; + } + coordinator.pkceFrozen = true; + coordinator.pkceSavePromise = Promise.resolve(innerSave(codeVerifier)); + return coordinator.pkceSavePromise; + }; + } + context.authProvider.redirectToAuthorization = async (authorizationUrl: URL) => { + const coordinator = context.coordinator; + const browserFlowBeforeRefresh = coordinator.browserFlow; + if (browserFlowBeforeRefresh !== undefined) return browserFlowBeforeRefresh.promptEmitted; + const refresh = coordinator.refreshInFlight; + if (refresh !== undefined && (await refresh)) { + unfreezePkce(coordinator); + return; + } + const concurrentBrowserFlow = coordinator.browserFlow ?? browserFlowBeforeRefresh; + if (concurrentBrowserFlow !== undefined) return concurrentBrowserFlow.promptEmitted; + if (!context.interactive) { + unfreezePkce(coordinator); + throw new Error("Authorization required but no interactive handler is available."); + } + + try { + const attempt = beginBrowserAuth(context); + const startPrompt = Promise.withResolvers(); + const promptEmitted = startPrompt.promise.then(() => inner(authorizationUrl)); + coordinator.browserFlow = { attempt, promptEmitted }; + startPrompt.resolve(undefined); + return promptEmitted; + } catch (err) { + unfreezePkce(coordinator); + throw err; + } + }; +} + /** * Fetch that always attaches the connect AbortSignal. SDK 403 upscoping calls * `auth()` with raw `_fetch` (no `requestInit.signal`); `_fetchWithInit` still @@ -104,20 +280,10 @@ function streamableHTTPTransportOptions( }; } -async function completeInteractiveAuth(context: HTTPAuthContext): Promise { - if (!context.interactive) - throw new Error("Authorization required but no interactive handler is available."); - const code = await context.callback.waitForCode(context.signal ?? new AbortController().signal); - await new StreamableHTTPClientTransport( - context.url, - streamableHTTPTransportOptions(context.authProvider, context.signal), - ).finishAuth(code); -} - /** - * Run interactive OAuth, retry the failed operation, and notify only when the - * retry itself succeeded — a failed re-auth must leave standing "needs auth" - * chrome alone. + * Run interactive OAuth, retry the failed operation, and notify when the retry + * succeeds or is aborted after auth completed — a failed re-auth must leave + * standing "needs auth" chrome alone. */ export async function retryAfterInteractiveAuth( completeAuth: () => Promise, @@ -125,47 +291,135 @@ export async function retryAfterInteractiveAuth( onAuthorized: (() => void) | undefined, ): Promise { await completeAuth(); - const value = await operation(); - onAuthorized?.(); - return value; + try { + const value = await operation(); + onAuthorized?.(); + return value; + } catch (err) { + if (isAbortError(err)) onAuthorized?.(); + throw err; + } } -async function recoverHTTPAuthorization( - err: unknown, - context: HTTPAuthContext | undefined, - operation: () => Promise, -): Promise { - if (context === undefined || !isRecoverableAuthError(err)) throw err; - let lastErr: unknown = err; - for (let attempt = 0; attempt < 2; attempt += 1) { - if (lastErr instanceof OAuthError) await context.authProvider.resetAuthorization(); - if (lastErr instanceof UnauthorizedError) { - return retryAfterInteractiveAuth( - () => completeInteractiveAuth(context), - operation, - context.onAuthorized === undefined - ? undefined - : () => context.onAuthorized?.(context.serverName), - ); - } +async function driveRecovery(err: UnauthorizedError | OAuthError, context: HTTPAuthContext) { + const coordinator = context.coordinator; + if (err instanceof OAuthError) await context.authProvider.resetAuthorization(); + if (err instanceof UnauthorizedError && coordinator.browserFlow === undefined) { + await getOrStartRefresh(context); + // SDK redirects waiting on this refresh must reserve the browser flow before the probe. + await Promise.resolve(); + } + + if (coordinator.browserFlow === undefined) { try { - return await operation(); - } catch (nextErr) { - if (!isRecoverableAuthError(nextErr)) throw nextErr; - lastErr = nextErr; + await coordinator.probe(); + return; + } catch (probeErr) { + if (!isRecoverableAuthError(probeErr)) throw probeErr; + if (coordinator.browserFlow === undefined) throw probeErr; } } - throw lastErr; + + const browserFlow = coordinator.browserFlow; + await browserFlow.promptEmitted; + const code = await waitForBrowserAuthCode(context); + await new StreamableHTTPClientTransport( + context.url, + streamableHTTPTransportOptions(context.authProvider, coordinator.lifecycle.signal), + ).finishAuth(code); + await coordinator.probe(); +} + +function completeVerifiedRecovery(context: HTTPAuthContext, generation: number): void { + const coordinator = context.coordinator; + if (coordinator.recoveryGeneration !== generation) return; + if (coordinator.notified) return; + coordinator.notified = true; + coordinator.pendingAttempt?.clear(); + delete coordinator.pendingAttempt; + browserAuthAttempts.delete(browserAuthKey(context)); + context.onAuthorized?.(context.serverName); +} + +function getOrStartRecovery( + err: UnauthorizedError | OAuthError, + context: HTTPAuthContext, +): Promise { + const coordinator = context.coordinator; + if (coordinator.inFlight !== undefined) return coordinator.inFlight; + coordinator.recoveryGeneration = (coordinator.recoveryGeneration ?? 0) + 1; + coordinator.notified = false; + delete coordinator.pendingAttempt; + const shared = Promise.resolve().then(() => driveRecovery(err, context)); + coordinator.inFlight = shared; + const clear = () => { + if (coordinator.inFlight !== shared) return; + if (coordinator.browserFlow?.attempt !== undefined) + coordinator.pendingAttempt = coordinator.browserFlow.attempt; + unfreezePkce(coordinator); + delete coordinator.inFlight; + delete coordinator.refreshInFlight; + delete coordinator.browserFlow; + }; + void shared.then(() => { + clear(); + if ((coordinator.waiters ?? 0) === 0) { + completeVerifiedRecovery(context, coordinator.recoveryGeneration ?? 0); + } + }, clear); + return shared; +} + +function abortReason(signal: AbortSignal): unknown { + return signal.reason ?? new DOMException("The operation was aborted", "AbortError"); +} + +function awaitRecovery(recovery: Promise, signal: AbortSignal | undefined): Promise { + if (signal === undefined) return recovery; + if (signal.aborted) return Promise.reject(abortReason(signal)); + return new Promise((resolve, reject) => { + const onAbort = () => { + signal.removeEventListener("abort", onAbort); + reject(abortReason(signal)); + }; + signal.addEventListener("abort", onAbort, { once: true }); + void recovery.then( + () => { + signal.removeEventListener("abort", onAbort); + resolve(); + }, + (err) => { + signal.removeEventListener("abort", onAbort); + reject(err); + }, + ); + }); } async function withHTTPAuthorizationRecovery( context: HTTPAuthContext | undefined, operation: () => Promise, + signal?: AbortSignal, ): Promise { try { return await operation(); } catch (err) { - return recoverHTTPAuthorization(err, context, operation); + if (context === undefined || !isRecoverableAuthError(err)) throw err; + const coordinator = context.coordinator; + coordinator.waiters = (coordinator.waiters ?? 0) + 1; + let generation = 0; + try { + return await retryAfterInteractiveAuth( + async () => { + await awaitRecovery(getOrStartRecovery(err, context), signal); + generation = coordinator.recoveryGeneration ?? 0; + }, + operation, + () => completeVerifiedRecovery(context, generation), + ); + } finally { + coordinator.waiters = (coordinator.waiters ?? 1) - 1; + } } } @@ -174,6 +428,7 @@ async function finishClient( serverName: string, authContext?: HTTPAuthContext, signal?: AbortSignal, + closeLifecycle?: () => void, ): Promise { const result = await withHTTPAuthorizationRecovery(authContext, () => signal === undefined ? client.listTools() : client.listTools(undefined, { signal }), @@ -192,13 +447,15 @@ async function finishClient( serverName, tools, async call(toolName, args, signal) { - const context = authContext === undefined ? undefined : { ...authContext, signal }; - const result = await withHTTPAuthorizationRecovery(context, () => - client.callTool({ name: toolName, arguments: args }, undefined, { signal }), + const result = await withHTTPAuthorizationRecovery( + authContext, + () => client.callTool({ name: toolName, arguments: args }, undefined, { signal }), + signal, ); return unwrapToolContent(result.content); }, async close() { + closeLifecycle?.(); authContext?.callback.close(); await client.close().catch(() => undefined); }, @@ -244,6 +501,14 @@ async function connectHttp( return { ok: false, serverName: config.name, error: "http MCP server requires a url" }; let callback: CallbackServer | undefined; let client: Client | undefined; + const lifecycle = new AbortController(); + const abortLifecycle = () => lifecycle.abort(options.signal?.reason); + const closeLifecycle = () => { + lifecycle.abort(); + options.signal?.removeEventListener("abort", abortLifecycle); + }; + if (options.signal?.aborted) abortLifecycle(); + else options.signal?.addEventListener("abort", abortLifecycle, { once: true }); try { const normalizedURL = normalizeMCPServerURL(config.url); const url = new URL(normalizedURL); @@ -263,35 +528,58 @@ async function connectHttp( redirectUrl: callback.redirectUrl, onAuthURL: (name, authUrl) => options.onAuthURL?.(name, authUrl), onAuthorizationState: callback.expectState, + fetchFn: fetchWithConnectAbort(lifecycle.signal), }); makeTransport = () => new StreamableHTTPClientTransport( url, - streamableHTTPTransportOptions(authProvider, options.signal), + streamableHTTPTransportOptions(authProvider, lifecycle.signal), ) as unknown as Transport; + const coordinator: HTTPAuthCoordinator = { + lifecycle, + probe: async () => { + const probeClient = new Client({ name: MCP_CLIENT_NAME, version: "1.0.0" }); + try { + await probeClient.connect(makeTransport(), { signal: lifecycle.signal }); + } finally { + await probeClient.close().catch(() => undefined); + } + }, + }; authContext = { url, authProvider, callback, + coordinator, interactive: options.onAuthURL !== undefined, serverName: config.name, ...(options.onAuthorized !== undefined ? { onAuthorized: options.onAuthorized } : {}), - ...(options.signal !== undefined ? { signal: options.signal } : {}), }; + gateRedirectToAuthorization(authContext); } const connectedClient = new Client({ name: MCP_CLIENT_NAME, version: "1.0.0" }); client = connectedClient; - await withHTTPAuthorizationRecovery(authContext, () => - connectedClient.connect( - makeTransport(), - options.signal === undefined ? undefined : { signal: options.signal }, - ), + await withHTTPAuthorizationRecovery( + authContext, + () => connectedClient.connect(makeTransport(), { signal: lifecycle.signal }), + lifecycle.signal, ); + if (authContext !== undefined) { + authContext.coordinator.probe = () => + connectedClient.listTools(undefined, { signal: lifecycle.signal }).then(() => undefined); + } return { ok: true, - client: await finishClient(connectedClient, config.name, authContext, options.signal), + client: await finishClient( + connectedClient, + config.name, + authContext, + options.signal, + closeLifecycle, + ), }; } catch (err) { + closeLifecycle(); await client?.close().catch(() => undefined); try { callback?.close(); diff --git a/src/mcp/oauth-provider.test.ts b/src/mcp/oauth-provider.test.ts index 0696f8c06..4328a401c 100644 --- a/src/mcp/oauth-provider.test.ts +++ b/src/mcp/oauth-provider.test.ts @@ -4,7 +4,9 @@ import { readFileSync } from "node:fs"; import { chmod, mkdir, mkdtemp, readFile, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"; import { authFilePath, loadAuthState, saveAuthState, deleteAuthState } from "./auth-store.js"; +import { fetchWithConnectAbort } from "./client.js"; import { createOAuthProvider } from "./oauth-provider.js"; async function tempHome(): Promise { @@ -603,4 +605,188 @@ describe("createOAuthProvider", () => { "scoped-secret", ); }); + + test("refreshToken posts grant_type=refresh_token with a resource and persists tokens", async () => { + const home = await tempHome(); + await saveAuthState( + linear, + { + clientInformation: clientInfo(1), + tokens: { + access_token: "stale", + token_type: "bearer", + refresh_token: "refresh-me", + }, + }, + home, + ); + + const tokenBodies: string[] = []; + const fetchFn = async (url: string | URL, init?: RequestInit): Promise => { + const href = String(url); + if (init?.method === "POST") { + tokenBodies.push(String(init.body)); + return new Response( + JSON.stringify({ + access_token: "fresh-access", + token_type: "bearer", + expires_in: 3600, + refresh_token: "fresh-refresh", + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + if (href.includes("oauth-protected-resource")) { + return new Response( + JSON.stringify({ + resource: "https://mcp.linear.app/mcp", + authorization_servers: ["https://mcp.linear.app"], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + if (href.includes("oauth-authorization-server") || href.includes("openid-configuration")) { + return new Response( + JSON.stringify({ + issuer: "https://mcp.linear.app", + authorization_endpoint: "https://mcp.linear.app/authorize", + token_endpoint: "https://mcp.linear.app/oauth/token", + response_types_supported: ["code"], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + return new Response(null, { status: 404 }); + }; + + const provider = await createOAuthProvider({ + serverName: "linear", + serverURL: linear.serverURL, + redirectUrl: "http://127.0.0.1:1/callback", + onAuthURL: () => undefined, + home, + fetchFn, + }); + + const tokens = await provider.refreshToken("refresh-me"); + expect(tokens.access_token).toBe("fresh-access"); + expect(tokenBodies).toHaveLength(1); + const body = tokenBodies[0] ?? ""; + expect(body).toContain("grant_type=refresh_token"); + expect(body).toContain("refresh_token=refresh-me"); + expect(body).toContain(`resource=${encodeURIComponent("https://mcp.linear.app/mcp")}`); + expect((await syncValue(provider.tokens()))?.access_token).toBe("fresh-access"); + expect((await loadAuthState(linear, home)).tokens?.access_token).toBe("fresh-access"); + }); + + test("refreshToken wraps failures as UnauthorizedError", async () => { + const home = await tempHome(); + await saveAuthState(linear, { clientInformation: clientInfo(1) }, home); + + const fetchFn = async (url: string | URL, init?: RequestInit): Promise => { + const href = String(url); + if (init?.method === "POST") { + return new Response(JSON.stringify({ error: "invalid_grant" }), { + status: 400, + headers: { "content-type": "application/json" }, + }); + } + if (href.includes("oauth-protected-resource")) { + return new Response( + JSON.stringify({ + resource: "https://mcp.linear.app/mcp", + authorization_servers: ["https://mcp.linear.app"], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + if (href.includes("oauth-authorization-server") || href.includes("openid-configuration")) { + return new Response( + JSON.stringify({ + issuer: "https://mcp.linear.app", + authorization_endpoint: "https://mcp.linear.app/authorize", + token_endpoint: "https://mcp.linear.app/oauth/token", + response_types_supported: ["code"], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + return new Response(null, { status: 404 }); + }; + + const provider = await createOAuthProvider({ + serverName: "linear", + serverURL: linear.serverURL, + redirectUrl: "http://127.0.0.1:1/callback", + onAuthURL: () => undefined, + home, + fetchFn, + }); + + await expect(provider.refreshToken("refresh-me")).rejects.toBeInstanceOf(UnauthorizedError); + await expect(provider.refreshToken("refresh-me")).rejects.toThrow( + "Token refresh failed for linear", + ); + }); + + test("refreshToken fetch honors the connect AbortSignal", async () => { + const home = await tempHome(); + await saveAuthState(linear, { clientInformation: clientInfo(1) }, home); + const abort = new AbortController(); + const seen: (AbortSignal | undefined)[] = []; + const fetchFn = fetchWithConnectAbort(abort.signal, (url, init) => { + seen.push(init?.signal ?? undefined); + const href = String(url); + if (init?.method === "POST") { + return new Promise((_resolve, reject) => { + const fail = (): void => { + reject(init.signal?.reason ?? new Error("aborted")); + }; + if (init.signal?.aborted === true) fail(); + else init.signal?.addEventListener("abort", fail, { once: true }); + }); + } + if (href.includes("oauth-protected-resource")) { + return Promise.resolve( + new Response( + JSON.stringify({ + resource: "https://mcp.linear.app/mcp", + authorization_servers: ["https://mcp.linear.app"], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ); + } + if (href.includes("oauth-authorization-server") || href.includes("openid-configuration")) { + return Promise.resolve( + new Response( + JSON.stringify({ + issuer: "https://mcp.linear.app", + authorization_endpoint: "https://mcp.linear.app/authorize", + token_endpoint: "https://mcp.linear.app/oauth/token", + response_types_supported: ["code"], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ); + } + return Promise.resolve(new Response(null, { status: 404 })); + }); + + const provider = await createOAuthProvider({ + serverName: "linear", + serverURL: linear.serverURL, + redirectUrl: "http://127.0.0.1:1/callback", + onAuthURL: () => undefined, + home, + fetchFn, + }); + + const pending = provider.refreshToken("refresh-me"); + while (seen.every((signal) => signal !== abort.signal)) await Promise.resolve(); + abort.abort(new DOMException("toolset disposed", "AbortError")); + await expect(pending).rejects.toThrow("toolset disposed"); + await expect(pending).rejects.not.toBeInstanceOf(UnauthorizedError); + expect(seen.some((signal) => signal === abort.signal)).toBe(true); + }); }); diff --git a/src/mcp/oauth-provider.ts b/src/mcp/oauth-provider.ts index 7c58b2b1a..615a3e488 100644 --- a/src/mcp/oauth-provider.ts +++ b/src/mcp/oauth-provider.ts @@ -1,12 +1,23 @@ import { statSync } from "node:fs"; import { homedir } from "node:os"; -import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js"; +import { + discoverAuthorizationServerMetadata, + discoverOAuthProtectedResourceMetadata, + refreshAuthorization, + selectResourceURL, + UnauthorizedError, + type OAuthClientProvider, +} from "@modelcontextprotocol/sdk/client/auth.js"; +import { resourceUrlFromServerUrl } from "@modelcontextprotocol/sdk/shared/auth-utils.js"; import type { + AuthorizationServerMetadata, OAuthClientInformationFull, OAuthClientInformationMixed, OAuthClientMetadata, + OAuthProtectedResourceMetadata, OAuthTokens, } from "@modelcontextprotocol/sdk/shared/auth.js"; +import type { FetchLike } from "@modelcontextprotocol/sdk/shared/transport.js"; import { authFilePath, tryLoadAuthStateSync, @@ -23,8 +34,12 @@ export interface OAuthProviderOptions { onAuthURL: (serverName: string, authorizationUrl: string) => void; onAuthorizationState?: (state: string) => void; home?: string; + fetchFn?: FetchLike; } -export type CorbitsOAuthProvider = OAuthClientProvider & { resetAuthorization(): Promise }; +export type CorbitsOAuthProvider = OAuthClientProvider & { + resetAuthorization(): Promise; + refreshToken(refreshToken: string): Promise; +}; function redirectUrisInclude( info: OAuthClientInformationFull | undefined, @@ -35,6 +50,10 @@ function redirectUrisInclude( return uris.includes(redirectUrl); } +function isAbortError(err: unknown): boolean { + return typeof err === "object" && err !== null && "name" in err && err.name === "AbortError"; +} + // Dynamic client registration bakes in the loopback redirect_uri (ephemeral port). // A later session that binds a new port cannot reuse that client_id for authorize // / token exchange — drop the stale registration when we have no refreshable @@ -144,7 +163,10 @@ export async function createOAuthProvider( }; let oauthState: string | undefined; - return { + let authorizationServerMetadata: AuthorizationServerMetadata | undefined; + let authorizationServerUrl: string | undefined; + let resource: URL | undefined; + const provider: CorbitsOAuthProvider = { get redirectUrl(): string { return opts.redirectUrl; }, @@ -208,5 +230,56 @@ export async function createOAuthProvider( }); delete stored.codeVerifier; }, + refreshToken: async (refreshToken: string): Promise => { + try { + const fetchFn = opts.fetchFn; + if (authorizationServerMetadata === undefined || authorizationServerUrl === undefined) { + let resourceMetadata: OAuthProtectedResourceMetadata | undefined; + try { + resourceMetadata = await discoverOAuthProtectedResourceMetadata( + opts.serverURL, + undefined, + fetchFn, + ); + } catch { + resourceMetadata = undefined; + } + const fromPrm = resourceMetadata?.authorization_servers?.[0]; + authorizationServerUrl = + fromPrm === undefined ? String(new URL("/", opts.serverURL)) : String(fromPrm); + authorizationServerMetadata = + (await discoverAuthorizationServerMetadata( + authorizationServerUrl, + fetchFn === undefined ? {} : { fetchFn }, + )) ?? undefined; + resource = + (await selectResourceURL(opts.serverURL, provider, resourceMetadata)) ?? + resourceUrlFromServerUrl(opts.serverURL); + } + if (authorizationServerMetadata === undefined) + throw new UnauthorizedError("authorization server metadata unavailable"); + const clientInformation = stored.clientInformation; + if (clientInformation === undefined) + throw new UnauthorizedError("no client registration to refresh"); + const resourceURL = resource ?? resourceUrlFromServerUrl(opts.serverURL); + const tokens = await refreshAuthorization(authorizationServerUrl, { + metadata: authorizationServerMetadata, + clientInformation, + refreshToken, + resource: resourceURL, + ...(fetchFn === undefined ? {} : { fetchFn }), + }); + await apply((state) => { + state.tokens = tokens; + }); + return tokens; + } catch (err) { + if (isAbortError(err) || err instanceof UnauthorizedError) throw err; + throw new UnauthorizedError( + `Token refresh failed for ${opts.serverName}: ${err instanceof Error ? err.message : String(err)}`, + ); + } + }, }; + return provider; }