diff --git a/src/services/mcp/McpOAuthClientProvider.ts b/src/services/mcp/McpOAuthClientProvider.ts index ee3277be4c..794f53dcda 100644 --- a/src/services/mcp/McpOAuthClientProvider.ts +++ b/src/services/mcp/McpOAuthClientProvider.ts @@ -93,6 +93,15 @@ export class McpOAuthClientProvider implements OAuthClientProvider { private _effectiveResourceIndicator: string | null /** Deduplicates concurrent _ensureCallbackServer() calls. */ private _ensureServerPromise: Promise | null = null + /** Callback promise most recently returned to a waitForAuthCode() caller. */ + private _awaitedAuthCodePromise: Promise | null = null + /** + * Rejected callback promise of a listener that failed while no caller was waiting. + * The next waitForAuthCode() call reports it once instead of starting another listener. + */ + private _unreportedAuthCodeFailure: Promise | null = null + /** Stops a detached listener. A replacement listener waits for it because it may reuse the port. */ + private _listenerCleanup: Promise | null = null private constructor( private readonly _serverUrl: string, @@ -243,6 +252,11 @@ export class McpOAuthClientProvider implements OAuthClientProvider { private async _doStartCallbackServer(): Promise { this._closed = false + // A new listener belongs to a new authorization attempt, so an earlier failure no longer applies. + this._unreportedAuthCodeFailure = null + if (this._listenerCleanup) { + await this._listenerCleanup + } const { port: configuredPort, host, path } = getCallbackAddress(this._configuredRedirectUri) const { server, port, result, cancel } = await startCallbackServer( configuredPort ?? this._port, @@ -261,18 +275,31 @@ export class McpOAuthClientProvider implements OAuthClientProvider { this._authCodePromise = authCodePromise // Static clients can start this promise before waitForAuthCode() observes it. - // Observe rejection now to prevent an unhandled rejection, but keep the original - // promise so waitForAuthCode() still receives the callback error. Stop only this - // promise's listener before clearing it so a later flow can bind a fresh listener. - void authCodePromise.catch(async () => { - if (this._authCodePromise === authCodePromise && this._server === server) { - await stopCallbackServer(server, cancel).catch(() => {}) - if (this._authCodePromise === authCodePromise && this._server === server) { - this._server = null - this._cancelCallbackServer = null - } + // Observe rejection now to prevent an unhandled rejection. A caller that is already + // waiting receives the error from this promise; otherwise keep the rejected promise so + // the next waitForAuthCode() call reports it instead of waiting on a new listener. + void authCodePromise.catch(() => { + // close() detaches this promise before it cancels the listener and stops it itself. + if (this._authCodePromise !== authCodePromise) return + if (this._awaitedAuthCodePromise !== authCodePromise) { + this._unreportedAuthCodeFailure = authCodePromise + } + // Detach only the failed listener so a later authorization attempt starts a fresh one. + // Keep the failed promise current for callers that are still resuming from startup. + this._server = null + this._cancelCallbackServer = null + this._stopDetachedListener(server, cancel) + }) + } + + /** Stops a listener that is no longer attached. A replacement listener waits because it may reuse the port. */ + private _stopDetachedListener(server: http.Server, cancel: () => void): void { + const cleanup = stopCallbackServerIgnoringErrors(server, cancel).finally(() => { + if (this._listenerCleanup === cleanup) { + this._listenerCleanup = null } }) + this._listenerCleanup = cleanup } state(): string { @@ -521,12 +548,24 @@ export class McpOAuthClientProvider implements OAuthClientProvider { * Resolves with the authorization code once the user completes the OAuth * browser flow and the local callback server receives the redirect. * Rejects on error or 5-minute timeout. + * + * If the listener failed before any caller waited, the first call rejects with + * that failure without starting another listener. A later call starts a fresh one. */ async waitForAuthCode(): Promise { - if (!this._authCodePromise || !this._server) { + if (!this._unreportedAuthCodeFailure && (!this._authCodePromise || !this._server)) { await this._ensureCallbackServer() } - return this._authCodePromise! + const unreportedFailure = this._unreportedAuthCodeFailure + if (unreportedFailure) { + this._unreportedAuthCodeFailure = null + return unreportedFailure + } + // A caller that waited for listener startup receives that listener's result, even when + // another caller already took the failure while this caller was resuming. + const authCodePromise = this._authCodePromise! + this._awaitedAuthCodePromise = authCodePromise + return authCodePromise } /** @@ -655,13 +694,30 @@ export class McpOAuthClientProvider implements OAuthClientProvider { if (this._ensureServerPromise) { await this._ensureServerPromise.catch(() => {}) } - if (!this._closed && this._server) { + // A closed flow must not report its failure to a later flow. + this._unreportedAuthCodeFailure = null + const server = this._server + if (!this._closed && server) { this._closed = true - await stopCallbackServer(this._server, this._cancelCallbackServer ?? (() => {})).catch(() => {}) + const cancel = this._cancelCallbackServer ?? (() => {}) + // Detach before cancelling so the rejection handler does not record the cancellation. this._server = null this._cancelCallbackServer = null this._authCodePromise = null + this._stopDetachedListener(server, cancel) } + if (this._listenerCleanup) { + await this._listenerCleanup + } + } +} + +/** Stops a callback listener. A cleanup error must not replace the callback error. */ +async function stopCallbackServerIgnoringErrors(server: http.Server, cancel: () => void): Promise { + try { + await stopCallbackServer(server, cancel) + } catch { + // The callback result, not the cleanup result, is the error that callers report. } } diff --git a/src/services/mcp/__tests__/McpHub.spec.ts b/src/services/mcp/__tests__/McpHub.spec.ts index 18ef524bcc..35ff10d64b 100644 --- a/src/services/mcp/__tests__/McpHub.spec.ts +++ b/src/services/mcp/__tests__/McpHub.spec.ts @@ -3211,6 +3211,29 @@ describe("McpHub", () => { ).resolves.toBeUndefined() }) + it("should record the callback error and disconnect when waitForAuthCode rejects", async () => { + vsc.window.showInformationMessage.mockResolvedValueOnce(t("mcp:oauth.flow.authenticateButton")) + // The provider reports a callback listener failure, such as a timeout before the user authenticated. + mockAuthProvider.waitForAuthCode.mockRejectedValueOnce(new Error("Callback timeout")) + + await mcpHub["_initiateOAuthFlow"]( + serverName, + source, + config, + mockAuthProvider, + mockTransport, + mockConnection, + ) + + expect(mockAuthProvider.openBrowser).toHaveBeenCalledTimes(1) + expect(mockAuthProvider.exchangeCodeForTokens).not.toHaveBeenCalled() + expect(mockAuthProvider.close).toHaveBeenCalled() + expect(mockConnection.server.status).toBe("disconnected") + expect(mcpHub["appendErrorMessage"]).toHaveBeenCalledTimes(1) + expect(mcpHub["appendErrorMessage"]).toHaveBeenCalledWith(mockConnection, "Callback timeout") + expect(mcpHub["connectToServer"]).not.toHaveBeenCalled() + }) + it("should disconnect and flag error when Cancel is pressed after Authenticate is clicked", async () => { // Simulate: user clicks Authenticate in the toast, then cancels via the progress bar // while waitForAuthCode is still pending. diff --git a/src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts b/src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts index 43e9e1ad18..185bf4c461 100644 --- a/src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts +++ b/src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts @@ -1,6 +1,7 @@ import { describe, it, expect, vi, beforeEach, afterAll } from "vitest" import type { Server } from "http" import type { OAuthClientInformationFull } from "@modelcontextprotocol/sdk/shared/auth.js" +import type { CallbackResult } from "../utils/callbackServer" // Mock vscode vi.mock("vscode", () => ({ @@ -112,9 +113,59 @@ function deferred() { return { promise, resolve, reject } } +/** Waits one macrotask so every pending promise reaction runs. */ +function flushPromises(): Promise { + return new Promise((resolve) => setTimeout(resolve, 0)) +} + +/** Runs a scenario and returns every unhandled rejection Node reports for it. */ +async function collectUnhandledRejections(run: () => Promise): Promise { + const unhandled: unknown[] = [] + const onUnhandledRejection = (reason: unknown) => unhandled.push(reason) + process.on("unhandledRejection", onUnhandledRejection) + try { + await run() + // Node reports an unhandled rejection only after the microtask queue drains. + await flushPromises() + } finally { + process.off("unhandledRejection", onUnhandledRejection) + } + return unhandled +} + +/** Queues callback listeners whose results each test settles explicitly. Listener ports start at 12345. */ +function mockCallbackListeners(count: number) { + return Array.from({ length: count }, (_, index) => { + const listener = { + // The provider only passes the server back to stopCallbackServer(), which is mocked. + server: { close: vi.fn((cb: () => void) => cb()) } as unknown as Server, + port: 12345 + index, + result: deferred(), + cancel: vi.fn(), + } + vi.mocked(startCallbackServer).mockResolvedValueOnce({ + server: listener.server, + port: listener.port, + result: listener.result.promise, + cancel: listener.cancel, + }) + return listener + }) +} + +/** A static client without a redirect URI starts its callback listener in create(). */ +function createStaticProvider(): Promise { + return McpOAuthClientProvider.create("https://example.com/mcp", createMockSecretStorage(), undefined, { + clientId: "registered-client", + }) +} + describe("McpOAuthClientProvider", () => { beforeEach(() => { vi.clearAllMocks() + // clearAllMocks() keeps queued one-time listeners. Reset them so a failed test cannot leak them. + vi.mocked(startCallbackServer).mockReset() + vi.mocked(stopCallbackServer).mockReset().mockResolvedValue(undefined) McpOAuthClientProvider.clearNonOAuthCache() }) @@ -1300,49 +1351,158 @@ describe("McpOAuthClientProvider", () => { await provider.close() }) - it("should close a failed listener, preserve its error, and start a fresh listener", async () => { - const firstResult = deferred<{ code?: string }>() - const secondResult = deferred<{ code?: string }>() - const cleanupFinished = deferred() - const firstServer = { close: vi.fn((cb: () => void) => cb()) } - const secondServer = { close: vi.fn((cb: () => void) => cb()) } - const firstCancel = vi.fn() - vi.mocked(stopCallbackServer).mockReturnValueOnce(cleanupFinished.promise) - vi.mocked(startCallbackServer) - .mockResolvedValueOnce({ - server: firstServer as unknown as Server, - port: 12345, - result: firstResult.promise, - cancel: firstCancel, - }) - .mockResolvedValueOnce({ - server: secondServer as unknown as Server, - port: 12346, - result: secondResult.promise, - cancel: vi.fn(), + it.each([ + { timing: "before the first caller waits", waitBeforeFailure: false }, + { timing: "while a caller is waiting", waitBeforeFailure: true }, + ])( + "should report a listener failure $timing once, close the listener, and retry on a fresh listener", + async ({ waitBeforeFailure }) => { + const [first, second] = mockCallbackListeners(2) + const callbackError = new Error("Callback listener failed") + vi.mocked(stopCallbackServer).mockRejectedValueOnce(new Error("Listener close failed")) + const provider = await createStaticProvider() + + const unhandled = await collectUnhandledRejections(async () => { + // Attach the expectation before the failure so a waiting caller is never unhandled. + const waitingRejection = waitBeforeFailure + ? expect(provider.waitForAuthCode()).rejects.toBe(callbackError) + : undefined + first.result.reject(callbackError) + await flushPromises() + + const report = waitingRejection ?? expect(provider.waitForAuthCode()).rejects.toBe(callbackError) + await flushPromises() + // Reporting the failure must not start another listener and wait on it instead. + expect(startCallbackServer).toHaveBeenCalledTimes(1) + // The cleanup error must not replace the callback error. + await report + expect(stopCallbackServer).toHaveBeenCalledWith(first.server, first.cancel) + + const retry = provider.waitForAuthCode() + await vi.waitFor(() => expect(startCallbackServer).toHaveBeenCalledTimes(2)) + expect(provider.redirectUrl).toBe("http://127.0.0.1:12346/callback") + second.result.resolve({ code: "second-code" }) + await expect(retry).resolves.toBe("second-code") }) - const provider = await McpOAuthClientProvider.create( - "https://example.com/mcp", - createMockSecretStorage(), - undefined, - { clientId: "registered-client" }, - ) - const firstWait = provider.waitForAuthCode() - const firstRejection = expect(firstWait).rejects.toThrow("Callback listener failed") - firstResult.reject(new Error("Callback listener failed")) - await firstRejection - await vi.waitFor(() => expect(stopCallbackServer).toHaveBeenCalledWith(firstServer, firstCancel)) + expect(unhandled).toEqual([]) + await provider.close() + }, + ) + + it.each([ + { timing: "before the first caller waits", waitBeforeCallback: false }, + { timing: "while a caller is waiting", waitBeforeCallback: true }, + ])( + "should return a code received $timing without starting another listener", + async ({ waitBeforeCallback }) => { + const [first] = mockCallbackListeners(1) + const provider = await createStaticProvider() + + const waiting = waitBeforeCallback ? provider.waitForAuthCode() : undefined + first.result.resolve({ code: "first-code" }) + await flushPromises() + + await expect(waiting ?? provider.waitForAuthCode()).resolves.toBe("first-code") + expect(startCallbackServer).toHaveBeenCalledTimes(1) + // A successful callback is not a failure, so the listener stays attached until close(). + expect(stopCallbackServer).not.toHaveBeenCalled() + await provider.close() + expect(stopCallbackServer).toHaveBeenCalledWith(first.server, first.cancel) + }, + ) + + it("should report an immediate listener failure to every caller that waited for the listener to start", async () => { + vi.mocked(startCallbackServer).mockResolvedValueOnce({ + // The provider only passes the server back to stopCallbackServer(), which is mocked. + server: { close: vi.fn((cb: () => void) => cb()) } as unknown as Server, + port: 12345, + result: Promise.resolve({ error: "access_denied" }), + cancel: vi.fn(), + }) + const provider = await McpOAuthClientProvider.create("https://example.com/mcp", createMockSecretStorage()) + + const results = await Promise.allSettled([provider.waitForAuthCode(), provider.waitForAuthCode()]) + + const expected = { status: "rejected", reason: new Error("OAuth authorization failed: access_denied") } + expect(results).toEqual([expected, expected]) expect(startCallbackServer).toHaveBeenCalledTimes(1) - cleanupFinished.resolve(undefined) - await cleanupFinished.promise - await Promise.resolve() + await provider.close() + }) + + it.each([ + { + action: "redirectToAuthorization() starts a replacement listener", + endFailedAttempt: (provider: McpOAuthClientProvider) => + provider.redirectToAuthorization(new URL("https://auth.example.com/authorize")), + }, + { + action: "close() ends the failed flow", + endFailedAttempt: (provider: McpOAuthClientProvider) => provider.close(), + }, + ])("should discard an unreported listener failure after $action", async ({ endFailedAttempt }) => { + const [first, second] = mockCallbackListeners(2) + const provider = await createStaticProvider() + first.result.reject(new Error("Callback listener failed")) + await flushPromises() + + await endFailedAttempt(provider) + const wait = provider.waitForAuthCode() + await vi.waitFor(() => expect(startCallbackServer).toHaveBeenCalledTimes(2)) + second.result.resolve({ code: "second-code" }) + + await expect(wait).resolves.toBe("second-code") + await provider.close() + }) - const secondWait = provider.waitForAuthCode() + it("should report a replacement listener failure that happens before its first caller waits", async () => { + const [first, second] = mockCallbackListeners(2) + const provider = await createStaticProvider() + const firstError = new Error("First listener failed") + const firstReport = expect(provider.waitForAuthCode()).rejects.toBe(firstError) + first.result.reject(firstError) + await firstReport + + // The SDK starts the next attempt's listener before the user can authenticate. + await provider.redirectToAuthorization(new URL("https://auth.example.com/authorize")) + expect(startCallbackServer).toHaveBeenCalledTimes(2) + const secondError = new Error("Callback timeout") + second.result.reject(secondError) + await flushPromises() + + // The earlier caller waited on the first listener, not this one, so this failure is still unreported. + const secondReport = expect(provider.waitForAuthCode()).rejects.toBe(secondError) + await flushPromises() + expect(startCallbackServer).toHaveBeenCalledTimes(2) + await secondReport + await provider.close() + }) + + it("should start a replacement listener only after the failed listener stops", async () => { + const [first, second] = mockCallbackListeners(2) + const cleanup = deferred() + vi.mocked(stopCallbackServer).mockReturnValueOnce(cleanup.promise) + const provider = await createStaticProvider() + const callbackError = new Error("Callback listener failed") + first.result.reject(callbackError) + await flushPromises() + await expect(provider.waitForAuthCode()).rejects.toBe(callbackError) + + const retry = provider.waitForAuthCode() + await flushPromises() + // The replacement reuses the failed listener's port, so it must wait until that listener closes. + expect(startCallbackServer).toHaveBeenCalledTimes(1) + + cleanup.resolve() await vi.waitFor(() => expect(startCallbackServer).toHaveBeenCalledTimes(2)) - expect(provider.redirectUrl).toBe("http://127.0.0.1:12346/callback") - secondResult.resolve({ code: "second-code" }) - await expect(secondWait).resolves.toBe("second-code") + expect(startCallbackServer).toHaveBeenLastCalledWith( + first.port, + expect.any(String), + "127.0.0.1", + "/callback", + ) + second.result.resolve({ code: "second-code" }) + await expect(retry).resolves.toBe("second-code") await provider.close() }) }) @@ -1382,6 +1542,61 @@ describe("McpOAuthClientProvider", () => { expect(stopCallbackServer).not.toHaveBeenCalled() }) + it("should wait for a failed listener to stop", async () => { + const [first] = mockCallbackListeners(1) + const cleanup = deferred() + vi.mocked(stopCallbackServer).mockReturnValueOnce(cleanup.promise) + const provider = await createStaticProvider() + first.result.reject(new Error("Callback listener failed")) + await flushPromises() + + let closed = false + const closing = provider.close().then(() => { + closed = true + }) + await flushPromises() + expect(closed).toBe(false) + + cleanup.resolve() + await closing + expect(stopCallbackServer).toHaveBeenCalledTimes(1) + }) + + it.each([ + { outcome: "succeeds", stopListener: async (_server: Server, cancel: () => void) => cancel() }, + { + outcome: "fails", + stopListener: async (_server: Server, cancel: () => void) => { + cancel() + throw new Error("Listener close failed") + }, + }, + ])( + "should start a fresh listener after close() when stopping the old listener $outcome", + async ({ stopListener }) => { + const [first, second] = mockCallbackListeners(2) + // Like the real listener, cancelling rejects the pending callback result. + first.cancel.mockImplementation(() => first.result.reject(new Error("Callback cancelled"))) + vi.mocked(stopCallbackServer).mockImplementationOnce(stopListener) + const provider = await createStaticProvider() + + const unhandled = await collectUnhandledRejections(async () => { + await provider.close() + expect(stopCallbackServer).toHaveBeenCalledTimes(1) + expect(first.cancel).toHaveBeenCalledTimes(1) + + // The cancellation belongs to the closed flow and must not be reported to the next attempt. + const wait = provider.waitForAuthCode() + await vi.waitFor(() => expect(startCallbackServer).toHaveBeenCalledTimes(2)) + second.result.resolve({ code: "second-code" }) + await expect(wait).resolves.toBe("second-code") + }) + + expect(unhandled).toEqual([]) + await provider.close() + }, + ) + it("should handle cancellation of an eagerly started callback without an unhandled rejection", async () => { const result = deferred<{ code?: string }>() const cancel = vi.fn(() => result.reject(new Error("Callback cancelled"))) @@ -1432,6 +1647,7 @@ describe("McpOAuthClientProvider", () => { }) it("should skip dynamic registration for a configured static client", async () => { + setupCallbackServerMock() const provider = await McpOAuthClientProvider.create( "https://example.com/mcp", createMockSecretStorage(), diff --git a/src/services/mcp/utils/__tests__/callbackServer.listener.spec.ts b/src/services/mcp/utils/__tests__/callbackServer.listener.spec.ts new file mode 100644 index 0000000000..119beb6140 --- /dev/null +++ b/src/services/mcp/utils/__tests__/callbackServer.listener.spec.ts @@ -0,0 +1,79 @@ +import * as net from "net" + +import { startCallbackServer, stopCallbackServer } from "../callbackServer" + +// This spec uses a real loopback listener. callbackServer.spec.ts mocks the http module. +const STATE = "listener-state" +const RESPONSE_WAIT_MS = 2_000 + +/** Opens a raw keep-alive connection, so that a request can reach the listener after the callback. */ +async function connectClient(port: number) { + const socket = net.connect(port, "127.0.0.1") + let received = "" + socket.setEncoding("utf8") + socket.on("data", (chunk: string) => { + received += chunk + }) + socket.on("error", () => { + // The server can close the connection when it stops. The assertions check the responses. + }) + await new Promise((resolve, reject) => { + socket.once("connect", resolve) + socket.once("error", reject) + }) + return { + send(path: string) { + socket.write(`GET ${path} HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: keep-alive\r\n\r\n`) + }, + /** Waits until the client has received one response, and returns its status line. */ + async waitForResponse(): Promise { + let status: string | undefined + await vi.waitFor( + () => { + status = received.match(/^HTTP\/1\.1 \d{3}/m)?.[0] + expect(status).toBeDefined() + }, + { timeout: RESPONSE_WAIT_MS }, + ) + return status + }, + destroy: () => socket.destroy(), + } +} + +describe("startCallbackServer with a real listener", () => { + it("should answer a request that arrives after the callback wait ends, so that the server can stop", async () => { + const { server, port, result, cancel } = await startCallbackServer(0, STATE) + const outcome = result.catch((error: unknown) => error) + const client = await connectClient(port) + let stopped = false + try { + // The listener no longer waits for a callback, but the connection is still open. + cancel() + expect(await outcome).toEqual(new Error("Callback cancelled")) + + client.send(`/callback?code=late-code&state=${STATE}`) + + expect(await client.waitForResponse()).toBe("HTTP/1.1 404") + await expect(stopCallbackServer(server, cancel)).resolves.toBeUndefined() + stopped = true + } finally { + client.destroy() + if (!stopped) server.close() + } + }) + + it("should accept a callback and then stop", async () => { + const { server, port, result, cancel } = await startCallbackServer(0, STATE) + const client = await connectClient(port) + try { + client.send(`/callback?code=first-code&state=${STATE}`) + + await expect(result).resolves.toMatchObject({ code: "first-code", state: STATE }) + expect(await client.waitForResponse()).toBe("HTTP/1.1 200") + await expect(stopCallbackServer(server, cancel)).resolves.toBeUndefined() + } finally { + client.destroy() + } + }) +}) diff --git a/src/services/mcp/utils/callbackServer.ts b/src/services/mcp/utils/callbackServer.ts index 22e98d337e..2ef72a3159 100644 --- a/src/services/mcp/utils/callbackServer.ts +++ b/src/services/mcp/utils/callbackServer.ts @@ -84,7 +84,12 @@ export function startCallbackServer( } server.on("request", (req: any, res: any) => { - if (resolved) return + if (resolved) { + // Answer late requests so that their connections can close when the server stops. + res.writeHead(404) + res.end("Not found") + return + } const url = new URL(req.url || "", `http://localhost:${actualPort}`) const pathname = url.pathname