fix: settle cancelled MCP OAuth callbacks (#2899)

Co-authored-by: yuchengzhen <yuchengzhen@moonshot.cn>
2026-08-16 12:15:29 +00:00 · 2026-08-13 22:56:06 +08:00 · 2026-08-13 22:56:06 +08:00 · 102984aa66
commit 102984aa66
parent 1414d46028
3 changed files with 156 additions and 34 deletions
--- a/.changeset/fix-mcp-oauth-cancel.md
+++ b/.changeset/fix-mcp-oauth-cancel.md
@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Fix MCP OAuth cancellation leaving an in-flight authorization waiting for its callback timeout.
--- a/packages/agent-core/src/mcp/oauth/callback-server.ts
+++ b/packages/agent-core/src/mcp/oauth/callback-server.ts
@ -24,11 +24,19 @@ export interface CallbackServer {
   *  - `signal` aborts → AbortError
   *  - `timeoutMs` elapses → Error('OAuth callback timed out')
   *  - the user's authorization server returns an error → Error('OAuth error: <code>')
+   *  - `close()` is called → OAuthCallbackClosedError
   */
  waitForCode(opts: { signal?: AbortSignal; timeoutMs?: number }): Promise<CallbackResult>;
  close(): Promise<void>;
 }

+export class OAuthCallbackClosedError extends Error {
+  constructor() {
+    super('OAuth callback listener closed');
+    this.name = 'OAuthCallbackClosedError';
+  }
+}
+
 const SUCCESS_HTML =
  '<!doctype html><html><head><meta charset="utf-8"><title>Authorized</title></head>' +
  '<body style="font-family:system-ui,sans-serif;padding:2rem;">' +
@ -46,12 +54,29 @@ const ERROR_HTML =
 export async function startCallbackServer(): Promise<CallbackServer> {
  let resolveCode: ((value: CallbackResult) => void) | undefined;
  let rejectCode: ((reason: Error) => void) | undefined;
-  let settled = false;
+  let cleanupWait: (() => void) | undefined;
+  let outcome:
+    | { readonly status: 'pending' }
+    | { readonly status: 'resolved'; readonly value: CallbackResult }
+    | { readonly status: 'rejected'; readonly reason: Error } = { status: 'pending' };

-  const settle = (fn: () => void) => {
-    if (settled) return;
-    settled = true;
-    fn();
+  const settle = (
+    next:
+      | { readonly status: 'resolved'; readonly value: CallbackResult }
+      | { readonly status: 'rejected'; readonly reason: Error },
+  ) => {
+    if (outcome.status !== 'pending') return;
+    outcome = next;
+    cleanupWait?.();
+    cleanupWait = undefined;
+    if (next.status === 'resolved') {
+      resolveCode?.(next.value);
+    } else {
+      rejectCode?.(next.reason);
+    }
+    resolveCode = undefined;
+    rejectCode = undefined;
+    void closeServer();
  };

  const server: Server = createServer((req, res) => {
@ -78,26 +103,26 @@ export async function startCallbackServer(): Promise<CallbackServer> {
    if (errorParam !== null) {
      const description = url.searchParams.get('error_description') ?? '';
      res.writeHead(400, { 'content-type': 'text/html; charset=utf-8' }).end(ERROR_HTML);
-      settle(() => {
-        rejectCode?.(
-          new Error(`OAuth error: ${errorParam}${description ? ` — ${description}` : ''}`),
-        );
+      settle({
+        status: 'rejected',
+        reason: new Error(
+          `OAuth error: ${errorParam}${description ? ` — ${description}` : ''}`,
+        ),
      });
      return;
    }
    const code = url.searchParams.get('code');
    if (code === null || code.length === 0) {
      res.writeHead(400, { 'content-type': 'text/html; charset=utf-8' }).end(ERROR_HTML);
-      settle(() => {
-        rejectCode?.(new Error('OAuth callback missing authorization code'));
+      settle({
+        status: 'rejected',
+        reason: new Error('OAuth callback missing authorization code'),
      });
      return;
    }
    const state = url.searchParams.get('state') ?? undefined;
    res.writeHead(200, { 'content-type': 'text/html; charset=utf-8' }).end(SUCCESS_HTML);
-    settle(() => {
-      resolveCode?.({ code, state });
-    });
+    settle({ status: 'resolved', value: { code, state } });
  }

  await new Promise<void>((resolve, reject) => {
@ -110,44 +135,49 @@ export async function startCallbackServer(): Promise<CallbackServer> {
  const port = (server.address() as AddressInfo).port;
  const redirectUri = `http://127.0.0.1:${port}/callback`;

-  let closed = false;
-  const close = async () => {
-    if (closed) return;
-    closed = true;
-    await new Promise<void>((resolve) => {
+  let closeServerPromise: Promise<void> | undefined;
+  const closeServer = (): Promise<void> => {
+    closeServerPromise ??= new Promise<void>((resolve) => {
      server.close(() => {
        resolve();
      });
    });
+    return closeServerPromise;
+  };
+  const close = async () => {
+    settle({ status: 'rejected', reason: new OAuthCallbackClosedError() });
+    await closeServer();
  };

  const waitForCode: CallbackServer['waitForCode'] = ({ signal, timeoutMs } = {}) => {
    return new Promise<CallbackResult>((resolve, reject) => {
+      if (outcome.status === 'resolved') {
+        resolve(outcome.value);
+        return;
+      }
+      if (outcome.status === 'rejected') {
+        reject(outcome.reason);
+        return;
+      }
+
      let timer: NodeJS.Timeout | undefined;
      const onAbort = () => {
-        settle(() =>
-          rejectCode?.(
+        settle({
+          status: 'rejected',
+          reason:
            signal?.reason instanceof Error ? signal.reason : new Error('OAuth flow aborted'),
-          ),
-        );
+        });
      };
      const cleanup = () => {
        if (timer !== undefined) clearTimeout(timer);
        signal?.removeEventListener('abort', onAbort);
      };
-      resolveCode = (value) => {
-        cleanup();
-        void close();
-        resolve(value);
-      };
-      rejectCode = (reason) => {
-        cleanup();
-        void close();
-        reject(reason);
-      };
+      cleanupWait = cleanup;
+      resolveCode = resolve;
+      rejectCode = reject;
      if (timeoutMs !== undefined) {
        timer = setTimeout(() => {
-          settle(() => rejectCode?.(new Error('OAuth callback timed out')));
+          settle({ status: 'rejected', reason: new Error('OAuth callback timed out') });
        }, timeoutMs);
      }
      if (signal !== undefined) {
--- a/packages/agent-core/test/mcp/oauth-callback-server.test.ts
+++ b/packages/agent-core/test/mcp/oauth-callback-server.test.ts
@ -0,0 +1,87 @@
+/**
+ * Scenario: lifecycle completion for the localhost MCP OAuth callback listener.
+ * Responsibilities: closing rejects pending waits, successful callbacks survive cleanup, and
+ * service cancellation settles in-flight completion. The listener and service are real; only the
+ * external MCP SDK authorization boundary is mocked.
+ * Run: pnpm --filter @moonshot-ai/agent-core exec vitest run test/mcp/oauth-callback-server.test.ts
+ */
+
+import { mkdtemp, rm } from 'node:fs/promises';
+import { tmpdir } from 'node:os';
+import { join } from 'pathe';
+
+import { auth } from '@modelcontextprotocol/sdk/client/auth.js';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import {
+  type BeginAuthorizationResult,
+  type CallbackServer,
+  JsonFileStore,
+  McpOAuthService,
+  OAuthCallbackClosedError,
+  startCallbackServer,
+} from '../../src/mcp/oauth';
+
+vi.mock('@modelcontextprotocol/sdk/client/auth.js', async (importOriginal) => ({
+  ...(await importOriginal<typeof import('@modelcontextprotocol/sdk/client/auth.js')>()),
+  auth: vi.fn(),
+}));
+
+describe('OAuth callback server', () => {
+  let server: CallbackServer | undefined;
+
+  afterEach(async () => {
+    await server?.close();
+    server = undefined;
+  });
+
+  it('rejects a pending callback wait with a closed error when explicitly closed', async () => {
+    server = await startCallbackServer();
+    const pending = server.waitForCode({ timeoutMs: 60_000 });
+    const rejection = expect(pending).rejects.toBeInstanceOf(OAuthCallbackClosedError);
+
+    await server.close();
+
+    await rejection;
+  });
+
+  it('delivers the callback payload when success closes the listener', async () => {
+    server = await startCallbackServer();
+    const pending = server.waitForCode({ timeoutMs: 60_000 });
+
+    await fetch(`${server.redirectUri}?code=code-1&state=state-1`);
+
+    await expect(pending).resolves.toEqual({ code: 'code-1', state: 'state-1' });
+  });
+});
+
+describe('McpOAuthService cancellation', () => {
+  let dir: string;
+  let flow: BeginAuthorizationResult | undefined;
+
+  beforeEach(async () => {
+    dir = await mkdtemp(join(tmpdir(), 'kimi-mcp-oauth-cancel-'));
+    vi.mocked(auth).mockImplementation(async (provider) => {
+      await provider.redirectToAuthorization(new URL('https://auth.example.test/authorize'));
+      return 'REDIRECT';
+    });
+  });
+
+  afterEach(async () => {
+    await flow?.cancel();
+    flow = undefined;
+    await rm(dir, { recursive: true, force: true });
+    vi.clearAllMocks();
+  });
+
+  it('rejects an in-flight completion when the authorization flow is cancelled', async () => {
+    const service = new McpOAuthService({ store: new JsonFileStore(dir) });
+    flow = await service.beginAuthorization('example', 'https://mcp.example.test/rpc');
+    const completion = flow.complete({ timeoutMs: 60_000 });
+    const rejection = expect(completion).rejects.toThrow('OAuth callback listener closed');
+
+    await flow.cancel();
+
+    await rejection;
+  });
+});