mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-08-24 08:06:38 +00:00
fix(agent-core-v2): settle early and cancelled MCP OAuth callbacks
This commit is contained in:
7Sageer
parent
29ca69a42d
commit
07b03eaa74
2 changed files with 204 additions and 34 deletions
|
|
@ -8,10 +8,24 @@ export interface CallbackResult {
|
|||
|
||||
export interface CallbackServer {
|
||||
readonly redirectUri: string;
|
||||
/**
|
||||
* Resolves with the OAuth callback payload, or rejects when:
|
||||
* - `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;">' +
|
||||
|
|
@ -29,12 +43,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) => {
|
||||
|
|
@ -61,26 +92,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) => {
|
||||
|
|
@ -93,44 +124,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) {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,134 @@
|
|||
import { createServer as createHttpServer, type Server as HttpServer } from 'node:http';
|
||||
import type { AddressInfo as HttpAddress } from 'node:net';
|
||||
|
||||
import { afterEach, describe, expect, it } from 'vitest';
|
||||
|
||||
import {
|
||||
type CallbackServer,
|
||||
OAuthCallbackClosedError,
|
||||
startCallbackServer,
|
||||
} from '#/mcpCore/oauth/callback-server';
|
||||
import { McpOAuthService } from '#/mcpCore/oauth/service';
|
||||
|
||||
import { createMemoryMcpOAuthStore } from '../stubs';
|
||||
|
||||
const cleanups: Array<() => Promise<void> | void> = [];
|
||||
afterEach(async () => {
|
||||
while (cleanups.length > 0) {
|
||||
await cleanups.pop()?.();
|
||||
}
|
||||
});
|
||||
|
||||
function trackServer(server: CallbackServer): void {
|
||||
cleanups.push(() => server.close());
|
||||
}
|
||||
|
||||
async function startRegistrationServer(): Promise<{ readonly url: string }> {
|
||||
const httpServer: HttpServer = createHttpServer((req, res) => {
|
||||
if (req.method !== 'POST' || req.url !== '/register') {
|
||||
res.writeHead(404).end();
|
||||
return;
|
||||
}
|
||||
let body = '';
|
||||
req.on('data', (chunk: Buffer) => {
|
||||
body += chunk.toString('utf-8');
|
||||
});
|
||||
req.on('end', () => {
|
||||
const metadata = JSON.parse(body) as Record<string, unknown>;
|
||||
res.writeHead(201, { 'content-type': 'application/json' });
|
||||
res.end(JSON.stringify({ ...metadata, client_id: 'test-client' }));
|
||||
});
|
||||
});
|
||||
await new Promise<void>((resolve) => httpServer.listen(0, '127.0.0.1', resolve));
|
||||
cleanups.push(
|
||||
() =>
|
||||
new Promise<void>((resolve, reject) => {
|
||||
httpServer.close((err) => {
|
||||
if (err) {
|
||||
reject(err);
|
||||
return;
|
||||
}
|
||||
resolve();
|
||||
});
|
||||
}),
|
||||
);
|
||||
const port = (httpServer.address() as HttpAddress).port;
|
||||
return { url: `http://127.0.0.1:${port}` };
|
||||
}
|
||||
|
||||
describe('startCallbackServer', () => {
|
||||
it('resolves a late waitForCode with a callback that arrived before it', async () => {
|
||||
const server = await startCallbackServer();
|
||||
trackServer(server);
|
||||
|
||||
const response = await fetch(`${server.redirectUri}?code=early-code&state=early-state`);
|
||||
expect(response.status).toBe(200);
|
||||
|
||||
await expect(server.waitForCode({ timeoutMs: 10_000 })).resolves.toEqual({
|
||||
code: 'early-code',
|
||||
state: 'early-state',
|
||||
});
|
||||
});
|
||||
|
||||
it('delivers the callback payload to a pending wait', async () => {
|
||||
const server = await startCallbackServer();
|
||||
trackServer(server);
|
||||
const pending = server.waitForCode({ timeoutMs: 10_000 });
|
||||
|
||||
await fetch(`${server.redirectUri}?code=code-1&state=state-1`);
|
||||
|
||||
await expect(pending).resolves.toEqual({ code: 'code-1', state: 'state-1' });
|
||||
});
|
||||
|
||||
it('rejects a pending wait with a closed error when explicitly closed', async () => {
|
||||
const server = await startCallbackServer();
|
||||
trackServer(server);
|
||||
const pending = server.waitForCode({ timeoutMs: 10_000 });
|
||||
const rejection = expect(pending).rejects.toBeInstanceOf(OAuthCallbackClosedError);
|
||||
|
||||
await server.close();
|
||||
|
||||
await rejection;
|
||||
});
|
||||
|
||||
it('rejects a late waitForCode after close', async () => {
|
||||
const server = await startCallbackServer();
|
||||
trackServer(server);
|
||||
|
||||
await server.close();
|
||||
|
||||
await expect(server.waitForCode({ timeoutMs: 10_000 })).rejects.toBeInstanceOf(
|
||||
OAuthCallbackClosedError,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('McpOAuthService cancellation', () => {
|
||||
it('rejects an in-flight completion when the authorization flow is cancelled', async () => {
|
||||
const service = new McpOAuthService({ store: createMemoryMcpOAuthStore() });
|
||||
cleanups.push(() => service.dispose());
|
||||
const registrationServer = await startRegistrationServer();
|
||||
const provider = service.getProvider('example', 'https://mcp.example.test/rpc');
|
||||
await provider.ready;
|
||||
await provider.saveDiscoveryState({
|
||||
authorizationServerUrl: registrationServer.url,
|
||||
authorizationServerMetadata: {
|
||||
issuer: registrationServer.url,
|
||||
authorization_endpoint: `${registrationServer.url}/authorize`,
|
||||
token_endpoint: `${registrationServer.url}/token`,
|
||||
registration_endpoint: `${registrationServer.url}/register`,
|
||||
response_types_supported: ['code'],
|
||||
grant_types_supported: ['authorization_code'],
|
||||
token_endpoint_auth_methods_supported: ['none'],
|
||||
},
|
||||
});
|
||||
|
||||
const flow = await service.beginAuthorization('example', 'https://mcp.example.test/rpc');
|
||||
const completion = flow.complete({ timeoutMs: 10_000 });
|
||||
const rejection = expect(completion).rejects.toThrow('OAuth callback listener closed');
|
||||
|
||||
await flow.cancel();
|
||||
|
||||
await rejection;
|
||||
}, 15000);
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue