fix(ai-chat): bound MCP connect + guard turn setup so a hung handshake can't wedge every run
A transient network blip during an external-MCP handshake left createMCPClient pending forever (@ai-sdk/mcp does not settle on abort). getOrBuildEntry caches the per-workspace build PROMISE, so the never-settling connect poisoned the cache and every later turn hung at step_count=0 before streamText — the run never finalized, the row stayed 'running', and the chat was permanently blocked with A_RUN_ALREADY_ACTIVE (an explicit Stop could not interrupt the un-signalled setup). - mcp-clients: wrap connect() in a settling timeout (connectWithTimeout) that closes a late-arriving client, so a hung handshake rejects instead of poisoning the build cache; the bad server is skipped and the build completes. - mcp-clients: close a connected-but-unregistered client when tools() fails, fixing a pre-existing transport leak in buildEntry. - ai-chat.service: bound the toolsFor build with the run's abort signal AND a deadline (raceAgainstAbortAndTimeout); re-throw an explicit Stop only when a run exists (runId) so the run finalizes as 'aborted', and keep the legacy socket-bound path unchanged; settle the run 'aborted' vs 'error' accordingly. - tests: cache-not-poisoned + orphan-client-close-once + Stop-during-setup finalizes the run once. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,121 @@
|
||||
import { Logger } from '@nestjs/common';
|
||||
|
||||
// Mock the AI SDK: the turn we drive is STOPPED during the pre-streamText setup
|
||||
// phase, so no provider call must ever be made. convertToModelMessages is reached
|
||||
// (before toolsFor) so it is stubbed to an empty transcript.
|
||||
jest.mock('ai', () => ({
|
||||
streamText: jest.fn(),
|
||||
generateText: jest.fn(),
|
||||
convertToModelMessages: jest.fn(async () => []),
|
||||
stepCountIs: jest.fn(() => () => false),
|
||||
}));
|
||||
|
||||
import { streamText } from 'ai';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
|
||||
/**
|
||||
* D2 — an explicit Stop DURING the external-MCP toolset build (the pre-streamText
|
||||
* setup phase) must:
|
||||
* (a) unwedge the turn (stream() rejects instead of hanging at step 0), and
|
||||
* (b) finalize the run as 'aborted' via the outer catch's onSettled — never leak
|
||||
* the run row as 'running' (which would 409 every later turn in this chat).
|
||||
*
|
||||
* The setup phase does NOT yet observe streamText's terminal callbacks, so before
|
||||
* the fix a hung `toolsFor` ignored the run's abort signal and never finalized.
|
||||
* `raceAgainstAbortAndTimeout(toolsFor, effectiveSignal, ...)` now rejects the
|
||||
* moment the run's signal aborts; the catch re-throws (signal aborted), and the
|
||||
* outer catch settles the run 'aborted'.
|
||||
*/
|
||||
describe('AiChatService.stream — abort during external-MCP setup finalizes the run (D2)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
|
||||
function makeService(mcpClients: { toolsFor: jest.Mock }) {
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
aiSettings as never,
|
||||
tools as never,
|
||||
mcpClients as never,
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo (openPage undefined -> never touched)
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc, tools };
|
||||
}
|
||||
|
||||
const body = {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
jest.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined as never);
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined as never);
|
||||
jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined as never);
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('stops the hung toolset build, rejects, and settles the run "aborted" — never reaching streamText', async () => {
|
||||
const runController = new AbortController();
|
||||
// The build hangs (never resolves); the run is STOPPED mid-build. Aborting on a
|
||||
// macrotask exercises the abort-listener path (a real user Stop during setup).
|
||||
const toolsFor = jest.fn(() => {
|
||||
setTimeout(() => runController.abort(new Error('user stop')), 0);
|
||||
return new Promise(() => {}); // never settles — models a hung MCP build
|
||||
});
|
||||
const { svc } = makeService({ toolsFor });
|
||||
|
||||
const onSettled = jest.fn();
|
||||
const begin = jest.fn(async () => ({
|
||||
runId: 'run-1',
|
||||
signal: runController.signal,
|
||||
}));
|
||||
|
||||
const promise = svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: { raw: {} } as never,
|
||||
signal: new AbortController().signal, // socket signal (distinct from the run)
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: {
|
||||
begin,
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled,
|
||||
} as never,
|
||||
});
|
||||
|
||||
// (a) The turn is UNWEDGED: it rejects (with the stop reason) instead of hanging.
|
||||
await expect(promise).rejects.toThrow('user stop');
|
||||
|
||||
// (b) The run is finalized as 'aborted' with NO error message (a Stop, not a
|
||||
// failure) — so the run row never leaks 'running'.
|
||||
expect(onSettled).toHaveBeenCalledTimes(1);
|
||||
expect(onSettled).toHaveBeenCalledWith('run-1', 'aborted', undefined);
|
||||
|
||||
// The build was reached, but the provider call was NEVER made (stopped at setup).
|
||||
expect(toolsFor).toHaveBeenCalledTimes(1);
|
||||
expect(streamTextMock).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -58,6 +58,15 @@ import {
|
||||
// multi-search research questions are not cut off mid-investigation.
|
||||
const MAX_AGENT_STEPS = 20;
|
||||
|
||||
// Wall-clock ceiling for building the external MCP toolset during the per-turn
|
||||
// setup phase (before streamText owns the lifecycle). Defense-in-depth ABOVE the
|
||||
// per-server connect bound in mcp-clients.service (CONNECT_TIMEOUT_MS): even if
|
||||
// that per-server timeout regressed, this outer deadline — together with the run's
|
||||
// abort signal — guarantees the setup phase can never wedge a turn at step 0 (the
|
||||
// production hang) and the run always finalizes. Generous: the per-server bound
|
||||
// should fire first in practice; this only backstops a total build stall.
|
||||
const MCP_TOOLSET_BUILD_DEADLINE_MS = 60_000;
|
||||
|
||||
// System-prompt addendum injected ONLY on the final step (see prepareAgentStep).
|
||||
// It forbids further tool calls and tells the model to synthesize the best
|
||||
// answer it can from what it already gathered, so a tool-heavy turn never ends
|
||||
@@ -176,6 +185,78 @@ export function sameInstant(
|
||||
return ta === tb;
|
||||
}
|
||||
|
||||
/**
|
||||
* Race `work` against an abort signal AND a wall-clock deadline, so a hung
|
||||
* external-MCP toolset build during the pre-streamText setup phase can NEITHER
|
||||
* wedge the turn NOR make it un-stoppable, and the run always finalizes. It
|
||||
* - resolves with `work`'s value when it settles first;
|
||||
* - REJECTS EARLY if `signal` aborts (with `signal.reason` when that is an Error,
|
||||
* else a generic `Error('aborted')`) — so an explicit Stop is honored mid-setup;
|
||||
* - REJECTS EARLY if `deadlineMs` elapses (defense-in-depth backstop);
|
||||
* - invokes `onLateResolve(value)` when `work` settles AFTER the race was already
|
||||
* lost, so the caller can release any resources that abandoned value owns
|
||||
* (e.g. close leased MCP clients that would otherwise leak their sockets).
|
||||
*
|
||||
* A rejection handler is attached to `work` so a late rejection is never an
|
||||
* unhandledRejection; the timer is unref'd and cleared once the race settles.
|
||||
*/
|
||||
export function raceAgainstAbortAndTimeout<T>(
|
||||
work: Promise<T>,
|
||||
signal: AbortSignal,
|
||||
deadlineMs: number,
|
||||
onLateResolve?: (value: T) => void,
|
||||
): Promise<T> {
|
||||
return new Promise<T>((resolve, reject) => {
|
||||
let settled = false;
|
||||
const cleanup = () => {
|
||||
clearTimeout(timer);
|
||||
signal.removeEventListener('abort', onAbort);
|
||||
};
|
||||
const onAbort = () => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
reject(
|
||||
signal.reason instanceof Error ? signal.reason : new Error('aborted'),
|
||||
);
|
||||
};
|
||||
const timer = setTimeout(() => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
reject(new Error(`setup timed out after ${deadlineMs}ms`));
|
||||
}, deadlineMs);
|
||||
// Do not keep the process alive just for this setup-deadline timer.
|
||||
timer.unref?.();
|
||||
if (signal.aborted) {
|
||||
onAbort();
|
||||
} else {
|
||||
signal.addEventListener('abort', onAbort, { once: true });
|
||||
}
|
||||
work.then(
|
||||
(value) => {
|
||||
if (settled) {
|
||||
// The race was already lost (abort/deadline): hand the abandoned value to
|
||||
// the caller so it can release the resources that value owns.
|
||||
onLateResolve?.(value);
|
||||
return;
|
||||
}
|
||||
settled = true;
|
||||
cleanup();
|
||||
resolve(value);
|
||||
},
|
||||
(err: unknown) => {
|
||||
// A late rejection after the race is already handled — swallow so it is
|
||||
// never an unhandledRejection.
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
reject(err instanceof Error ? err : new Error(String(err)));
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Payload accepted from the client `useChat` POST body. We do NOT bind a strict
|
||||
* DTO (the global ValidationPipe whitelist would strip the useChat-specific
|
||||
@@ -704,10 +785,40 @@ export class AiChatService implements OnModuleInit {
|
||||
instructions: [],
|
||||
};
|
||||
try {
|
||||
external = await this.mcpClients.toolsFor(workspace.id);
|
||||
// Bound the external-MCP toolset build by BOTH the run's abort signal and
|
||||
// a generous wall-clock deadline. This is the pre-streamText setup phase,
|
||||
// which streamText's terminal callbacks do NOT yet govern — so without this
|
||||
// a hung build would hang the turn at step 0 forever (the production hang),
|
||||
// unobservant of an explicit Stop. The deadline is defense-in-depth ABOVE
|
||||
// the per-server connect bound in mcp-clients.service. On a LATE resolve
|
||||
// (the race was already lost) close the abandoned toolset's leased clients
|
||||
// so their transports are not leaked.
|
||||
external = await raceAgainstAbortAndTimeout(
|
||||
this.mcpClients.toolsFor(workspace.id),
|
||||
effectiveSignal,
|
||||
MCP_TOOLSET_BUILD_DEADLINE_MS,
|
||||
(late) => {
|
||||
void Promise.all(
|
||||
late.clients.map((c) => c.close().catch(() => undefined)),
|
||||
);
|
||||
},
|
||||
);
|
||||
} catch (err) {
|
||||
// Building the external toolset must never break the turn; proceed with
|
||||
// Docmost-only tools. Never log URLs/headers — short message only.
|
||||
// An explicit Stop reached the RUN's signal DURING setup: re-throw so the
|
||||
// outer catch finalizes the run as aborted — never swallow a Stop. Gated on
|
||||
// `runId`: the re-throw exists ONLY to finalize the run, which exists only
|
||||
// in autonomous mode. On the legacy path (no runId) `effectiveSignal` is the
|
||||
// SOCKET signal (it aborts on a client disconnect); re-throwing there would
|
||||
// change prior behavior and make the controller write JSON to an already-
|
||||
// closed socket (it only attaches res.raw.on('error') in autonomous mode).
|
||||
// So legacy keeps its prior behavior — warn + proceed, and streamText then
|
||||
// observes the aborted socket signal.
|
||||
if (runId && effectiveSignal.aborted) {
|
||||
throw err;
|
||||
}
|
||||
// Otherwise a down/slow server (build timeout or other fault) must never
|
||||
// break the turn: proceed with Docmost-only tools. Never log URLs/headers —
|
||||
// short message only.
|
||||
this.logger.warn(
|
||||
`External MCP toolset unavailable: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
@@ -1307,12 +1418,19 @@ export class AiChatService implements OnModuleInit {
|
||||
if (this.environment?.isAiChatResumableStreamEnabled?.()) {
|
||||
this.streamRegistry?.abortEntry(chatId, runId);
|
||||
}
|
||||
// Distinguish an explicit Stop (the run's signal aborted during setup) from
|
||||
// a real failure, so the run settles with the correct terminal status
|
||||
// instead of always 'error'. onSettled/finalizeRun is idempotent, so this
|
||||
// is safe even if a streamText callback also settles the run.
|
||||
const settleStatus = effectiveSignal.aborted ? 'aborted' : 'error';
|
||||
await runHooks?.onSettled?.(
|
||||
runId,
|
||||
'error',
|
||||
err instanceof Error
|
||||
? err.message
|
||||
: 'Agent run failed before streaming started',
|
||||
settleStatus,
|
||||
settleStatus === 'aborted'
|
||||
? undefined
|
||||
: err instanceof Error
|
||||
? err.message
|
||||
: 'Agent run failed before streaming started',
|
||||
);
|
||||
}
|
||||
throw err;
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
import { McpClientsService } from './mcp-clients.service';
|
||||
|
||||
/**
|
||||
* D1 — a HUNG MCP handshake must not POISON the per-workspace build cache.
|
||||
*
|
||||
* THE BUG (production hang): `createMCPClient` (inside the private `connect`) is
|
||||
* NOT bounded by a timeout and — like @ai-sdk/mcp's tool calls — its promise does
|
||||
* NOT settle on abort. A transient network blip mid-handshake made connect hang
|
||||
* FOREVER. Because getOrBuildEntry caches the build PROMISE, that never-settling
|
||||
* connect wedged EVERY later turn for the workspace (each awaited the same pending
|
||||
* build) — step_count stuck at 0, run row leaking 'running', chat 409ing forever.
|
||||
*
|
||||
* THE FIX: `connectWithTimeout` races `connect` against a SETTLING timeout
|
||||
* (CONNECT_TIMEOUT_MS). On timeout it REJECTS, so buildEntry catches it, records
|
||||
* the server `ok:false`, and the build COMPLETES with that server skipped — the
|
||||
* cache is never poisoned and a subsequent `toolsFor` returns instead of hanging.
|
||||
*
|
||||
* REACHABILITY NOTE: the smallest network-free path that exercises the fix is to
|
||||
* spy on the private `connect` (the same harness the namespacing spec uses) —
|
||||
* `connectWithTimeout` wraps exactly that call, so a never-resolving `connect`
|
||||
* models a never-settling `createMCPClient` precisely, without DNS/sockets.
|
||||
*
|
||||
* Fake timers prove the timeout fires WITHOUT real waiting.
|
||||
*/
|
||||
|
||||
// Mirrors the private CONNECT_TIMEOUT_MS constant in mcp-clients.service.ts.
|
||||
const CONNECT_TIMEOUT_MS = 5000;
|
||||
|
||||
interface FakeServer {
|
||||
id: string;
|
||||
name: string;
|
||||
transport: string;
|
||||
url: string;
|
||||
headersEnc: string | null;
|
||||
toolAllowlist: string[] | null;
|
||||
instructions?: string | null;
|
||||
}
|
||||
|
||||
function server(
|
||||
over: Partial<FakeServer> & { id: string; name: string },
|
||||
): FakeServer {
|
||||
return {
|
||||
transport: 'http',
|
||||
url: 'https://example.com/mcp',
|
||||
headersEnc: null,
|
||||
toolAllowlist: null,
|
||||
...over,
|
||||
};
|
||||
}
|
||||
|
||||
function buildService(servers: FakeServer[]) {
|
||||
const repoStub = { listEnabled: jest.fn().mockResolvedValue(servers) };
|
||||
const service = new McpClientsService(repoStub as never, {} as never);
|
||||
// Silence the expected "server unavailable" warning.
|
||||
jest
|
||||
.spyOn(
|
||||
(service as unknown as { logger: { warn: (...a: unknown[]) => void } })
|
||||
.logger,
|
||||
'warn',
|
||||
)
|
||||
.mockImplementation(() => undefined);
|
||||
return service;
|
||||
}
|
||||
|
||||
// Spy on the private `connect` with a per-server implementation.
|
||||
function stubConnect(
|
||||
service: McpClientsService,
|
||||
impl: (s: FakeServer) => Promise<unknown>,
|
||||
) {
|
||||
return jest
|
||||
.spyOn(
|
||||
service as unknown as { connect: (s: FakeServer) => Promise<unknown> },
|
||||
'connect',
|
||||
)
|
||||
.mockImplementation(impl);
|
||||
}
|
||||
|
||||
describe('McpClientsService.connectWithTimeout — hung connect does not poison the cache (D1)', () => {
|
||||
beforeEach(() => jest.useFakeTimers());
|
||||
afterEach(() => {
|
||||
jest.clearAllTimers();
|
||||
jest.useRealTimers();
|
||||
jest.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('buildEntry completes (server recorded ok:false) when connect never settles, and toolsFor does not hang', async () => {
|
||||
const svc = buildService([server({ id: 'id-hung', name: 'hung' })]);
|
||||
// connect NEVER settles — models a wedged createMCPClient handshake.
|
||||
stubConnect(svc, () => new Promise<never>(() => {}));
|
||||
|
||||
const toolsetPromise = svc.toolsFor('ws-1');
|
||||
// Drive fake time past the connect bound so connectWithTimeout rejects and
|
||||
// buildEntry catches it (records ok:false) — flushing the microtasks.
|
||||
await jest.advanceTimersByTimeAsync(CONNECT_TIMEOUT_MS + 1);
|
||||
|
||||
const toolset = await toolsetPromise;
|
||||
// The build COMPLETED with the bad server skipped (no tools, ok:false).
|
||||
expect(Object.keys(toolset.tools)).toHaveLength(0);
|
||||
expect(toolset.outcomes).toEqual([
|
||||
{ name: 'hung', ok: false, reason: 'MCP connect timed out after 5000ms' },
|
||||
]);
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
|
||||
// The cache is NOT poisoned: a subsequent turn returns (served from the warm
|
||||
// cached entry) instead of awaiting a never-settling build.
|
||||
const again = await svc.toolsFor('ws-1');
|
||||
expect(Object.keys(again.tools)).toHaveLength(0);
|
||||
await Promise.all(again.clients.map((c) => c.close()));
|
||||
});
|
||||
|
||||
it('a hung server is skipped but a healthy server in the SAME build still contributes its tools', async () => {
|
||||
const svc = buildService([
|
||||
server({ id: 'id-hung', name: 'hung' }),
|
||||
server({ id: 'id-ok', name: 'ok' }),
|
||||
]);
|
||||
const okClient = {
|
||||
tools: () => Promise.resolve({ search: { description: 'x' } }),
|
||||
close: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
stubConnect(svc, (s) =>
|
||||
s.id === 'id-hung'
|
||||
? new Promise<never>(() => {})
|
||||
: Promise.resolve(okClient),
|
||||
);
|
||||
|
||||
const toolsetPromise = svc.toolsFor('ws-2');
|
||||
await jest.advanceTimersByTimeAsync(CONNECT_TIMEOUT_MS + 1);
|
||||
|
||||
const toolset = await toolsetPromise;
|
||||
// Healthy server's tool survives (namespaced); hung server recorded ok:false.
|
||||
expect(Object.keys(toolset.tools)).toEqual(['ok_search']);
|
||||
expect(toolset.outcomes).toEqual([
|
||||
{ name: 'hung', ok: false, reason: 'MCP connect timed out after 5000ms' },
|
||||
{ name: 'ok', ok: true },
|
||||
]);
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
});
|
||||
|
||||
it('closes the ORPHANED client when connect resolves LATE (after the timeout)', async () => {
|
||||
const svc = buildService([server({ id: 'id-late', name: 'late' })]);
|
||||
const lateClient = {
|
||||
tools: () => Promise.resolve({}),
|
||||
close: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
// connect resolves only AFTER the connect bound has already elapsed, so
|
||||
// connectWithTimeout has already rejected and must close this orphan.
|
||||
stubConnect(
|
||||
svc,
|
||||
() =>
|
||||
new Promise((resolve) => {
|
||||
setTimeout(() => resolve(lateClient), CONNECT_TIMEOUT_MS * 2);
|
||||
}),
|
||||
);
|
||||
|
||||
const toolsetPromise = svc.toolsFor('ws-3');
|
||||
// Fire the timeout: the build completes with the server skipped.
|
||||
await jest.advanceTimersByTimeAsync(CONNECT_TIMEOUT_MS + 1);
|
||||
const toolset = await toolsetPromise;
|
||||
expect(toolset.outcomes[0]?.ok).toBe(false);
|
||||
expect(lateClient.close).not.toHaveBeenCalled();
|
||||
|
||||
// Now let the late connect resolve — the orphan must be closed, not leaked.
|
||||
await jest.advanceTimersByTimeAsync(CONNECT_TIMEOUT_MS * 2);
|
||||
expect(lateClient.close).toHaveBeenCalledTimes(1);
|
||||
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
});
|
||||
});
|
||||
|
||||
describe('McpClientsService.buildEntry — closes a connected client whose tools() fails (leak fix)', () => {
|
||||
beforeEach(() => jest.useFakeTimers());
|
||||
afterEach(() => {
|
||||
jest.clearAllTimers();
|
||||
jest.useRealTimers();
|
||||
jest.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('connect succeeds but tools() REJECTS: the client is close()d exactly once and the server is skipped, while a healthy server still contributes', async () => {
|
||||
const svc = buildService([
|
||||
server({ id: 'id-bad', name: 'bad' }),
|
||||
server({ id: 'id-ok', name: 'ok' }),
|
||||
]);
|
||||
// The bad server connects fine, then tools() rejects — the client would leak if
|
||||
// buildEntry did not close it in the per-server catch (it was never registered).
|
||||
const badClient = {
|
||||
tools: () => Promise.reject(new Error('tools listing failed')),
|
||||
close: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
const okClient = {
|
||||
tools: () => Promise.resolve({ search: { description: 'x' } }),
|
||||
close: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
stubConnect(svc, (s) =>
|
||||
s.id === 'id-bad' ? Promise.resolve(badClient) : Promise.resolve(okClient),
|
||||
);
|
||||
|
||||
const toolset = await svc.toolsFor('ws-4');
|
||||
|
||||
// The orphaned (never-registered) client is closed exactly once — no leak.
|
||||
expect(badClient.close).toHaveBeenCalledTimes(1);
|
||||
// Healthy server survives; bad server recorded ok:false and skipped.
|
||||
expect(Object.keys(toolset.tools)).toEqual(['ok_search']);
|
||||
expect(toolset.outcomes).toEqual([
|
||||
{ name: 'bad', ok: false, reason: 'tools listing failed' },
|
||||
{ name: 'ok', ok: true },
|
||||
]);
|
||||
|
||||
// The healthy (registered) client is NOT closed by the loop — it is owned by the
|
||||
// cache entry and stays warm (closed only on eviction/teardown, not on lease
|
||||
// release). Releasing the lease keeps it warm since the entry is not evicted.
|
||||
expect(okClient.close).not.toHaveBeenCalled();
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
expect(okClient.close).not.toHaveBeenCalled();
|
||||
// The failed client is never double-closed.
|
||||
expect(badClient.close).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('connect succeeds but tools() HANGS (times out): the client is close()d once and the server is skipped', async () => {
|
||||
const svc = buildService([server({ id: 'id-slow', name: 'slow' })]);
|
||||
const slowClient = {
|
||||
// tools() never settles -> withTimeout rejects after CONNECT_TIMEOUT_MS.
|
||||
tools: () => new Promise<Record<string, never>>(() => {}),
|
||||
close: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
stubConnect(svc, () => Promise.resolve(slowClient));
|
||||
|
||||
const toolsetPromise = svc.toolsFor('ws-5');
|
||||
// Drive fake time past the tools() bound so withTimeout rejects.
|
||||
await jest.advanceTimersByTimeAsync(CONNECT_TIMEOUT_MS + 1);
|
||||
const toolset = await toolsetPromise;
|
||||
|
||||
expect(slowClient.close).toHaveBeenCalledTimes(1);
|
||||
expect(Object.keys(toolset.tools)).toHaveLength(0);
|
||||
expect(toolset.outcomes[0]?.ok).toBe(false);
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
});
|
||||
});
|
||||
@@ -195,7 +195,7 @@ export class McpClientsService {
|
||||
): Promise<{ ok: true; tools: string[] } | { ok: false; error: string }> {
|
||||
let client: McpClient | undefined;
|
||||
try {
|
||||
client = await this.connect(server);
|
||||
client = await this.connectWithTimeout(server, CONNECT_TIMEOUT_MS);
|
||||
const raw = await withTimeout(client.tools(), CONNECT_TIMEOUT_MS);
|
||||
return { ok: true, tools: Object.keys(raw) };
|
||||
} catch (err) {
|
||||
@@ -256,10 +256,18 @@ export class McpClientsService {
|
||||
const instructions: McpServerInstruction[] = [];
|
||||
|
||||
for (const server of servers) {
|
||||
// Track the connected client OUTSIDE the try so the catch can close it when
|
||||
// it was obtained but not yet registered in `clients` (e.g. tools() threw or
|
||||
// timed out after a successful connect). `registered` flips only once the
|
||||
// client is owned by `clients` (closed at entry teardown), so the catch never
|
||||
// double-closes a registered client.
|
||||
let client: McpClient | undefined;
|
||||
let registered = false;
|
||||
try {
|
||||
const client = await this.connect(server);
|
||||
client = await this.connectWithTimeout(server, CONNECT_TIMEOUT_MS);
|
||||
const raw = await withTimeout(client.tools(), CONNECT_TIMEOUT_MS);
|
||||
clients.push(client);
|
||||
registered = true;
|
||||
const allow = server.toolAllowlist;
|
||||
const picked =
|
||||
Array.isArray(allow) && allow.length > 0 ? pick(raw, allow) : raw;
|
||||
@@ -290,9 +298,15 @@ export class McpClientsService {
|
||||
});
|
||||
}
|
||||
} catch (err) {
|
||||
// A failed server is skipped — the turn proceeds with the rest. Log a
|
||||
// short warning (never the URL/headers) so ops can see degradation, and
|
||||
// record the outcome so the UI can show "tool X unavailable".
|
||||
// A failed server is skipped — the turn proceeds with the rest. If connect
|
||||
// returned a live client but a later step (tools()) threw, that client was
|
||||
// never registered in `clients`, so close it here or its transport/socket
|
||||
// leaks (compounding every 60s cache rebuild during a flaky-server outage).
|
||||
if (client && !registered) {
|
||||
void client.close().catch(() => undefined);
|
||||
}
|
||||
// Log a short warning (never the URL/headers) so ops can see degradation,
|
||||
// and record the outcome so the UI can show "tool X unavailable".
|
||||
const reason = shortError(err);
|
||||
this.logger.warn(
|
||||
`External MCP server "${server.name}" unavailable: ${reason}`,
|
||||
@@ -383,6 +397,55 @@ export class McpClientsService {
|
||||
return client;
|
||||
}
|
||||
|
||||
/**
|
||||
* Race {@link connect} against a SETTLING timeout so a hung MCP handshake can
|
||||
* never POISON the per-workspace build cache. `createMCPClient` (inside connect)
|
||||
* is NOT bounded internally, and — exactly like @ai-sdk/mcp's tool calls
|
||||
* (see wrapToolWithCallTimeout) — its promise does NOT settle on abort. So a
|
||||
* transient network blip mid-handshake can make connect hang FOREVER. Because
|
||||
* getOrBuildEntry caches the build PROMISE, a never-settling connect would then
|
||||
* wedge EVERY later turn for the workspace (each awaits the same pending build,
|
||||
* step_count stuck at 0, run row leaks 'running', chat 409s forever). Bounding
|
||||
* connect here guarantees buildEntry always gets a client OR a rejection within
|
||||
* `ms` — so the build completes (bad server skipped) and the cache stays clean.
|
||||
*
|
||||
* If connect resolves LATE (after we already rejected on the timeout), we close
|
||||
* the orphaned client so its transport/socket is not leaked.
|
||||
*/
|
||||
private connectWithTimeout(
|
||||
server: Pick<AiMcpServer, 'transport' | 'url' | 'headersEnc'>,
|
||||
ms: number,
|
||||
): Promise<McpClient> {
|
||||
return new Promise<McpClient>((resolve, reject) => {
|
||||
let settled = false;
|
||||
const timer = setTimeout(() => {
|
||||
settled = true;
|
||||
reject(new Error(`MCP connect timed out after ${ms}ms`));
|
||||
}, ms);
|
||||
// Do not keep the process alive just for this connect-timeout timer.
|
||||
timer.unref?.();
|
||||
this.connect(server).then(
|
||||
(client) => {
|
||||
if (settled) {
|
||||
// The race was already lost to the timeout: close the orphaned client
|
||||
// so its socket is not leaked, and drop the late result.
|
||||
void client.close().catch(() => undefined);
|
||||
return;
|
||||
}
|
||||
clearTimeout(timer);
|
||||
settled = true;
|
||||
resolve(client);
|
||||
},
|
||||
(err: unknown) => {
|
||||
if (settled) return; // late rejection after the timeout — already handled
|
||||
clearTimeout(timer);
|
||||
settled = true;
|
||||
reject(err instanceof Error ? err : new Error(String(err)));
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Decrypt the stored auth headers. Returns undefined when none are set. The
|
||||
* plaintext headers live only in this returned object and are passed straight
|
||||
|
||||
Reference in New Issue
Block a user