Skip to content

Commit dc4ca21

Browse files
cdcabreracdcabrera
committed
test: review updates
1 parent a655cd2 commit dc4ca21

3 files changed

Lines changed: 226 additions & 44 deletions

File tree

‎src/__tests__/__snapshots__/server.tools.test.ts.snap‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,23 @@ exports[`makeProxyCreators should attempt to return proxy creators, a function w
313313
}
314314
`;
315315

316+
exports[`sendToolsHostShutdown should attempt force shutdown of child and fail NaN`] = `
317+
[
318+
[
319+
"Failed to send shutdown signal to Tools Host child process: Error: Mock send failure",
320+
],
321+
[
322+
"Failed to force-kill Tools Host child process: Error: Mock failed to kill child process",
323+
],
324+
[
325+
"Failed to close Tools Host stderr reader: Error: Mock close failure 1",
326+
],
327+
[
328+
"Slow shutdown response. Primary fallback force-killing Tools Host child process.",
329+
],
330+
]
331+
`;
332+
316333
exports[`spawnToolsHost attempt to spawn the Tools Host, with no pluginIsolation, node 24: spawn 1`] = `
317334
{
318335
"spawn": [

‎src/__tests__/server.tools.test.ts‎

Lines changed: 138 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -389,9 +389,144 @@ describe('makeProxyCreators', () => {
389389
});
390390

391391
describe('sendToolsHostShutdown', () => {
392-
it('should exist', () => {
393-
// placeholder test
394-
expect(sendToolsHostShutdown).toBeDefined();
392+
const MockLog = jest.mocked(log);
393+
const MockSend = jest.mocked(send);
394+
let mapGetSpy: jest.SpyInstance;
395+
let mapDeleteSpy: jest.SpyInstance;
396+
397+
beforeEach(() => {
398+
jest.clearAllMocks();
399+
jest.useFakeTimers();
400+
mapGetSpy = jest.spyOn(Map.prototype, 'get');
401+
mapDeleteSpy = jest.spyOn(Map.prototype, 'delete');
402+
});
403+
404+
afterEach(() => {
405+
jest.useRealTimers();
406+
mapGetSpy.mockRestore();
407+
mapDeleteSpy.mockRestore();
408+
});
409+
410+
it('should attempt graceful shutdown of child', async () => {
411+
const onceHandlers: Record<string, any> = {};
412+
const child = {
413+
kill: jest.fn(),
414+
killed: false,
415+
once: jest.fn((event: string, handler: any) => {
416+
onceHandlers[event] = handler;
417+
}),
418+
off: jest.fn(),
419+
stderr: {
420+
on: jest.fn(),
421+
off: jest.fn()
422+
}
423+
};
424+
const handle = { child, closeStderr: jest.fn() };
425+
const sessionId = 'test-session-id';
426+
427+
mapGetSpy.mockReturnValue(handle);
428+
429+
const promise = sendToolsHostShutdown({ pluginHost: { gracePeriodMs: 10 } } as any, { sessionId } as any);
430+
431+
onceHandlers['disconnect']();
432+
433+
await promise;
434+
435+
expect(MockSend).toHaveBeenCalledTimes(1);
436+
expect(child.once).toHaveBeenCalledTimes(2);
437+
expect(child.off).toHaveBeenCalledWith('exit', onceHandlers['exit']);
438+
expect(child.off).toHaveBeenCalledWith('disconnect', onceHandlers['disconnect']);
439+
expect(handle.closeStderr).toHaveBeenCalledTimes(1);
440+
expect(mapDeleteSpy).toHaveBeenCalledWith(sessionId);
441+
442+
jest.advanceTimersByTime(220);
443+
expect(child.kill).not.toHaveBeenCalled();
444+
});
445+
446+
it('should attempt force shutdown of child', async () => {
447+
const child = {
448+
// eslint-disable-next-line func-names
449+
kill: jest.fn(function (this: any) {
450+
this.killed = true;
451+
452+
return true;
453+
}),
454+
killed: false,
455+
once: jest.fn(),
456+
off: jest.fn(),
457+
stderr: {
458+
on: jest.fn(),
459+
off: jest.fn()
460+
}
461+
};
462+
const handle = { child, closeStderr: jest.fn() };
463+
const sessionId = 'test-session-id';
464+
465+
mapGetSpy.mockReturnValue(handle);
466+
467+
const promise = sendToolsHostShutdown({ pluginHost: { gracePeriodMs: 10 } } as any, { sessionId } as any);
468+
469+
jest.advanceTimersByTime(20);
470+
await promise;
471+
472+
jest.advanceTimersByTime(220);
473+
474+
expect(MockSend).toHaveBeenCalledTimes(1);
475+
expect(child.once).toHaveBeenCalledTimes(2);
476+
expect(child.kill).toHaveBeenCalledTimes(1);
477+
expect(child.kill).toHaveBeenCalledWith('SIGKILL');
478+
expect(child.killed).toBe(true);
479+
expect(child.off).toHaveBeenCalledTimes(2);
480+
expect(handle.closeStderr).toHaveBeenCalledTimes(1);
481+
expect(mapDeleteSpy).toHaveBeenCalledWith(sessionId);
482+
});
483+
484+
it('should attempt force shutdown of child and fail', async () => {
485+
const child = {
486+
kill: jest.fn()
487+
.mockImplementationOnce(() => {
488+
throw new Error('Mock failed to kill child process');
489+
})
490+
// eslint-disable-next-line func-names
491+
.mockImplementationOnce(function (this: any) {
492+
this.killed = true;
493+
494+
return true;
495+
}),
496+
killed: false,
497+
once: jest.fn(),
498+
off: jest.fn(),
499+
stderr: {
500+
on: jest.fn(),
501+
off: jest.fn()
502+
}
503+
};
504+
const handle = {
505+
child,
506+
closeStderr: jest.fn()
507+
.mockImplementationOnce(() => {
508+
throw new Error('Mock close failure 1');
509+
})
510+
.mockImplementationOnce(() => {})
511+
};
512+
const sessionId = 'test-session-id';
513+
514+
MockSend.mockImplementationOnce(() => {
515+
throw new Error('Mock send failure');
516+
});
517+
518+
mapGetSpy.mockReturnValue(handle);
519+
520+
const promise = sendToolsHostShutdown({ pluginHost: { gracePeriodMs: 10 } } as any, { sessionId } as any);
521+
522+
jest.advanceTimersByTime(10);
523+
await promise;
524+
525+
jest.advanceTimersByTime(220);
526+
527+
expect(child.kill).toHaveBeenCalledWith('SIGKILL');
528+
expect(child.killed).toBe(false);
529+
expect([...MockLog.error.mock.calls, ...MockLog.warn.mock.calls, ...MockLog.info.mock.calls]).toMatchSnapshot();
395530
});
396531
});
397532

‎src/server.tools.ts‎

Lines changed: 71 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,7 @@ const makeProxyCreators = (
288288
{ pluginHost }: GlobalOptions = getOptions()
289289
): McpToolCreator[] => handle.tools.map((tool): McpToolCreator => () => {
290290
const name = tool.name;
291+
const invokeTimeoutMs = Math.max(0, Number(pluginHost?.invokeTimeoutMs) || 0);
291292

292293
// Rebuild Zod schema from serialized JSON.
293294
const zodSchemaStrict = jsonSchemaToZod(tool.inputSchema);
@@ -330,7 +331,7 @@ const makeProxyCreators = (
330331
const response = await awaitIpc(
331332
handle.child,
332333
isInvokeResult(requestId),
333-
pluginHost.invokeTimeoutMs
334+
invokeTimeoutMs
334335
);
335336

336337
if ('ok' in response && response.ok === false) {
@@ -367,7 +368,6 @@ const makeProxyCreators = (
367368
*
368369
* @param {GlobalOptions} options - Global options.
369370
* @param {AppSession} sessionOptions - Session options.
370-
* @returns {Promise<void>} Promise that resolves when the host is stopped or noop.
371371
*/
372372
const sendToolsHostShutdown = async (
373373
{ pluginHost }: GlobalOptions = getOptions(),
@@ -379,64 +379,94 @@ const sendToolsHostShutdown = async (
379379
return;
380380
}
381381

382-
const gracePeriodMs = (Number.isInteger(pluginHost?.gracePeriodMs) && pluginHost.gracePeriodMs) || 0;
382+
const gracePeriodMs = Math.max(0, Number(pluginHost?.gracePeriodMs) || 0);
383383
const fallbackGracePeriodMs = gracePeriodMs + 200;
384384

385385
const child = handle.child;
386386
let resolved = false;
387387
let forceKillPrimary: NodeJS.Timeout | undefined;
388-
let forceKillFallback: NodeJS.Timeout | undefined;
388+
let forceKillSecondary: NodeJS.Timeout | undefined;
389+
let resolveIt: ((value: PromiseLike<void> | void) => void) | undefined;
389390

390-
await new Promise<void>(resolve => {
391-
const resolveOnce = () => {
392-
if (resolved) {
393-
return;
394-
}
391+
// Attempt exit, disconnect, then remove from activeHostsBySession and finally resolve
392+
const shutdownChild = () => {
393+
if (resolved) {
394+
return;
395+
}
395396

396-
resolved = true;
397-
child.off('exit', resolveOnce);
398-
child.off('disconnect', resolveOnce);
397+
resolved = true;
398+
child.off('exit', shutdownChild);
399+
child.off('disconnect', shutdownChild);
399400

400-
if (forceKillPrimary) {
401-
clearTimeout(forceKillPrimary);
402-
}
401+
if (forceKillPrimary) {
402+
clearTimeout(forceKillPrimary);
403+
}
403404

404-
if (forceKillFallback) {
405-
clearTimeout(forceKillFallback);
406-
}
405+
if (forceKillSecondary) {
406+
clearTimeout(forceKillSecondary);
407+
}
408+
409+
try {
410+
(handle as any).closeStderr();
411+
log.info('Tools Host stderr reader closed.');
412+
} catch (error) {
413+
log.error(`Failed to close Tools Host stderr reader: ${formatUnknownError(error)}`);
414+
}
407415

408-
try {
409-
handle.closeStderr?.();
410-
} catch {}
416+
const confirmHandle = activeHostsBySession.get(sessionId);
411417

418+
if (confirmHandle?.child === child) {
412419
activeHostsBySession.delete(sessionId);
413-
resolve();
414-
};
420+
}
415421

416-
try {
417-
send(child, { t: 'shutdown', id: makeId() });
418-
} catch {}
422+
resolveIt?.();
423+
};
419424

420-
const shutdownChild = () => {
421-
try {
422-
if (!child?.killed) {
423-
child.kill('SIGKILL');
424-
}
425-
} finally {
426-
resolveOnce();
425+
// Forced shutdown.
426+
const sigkillChild = (isSecondaryFallback: boolean = false) => {
427+
try {
428+
if (!child?.killed) {
429+
log.warn(
430+
`${
431+
(resolved && 'Already attempted shutdown.') || 'Slow shutdown response.'
432+
} ${
433+
(isSecondaryFallback && 'Secondary') || 'Primary'
434+
} fallback force-killing Tools Host child process.`
435+
);
436+
child.kill('SIGKILL');
427437
}
428-
};
438+
} catch (error) {
439+
log.error(`Failed to force-kill Tools Host child process: ${formatUnknownError(error)}`);
440+
}
441+
};
442+
443+
// Start the shutdown process
444+
await new Promise<void>(resolve => {
445+
resolveIt = resolve;
446+
// Send a shutdown signal to child. We try/catch in case the process is already dead, and
447+
// since we're still following it up with a graceful shutdown, then force-kill.
448+
try {
449+
send(child, { t: 'shutdown', id: makeId() });
450+
} catch (error) {
451+
log.error(`Failed to send shutdown signal to Tools Host child process: ${formatUnknownError(error)}`);
452+
}
429453

430-
// Primary grace period
431-
forceKillPrimary = setTimeout(shutdownChild, gracePeriodMs);
454+
// Set primary timeout for force shutdown
455+
forceKillPrimary = setTimeout(() => {
456+
sigkillChild();
457+
shutdownChild();
458+
}, gracePeriodMs);
432459
forceKillPrimary?.unref?.();
433460

434-
// Fallback grace period
435-
forceKillFallback = setTimeout(shutdownChild, fallbackGracePeriodMs);
436-
forceKillFallback?.unref?.();
461+
// Set fallback timeout for force shutdown
462+
forceKillSecondary = setTimeout(() => {
463+
sigkillChild(true);
464+
}, fallbackGracePeriodMs);
465+
forceKillSecondary?.unref?.();
437466

438-
child.once('exit', resolveOnce);
439-
child.once('disconnect', resolveOnce);
467+
// Set up exit/disconnect handlers to resolve
468+
child.once('exit', shutdownChild);
469+
child.once('disconnect', shutdownChild);
440470
});
441471
};
442472

0 commit comments

Comments
 (0)