diff --git a/src/application/AnalysisProvider.ts b/src/application/AnalysisProvider.ts index 87a717795..cc17cac82 100644 --- a/src/application/AnalysisProvider.ts +++ b/src/application/AnalysisProvider.ts @@ -123,6 +123,12 @@ export interface ProviderRequestActivitySnapshot { /** Provider session bound to exactly one parsed artifact. */ export interface AnalysisClient extends AnalysisOperationPort { + /** Execute only if this operation can run without loading or starting a provider. */ + executeIfReady?( + operation: AnalysisOperation, + parameters: Readonly>, + options?: ExecutionOptions, + ): Promise | undefined>; /** Retained token-verified process snapshots for dynamic providers this client started. */ runtimeLineageSnapshots?(): readonly ProviderRuntimeLineageSnapshot[]; /** Current provider work that may outlive a timed-out or cancelled caller. */ diff --git a/src/application/binary/BinarySession.ts b/src/application/binary/BinarySession.ts index 9eff0e0b8..f7b2d3529 100644 --- a/src/application/binary/BinarySession.ts +++ b/src/application/binary/BinarySession.ts @@ -81,7 +81,9 @@ export class BinarySession | undefined; #transition: Promise = Promise.resolve(); #transitionGeneration = 0; + #transitionPending = 0; readonly #calls = new Map, number>(); + readonly #promptCompletionControllers = new Set(); readonly #providerRouter: SessionProviderRouter; readonly #runtimeUnavailability = new Map< string, @@ -532,6 +534,38 @@ export class BinarySession return call.finally(() => this.#calls.delete(call)); } + /** Run a prompt lookup only against an already connected provider child. */ + executeIfProviderReady( + name: Parameters[0], + arguments_: Readonly>, + options?: Pick, + ): Promise | undefined> { + const active = this.#active; + if ( + this.#transitionPending > 0 || + active === undefined || + active.client.executeIfReady === undefined + ) + return Promise.resolve(undefined); + + const controller = new AbortController(); + this.#promptCompletionControllers.add(controller); + const signal = + options?.signal === undefined + ? controller.signal + : AbortSignal.any([controller.signal, options.signal]); + const generation = this.#transitionGeneration; + const call = active.client.executeIfReady(name, arguments_, { + ...options, + signal, + }); + this.#calls.set(call, generation); + return call.finally(() => { + this.#calls.delete(call); + this.#promptCompletionControllers.delete(controller); + }); + } + /** Admit one composed request until its result and session bookkeeping settle. */ withAdmittedAnalysis( operationName: string, @@ -728,6 +762,9 @@ export class BinarySession #serialize( operation: (admittedThrough: number) => Promise, ): Promise { + for (const controller of this.#promptCompletionControllers) + controller.abort(); + this.#transitionPending += 1; // Later calls wait for this transition and must not be included in its drain. const admittedThrough = this.#transitionGeneration++; const result = this.#transition.then( @@ -738,6 +775,14 @@ export class BinarySession () => undefined, () => undefined, ); + void result.then( + () => { + this.#transitionPending -= 1; + }, + () => { + this.#transitionPending -= 1; + }, + ); return result; } diff --git a/src/application/binary/BinarySessionPort.ts b/src/application/binary/BinarySessionPort.ts index ee666f665..749e98b2f 100644 --- a/src/application/binary/BinarySessionPort.ts +++ b/src/application/binary/BinarySessionPort.ts @@ -14,6 +14,7 @@ import type { EvidenceIntegrityError } from "../../domain/evidenceErrors.js"; import type { WorkflowSnapshotRecordInput } from "./BinarySessionRecords.js"; import type { + AnalysisExecution, AnalysisOperation, AnalysisOperationPort, ExecutionOptions, @@ -35,6 +36,12 @@ export interface SavedAnalysisSnapshot { /** Target lifecycle used by CLI and MCP without exposing a concrete provider. */ export interface BinarySessionPort extends AnalysisOperationPort, InvestigationRecordPort { + /** Execute only against a positively ready, already-created provider client. */ + executeIfProviderReady( + operation: AnalysisOperation, + parameters: Readonly>, + options?: Pick, + ): Promise | undefined>; /** Hold the active target binding through a complete composed operation. */ withAdmittedAnalysis( operationName: string, diff --git a/src/application/binary/CompositeProvider.ts b/src/application/binary/CompositeProvider.ts index b0a7e7c13..05184d279 100644 --- a/src/application/binary/CompositeProvider.ts +++ b/src/application/binary/CompositeProvider.ts @@ -94,6 +94,15 @@ export class CompositeProvider implements AnalysisProvider { return outcomes.find((outcome) => !outcome.ok) ?? ok(null); }; return { + executeIfReady: (operation, parameters, options) => { + const provider = this.#providerByOperation.get(operation); + const client = + provider === undefined ? undefined : clients.get(provider); + return ( + client?.executeIfReady?.(operation, parameters, options) ?? + Promise.resolve(undefined) + ); + }, execute: (operation, parameters, options) => { if (operation === "health") return Promise.resolve( diff --git a/src/application/binary/LazyAnalysisProvider.ts b/src/application/binary/LazyAnalysisProvider.ts index b9127e48b..a87e77445 100644 --- a/src/application/binary/LazyAnalysisProvider.ts +++ b/src/application/binary/LazyAnalysisProvider.ts @@ -79,6 +79,14 @@ class LazyAnalysisClient implements AnalysisClient { #loading: Promise | undefined; #closed = false; + executeIfReady: NonNullable = ( + operation, + parameters, + options, + ) => + this.#client?.executeIfReady?.(operation, parameters, options) ?? + Promise.resolve(undefined); + constructor(loadClient: () => Promise, providerId: string) { this.#loadClient = loadClient; this.#providerId = providerId; diff --git a/src/ghidra/GhidraClient.ts b/src/ghidra/GhidraClient.ts index 7731cf872..d1dee3f01 100644 --- a/src/ghidra/GhidraClient.ts +++ b/src/ghidra/GhidraClient.ts @@ -115,6 +115,7 @@ export class GhidraClient { #lastDiagnostics: Readonly> = {}; #startupController: AbortController | undefined; #startPromise: Promise | undefined; + #ready = false; #closePromise: Promise | undefined; readonly #failure = bindGhidraSessionFailure( () => this.#diagnostics(), @@ -216,7 +217,8 @@ export class GhidraClient { void started.then( (result) => { signal?.removeEventListener("abort", onAbort); - if (!result.ok) reset(); + if (result.ok) this.#ready = true; + else reset(); }, (cause: unknown) => { this.#logger.debug( @@ -293,6 +295,20 @@ export class GhidraClient { ); } + /** Dispatch against the connected bridge without invoking startup or retry. */ + callToolIfReady( + operation: GhidraOperation, + parameters: Readonly>, + options: { readonly signal?: AbortSignal } = {}, + ): Promise | undefined> { + if (!this.isReady()) return Promise.resolve(undefined); + return this.#requestQueue.run( + operation, + parameters, + options.signal === undefined ? {} : { signal: options.signal }, + ); + } + /** Read the private immutable source for format parsers that require captured bytes. */ readTargetSnapshot( maximumBytes?: number, @@ -416,6 +432,16 @@ export class GhidraClient { } } + /** True only after the authenticated startup handshake has completed. */ + isReady(): boolean { + return ( + this.#ready && + this.#socket !== undefined && + !this.#socket.destroyed && + this.#closePromise === undefined + ); + } + #cleanupError(cause: unknown): AnalysisError { if (cause instanceof AnalysisError && cause.cleanupIncomplete) return cause; return new ProviderCleanupError( @@ -462,6 +488,7 @@ export class GhidraClient { } #requestClose(forceStop: boolean): Promise { + this.#ready = false; const starting = this.#startPromise; const controller = this.#startupController; controller?.abort(); diff --git a/src/ghidra/GhidraProviderClient.ts b/src/ghidra/GhidraProviderClient.ts index d2b963e52..fc0297789 100644 --- a/src/ghidra/GhidraProviderClient.ts +++ b/src/ghidra/GhidraProviderClient.ts @@ -108,6 +108,7 @@ export type GhidraProviderClientFactory = ( Partial< Pick< GhidraClient, + | "callToolIfReady" | "operationHealthSnapshot" | "runtimeLineage" | "readTargetSnapshot" @@ -667,6 +668,40 @@ export const createGhidraProviderClient = (input: { startupTimeoutMs: config.ghidraStartupTimeoutMs, }) : undefined; + const projectExecution = ( + operation: AnalysisOperation, + normalized: JsonValue, + rawResult: JsonValue, + observationLimitations: readonly string[] = [], + ) => { + const relationshipCoverage = relationshipCoverageFor(operation); + return createAnalysisExecution(normalized, committedProfile.provider, { + rawResult, + analysisProfile: committedProfile, + ...(relationshipCoverage === undefined ? {} : { relationshipCoverage }), + limitations: [ + ...limitationsFor(operation), + ...observationLimitations, + ...sessionLimitations(), + ], + }); + }; + const inventoryExecution = ( + operation: "list_documents" | "list_procedures", + called: Result, + ) => { + if (!called.ok) + return err( + projectSessionError( + operation, + called.error, + config.ghidraStartupTimeoutMs, + ), + ); + const parsed = parseGhidraInventoryResult(operation, called.value); + if (!parsed.ok) return parsed; + return ok(projectExecution(operation, parsed.value, called.value)); + }; return { execute: async (operation, parameters, options) => { if (extensionFailure !== undefined) return err(extensionFailure); @@ -724,6 +759,8 @@ export const createGhidraProviderClient = (input: { ? {} : { progress: options.progress }), }); + if (operation === "list_documents" || operation === "list_procedures") + return inventoryExecution(operation, called); if (!called.ok) return err( projectSessionError( @@ -827,22 +864,30 @@ export const createGhidraProviderClient = (input: { if (!attested.ok) return attested; normalized = attested.value; } - const relationshipCoverage = relationshipCoverageFor(operation); return ok( - createAnalysisExecution(normalized, committedProfile.provider, { - rawResult: called.value, - analysisProfile: committedProfile, - ...(relationshipCoverage === undefined - ? {} - : { relationshipCoverage }), - limitations: [ - ...limitationsFor(operation), - ...observationLimitations, - ...sessionLimitations(), - ], - }), + projectExecution( + operation, + normalized, + called.value, + observationLimitations, + ), ); }, + executeIfReady: async (operation, parameters, options) => { + if (operation !== "list_documents" && operation !== "list_procedures") + return undefined; + if (extensionFailure !== undefined) return err(extensionFailure); + const input = parseGhidraInventoryInput(operation, parameters); + if (!input.ok) return input; + if (client.callToolIfReady === undefined) return undefined; + const called = await client.callToolIfReady( + operation, + input.value, + options?.signal === undefined ? {} : { signal: options.signal }, + ); + if (called === undefined) return undefined; + return inventoryExecution(operation, called); + }, operationHealthSnapshot: () => client.operationHealthSnapshot?.(), runtimeLineageSnapshots: () => { const observation = client.runtimeLineage?.() ?? null; diff --git a/src/hopper/HopperClient.ts b/src/hopper/HopperClient.ts index 2e903c8d2..8196adcee 100644 --- a/src/hopper/HopperClient.ts +++ b/src/hopper/HopperClient.ts @@ -121,6 +121,7 @@ export class HopperClient { #launcherFailureDiagnostic: HopperStartupFailureDiagnostic | undefined; #operationFailure: OperationFailureRecord | undefined; #startupController: AbortController | undefined; + #ready = false; #startupFailureResult: AnalysisError | undefined; #runId: string | undefined; #startPromise: Promise> | undefined; @@ -257,6 +258,16 @@ export class HopperClient { }; } + /** True only after the authenticated health handshake has completed. */ + isReady(): boolean { + return ( + this.#ready && + this.#socket !== undefined && + !this.#socket.destroyed && + this.#closePromise === undefined + ); + } + /** Launch the bridge once and complete its authenticated health handshake. */ start( signal?: AbortSignal, @@ -292,7 +303,7 @@ export class HopperClient { if (!result.ok) { this.#rememberStartupFailure(result.error); reset(); - } + } else this.#ready = true; }, (cause: unknown) => { this.#logger.debug( @@ -466,7 +477,28 @@ export class HopperClient { ): Promise> { const started = await this.start(options.signal); if (!started.ok) return started; - const result = await this.#request(name, arguments_, options); + return this.#normalizeToolResult( + name, + await this.#request(name, arguments_, options), + ); + } + + /** Dispatch against the connected bridge without invoking startup or retry. */ + callToolIfReady( + name: string, + arguments_: Readonly> = {}, + options: { readonly signal?: AbortSignal } = {}, + ): Promise | undefined> { + if (!this.isReady()) return Promise.resolve(undefined); + return this.#request(name, arguments_, options).then((result) => + this.#normalizeToolResult(name, result), + ); + } + + #normalizeToolResult( + name: string, + result: Result, + ): Result { return !result.ok && result.error instanceof HopperRemoteError && result.error.diagnosticType === "capability_unavailable" @@ -579,6 +611,7 @@ export class HopperClient { options: HopperClientCloseOptions = {}, ): Promise> { this.#closing = true; + this.#ready = false; try { return await cleanupHopperSession({ socket: this.#socket, diff --git a/src/hopper/HopperProvider.ts b/src/hopper/HopperProvider.ts index 0b7316a19..c52eefe28 100644 --- a/src/hopper/HopperProvider.ts +++ b/src/hopper/HopperProvider.ts @@ -5,6 +5,7 @@ import { accessSync, constants, statSync } from "node:fs"; import { createAnalysisExecution } from "../application/AnalysisProvider.js"; import type { AnalysisClient, + AnalysisOperation, AnalysisClientContext, AnalysisProfileResolutionOptions, AnalysisProviderCandidate, @@ -16,6 +17,7 @@ import type { import type { AnalysisProfileCommitment } from "../domain/analysisProfile.js"; import type { AppConfig } from "../config/types.js"; import type { BinaryTarget } from "../domain/binaryTargetTypes.js"; +import type { JsonValue } from "../domain/jsonValue.js"; import type { Logger } from "pino"; import { AnalysisCapabilityUnavailableError } from "../domain/analysisErrorCore.js"; import { err, ok } from "../domain/result.js"; @@ -251,6 +253,40 @@ export class HopperProvider implements AnalysisProviderCandidate { this.logger.info(diagnostic, "Hopper bridge reported a diagnostic"), }); const regexSearch = new HopperRegexSearch(client); + const projectExecution = ( + operation: AnalysisOperation, + normalized: JsonValue, + rawResult: JsonValue | null = null, + ) => + createAnalysisExecution(normalized, executionProvider, { + // Only the file-offset mapping differs from Hopper's reply. + rawResult, + ...(profile === undefined ? {} : { analysisProfile: profile }), + ...(operation === "procedure_callers" || + operation === "procedure_callees" + ? { relationshipCoverage: { status: "complete" as const } } + : {}), + limitations: [ + ...(preparedImage === undefined + ? [] + : [ + "REA loaded a verified thin Mach-O slice from the original FAT64 source. Identity and file offsets refer to the original file. The owned temporary image and its Hopper document close together, including on MCP shutdown.", + ]), + ...(CAPABILITIES.find( + (descriptor) => descriptor.operation === operation, + )?.limitations ?? []), + ...(operation === "read_bytes" + ? [ + "Hopper byte reads stop at the containing segment's exclusive end; a readable prefix is returned with complete=false when the requested range crosses it.", + ] + : []), + ...(operation === "procedure_references" + ? [ + "Native CallReference classifications and unresolved targets are preserved. Detailed reference flags and coverage beyond reported CallReference objects remain unknown.", + ] + : []), + ], + }); return { execute: async (operation, parameters, options) => { const result = @@ -265,40 +301,26 @@ export class HopperProvider implements AnalysisProviderCandidate { ? await mapHopperFileOffset(target, result.value, options?.signal) : result; return mapped.ok - ? { - ok: true, - value: createAnalysisExecution(mapped.value, executionProvider, { - // Only the file-offset mapping differs from Hopper's reply. - rawResult: mapped === result ? null : result.value, - ...(profile === undefined ? {} : { analysisProfile: profile }), - ...(operation === "procedure_callers" || - operation === "procedure_callees" - ? { relationshipCoverage: { status: "complete" as const } } - : {}), - limitations: [ - ...(preparedImage === undefined - ? [] - : [ - "REA loaded a verified thin Mach-O slice from the original FAT64 source. Identity and file offsets refer to the original file. The owned temporary image and its Hopper document close together, including on MCP shutdown.", - ]), - ...(CAPABILITIES.find( - (descriptor) => descriptor.operation === operation, - )?.limitations ?? []), - ...(operation === "read_bytes" - ? [ - "Hopper byte reads stop at the containing segment's exclusive end; a readable prefix is returned with complete=false when the requested range crosses it.", - ] - : []), - ...(operation === "procedure_references" - ? [ - "Native CallReference classifications and unresolved targets are preserved. Detailed reference flags and coverage beyond reported CallReference objects remain unknown.", - ] - : []), - ], - }), - } + ? ok( + projectExecution( + operation, + mapped.value, + mapped === result ? null : result.value, + ), + ) : mapped; }, + executeIfReady: async (operation, parameters, options) => { + if (operation !== "list_documents" && operation !== "list_procedures") + return undefined; + const result = await client.callToolIfReady( + operation, + parameters, + options?.signal === undefined ? {} : { signal: options.signal }, + ); + if (result === undefined || !result.ok) return result; + return ok(projectExecution(operation, result.value)); + }, runtimeLineageSnapshots: () => { const observation = client.runtimeLineage(); return observation === null diff --git a/src/ida/IdaSessionClient.test.ts b/src/ida/IdaSessionClient.test.ts index a7d7be083..e43f18cf7 100644 --- a/src/ida/IdaSessionClient.test.ts +++ b/src/ida/IdaSessionClient.test.ts @@ -116,6 +116,86 @@ describe("headless IDA producer pagination", () => { }); describe("attached IDA observation semantics", () => { + it("returns cold and pending startup calls immediately, then cancels queued work before dispatch", async () => { + const { producer, client } = await fixture("headless"); + await expect(client.executeIfReady("list_procedures", {})).resolves.toBe( + undefined, + ); + expect(producer.connects).toBe(0); + + let enterOpen: (() => void) | undefined; + let finishOpen: (() => void) | undefined; + const openEntered = new Promise((resolve) => { + enterOpen = resolve; + }); + const openGate = new Promise((resolve) => { + finishOpen = resolve; + }); + producer.beforeCall = async (name) => { + if (name === "idb_open") { + enterOpen?.(); + await openGate; + } + }; + const startup = client.execute("list_procedures", {}); + await openEntered; + await expect(client.executeIfReady("list_strings", {})).resolves.toBe( + undefined, + ); + expect( + producer.calls.filter(({ name }) => name === "idb_open"), + ).toHaveLength(1); + finishOpen?.(); + const started = await startup; + if (!started.ok) throw started.error; + const warm = await client.executeIfReady("list_procedures", {}); + if (warm === undefined) throw new Error("expected ready IDA dispatch"); + if (!warm.ok) throw warm.error; + expect(warm.value.result).toEqual([ + expect.objectContaining({ address: "0x1000", value: "main" }), + ]); + + let enterRead: (() => void) | undefined; + let finishRead: (() => void) | undefined; + const readEntered = new Promise((resolve) => { + enterRead = resolve; + }); + const readGate = new Promise((resolve) => { + finishRead = resolve; + }); + producer.beforeCall = async (name) => { + if (name === "list_funcs") { + enterRead?.(); + await readGate; + } + }; + const active = client.execute("list_procedures", {}); + await readEntered; + const controller = new AbortController(); + const queued = client.executeIfReady( + "list_strings", + {}, + { + signal: controller.signal, + }, + ); + controller.abort(); + await expect(queued).resolves.toMatchObject({ + ok: false, + error: { _tag: "AnalysisCancelledError" }, + }); + finishRead?.(); + const activeResult = await active; + if (!activeResult.ok) throw activeResult.error; + expect( + producer.calls.filter(({ name }) => name === "list_funcs"), + ).toHaveLength(3); + expect(producer.calls.some(({ name }) => name === "find_regex")).toBe( + false, + ); + await client.close(); + }); + it("accepts a normalized spelling of the bound document path", async () => { const { target, producer, client } = await fixture(); const equivalent = diff --git a/src/ida/IdaSessionClient.ts b/src/ida/IdaSessionClient.ts index 0633d9f68..f2a624569 100644 --- a/src/ida/IdaSessionClient.ts +++ b/src/ida/IdaSessionClient.ts @@ -1,4 +1,8 @@ import { z } from "zod"; +import { + ABORTED, + waitForAbortable, +} from "../application/binary/AbortablePromise.js"; import { createAnalysisExecution, type AnalysisClient, @@ -141,10 +145,35 @@ export class IdaSessionClient implements AnalysisClient { return call; } + async executeIfReady( + operation: AnalysisOperation, + parameters: Readonly>, + options?: ExecutionOptions, + ) { + if (!this.#openCompleted || this.#closed || this.#closing) + return Promise.resolve(undefined); + const call = this.#tail.then(() => { + if (!this.#openCompleted || this.#closed || this.#closing) + return undefined; + if (wasCancelled(options)) + return err(new AnalysisCancelledError(operation)); + return this.#execute(operation, parameters, options, true); + }); + this.#tail = call.then( + () => undefined, + () => undefined, + ); + const result = await waitForAbortable(call, options?.signal); + return result === ABORTED + ? err(new AnalysisCancelledError(operation)) + : result; + } + async #execute( operation: AnalysisOperation, parameters: Readonly>, options?: ExecutionOptions, + readyOnly = false, ) { if (this.#closed || this.#closing) return err( @@ -202,7 +231,7 @@ export class IdaSessionClient implements AnalysisClient { parameters, }; try { - await this.#start(); + if (!readyOnly) await this.#start(); const before = await this.#observe(options); failureContext.metadataBefore = before; if (wasCancelled(options)) throw new AnalysisCancelledError(operation); diff --git a/src/server/EvidenceMcpServer.ts b/src/server/EvidenceMcpServer.ts index d5f5e0598..dda6d17f3 100644 --- a/src/server/EvidenceMcpServer.ts +++ b/src/server/EvidenceMcpServer.ts @@ -1,3 +1,5 @@ +import { AsyncLocalStorage } from "node:async_hooks"; + import { McpServer, isCallToolResult, @@ -34,6 +36,12 @@ export interface CompactInputSchemaPresentation { readonly budgetBytes: number; } +const currentRequestSignal = new AsyncLocalStorage(); + +/** Signal for the completion request currently being dispatched by the SDK. */ +export const currentMcpRequestSignal = (): AbortSignal | undefined => + currentRequestSignal.getStore(); + /** A Zod schema whose standard interface exposes a JSON Schema projection. */ const hasAdvertisedInputJsonSchema = (value: unknown): boolean => { if (typeof value !== "object" || value === null) return false; @@ -124,6 +132,7 @@ class ErrorEvidenceTransport implements Transport { onmessage: Transport["onmessage"]; readonly hasPerRequestStream?: boolean; private readonly pendingTools = new Map(); + private readonly pendingCompletions = new Map(); constructor( private readonly transport: Transport, @@ -141,6 +150,8 @@ class ErrorEvidenceTransport implements Transport { typeof message.params?.["name"] === "string" ) this.pendingTools.set(message.id, message.params["name"]); + if (isJSONRPCRequest(message) && message.method === "completion/complete") + this.pendingCompletions.set(message.id, new AbortController()); // The SDK suppresses replies after cancellation, so send() cannot release // this bookkeeping. Keep transport authentication outside retained data. if ( @@ -148,14 +159,28 @@ class ErrorEvidenceTransport implements Transport { message.method === "notifications/cancelled" ) { const id = message.params?.["requestId"]; - if (typeof id === "string" || typeof id === "number") + if (typeof id === "string" || typeof id === "number") { this.pendingTools.delete(id); + this.pendingCompletions.get(id)?.abort(); + this.pendingCompletions.delete(id); + } } - this.onmessage?.(message, extra); + const completion = + isJSONRPCRequest(message) && message.method === "completion/complete" + ? this.pendingCompletions.get(message.id) + : undefined; + if (completion === undefined) this.onmessage?.(message, extra); + else + currentRequestSignal.run(completion.signal, () => + this.onmessage?.(message, extra), + ); }; transport.onerror = (error) => this.onerror?.(error); transport.onclose = () => { this.pendingTools.clear(); + for (const completion of this.pendingCompletions.values()) + completion.abort(); + this.pendingCompletions.clear(); this.onclose?.(); }; } @@ -176,6 +201,9 @@ class ErrorEvidenceTransport implements Transport { await this.transport.close(); } finally { this.pendingTools.clear(); + for (const completion of this.pendingCompletions.values()) + completion.abort(); + this.pendingCompletions.clear(); } } @@ -211,8 +239,10 @@ class ErrorEvidenceTransport implements Transport { await this.transport.send({ ...recovered, result }, options); } else await this.transport.send(recovered, options); } finally { - if (!isJSONRPCRequest(message) && id !== undefined) + if (!isJSONRPCRequest(message) && id !== undefined) { this.pendingTools.delete(id); + this.pendingCompletions.delete(id); + } } } diff --git a/src/server/promptCompletion.ts b/src/server/promptCompletion.ts index bc37e8fda..c4a4c028a 100644 --- a/src/server/promptCompletion.ts +++ b/src/server/promptCompletion.ts @@ -2,9 +2,11 @@ import type { EvidenceReader } from "../application/investigation/InvestigationR import { z } from "zod"; import { uniqueSorted } from "../domain/canonicalOrdering.js"; +import type { JsonValue } from "../domain/jsonValue.js"; import type { AnalysisOperationPort } from "../application/AnalysisProvider.js"; import type { BinarySessionPort } from "../application/binary/BinarySessionPort.js"; import type { PromptCompletionKind } from "../contracts/promptContracts.js"; +import { currentMcpRequestSignal } from "./EvidenceMcpServer.js"; import { artifactInventoryResultSchema } from "../domain/artifactGraph.js"; import { artifactInspectionResultSchema } from "../domain/artifactInspection.js"; import { parseProcessCapture } from "../domain/process/processCaptureParsing.js"; @@ -57,9 +59,13 @@ const completionCandidates = async ( ): Promise => { switch (kind) { case "document": - return documentCandidates(analysis); + return documentCandidates(analysis, session); case "procedure": - return procedureCandidates(analysis, context?.arguments?.document); + return procedureCandidates( + analysis, + session, + context?.arguments?.document, + ); case "provider": return providerCandidates(session); case "evidence": @@ -74,8 +80,15 @@ const completionCandidates = async ( const documentCandidates = async ( analysis: AnalysisOperationPort, + session: BinarySessionPort | undefined, ): Promise => { - const result = await analysis.execute("list_documents", {}); + const result = await completionInventoryLookup( + analysis, + session, + "list_documents", + {}, + ); + if (result === undefined) return []; if (!result.ok) return []; const parsed = documentListSchema.safeParse(result.value.result); return parsed.success ? parsed.data : []; @@ -83,12 +96,19 @@ const documentCandidates = async ( const procedureCandidates = async ( analysis: AnalysisOperationPort, + session: BinarySessionPort | undefined, document: string | undefined, ): Promise => { const byAddress = new Map(); const input = document === undefined || document.length === 0 ? {} : { document }; - const result = await analysis.execute("list_procedures", input); + const result = await completionInventoryLookup( + analysis, + session, + "list_procedures", + input, + ); + if (result === undefined) return []; if (result.ok) { const parsed = z .array(z.object({ address: z.string(), value: z.string() })) @@ -111,6 +131,21 @@ const procedureCandidates = async ( return candidates; }; +const completionInventoryLookup = async ( + analysis: AnalysisOperationPort, + session: BinarySessionPort | undefined, + operation: "list_documents" | "list_procedures", + parameters: Readonly>, +) => { + const signal = currentMcpRequestSignal(); + const options = signal === undefined ? {} : { signal }; + const result = + session === undefined + ? await analysis.execute(operation, parameters, options) + : await session.executeIfProviderReady(operation, parameters, options); + return signal?.aborted === true ? undefined : result; +}; + const providerCandidates = ( session: BinarySessionPort | undefined, ): readonly string[] => { diff --git a/tests/boundary/mcp/promptMcp.test.ts b/tests/boundary/mcp/promptMcp.test.ts index 4e0072e8f..9138a102d 100644 --- a/tests/boundary/mcp/promptMcp.test.ts +++ b/tests/boundary/mcp/promptMcp.test.ts @@ -11,7 +11,22 @@ import { afterEach, describe, expect, it } from "vitest"; import { createTestTempDirectory } from "../../fixtures/temporaryDirectory.js"; -import type { AnalysisClient } from "../../../src/application/AnalysisProvider.js"; +import { + createAnalysisExecution, + type AnalysisClient, +} from "../../../src/application/AnalysisProvider.js"; +import type { JsonValue } from "../../../src/domain/jsonValue.js"; +import type { + AnalysisProvider, + AnalysisProviderCandidate, + CapabilityDescriptor, + ProviderIdentity, +} from "../../../src/application/AnalysisProvider.js"; +import { AnalysisProviderRegistry } from "../../../src/application/binary/AnalysisProviderRegistry.js"; +import { composeBinarySession } from "../../../src/application/binary/BinarySessionComposition.js"; +import { createAnalysisProfile } from "../../../src/domain/analysisProfile.js"; +import { AnalysisCancelledError } from "../../../src/domain/analysisErrorCore.js"; +import { err, ok } from "../../../src/domain/result.js"; import { createCacheProvider, createTestBinarySession, @@ -304,6 +319,67 @@ describe("guided prompt completion registry", () => { }); }); +it("guided prompt completion avoids cold startup and cancels warm lookup on close", async () => { + const directory = await createTestTempDirectory("rea-prompt-warm-only-"); + temporaryDirectories.push(directory); + const target = join(directory, "target.hop"); + await writeFile(target, "fixture target"); + const fixture = await createWarmCompletionFixture(target); + const { client } = fixture; + + const opened = await client.callTool({ + name: "open_binary", + arguments: { path: target }, + }); + expect(opened.isError).not.toBe(true); + expect(fixture.childrenCreated()).toBe(0); + expect(await complete(client, "document", "")).toEqual([]); + expect(await complete(client, "procedure", "")).toEqual([]); + expect(fixture.childrenCreated()).toBe(0); + + const startup = fixture.holdStartup(); + const warming = client.callTool({ name: "list_documents", arguments: {} }); + await startup.started; + expect(await complete(client, "document", "")).toEqual([]); + expect(await complete(client, "procedure", "")).toEqual([]); + expect(fixture.childrenCreated()).toBe(1); + startup.finish(); + const warmed = await warming; + expect(warmed.isError, JSON.stringify(warmed)).not.toBe(true); + expect(fixture.childrenCreated()).toBe(1); + expect(await complete(client, "document", "fixture")).toEqual([ + "fixture-document", + ]); + + const cancellation = fixture.holdCompletion(); + const controller = new AbortController(); + const cancelled = client.complete( + { + ref: { type: "ref/prompt", name: "investigate_feature" }, + argument: { name: "document", value: "" }, + }, + { signal: controller.signal }, + ); + await cancellation.started; + controller.abort(); + await expect(cancelled).rejects.toThrow(); + await cancellation.aborted; + + const lifecycle = fixture.holdCompletion(); + const pending = client.complete({ + ref: { type: "ref/prompt", name: "investigate_feature" }, + argument: { name: "document", value: "" }, + }); + await lifecycle.started; + const closed = await client.callTool({ + name: "close_binary", + arguments: {}, + }); + expect(closed.isError).not.toBe(true); + await lifecycle.aborted; + expect((await pending).completion.values).toEqual([]); +}, 10_000); + describe("guided prompt completion lifecycle", () => { it("survives a cancelled completion and ignores its late result", async () => { let calls = 0; @@ -410,6 +486,181 @@ describe("guided prompt completion lifecycle", () => { }, 10_000); }); +const createWarmCompletionFixture = async (target: string) => { + const identity = { + id: "deep-fixture", + name: "Deep fixture provider", + version: "1", + }; + const analysisProfile = createAnalysisProfile(identity, { fixture: true }); + const observeCompletion = (result: JsonValue) => + resultOk(createAnalysisExecution(result, identity, { analysisProfile })); + let childrenCreated = 0; + let ready = false; + let startupGate: + | { + readonly started: () => void; + readonly release: Promise; + } + | undefined; + let completionGate: + | { + readonly started: () => void; + readonly aborted: () => void; + } + | undefined; + const deep = deepProvider(identity, { + createClient() { + childrenCreated += 1; + return { + execute: async (operation) => { + if (operation !== "health") { + const gate = startupGate; + if (gate !== undefined) { + startupGate = undefined; + gate.started(); + await gate.release; + } + ready = true; + } + return observeCompletion( + operation === "list_documents" + ? ["fixture-document"] + : operation === "list_procedures" + ? [{ address: "0x1000", value: "fixture-procedure" }] + : null, + ); + }, + executeIfReady(operation, _parameters, options) { + if ( + !ready || + (operation !== "list_documents" && operation !== "list_procedures") + ) + return Promise.resolve(undefined); + const gate = completionGate; + if (gate === undefined) + return Promise.resolve( + observeCompletion( + operation === "list_documents" + ? ["fixture-document"] + : [{ address: "0x1000", value: "fixture-procedure" }], + ), + ); + completionGate = undefined; + gate.started(); + return new Promise((resolve) => { + const signal = options?.signal; + const cancel = (): void => { + gate.aborted(); + resolve(err(new AnalysisCancelledError(operation))); + }; + if (signal?.aborted === true) cancel(); + else signal?.addEventListener("abort", cancel, { once: true }); + }); + }, + close: () => Promise.resolve(resultOk(null)), + }; + }, + }); + const auxiliaryIdentity: ProviderIdentity = { + id: "aux-fixture", + name: "Aux fixture", + version: "1", + }; + const auxiliary: AnalysisProvider = { + identity: () => auxiliaryIdentity, + capabilities: () => [ + promptCapability(auxiliaryIdentity, "inspect_artifact"), + ], + createClient: () => ({ + execute: () => Promise.resolve(observed(null)), + close: () => Promise.resolve(resultOk(null)), + }), + }; + const session = composeBinarySession(new AnalysisProviderRegistry([deep]), [ + auxiliary, + ]); + const client = await connect(createServer({ kind: "session", session })); + resources.push(session); + return { + client, + childrenCreated: () => childrenCreated, + holdStartup: () => { + let started = (): void => undefined; + let finish = (): void => undefined; + const observedStart = new Promise((resolve) => { + started = resolve; + }); + const release = new Promise((resolve) => { + finish = resolve; + }); + startupGate = { started, release }; + return { started: observedStart, finish }; + }, + holdCompletion: () => { + let started = (): void => undefined; + let aborted = (): void => undefined; + const observedStart = new Promise((resolve) => { + started = resolve; + }); + const observedAbort = new Promise((resolve) => { + aborted = resolve; + }); + completionGate = { started, aborted }; + return { started: observedStart, aborted: observedAbort }; + }, + }; +}; + +const deepProvider = ( + identity: Parameters[0], + options: { + readonly createClient: AnalysisProviderCandidate["createClient"]; + }, +): AnalysisProviderCandidate => ({ + identity: () => identity, + capabilities: () => [ + promptCapability(identity, "list_documents"), + promptCapability(identity, "list_procedures"), + ], + inspectAvailability: () => ({ + status: "available", + code: null, + reason: null, + diagnostics: {}, + }), + inspectTargetSupport: () => ({ + status: "supported", + code: null, + reason: null, + diagnostics: {}, + }), + resolveAnalysisProfile: async () => + ok({ profile: createAnalysisProfile(identity, { fixture: true }) }), + createClient: options.createClient, +}); + +const promptCapability = ( + provider: ProviderIdentity, + operation: CapabilityDescriptor["operation"], +): CapabilityDescriptor => ({ + provider, + operation, + available: true, + reason: null, + cachePolicy: "live", + effects: { + mutatesArtifact: false, + launchesProcess: false, + mayShowUi: false, + mayAccessNetwork: false, + mayWriteFilesystem: false, + changesPermissions: false, + requiresRoot: false, + }, + limitations: [], +}); + const fixtureProvider = { id: "fixture", name: "Fixture", version: "1" }; const fixtureSession = ( @@ -461,6 +712,22 @@ const lifecycleClient = ( } return Promise.resolve(observed(null)); }, + executeIfReady(operation, parameters) { + if (operation === "list_documents") + return Promise.resolve(observed([`${target}-document`])); + if (operation === "list_procedures") { + requests.push(parameters); + return Promise.resolve( + observed([ + { + address: target === "first" ? "0x1000" : "0x2000", + value: `${target}-procedure`, + }, + ]), + ); + } + return Promise.resolve(undefined); + }, close: () => Promise.resolve(resultOk(null)), }); diff --git a/tests/boundary/providers/ghidra/ghidraClient.test.ts b/tests/boundary/providers/ghidra/ghidraClient.test.ts index 43abedb3d..9ce3f77ce 100644 --- a/tests/boundary/providers/ghidra/ghidraClient.test.ts +++ b/tests/boundary/providers/ghidra/ghidraClient.test.ts @@ -150,6 +150,26 @@ afterEach(async () => { await Promise.all(clients.splice(0).map((client) => client.close())); }); +describe("GhidraClient ready-only dispatch", () => { + it("does not start cold and preserves the normal result once ready", async () => { + const launcher = new FixtureLauncher(); + const client = clientFor(launcher); + + await expect(client.callToolIfReady("list_strings", {})).resolves.toBe( + undefined, + ); + expect(launcher.processes).toEqual([]); + + const started = await client.start(); + if (!started.ok) throw started.error; + await expect(client.callToolIfReady("list_strings", {})).resolves.toEqual({ + ok: true, + value: [expect.objectContaining({ value: "fixture value" })], + }); + expect(launcher.processes).toHaveLength(1); + }); +}); + it("reports unknown startup work and stops notifications when cancellation releases the process", async () => { vi.useFakeTimers({ toFake: ["setInterval", "clearInterval"] }); const launcher = new FixtureLauncher("silent"); diff --git a/tests/boundary/providers/hopper/hopperClient.protocol.test.ts b/tests/boundary/providers/hopper/hopperClient.protocol.test.ts index 007f9d336..02a6ae73b 100644 --- a/tests/boundary/providers/hopper/hopperClient.protocol.test.ts +++ b/tests/boundary/providers/hopper/hopperClient.protocol.test.ts @@ -15,6 +15,26 @@ import { } from "./hopperClient.fixture.js"; describe("HopperClient protocol", () => { + it("does not start cold and preserves normal dispatch once ready", async () => { + const launcher = new FixtureLauncher(); + const client = new HopperClient({ launcher, startupTimeoutMs: 1_000 }); + onTestFinished(async () => { + await client.close(); + }); + + await expect( + client.callToolIfReady("echo", { label: "cold" }), + ).resolves.toBe(undefined); + expect(launcher.processes).toEqual([]); + + const started = await client.start(); + if (!started.ok) throw started.error; + await expect( + client.callToolIfReady("echo", { label: "warm" }), + ).resolves.toEqual({ ok: true, value: { label: "warm" } }); + expect(launcher.processes).toHaveLength(1); + }); + it("sanitizes unexpected owned-cleanup exceptions", () => { expect( providerCleanupFailure(