From d308e7f921e7c53a8fb1d26dacd5d0c7a226126f Mon Sep 17 00:00:00 2001 From: Rohit Ghumare Date: Tue, 13 Oct 2026 00:14:17 +0530 Subject: [PATCH] fix(evidence): admit bundle imports across session transitions --- src/server/registerSessionRecordTools.ts | 78 +++++---- .../mcp/evidenceImportLifecycle.test.ts | 155 ++++++++++++++++++ 2 files changed, 202 insertions(+), 31 deletions(-) create mode 100644 tests/boundary/mcp/evidenceImportLifecycle.test.ts diff --git a/src/server/registerSessionRecordTools.ts b/src/server/registerSessionRecordTools.ts index 41efe981e74..445345fbfbb 100644 --- a/src/server/registerSessionRecordTools.ts +++ b/src/server/registerSessionRecordTools.ts @@ -12,6 +12,8 @@ import { err, ok } from "../domain/result.js"; import { AnalysisCancelledError } from "../domain/analysisErrorCore.js"; import { inspectEvidenceBundle } from "../application/investigation/InspectEvidenceBundle.js"; import { toolRegistrationOptions } from "./toolRegistrationOptions.js"; +import { withAdmittedAnalysis } from "./analysisAdmission.js"; +import { runAdmittedToolOperation } from "./admittedToolOperation.js"; interface EvidenceToolRegistration { readonly server: EvidenceMcpServer; @@ -92,40 +94,54 @@ const registerImportEvidenceTool = ({ resources, importContract, }: EvidenceToolRegistration): void => { + const admission = withAdmittedAnalysis({ kind: "session", session }); server.registerTool( importContract.name, toolRegistrationOptions(importContract), - async (input, context) => { - const path = input.path; - const signal = context.mcpReq.signal; - const loaded = await readImmutableEvidenceBundle(path, signal, resources); - if (!loaded.ok) - return server.delivery.toCallToolResult(loaded, importContract); - const retainedUnknownRevisions = new Set( - bundleForSerialization(session).unknowns.map((unknown) => - unknownRevisionKey(unknown), - ), - ); - if (signal.aborted) - return server.delivery.toCallToolResult( - err(new AnalysisCancelledError("import_evidence_bundle")), - importContract, - ); - const imported = session.importEvidenceBundle(loaded.value); - return imported.ok - ? server.delivery.toCallToolResult( - ok({ - imported: imported.value, - unknowns_added: loaded.value.unknowns.filter( - (unknown) => - !retainedUnknownRevisions.has(unknownRevisionKey(unknown)), - ).length, - total: bundleForSerialization(session).records.length, - }), - importContract, - ) - : server.delivery.toCallToolResult(imported, importContract); - }, + async (input, context) => + runAdmittedToolOperation( + server, + admission, + importContract.name, + context.mcpReq.signal, + async () => { + const path = input.path; + const signal = context.mcpReq.signal; + const loaded = await readImmutableEvidenceBundle( + path, + signal, + resources, + ); + if (!loaded.ok) + return server.delivery.toCallToolResult(loaded, importContract); + const retainedUnknownRevisions = new Set( + bundleForSerialization(session).unknowns.map((unknown) => + unknownRevisionKey(unknown), + ), + ); + if (signal.aborted) + return server.delivery.toCallToolResult( + err(new AnalysisCancelledError("import_evidence_bundle")), + importContract, + ); + const imported = session.importEvidenceBundle(loaded.value); + return imported.ok + ? server.delivery.toCallToolResult( + ok({ + imported: imported.value, + unknowns_added: loaded.value.unknowns.filter( + (unknown) => + !retainedUnknownRevisions.has( + unknownRevisionKey(unknown), + ), + ).length, + total: bundleForSerialization(session).records.length, + }), + importContract, + ) + : server.delivery.toCallToolResult(imported, importContract); + }, + ), ); }; diff --git a/tests/boundary/mcp/evidenceImportLifecycle.test.ts b/tests/boundary/mcp/evidenceImportLifecycle.test.ts new file mode 100644 index 00000000000..f0371ef897f --- /dev/null +++ b/tests/boundary/mcp/evidenceImportLifecycle.test.ts @@ -0,0 +1,155 @@ +import { Client, InMemoryTransport } from "@modelcontextprotocol/client"; +import { writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import { expect, it, onTestFinished, vi } from "vitest"; + +import { createEvidence } from "../../../src/domain/evidence.js"; +import { createEvidenceBundle } from "../../../src/domain/evidenceBundle.js"; +import { OwnedFileHandle } from "../../../src/filesystem/OwnedFileHandle.js"; +import { createServer } from "../../../src/server/createServer.js"; +import { + createDeferred, + createTestBinarySession, +} from "../../fixtures/binarySession.js"; +import { createTestTempDirectory } from "../../fixtures/temporaryDirectory.js"; + +const connectImportSession = async () => { + const root = await createTestTempDirectory("rea-mcp-import-lifecycle-"); + const session = createTestBinarySession(() => { + throw new Error("Evidence import must not start a binary provider"); + }); + const server = createServer({ kind: "session", session }); + const client = new Client({ + name: "evidence-import-lifecycle", + version: "1", + }); + const readCompleted = createDeferred(); + const releaseRead = createDeferred(); + const originalFileClose = OwnedFileHandle.prototype.close; + vi.spyOn(OwnedFileHandle.prototype, "close").mockImplementationOnce( + async function (this: OwnedFileHandle) { + await originalFileClose.call(this); + readCompleted.resolve(); + await releaseRead.promise; + }, + ); + const closeStarted = createDeferred(); + const originalSessionClose = session.close.bind(session); + const events: string[] = []; + vi.spyOn(session, "close").mockImplementation((...args) => { + const closing = originalSessionClose(...args); + closeStarted.resolve(); + return closing.then((result) => { + events.push("close"); + return result; + }); + }); + const originalImport = session.importEvidenceBundle.bind(session); + vi.spyOn(session, "importEvidenceBundle").mockImplementation((bundle) => { + const imported = originalImport(bundle); + events.push("import"); + return imported; + }); + onTestFinished(async () => { + releaseRead.resolve(); + vi.restoreAllMocks(); + await client.close(); + await server.close(); + await session.close(); + }); + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + const laterImportDispatched = createDeferred(); + let importsDispatched = 0; + const send = clientTransport.send.bind(clientTransport); + vi.spyOn(clientTransport, "send").mockImplementation( + async (message, ...args) => { + await send(message, ...args); + if ( + "method" in message && + message.method === "tools/call" && + message.params?.name === "import_evidence_bundle" + ) { + importsDispatched += 1; + if (importsDispatched === 2) laterImportDispatched.resolve(); + } + }, + ); + await server.connect(serverTransport); + await client.connect(clientTransport); + const importBundle = async (operation: string) => { + const record = createEvidence( + undefined, + { id: "fixture", name: "Fixture", version: "1" }, + { operation, parameters: {}, result: { retained: operation } }, + ); + const path = join(root, `${operation}.json`); + await writeFile(path, JSON.stringify(createEvidenceBundle([record]))); + return client.callTool({ + name: "import_evidence_bundle", + arguments: { path }, + }); + }; + return { + client, + session, + events, + importBundle, + readCompleted: readCompleted.promise, + closeStarted: closeStarted.promise, + laterImportDispatched: laterImportDispatched.promise, + releaseRead: () => releaseRead.resolve(), + }; +}; + +it("drains a file import before close_binary clears the session", async () => { + const fixture = await connectImportSession(); + const importing = fixture.importBundle("before_close"); + await fixture.readCompleted; + const closing = fixture.client.callTool({ + name: "close_binary", + arguments: {}, + }); + await fixture.closeStarted; + await fixture.client.ping(); + fixture.releaseRead(); + const [imported, closed] = await Promise.all([importing, closing]); + + expect(imported.isError, JSON.stringify(imported)).not.toBe(true); + expect(imported.structuredContent).toMatchObject({ + result: { imported: 1, unknowns_added: 0, total: 1 }, + }); + expect(closed.isError, JSON.stringify(closed)).not.toBe(true); + expect(fixture.events).toEqual(["import", "close"]); + expect(fixture.session.exportEvidenceBundle().records).toEqual([]); + await fixture.client.ping(); +}); + +it("admits a later import after close_binary without retaining the earlier bundle", async () => { + const fixture = await connectImportSession(); + const earlier = fixture.importBundle("before_close"); + await fixture.readCompleted; + const closing = fixture.client.callTool({ + name: "close_binary", + arguments: {}, + }); + await fixture.closeStarted; + const later = fixture.importBundle("after_close"); + await fixture.laterImportDispatched; + await fixture.client.ping(); + fixture.releaseRead(); + const results = await Promise.all([earlier, closing, later]); + + for (const result of results) + expect(result.isError, JSON.stringify(result)).not.toBe(true); + expect(results[2]?.structuredContent).toMatchObject({ + result: { imported: 1, unknowns_added: 0, total: 1 }, + }); + expect(fixture.events).toEqual(["import", "close", "import"]); + expect( + fixture.session + .exportEvidenceBundle() + .records.map((record) => record.operation), + ).toEqual(["after_close"]); + await fixture.client.ping(); +});