diff --git a/.changeset/fix-mcp-oauth-cancel.md b/.changeset/fix-mcp-oauth-cancel.md new file mode 100644 index 000000000..07f295679 --- /dev/null +++ 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. diff --git a/packages/agent-core/src/mcp/oauth/callback-server.ts b/packages/agent-core/src/mcp/oauth/callback-server.ts index a1c902e28..5b8a0c0cf 100644 --- 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: ') + * - `close()` is called → OAuthCallbackClosedError */ waitForCode(opts: { signal?: AbortSignal; timeoutMs?: number }): Promise; close(): Promise; } +export class OAuthCallbackClosedError extends Error { + constructor() { + super('OAuth callback listener closed'); + this.name = 'OAuthCallbackClosedError'; + } +} + const SUCCESS_HTML = 'Authorized' + '' + @@ -46,12 +54,29 @@ const ERROR_HTML = export async function startCallbackServer(): Promise { 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 { 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((resolve, reject) => { @@ -110,44 +135,49 @@ export async function startCallbackServer(): Promise { 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((resolve) => { + let closeServerPromise: Promise | undefined; + const closeServer = (): Promise => { + closeServerPromise ??= new Promise((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((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) { diff --git a/packages/agent-core/test/mcp/oauth-callback-server.test.ts b/packages/agent-core/test/mcp/oauth-callback-server.test.ts new file mode 100644 index 000000000..19d6562d4 --- /dev/null +++ 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()), + 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; + }); +});