mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-08-25 08:34:39 +00:00
fix(mcp): abort the auth::complete long poll on client disconnect
This commit is contained in:
parent
cd8a10d9aa
commit
3cf106152b
4 changed files with 61 additions and 2 deletions
|
|
@ -1,3 +1,5 @@
|
|||
import type { ServerResponse } from 'node:http';
|
||||
|
||||
import {
|
||||
ErrorCodes,
|
||||
IConfigService,
|
||||
|
|
@ -97,7 +99,12 @@ const inspectServersBodySchema = z.object({
|
|||
|
||||
const authCompleteBodySchema = z.object({
|
||||
flowId: z.string().min(1),
|
||||
timeoutMs: z.number().int().min(1).optional(),
|
||||
timeoutMs: z
|
||||
.number()
|
||||
.int()
|
||||
.min(1)
|
||||
.max(2 ** 31 - 1)
|
||||
.optional(),
|
||||
});
|
||||
|
||||
const authCancelBodySchema = z.object({ flowId: z.string().min(1) });
|
||||
|
|
@ -484,11 +491,20 @@ export function registerV2McpRoutes(app: V2McpRouteHost, core: Scope): void {
|
|||
tags: ['v2-mcp'],
|
||||
},
|
||||
async (req, reply) => {
|
||||
const { raw } = reply as unknown as { raw: ServerResponse };
|
||||
const disconnect = new AbortController();
|
||||
const onClose = (): void => {
|
||||
if (raw.writableFinished) return;
|
||||
disconnect.abort();
|
||||
};
|
||||
raw.once('close', onClose);
|
||||
try {
|
||||
await management().completeServerAuth(req.body);
|
||||
await management().completeServerAuth(req.body, { signal: disconnect.signal });
|
||||
reply.send(okEnvelope(null, req.id));
|
||||
} catch (err) {
|
||||
sendMappedError(reply, req.id, err);
|
||||
} finally {
|
||||
raw.off('close', onClose);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
|
|
|||
|
|
@ -301,6 +301,7 @@ export async function startServer(opts: ServerStartOptions): Promise<RunningServ
|
|||
disableRequestLogging: true,
|
||||
genReqId: (req) => resolveRequestId(req.headers),
|
||||
}) as unknown as FastifyInstance;
|
||||
app.server.requestTimeout = 0;
|
||||
registerRequestLogging(app);
|
||||
app.setValidatorCompiler(() => () => true);
|
||||
app.setSerializerCompiler(() => (data) => JSON.stringify(data));
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ interface FetchOptions {
|
|||
readonly method?: string;
|
||||
readonly headers?: HeaderMap;
|
||||
readonly body?: string;
|
||||
readonly signal?: AbortSignal;
|
||||
}
|
||||
|
||||
export function bearerToken(server: RunningServer): string {
|
||||
|
|
|
|||
|
|
@ -461,5 +461,46 @@ describe('server /api/v2/mcp', () => {
|
|||
expect(reset.body).toMatchObject({ code: 0, data: null });
|
||||
expect(stub.state.lastResetLocator).toEqual({ source: 'plugin', pluginId: 'p', serverName: 's' });
|
||||
});
|
||||
|
||||
it('rejects an overflowing auth:complete timeoutMs with 40001', async () => {
|
||||
const stub = makeMcpStub();
|
||||
await boot(stub);
|
||||
|
||||
const res = await call('POST', '/api/v2/mcp/auth:complete', {
|
||||
flowId: 'flow-1',
|
||||
timeoutMs: 2 ** 31,
|
||||
});
|
||||
|
||||
expect(res.body.code).toBe(40001);
|
||||
expect(stub.calls).toEqual([]);
|
||||
});
|
||||
|
||||
it('aborts the engine wait when the client disconnects mid-complete', async () => {
|
||||
const stub = makeMcpStub();
|
||||
let seenSignal: AbortSignal | undefined;
|
||||
let reached = false;
|
||||
stub.service.completeServerAuth = async (_handle, options) => {
|
||||
seenSignal = options?.signal;
|
||||
reached = true;
|
||||
await new Promise<void>((resolve) => {
|
||||
options?.signal?.addEventListener('abort', () => resolve(), { once: true });
|
||||
});
|
||||
};
|
||||
await boot(stub);
|
||||
|
||||
const controller = new AbortController();
|
||||
const pending = authedFetch(server as RunningServer, base, '/api/v2/mcp/auth:complete', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({ flowId: 'flow-1' }),
|
||||
signal: controller.signal,
|
||||
});
|
||||
await vi.waitFor(() => expect(reached).toBe(true));
|
||||
|
||||
controller.abort();
|
||||
|
||||
await expect(pending).rejects.toThrow();
|
||||
await vi.waitFor(() => expect(seenSignal?.aborted).toBe(true));
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue