Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 47 additions & 31 deletions src/server/registerSessionRecordTools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
},
),
);
};

Expand Down
155 changes: 155 additions & 0 deletions tests/boundary/mcp/evidenceImportLifecycle.test.ts
Original file line number Diff line number Diff line change
@@ -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();
});