diff --git a/src/application/DirectAnalysis.ts b/src/application/DirectAnalysis.ts index ea7ed8f6cd..a798a835b5 100644 --- a/src/application/DirectAnalysis.ts +++ b/src/application/DirectAnalysis.ts @@ -30,11 +30,7 @@ import { snapshotEvidenceForQuery, snapshotMatchesTarget, } from "../domain/analysisSnapshot.js"; -import { - analysisProfileSchema, - committedProviderSchema, - type AnalysisProfileCommitment, -} from "../domain/analysisProfile.js"; +import type { AnalysisProfileCommitment } from "../domain/analysisProfile.js"; import { err, ok, type Result } from "../domain/result.js"; import type { AnalysisSnapshot } from "../domain/analysisSnapshot.js"; import type { @@ -45,6 +41,10 @@ import { REA_WORKFLOW_PROVIDER, workflowAnalysisProfile, } from "./InvestigationProviders.js"; +import { + createWorkflowEvidence, + workflowSnapshotRecord, +} from "./WorkflowEvidence.js"; import type { AnalysisProviderSelector } from "../contracts/providerSelection.js"; import { artifactInspectionResultSchema } from "../domain/artifactInspection.js"; @@ -268,35 +268,13 @@ const runAnalysis = async ( tool !== "trace_native_ui_action" && snapshotPath !== undefined && evidence !== undefined && - "analysis_profile" in evidence && session.allowsSnapshotReplay(tool) ) { - const recorded = session.recordWorkflowSnapshot({ - operation: tool, - parameters: arguments_, - execution: { - result: evidence.normalized_result, - rawResult: evidence.raw_result, - provider: committedProviderSchema.parse(evidence.provider), - analysisProfile: analysisProfileSchema.parse( - evidence.analysis_profile, - ), - limitations: evidence.limitations, - locations: evidence.locations, - subject: - evidence.subject === null - ? null - : { - path: evidence.subject.local_path, - sha256: evidence.subject.digest.sha256, - format: evidence.subject.format, - ...(evidence.subject.architecture === null - ? {} - : { architecture: evidence.subject.architecture }), - }, - }, - }); - if (!recorded.ok) return cliError(recorded.error); + const workflowRecord = workflowSnapshotRecord(evidence, tool); + if (workflowRecord !== undefined) { + const recorded = session.recordWorkflowSnapshot(workflowRecord); + if (!recorded.ok) return cliError(recorded.error); + } } if ( tool !== "trace_native_ui_action" && @@ -356,15 +334,12 @@ const executeAnalysisTool = async (input: { signal, ); if (!result.ok) return { output: cliError(result.error) }; - const evidence = createEvidence(input.openedTarget, REA_WORKFLOW_PROVIDER, { + const evidence = createWorkflowEvidence({ + target: input.openedTarget, operation: tool, parameters: input.arguments, result: result.value, - ...(evidenceProfile === undefined - ? {} - : { analysisProfile: evidenceProfile }), - confidence: "derived", - limitations: ["Derived by an REA composed workflow."], + upstreamProfile: evidenceProfile, }); return { output: evidence, evidence }; } @@ -493,8 +468,7 @@ const analysisProfileForEvidence = ( | DirectAnalysisTool, ): AnalysisProfileCommitment | undefined => { if (!isWorkflowEvidenceTool(tool)) return session.analysisProfile(tool); - const upstream = session.analysisProfile(); - return upstream === undefined ? undefined : workflowAnalysisProfile(upstream); + return session.analysisProfile(); }; const isWorkflowEvidenceTool = ( diff --git a/src/application/WorkflowEvidence.ts b/src/application/WorkflowEvidence.ts new file mode 100644 index 0000000000..7944206a93 --- /dev/null +++ b/src/application/WorkflowEvidence.ts @@ -0,0 +1,64 @@ +import { + REA_WORKFLOW_PROVIDER, + workflowAnalysisProfile, +} from "./InvestigationProviders.js"; +import type { AnalysisOperation } from "./AnalysisProvider.js"; +import type { AnalysisProfileCommitment } from "../domain/analysisProfile.js"; +import type { WorkflowSnapshotRecordInput } from "./binary/BinarySessionRecords.js"; +import { + analysisProfileSchema, + committedProviderSchema, +} from "../domain/analysisProfile.js"; +import type { BinaryTarget } from "../domain/binaryTarget.js"; +import { createEvidence, type Evidence } from "../domain/evidence.js"; +import type { JsonValue } from "../domain/jsonValue.js"; + +/** Build the shared Evidence record for a successful composed binary workflow. */ +export const createWorkflowEvidence = (input: { + readonly target: BinaryTarget | undefined; + readonly operation: AnalysisOperation; + readonly parameters: Readonly>; + readonly result: JsonValue; + readonly upstreamProfile: AnalysisProfileCommitment | undefined; +}): Evidence => + createEvidence(input.target, REA_WORKFLOW_PROVIDER, { + operation: input.operation, + parameters: input.parameters, + result: input.result, + ...(input.upstreamProfile === undefined + ? {} + : { analysisProfile: workflowAnalysisProfile(input.upstreamProfile) }), + confidence: "derived", + limitations: ["Derived by an REA composed workflow."], + }); + +/** Convert workflow Evidence into its exact replay binding, when it has a profile. */ +export const workflowSnapshotRecord = ( + evidence: Evidence, + operation: AnalysisOperation, +): WorkflowSnapshotRecordInput | undefined => { + if (!("analysis_profile" in evidence)) return undefined; + return { + operation, + parameters: evidence.parameters, + execution: { + result: evidence.normalized_result, + rawResult: evidence.raw_result, + provider: committedProviderSchema.parse(evidence.provider), + analysisProfile: analysisProfileSchema.parse(evidence.analysis_profile), + limitations: evidence.limitations, + locations: evidence.locations, + subject: + evidence.subject === null + ? null + : { + path: evidence.subject.local_path, + sha256: evidence.subject.digest.sha256, + format: evidence.subject.format, + ...(evidence.subject.architecture === null + ? {} + : { architecture: evidence.subject.architecture }), + }, + }, + }; +}; diff --git a/src/application/binary/BinarySessionPort.ts b/src/application/binary/BinarySessionPort.ts index a8bd37a043..fb7317f000 100644 --- a/src/application/binary/BinarySessionPort.ts +++ b/src/application/binary/BinarySessionPort.ts @@ -7,6 +7,8 @@ import type { AnalysisProfileCommitment } from "../../domain/analysisProfile.js" import type { AnalysisProviderSelector } from "../../contracts/providerSelection.js"; import type { AnalysisError } from "../../domain/analysisErrorBase.js"; import type { Result } from "../../domain/result.js"; +import type { EvidenceIntegrityError } from "../../domain/evidenceErrors.js"; +import type { WorkflowSnapshotRecordInput } from "./BinarySessionRecords.js"; import type { AnalysisOperation, @@ -43,6 +45,10 @@ export interface BinarySessionPort analysisProfile( operation?: AnalysisOperation, ): AnalysisProfileCommitment | undefined; + allowsSnapshotReplay(operation: AnalysisOperation): boolean; + recordWorkflowSnapshot( + input: WorkflowSnapshotRecordInput, + ): Result; openCompatibility(): Readonly>; onAvailabilityChanged?(listener: () => void | Promise): () => void; onAnalysisSnapshotChanged?(listener: () => void | Promise): () => void; diff --git a/src/application/binary/BinarySessionRecords.ts b/src/application/binary/BinarySessionRecords.ts index ee97d88492..581f45b846 100644 --- a/src/application/binary/BinarySessionRecords.ts +++ b/src/application/binary/BinarySessionRecords.ts @@ -3,6 +3,7 @@ import type { AnalysisSnapshot } from "../../domain/analysisSnapshot.js"; import type { BinaryTarget } from "../../domain/binaryTarget.js"; import type { Evidence } from "../../domain/evidence.js"; import type { EvidenceBundle } from "../../domain/evidenceBundle.js"; +import type { JsonValue } from "../../domain/jsonValue.js"; import { evidenceBundleForTarget } from "../../domain/evidenceBundle.js"; import { EvidenceIntegrityError } from "../../domain/evidenceErrors.js"; import { type AnalysisError } from "../../domain/analysisErrorBase.js"; @@ -26,6 +27,15 @@ export interface ActiveAnalysisBinding { readonly profile: AnalysisProfileCommitment | null; } +/** Validated composed-workflow payload bound to the active target and profile. */ +export interface WorkflowSnapshotRecordInput { + readonly operation: AnalysisOperation; + readonly parameters: Readonly>; + readonly execution: Parameters< + AnalysisSnapshotCache["recordWorkflow"] + >[0]["execution"]; +} + /** Binary snapshot owner and compatibility facade for composed investigation records. */ export abstract class BinarySessionRecords { readonly #records: InvestigationRecords; @@ -161,15 +171,9 @@ export abstract class BinarySessionRecords { } /** Retain one derived workflow result alongside its provider cache entries. */ - recordWorkflowSnapshot(input: { - readonly operation: string; - readonly parameters: Readonly< - Record - >; - readonly execution: Parameters< - AnalysisSnapshotCache["recordWorkflow"] - >[0]["execution"]; - }): Result { + recordWorkflowSnapshot( + input: WorkflowSnapshotRecordInput, + ): Result { const active = this.activeAnalysisBinding(); if (active === undefined || active.profile === null) return err( diff --git a/src/server/createServer.ts b/src/server/createServer.ts index d011e5ac85..fd71d210cd 100644 --- a/src/server/createServer.ts +++ b/src/server/createServer.ts @@ -349,6 +349,14 @@ const registerBinaryAnalysisTools = ({ ...analysisOptions, analysisProfile: session === undefined ? undefined : () => session.analysisProfile(), + allowsSnapshotReplay: + session === undefined + ? undefined + : (operation) => session.allowsSnapshotReplay(operation), + recordWorkflowSnapshot: + session === undefined + ? undefined + : (input) => session.recordWorkflowSnapshot(input), }); const evidenceOptions = { logger, activeTarget, recordEvidence }; registerNativeTools(server, analysis, evidenceOptions); diff --git a/src/server/registerEnhancedTools.ts b/src/server/registerEnhancedTools.ts index 3a6c156afc..4797998a9e 100644 --- a/src/server/registerEnhancedTools.ts +++ b/src/server/registerEnhancedTools.ts @@ -9,19 +9,19 @@ import type { } from "@modelcontextprotocol/server"; import type { AnalysisOperationPort } from "../application/AnalysisProvider.js"; +import type { BinarySessionPort } from "../application/binary/BinarySessionPort.js"; import { EnhancedTools, type ValidatedEnhancedCall, } from "../application/EnhancedTools.js"; import { - REA_WORKFLOW_PROVIDER, - workflowAnalysisProfile, -} from "../application/InvestigationProviders.js"; + createWorkflowEvidence, + workflowSnapshotRecord, +} from "../application/WorkflowEvidence.js"; import { toolContract, type ToolContract } from "../contracts/toolContracts.js"; import type { AnalysisProfileCommitment } from "../domain/analysisProfile.js"; import type { BinaryTarget } from "../domain/binaryTarget.js"; import { UnknownRegistryError } from "../domain/unknownRegistryError.js"; -import { createEvidence } from "../domain/evidence.js"; import type { JsonValue } from "../domain/jsonValue.js"; import type { Logger } from "../logger.js"; import { mcpProgressReporter } from "./mcpProgress.js"; @@ -39,6 +39,14 @@ export interface EnhancedToolRegistration { | undefined; readonly recordEvidence: EvidenceWriter["recordEvidence"] | undefined; readonly recordUnknown: UnknownRegistryPort["recordUnknown"] | undefined; + readonly allowsSnapshotReplay: + | ((operation: ValidatedEnhancedCall["name"]) => boolean) + | undefined; + readonly recordWorkflowSnapshot: + | (( + input: Parameters[0], + ) => ReturnType) + | undefined; } /** Register composed workflows against the same port as direct bridge tools. */ @@ -235,28 +243,27 @@ const executeEnhancedTool = async ( terminal: true, }); if (result.ok) { - const upstreamProfile = registration.analysisProfile?.(); - const evidence = createEvidence( - registration.activeTarget?.(), - REA_WORKFLOW_PROVIDER, - { - operation: name, - parameters, - result: result.value, - ...(upstreamProfile === undefined - ? {} - : { - analysisProfile: workflowAnalysisProfile(upstreamProfile), - }), - confidence: "derived", - limitations: [ - "Derived by an REA workflow from one or more provider observations.", - ], - }, - ); + const evidence = createWorkflowEvidence({ + target: registration.activeTarget?.(), + operation: name, + parameters, + result: result.value, + upstreamProfile: registration.analysisProfile?.(), + }); const recorded = registration.recordEvidence?.(evidence); if (recorded !== undefined && !recorded.ok) return toCallToolResult(recorded, contract); + if ( + name !== "trace_native_ui_action" && + registration.allowsSnapshotReplay?.(name) === true + ) { + const workflowRecord = workflowSnapshotRecord(evidence, name); + if (workflowRecord !== undefined) { + const snapshot = registration.recordWorkflowSnapshot?.(workflowRecord); + if (snapshot !== undefined && !snapshot.ok) + return toCallToolResult(snapshot, contract); + } + } const unknowns = recordWorkflowUnknowns({ name, result: result.value, diff --git a/tests/boundary/mcp/workflowSnapshotReplayMcp.test.ts b/tests/boundary/mcp/workflowSnapshotReplayMcp.test.ts new file mode 100644 index 0000000000..2a985d61e2 --- /dev/null +++ b/tests/boundary/mcp/workflowSnapshotReplayMcp.test.ts @@ -0,0 +1,230 @@ +import { writeFile } from "node:fs/promises"; +import { join } from "node:path"; + +import { Client, InMemoryTransport } from "@modelcontextprotocol/client"; +import { afterEach, describe, expect, it } from "vitest"; +import { z } from "zod"; + +import { + createAnalysisExecution, + type AnalysisProvider, + type CapabilityDescriptor, +} from "../../../src/application/AnalysisProvider.js"; +import type { DirectAnalysisDependencies } from "../../../src/application/DirectAnalysisDependencies.js"; +import { runDirectAnalysis } from "../../../src/application/DirectAnalysis.js"; +import { createTestBinarySession } from "../../fixtures/binarySession.js"; +import { createTestTempDirectory } from "../../fixtures/temporaryDirectory.js"; +import { createAnalysisProfile } from "../../../src/domain/analysisProfile.js"; +import { readAnalysisSnapshot } from "../../../src/application/binary/AnalysisSnapshotFiles.js"; +import { ok } from "../../../src/domain/result.js"; +import { silentLogger } from "../../../src/logger.js"; +import { createServer } from "../../../src/server/createServer.js"; +import { toolContract } from "../../../src/contracts/toolContracts.js"; + +const identity = { id: "fixture", name: "Fixture", version: "1" } as const; +const profile = createAnalysisProfile(identity, { fixture: true }); +const operations = [ + "list_segments", + "list_documents", + "list_procedures", + "list_strings", +] as const; + +const resources: Array<{ close(): Promise }> = []; + +afterEach(async () => { + await Promise.all(resources.splice(0).map((resource) => resource.close())); +}); + +describe("MCP composed workflow snapshot replay", () => { + it("exports an MCP workflow binding and replays its Evidence without provider startup or calls", async () => { + const directory = await createTestTempDirectory( + "rea-mcp-workflow-snapshot-", + ); + const targetPath = join(directory, "fixture.hop"); + const snapshotPath = join(directory, "analysis.json"); + await writeFile(targetPath, "fixture binary"); + const starts: string[] = []; + const calls: string[] = []; + const provider = makeProvider(starts, calls); + const session = createTestBinarySession(provider, { + resolveAnalysisProfile: () => + Promise.resolve(ok({ profile, compatibility: {} })), + }); + const server = createServer(session, session, { logger: silentLogger }); + const mcp = new Client({ name: "workflow-snapshot", version: "1.0.0" }); + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + resources.push(mcp, server); + await server.connect(serverTransport); + await mcp.connect(clientTransport); + + const opened = await mcp.callTool({ + name: "open_binary", + arguments: { path: targetPath }, + }); + expect(opened.isError).not.toBe(true); + const analyzed = await mcp.callTool({ + name: "binary_overview", + arguments: {}, + }); + expect(analyzed.isError, JSON.stringify(analyzed.content)).not.toBe(true); + const output = toolContract("binary_overview").outputSchema.parse( + analyzed.structuredContent, + ); + const workflowEvidence = z + .object({ + operation: z.literal("binary_overview"), + analysis_profile: z.object({ digest: z.string() }), + }) + .passthrough() + .parse(output.evidence); + const afterAnalysisCalls = [...calls]; + const afterAnalysisStarts = [...starts]; + + const closed = await mcp.callTool({ + name: "close_binary", + arguments: { snapshot_path: snapshotPath }, + }); + expect(closed.isError, JSON.stringify(closed.content)).not.toBe(true); + const snapshot = await readAnalysisSnapshot(snapshotPath); + if (!snapshot.ok) throw snapshot.error; + expect(snapshot.value.workflow_entries).toHaveLength(1); + expect(snapshot.value.evidence_bundle.records).toContainEqual( + expect.objectContaining({ + evidence_id: workflowEvidence.evidence_id, + operation: "binary_overview", + }), + ); + + const dependencies: DirectAnalysisDependencies = { + createBinarySession: () => + createTestBinarySession(provider, { + resolveAnalysisProfile: () => + Promise.resolve(ok({ profile, compatibility: {} })), + }), + createManagedBinarySession: () => + createTestBinarySession(provider, { + resolveAnalysisProfile: () => + Promise.resolve(ok({ profile, compatibility: {} })), + }), + }; + const replay = await runDirectAnalysis( + dependencies, + targetPath, + "binary_overview", + {}, + { snapshotPath }, + ); + + expect(replay).toEqual(output.evidence); + expect(calls).toEqual(afterAnalysisCalls); + expect(starts).toEqual(afterAnalysisStarts); + }); + + it("does not bind a workflow when an upstream operation is live", async () => { + const directory = await createTestTempDirectory( + "rea-mcp-live-workflow-snapshot-", + ); + const targetPath = join(directory, "fixture.hop"); + const snapshotPath = join(directory, "analysis.json"); + await writeFile(targetPath, "fixture binary"); + const starts: string[] = []; + const calls: string[] = []; + const session = createTestBinarySession( + makeProvider(starts, calls, "list_strings"), + { + resolveAnalysisProfile: () => + Promise.resolve(ok({ profile, compatibility: {} })), + }, + ); + const server = createServer(session, session, { logger: silentLogger }); + const mcp = new Client({ + name: "live-workflow-snapshot", + version: "1.0.0", + }); + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + resources.push(mcp, server); + await server.connect(serverTransport); + await mcp.connect(clientTransport); + + expect( + ( + await mcp.callTool({ + name: "open_binary", + arguments: { path: targetPath }, + }) + ).isError, + ).not.toBe(true); + const analyzed = await mcp.callTool({ + name: "binary_overview", + arguments: {}, + }); + expect(analyzed.isError, JSON.stringify(analyzed.content)).not.toBe(true); + expect( + toolContract("binary_overview").outputSchema.parse( + analyzed.structuredContent, + ).evidence, + ).toMatchObject({ operation: "binary_overview" }); + const closed = await mcp.callTool({ + name: "close_binary", + arguments: { snapshot_path: snapshotPath }, + }); + expect(closed.isError, JSON.stringify(closed.content)).not.toBe(true); + + const snapshot = await readAnalysisSnapshot(snapshotPath); + if (!snapshot.ok) throw snapshot.error; + expect(snapshot.value.workflow_entries ?? []).toEqual([]); + }); +}); + +const makeProvider = ( + starts: string[], + calls: string[], + liveOperation?: (typeof operations)[number], +): AnalysisProvider => { + const capabilities: CapabilityDescriptor[] = operations.map((operation) => ({ + operation, + provider: identity, + available: true, + reason: null, + cachePolicy: operation === liveOperation ? "live" : "snapshot", + effects: { + mutatesArtifact: false, + launchesProcess: true, + mayShowUi: false, + mayAccessNetwork: false, + mayWriteFilesystem: false, + changesPermissions: false, + requiresRoot: false, + }, + limitations: [], + })); + return { + identity: () => identity, + capabilities: () => capabilities, + resolveAnalysisProfile: () => + Promise.resolve(ok({ profile, compatibility: {} })), + createClient: () => { + starts.push("start"); + return { + execute: async (operation) => { + calls.push(operation); + const result = + operation === "health" + ? null + : operation === "list_segments" + ? [{ name: "__TEXT", start: "0x1000", end: "0x2000" }] + : operation === "list_documents" + ? ["fixture"] + : operation === "list_procedures" + ? ["0x1000"] + : { "0x1000": "fixture" }; + return ok(createAnalysisExecution(result, identity)); + }, + close: () => Promise.resolve(), + }; + }, + }; +};