Skip to content

Commit cca90fa

Browse files
Report an MCP OAuth callback failure before a retry
A callback listener can fail before waitForAuthCode() starts to wait. The provider then started a new listener and hid the original error for up to five minutes. - Keep a failure that occurs while no caller waits. - Report that failure once on the next waitForAuthCode() call. - Start a new listener on the call after that. - Close each failed listener. Ignore errors from that cleanup. - Start a new listener only after the old listener closes. - Discard a kept failure when a new listener starts or close() runs.
1 parent 0e208c4 commit cca90fa

3 files changed

Lines changed: 342 additions & 52 deletions

File tree

‎src/services/mcp/McpOAuthClientProvider.ts‎

Lines changed: 70 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,15 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
9393
private _effectiveResourceIndicator: string | null
9494
/** Deduplicates concurrent _ensureCallbackServer() calls. */
9595
private _ensureServerPromise: Promise<void> | null = null
96+
/** Callback promise most recently returned to a waitForAuthCode() caller. */
97+
private _awaitedAuthCodePromise: Promise<string> | null = null
98+
/**
99+
* Rejected callback promise of a listener that failed while no caller was waiting.
100+
* The next waitForAuthCode() call reports it once instead of starting another listener.
101+
*/
102+
private _unreportedAuthCodeFailure: Promise<string> | null = null
103+
/** Stops a detached listener. A replacement listener waits for it because it may reuse the port. */
104+
private _listenerCleanup: Promise<void> | null = null
96105

97106
private constructor(
98107
private readonly _serverUrl: string,
@@ -243,6 +252,11 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
243252

244253
private async _doStartCallbackServer(): Promise<void> {
245254
this._closed = false
255+
// A new listener belongs to a new authorization attempt, so an earlier failure no longer applies.
256+
this._unreportedAuthCodeFailure = null
257+
if (this._listenerCleanup) {
258+
await this._listenerCleanup
259+
}
246260
const { port: configuredPort, host } = getCallbackAddress(this._configuredRedirectUri)
247261
const { server, port, result, cancel } = await startCallbackServer(
248262
configuredPort ?? this._port,
@@ -260,18 +274,31 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
260274
this._authCodePromise = authCodePromise
261275

262276
// Static clients can start this promise before waitForAuthCode() observes it.
263-
// Observe rejection now to prevent an unhandled rejection, but keep the original
264-
// promise so waitForAuthCode() still receives the callback error. Stop only this
265-
// promise's listener before clearing it so a later flow can bind a fresh listener.
266-
void authCodePromise.catch(async () => {
267-
if (this._authCodePromise === authCodePromise && this._server === server) {
268-
await stopCallbackServer(server, cancel).catch(() => {})
269-
if (this._authCodePromise === authCodePromise && this._server === server) {
270-
this._server = null
271-
this._cancelCallbackServer = null
272-
}
277+
// Observe rejection now to prevent an unhandled rejection. A caller that is already
278+
// waiting receives the error from this promise; otherwise keep the rejected promise so
279+
// the next waitForAuthCode() call reports it instead of waiting on a new listener.
280+
void authCodePromise.catch(() => {
281+
// close() detaches this promise before it cancels the listener and stops it itself.
282+
if (this._authCodePromise !== authCodePromise) return
283+
if (this._awaitedAuthCodePromise !== authCodePromise) {
284+
this._unreportedAuthCodeFailure = authCodePromise
285+
}
286+
// Detach only the failed listener so a later authorization attempt starts a fresh one.
287+
// Keep the failed promise current for callers that are still resuming from startup.
288+
this._server = null
289+
this._cancelCallbackServer = null
290+
this._stopDetachedListener(server, cancel)
291+
})
292+
}
293+
294+
/** Stops a listener that is no longer attached. A replacement listener waits because it may reuse the port. */
295+
private _stopDetachedListener(server: http.Server, cancel: () => void): void {
296+
const cleanup = stopCallbackServerIgnoringErrors(server, cancel).finally(() => {
297+
if (this._listenerCleanup === cleanup) {
298+
this._listenerCleanup = null
273299
}
274300
})
301+
this._listenerCleanup = cleanup
275302
}
276303

277304
state(): string {
@@ -520,12 +547,24 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
520547
* Resolves with the authorization code once the user completes the OAuth
521548
* browser flow and the local callback server receives the redirect.
522549
* Rejects on error or 5-minute timeout.
550+
*
551+
* If the listener failed before any caller waited, the first call rejects with
552+
* that failure without starting another listener. A later call starts a fresh one.
523553
*/
524554
async waitForAuthCode(): Promise<string> {
525-
if (!this._authCodePromise || !this._server) {
555+
if (!this._unreportedAuthCodeFailure && (!this._authCodePromise || !this._server)) {
526556
await this._ensureCallbackServer()
527557
}
528-
return this._authCodePromise!
558+
const unreportedFailure = this._unreportedAuthCodeFailure
559+
if (unreportedFailure) {
560+
this._unreportedAuthCodeFailure = null
561+
return unreportedFailure
562+
}
563+
// A caller that waited for listener startup receives that listener's result, even when
564+
// another caller already took the failure while this caller was resuming.
565+
const authCodePromise = this._authCodePromise!
566+
this._awaitedAuthCodePromise = authCodePromise
567+
return authCodePromise
529568
}
530569

531570
/**
@@ -654,13 +693,30 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
654693
if (this._ensureServerPromise) {
655694
await this._ensureServerPromise.catch(() => {})
656695
}
657-
if (!this._closed && this._server) {
696+
// A closed flow must not report its failure to a later flow.
697+
this._unreportedAuthCodeFailure = null
698+
const server = this._server
699+
if (!this._closed && server) {
658700
this._closed = true
659-
await stopCallbackServer(this._server, this._cancelCallbackServer ?? (() => {})).catch(() => {})
701+
const cancel = this._cancelCallbackServer ?? (() => {})
702+
// Detach before cancelling so the rejection handler does not record the cancellation.
660703
this._server = null
661704
this._cancelCallbackServer = null
662705
this._authCodePromise = null
706+
this._stopDetachedListener(server, cancel)
663707
}
708+
if (this._listenerCleanup) {
709+
await this._listenerCleanup
710+
}
711+
}
712+
}
713+
714+
/** Stops a callback listener. A cleanup error must not replace the callback error. */
715+
async function stopCallbackServerIgnoringErrors(server: http.Server, cancel: () => void): Promise<void> {
716+
try {
717+
await stopCallbackServer(server, cancel)
718+
} catch {
719+
// The callback result, not the cleanup result, is the error that callers report.
664720
}
665721
}
666722

‎src/services/mcp/__tests__/McpHub.spec.ts‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3211,6 +3211,29 @@ describe("McpHub", () => {
32113211
).resolves.toBeUndefined()
32123212
})
32133213

3214+
it("should record the callback error and disconnect when waitForAuthCode rejects", async () => {
3215+
vsc.window.showInformationMessage.mockResolvedValueOnce(t("mcp:oauth.flow.authenticateButton"))
3216+
// The provider reports a callback listener failure, such as a timeout before the user authenticated.
3217+
mockAuthProvider.waitForAuthCode.mockRejectedValueOnce(new Error("Callback timeout"))
3218+
3219+
await mcpHub["_initiateOAuthFlow"](
3220+
serverName,
3221+
source,
3222+
config,
3223+
mockAuthProvider,
3224+
mockTransport,
3225+
mockConnection,
3226+
)
3227+
3228+
expect(mockAuthProvider.openBrowser).toHaveBeenCalledTimes(1)
3229+
expect(mockAuthProvider.exchangeCodeForTokens).not.toHaveBeenCalled()
3230+
expect(mockAuthProvider.close).toHaveBeenCalled()
3231+
expect(mockConnection.server.status).toBe("disconnected")
3232+
expect(mcpHub["appendErrorMessage"]).toHaveBeenCalledTimes(1)
3233+
expect(mcpHub["appendErrorMessage"]).toHaveBeenCalledWith(mockConnection, "Callback timeout")
3234+
expect(mcpHub["connectToServer"]).not.toHaveBeenCalled()
3235+
})
3236+
32143237
it("should disconnect and flag error when Cancel is pressed after Authenticate is clicked", async () => {
32153238
// Simulate: user clicks Authenticate in the toast, then cancels via the progress bar
32163239
// while waitForAuthCode is still pending.

0 commit comments

Comments
 (0)