fix(evidence): admit bundle imports across session transitions (#2303)

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);
+        },
+      ),
   );
 };
 

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<void>();
+  const releaseRead = createDeferred<void>();
+  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<void>();
+  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<void>();
+  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();
+});