From 07b36c148946f46b99a4ce04e55726e97d0ee03a Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:27:51 +0800 Subject: [PATCH] fix(mcp): keep prompt completion within ready provider sessions Avoid lazy provider startup for document and procedure completion. Propagate MCP cancellation through completers, cancel lookups during session transitions, and share provider result projection with ordinary operations. Cover cold, starting, warm, queued, cancelled, and closed sessions through MCP composition and provider boundaries. --- src/application/AnalysisProvider.ts | 6 + src/application/binary/BinarySession.ts | 45 +++ src/application/binary/BinarySessionPort.ts | 7 + src/application/binary/CompositeProvider.ts | 9 + .../binary/LazyAnalysisProvider.ts | 8 + src/ghidra/GhidraClient.ts | 29 +- src/ghidra/GhidraProviderClient.ts | 71 ++++- src/hopper/HopperClient.ts | 37 ++- src/hopper/HopperProvider.ts | 86 +++--- src/ida/IdaSessionClient.test.ts | 80 ++++++ src/ida/IdaSessionClient.ts | 31 +- src/server/EvidenceMcpServer.ts | 36 ++- src/server/promptCompletion.ts | 43 ++- tests/boundary/mcp/promptMcp.test.ts | 269 +++++++++++++++++- .../providers/ghidra/ghidraClient.test.ts | 20 ++ .../hopper/hopperClient.protocol.test.ts | 20 ++ 16 files changed, 740 insertions(+), 57 deletions(-) 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(