From 41480bc44f270b925cb35e4a9915a58317b2a492 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Tue, 7 Jul 2026 18:05:10 +0300 Subject: [PATCH] =?UTF-8?q?test(ai-chat):=20harden=20the=20MCP-hang=20fix?= =?UTF-8?q?=20=E2=80=94=20cover=20defense-in-depth=20paths=20+=20paralleli?= =?UTF-8?q?ze=20build=20(#397)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the merged MCP-hang fix (16b476a2); post-hoc review DO. The fix itself is unchanged and prod-working — this adds the coverage + one coherence fix the review asked for. Tests (ai-chat.service.setup-abort.spec.ts): - onLateResolve late-close: a toolset that resolves AFTER the setup race was lost has its leased clients released (close spy asserted). - pure 60s deadline (signal NOT aborted): the turn proceeds Docmost-only (reaches streamText, run NOT finalized 'aborted') — the defense-in-depth backstop, previously untested. - legacy no-runId: a setup abort does NOT re-throw (the `runId &&` guard); together with the deadline test this locks both halves of the catch guard. Coherence (external-mcp/mcp-clients.service.ts buildEntry): - The per-server connects now run via Promise.all instead of a sequential for-await, so total build time is bounded by the slowest single server (~2×CONNECT_TIMEOUT_MS) instead of the sum. At >6 all-timing-out servers the sequential build exceeded the outer 60s deadline, inverting the "per-server bound primary, outer deadline is backstop" invariant. Merge is done in a sequential post-Promise.all loop in server order, so tool-key precedence/disambiguation, outcomes and client ordering are byte-identical to the sequential build; each server keeps its own timeout + failure-close. Comments: the onLateResolve note now says it releases the lease (refcount), not force-closes transports (the cache owns them); the invariant comment reflects the parallel build. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ai-chat.service.setup-abort.spec.ts | 176 +++++++++++++++++- .../src/core/ai-chat/ai-chat.service.ts | 15 +- .../external-mcp/mcp-clients.service.ts | 101 ++++++---- 3 files changed, 253 insertions(+), 39 deletions(-) diff --git a/apps/server/src/core/ai-chat/ai-chat.service.setup-abort.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.setup-abort.spec.ts index 780f960a..c04057c9 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.setup-abort.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.setup-abort.spec.ts @@ -64,6 +64,32 @@ describe('AiChatService.stream — abort during external-MCP setup finalizes the ], }; + // A minimal raw ServerResponse stand-in for the turns that PROCEED past setup + // and reach streamText (the deadline + legacy paths). The setup-only abort test + // never wires the stream, so it keeps using `{ raw: {} }`. + function makeRawRes() { + return { + raw: { + writeHead: jest.fn(function writeHead(this: unknown) { + return this; + }), + write: jest.fn(), + once: jest.fn(), + flushHeaders: jest.fn(), + }, + }; + } + + // A fake streamText result: the service only calls consumeStream() and + // pipeUIMessageStreamToResponse() on it (both fire-and-forget). Its terminal + // callbacks are never invoked, so the run is not finalized through them. + function makeStreamResult() { + return { + consumeStream: jest.fn(), + pipeUIMessageStreamToResponse: jest.fn(), + }; + } + beforeEach(() => { streamTextMock.mockReset(); jest.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined as never); @@ -71,7 +97,10 @@ describe('AiChatService.stream — abort during external-MCP setup finalizes the jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined as never); }); - afterEach(() => jest.restoreAllMocks()); + afterEach(() => { + jest.restoreAllMocks(); + jest.useRealTimers(); + }); it('stops the hung toolset build, rejects, and settles the run "aborted" — never reaching streamText', async () => { const runController = new AbortController(); @@ -118,4 +147,149 @@ describe('AiChatService.stream — abort during external-MCP setup finalizes the expect(toolsFor).toHaveBeenCalledTimes(1); expect(streamTextMock).not.toHaveBeenCalled(); }); + + // Item 1 — the onLateResolve leg of raceAgainstAbortAndTimeout. When `toolsFor` + // loses the race (abort) but RESOLVES LATER with a leased toolset, the setup site + // must release that abandoned toolset's leases (call close() on its client + // handles) so their lease refcount is not pinned forever by a toolset nobody + // consumes. Nothing else exercises this path. + it('releases the leases of a toolset that resolves AFTER the race was already lost (onLateResolve)', async () => { + const runController = new AbortController(); + // A controllable build: it hangs until we resolve it by hand, and the run is + // stopped mid-build so the race rejects BEFORE the build settles. + let resolveTools: (v: unknown) => void = () => undefined; + const toolsForPromise = new Promise((resolve) => { + resolveTools = resolve; + }); + const toolsFor = jest.fn(() => { + setTimeout(() => runController.abort(new Error('user stop')), 0); + return toolsForPromise; + }); + const { svc } = makeService({ toolsFor }); + + 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, + model: {} as never, + role: null, + runHooks: { + begin, + onAssistantSeeded: jest.fn(), + onStep: jest.fn(), + onSettled: jest.fn(), + } as never, + }); + + // The race is lost to the abort: the turn rejects with the stop reason. + await expect(promise).rejects.toThrow('user stop'); + + // NOW the abandoned build resolves late with a leased client. onLateResolve must + // release it (call close on the lease handle). + const close = jest.fn().mockResolvedValue(undefined); + resolveTools({ + tools: {}, + clients: [{ close }], + outcomes: [], + instructions: [], + }); + // Flush the microtasks so work.then -> onLateResolve -> Promise.all(close) runs. + await new Promise((r) => setImmediate(r)); + + expect(close).toHaveBeenCalledTimes(1); + }); + + // Item 2 — the PURE DEADLINE branch (MCP_TOOLSET_BUILD_DEADLINE_MS). `toolsFor` + // never settles and the run's signal is NOT aborted: the race rejects with a + // "setup timed out" error, the catch does NOT re-throw (runId set but signal not + // aborted), and the turn PROCEEDS Docmost-only. It must reach streamText (the turn + // continues, not wedged) and must NOT be finalized 'aborted'. + it('proceeds Docmost-only (reaches streamText) when the build hits the deadline without an abort', async () => { + jest.useFakeTimers(); + + // The build hangs forever; the run's signal is never aborted. + const toolsFor = jest.fn(() => new Promise(() => {})); + const { svc } = makeService({ toolsFor }); + + streamTextMock.mockReturnValue(makeStreamResult() as never); + + const onSettled = jest.fn(); + const runSignal = new AbortController().signal; // never aborts + const begin = jest.fn(async () => ({ runId: 'run-1', signal: runSignal })); + + const promise = svc.stream({ + user: { id: 'user-1' } as never, + workspace: { id: 'ws-1' } as never, + sessionId: 'sess-1', + body: body as never, + res: makeRawRes() as never, + signal: new AbortController().signal, + model: {} as never, + role: null, + runHooks: { + begin, + onAssistantSeeded: jest.fn(), + onStep: jest.fn(), + onSettled, + } as never, + }); + + // Advance past the 60s build deadline; advanceTimersByTimeAsync flushes the + // promise microtasks between timer fires so the whole setup chain runs. + await jest.advanceTimersByTimeAsync(60_001); + // The turn does not throw out of setup — it continues to stream. + await expect(promise).resolves.toBeUndefined(); + + // The turn CONTINUED: streamText was reached (Docmost-only), not wedged. + expect(toolsFor).toHaveBeenCalledTimes(1); + expect(streamTextMock).toHaveBeenCalledTimes(1); + // The run was NOT finalized as aborted (the deadline is not a Stop) — the setup + // catch settle path never ran, so onSettled is left to streamText's callbacks. + expect(onSettled).not.toHaveBeenCalled(); + }); + + // Item 3 — the LEGACY no-runId path. The catch's re-throw is gated on + // `runId && effectiveSignal.aborted`. With NO runId (no runHooks) an abort during + // setup must NOT re-throw (runId falsy) — the turn warns + proceeds Docmost-only + // and streams, and is never finalized 'aborted' via the re-throw. Locks the + // `runId &&` half of the guard. + it('does NOT re-throw on a setup abort when there is no runId (legacy path proceeds Docmost-only)', async () => { + const socketController = new AbortController(); + // The build hangs; the SOCKET signal (legacy effectiveSignal) aborts mid-build. + const toolsFor = jest.fn(() => { + setTimeout(() => socketController.abort(new Error('socket closed')), 0); + return new Promise(() => {}); + }); + const { svc } = makeService({ toolsFor }); + + streamTextMock.mockReturnValue(makeStreamResult() as never); + + // No runHooks => runId undefined, effectiveSignal === the socket 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: makeRawRes() as never, + signal: socketController.signal, + model: {} as never, + role: null, + }); + + // The turn does NOT reject out of setup (no re-throw on the legacy path). + await expect(promise).resolves.toBeUndefined(); + + // It proceeded Docmost-only and reached streamText — streamText then observes + // the already-aborted socket signal via its own abortSignal. + expect(toolsFor).toHaveBeenCalledTimes(1); + expect(streamTextMock).toHaveBeenCalledTimes(1); + }); }); diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index f0793867..9986c2a0 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -63,8 +63,12 @@ const MAX_AGENT_STEPS = 20; // 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. +// production hang) and the run always finalizes. It stays a TRUE backstop because +// buildEntry connects to the servers CONCURRENTLY, so the total build time is +// bounded by the SLOWEST single server (~2×CONNECT_TIMEOUT_MS), not the SUM across +// them — the per-server bound fires first no matter how many servers are enabled, +// and this outer deadline only catches a total build stall the per-server bound +// somehow missed. const MCP_TOOLSET_BUILD_DEADLINE_MS = 60_000; // System-prompt addendum injected ONLY on the final step (see prepareAgentStep). @@ -791,8 +795,11 @@ export class AiChatService implements OnModuleInit { // 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. + // (the race was already lost) RELEASE the abandoned toolset's leases — + // c.close() here is the lease handle, so it decrements the cache entry's + // refcount; it does NOT force-close the transports (the cache OWNS the + // clients and closes them on TTL/evict). This just prevents the lease + // refcount from being pinned >=1 forever by a toolset nobody will consume. external = await raceAgainstAbortAndTimeout( this.mcpClients.toolsFor(workspace.id), effectiveSignal, diff --git a/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts b/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts index 762846f2..1e688974 100644 --- a/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts +++ b/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts @@ -255,54 +255,49 @@ export class McpClientsService { const callTimeoutMs = mcpCallTimeoutMs(); 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. + // Per-server connect+tools result, still tagged with its server so the merge + // below can be applied in the SAME order as `servers` (see the parallel note). + type PerServerResult = + | { ok: true; client: McpClient; guarded: Record } + | { ok: false; reason: string }; + + // Connect to (and list tools for) every enabled server CONCURRENTLY, so the + // total build time is bounded by the SLOWEST single server (~2× + // CONNECT_TIMEOUT_MS: connect + tools), NOT the SUM across servers. The + // sequential loop this replaced summed those bounds, so with enough all-timing- + // out servers the outer MCP_TOOLSET_BUILD_DEADLINE_MS could fire before the + // per-server bounds, dropping ALL external tools and inverting the "per-server + // bound is primary, outer is a backstop" invariant. Each server keeps its OWN + // try/catch + connectWithTimeout/withTimeout bound + close-on-failure logic; a + // failed server is skipped, never fatal. Nothing here mutates the shared + // arrays — every result is merged IN SERVER ORDER after Promise.all, so tool- + // key precedence/disambiguation, `outcomes`, `instructions` and `clients` + // ordering all match the previous sequential behavior exactly. + const perServer = async ( + server: (typeof servers)[number], + ): Promise => { + // Track the connected client so the catch can close it when it was obtained + // but tools() then threw/timed out (connectWithTimeout closes its OWN orphan + // on a connect timeout, so `client` stays undefined on that path). On success + // the client is handed back and registered by the merge below (owned by the + // entry, closed at teardown) — so it is never double-closed. let client: McpClient | undefined; - let registered = false; try { 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; // Bound each tool's execute with a per-call total-timeout guard before // merging, so a single chatty-but-stuck call is aborted after the cap. const guarded = wrapToolsWithCallTimeout(picked, callTimeoutMs); - // Namespace each tool with the sanitized server name AND disambiguate - // against names already merged from earlier servers, so no external - // tool is silently overwritten on collision. The returned count drives - // whether this server's prompt guidance is included (≥1 tool merged). - const merged = this.mergeNamespaced( - tools, - guarded, - server.name, - server.id, - ); - outcomes.push({ name: server.name, ok: true }); - // Include this server's guidance ONLY when it actually contributed at - // least one tool the agent can call (allowlist may have filtered all of - // them out) AND the admin authored non-blank instructions. The header - // prefix is the sanitized server name (= the tool namespace prefix). - const guide = server.instructions?.trim(); - if (merged.count > 0 && guide) { - instructions.push({ - serverName: server.name, - toolPrefix: merged.prefix, - instructions: guide, - }); - } + return { ok: true, client, guarded }; } catch (err) { // 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) { + if (client) { void client.close().catch(() => undefined); } // Log a short warning (never the URL/headers) so ops can see degradation, @@ -311,7 +306,45 @@ export class McpClientsService { this.logger.warn( `External MCP server "${server.name}" unavailable: ${reason}`, ); - outcomes.push({ name: server.name, ok: false, reason }); + return { ok: false, reason }; + } + }; + + // Promise.all preserves array order regardless of settle order, so `results[i]` + // is `servers[i]`'s outcome — the merge below stays deterministic and matches + // the old sequential order (later servers still override/disambiguate against + // earlier ones on a tool-key clash). + const results = await Promise.all(servers.map(perServer)); + for (let i = 0; i < servers.length; i += 1) { + const server = servers[i]; + const result = results[i]; + if (result.ok !== true) { + outcomes.push({ name: server.name, ok: false, reason: result.reason }); + continue; + } + clients.push(result.client); + // Namespace each tool with the sanitized server name AND disambiguate + // against names already merged from earlier servers, so no external + // tool is silently overwritten on collision. The returned count drives + // whether this server's prompt guidance is included (≥1 tool merged). + const merged = this.mergeNamespaced( + tools, + result.guarded, + server.name, + server.id, + ); + outcomes.push({ name: server.name, ok: true }); + // Include this server's guidance ONLY when it actually contributed at + // least one tool the agent can call (allowlist may have filtered all of + // them out) AND the admin authored non-blank instructions. The header + // prefix is the sanitized server name (= the tool namespace prefix). + const guide = server.instructions?.trim(); + if (merged.count > 0 && guide) { + instructions.push({ + serverName: server.name, + toolPrefix: merged.prefix, + instructions: guide, + }); } }