Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/fix-mcp-oauth-cancel.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@moonshot-ai/kimi-code": patch
---

Fix MCP OAuth cancellation leaving an in-flight authorization waiting for its callback timeout.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Port the OAuth cancellation fix to v2

This release note claims the Kimi Code MCP OAuth cancellation hang is fixed, but the patch only changes the legacy packages/agent-core callback server; the default Kimi Code/kap-server path uses packages/agent-core-v2/src/mcpCore/oauth/service.ts and its mcpCore/oauth/callback-server.ts still has close() only closing the listener without rejecting a pending waitForCode(). Users on the v2 path can still cancel an auth flow and wait until the callback timeout, so please port the same outcome/closed-error handling to agent-core-v2 or narrow this changeset to the legacy path.

AGENTS.md reference: AGENTS.md:L21-L22

Useful? React with 👍 / 👎.

98 changes: 64 additions & 34 deletions packages/agent-core/src/mcp/oauth/callback-server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;">' +
Expand All @@ -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) => {
Expand All @@ -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) => {
Expand All @@ -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) {
Expand Down
87 changes: 87 additions & 0 deletions packages/agent-core/test/mcp/oauth-callback-server.test.ts
Original file line number Diff line number Diff line change
@@ -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;
});
});
Loading