diff --git a/src/server/createServer.ts b/src/server/createServer.ts index bace9e2ff5..1c521b91ad 100644 --- a/src/server/createServer.ts +++ b/src/server/createServer.ts @@ -792,8 +792,12 @@ export const createServer = ( const closeServer = server.close.bind(server); server.close = async () => { const results = await Promise.allSettled([closeServer(), closeProviders()]); - for (const result of results) - if (result.status === "rejected") throw result.reason; + const failures = results.flatMap((result) => + result.status === "rejected" ? [result.reason as unknown] : [], + ); + if (failures.length === 1) throw failures[0]; + if (failures.length > 1) + throw new AggregateError(failures, "REA server cleanup failed"); }; registerCaptureTools( server, diff --git a/tests/boundary/mcp/createServerCleanup.test.ts b/tests/boundary/mcp/createServerCleanup.test.ts new file mode 100644 index 0000000000..909b0fcd67 --- /dev/null +++ b/tests/boundary/mcp/createServerCleanup.test.ts @@ -0,0 +1,49 @@ +import { afterEach, expect, it, vi } from "vitest"; +import { BinarySession } from "../../../src/application/binary/BinarySession.js"; +import { AnalysisProviderRegistry } from "../../../src/application/binary/AnalysisProviderRegistry.js"; +import { SessionProviderRouter } from "../../../src/application/binary/SessionProviderRouter.js"; +import { InvestigationRecords } from "../../../src/application/investigation/InvestigationRecords.js"; +import { CdpBrowserProvider } from "../../../src/browser/CdpBrowserProvider.js"; +import { EvidenceMcpServer } from "../../../src/server/EvidenceMcpServer.js"; +import { createServer } from "../../../src/server/createServer.js"; + +afterEach(() => vi.restoreAllMocks()); + +it("preserves both shutdown failures and retries provider cleanup after onclose", async () => { + const transportFailure = new Error("transport cleanup failed"); + const providerFailure = new Error("browser cleanup failed"); + const transportClose = vi + .spyOn(EvidenceMcpServer.prototype, "close") + .mockImplementation(async function (this: EvidenceMcpServer) { + this.server.onclose?.(); + throw transportFailure; + }); + const browser = new CdpBrowserProvider(); + const providerClose = vi + .spyOn(browser, "close") + .mockRejectedValueOnce(providerFailure) + .mockResolvedValue(undefined); + const session = new BinarySession( + SessionProviderRouter.selectable(new AnalysisProviderRegistry([]), []), + new InvestigationRecords(), + ); + const server = createServer( + { kind: "session", session }, + { + browserObservation: browser, + }, + ); + const failure: unknown = await server + .close() + .catch((cause: unknown) => cause); + expect(failure).toBeInstanceOf(AggregateError); + const causes = (cause: unknown): unknown[] => + cause instanceof AggregateError ? cause.errors.flatMap(causes) : [cause]; + expect(causes(failure)).toEqual([transportFailure, providerFailure]); + expect(providerClose).toHaveBeenCalledTimes(1); + transportClose.mockResolvedValue(undefined); + await expect(server.close()).resolves.toBeUndefined(); + expect(providerClose).toHaveBeenCalledTimes(2); + transportClose.mockRestore(); + await session.close(); +}); diff --git a/tests/boundary/mcp/processCleanupRetryMcp.test.ts b/tests/boundary/mcp/processCleanupRetryMcp.test.ts index fcfa70a581..0f3f6a5cb1 100644 --- a/tests/boundary/mcp/processCleanupRetryMcp.test.ts +++ b/tests/boundary/mcp/processCleanupRetryMcp.test.ts @@ -121,7 +121,7 @@ it("projects the retained process root through close_binary and resolves its ret const firstWire = JSON.stringify(failed); expect(failed.error).toMatchObject({ code: "cleanup_incomplete", - details: { resources: ["temporary_root"] }, + details: { resources: [root] }, }); expect(fault.removalAttempts).toBe(2); expect(fault.retryFailure).toBeInstanceOf(ProcessCaptureError); @@ -129,7 +129,7 @@ it("projects the retained process root through close_binary and resolves its ret throw new Error("Expected the process scope's typed retry error"); const retryError = fault.retryFailure; expect(retryError.cleanupIncomplete).toBe(true); - expect(retryError.cleanupResources).toEqual(["temporary_root"]); + expect(retryError.cleanupResources).toEqual([root]); denyRemoval = false; const recovered = await callClose(); @@ -185,7 +185,7 @@ it.each(["returned", "thrown", "synchronous"] as const)( expect(failed.error).toMatchObject({ code: "cleanup_incomplete", details: { - resources: ["temporary_root"], + resources: [root], diagnostics: { primary_error: { details: { @@ -207,7 +207,9 @@ it.each(["returned", "thrown", "synchronous"] as const)( throw new Error("Expected the process scope's typed retry error"); const retryError = fault.retryFailure; expect(retryError.cleanupIncomplete).toBe(true); + expect(retryError.cleanupResources).toEqual([root]); expect(composed.cleanupIncomplete).toBe(true); + expect(composed.cleanupResources).toEqual([root]); denyRemoval = false; const recovered = parseMcpToolError(await callClose());