From 5f30ea6ff2dd89aab7af6da6e8dfa5e508ed459b Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Sun, 11 Oct 2026 00:03:23 +0800 Subject: [PATCH] fix(engines): serialize JEB acquisition and respect IDA cancellation Validation: integrated changes passed formatting, typecheck, lint, resolved module-boundary verification and git diff --check. Tests and real-provider workflows were not run. --- src/ida/IdaOperationRunner.ts | 4 + src/ida/IdaSessionClient.ts | 3 + src/jeb/JebMcpConnection.ts | 177 +++++++++++++++++++++++++++++----- 3 files changed, 159 insertions(+), 25 deletions(-) diff --git a/src/ida/IdaOperationRunner.ts b/src/ida/IdaOperationRunner.ts index e775eb864..553053a31 100644 --- a/src/ida/IdaOperationRunner.ts +++ b/src/ida/IdaOperationRunner.ts @@ -2,6 +2,7 @@ import { z } from "zod"; import { analysisSearchInput } from "../contracts/analysisSearchContract.js"; import { AnalysisInputError, + AnalysisCancelledError, AnalysisProtocolError, } from "../domain/analysisErrorCore.js"; import { @@ -62,16 +63,19 @@ export class IdaOperationRunner { private readonly connection: IdaMcpConnection, private readonly protocol: "legacy" | "modern", private readonly database?: string, + private readonly signal?: AbortSignal, ) {} async call( tool: string, args: Readonly>, ): Promise { + if (this.signal?.aborted) throw new AnalysisCancelledError(tool); const scoped = this.database === undefined ? args : { ...args, database: this.database }; const result = await this.connection.call(tool, scoped); this.raw.push({ tool, arguments: scoped, result }); + if (this.signal?.aborted) throw new AnalysisCancelledError(tool); return result; } diff --git a/src/ida/IdaSessionClient.ts b/src/ida/IdaSessionClient.ts index db6301c0c..d4e51d26e 100644 --- a/src/ida/IdaSessionClient.ts +++ b/src/ida/IdaSessionClient.ts @@ -195,8 +195,11 @@ export class IdaSessionClient implements AnalysisClient { this.connection, this.config.mode === "attached" ? "legacy" : "modern", this.#database, + options?.signal, ); const result = await runner.run(operation, parameters); + if (wasCancelled(options)) + return err(new AnalysisCancelledError(operation)); const after = await this.#observe(); if (wasCancelled(options)) return err(new AnalysisCancelledError(operation)); diff --git a/src/jeb/JebMcpConnection.ts b/src/jeb/JebMcpConnection.ts index b70249ce2..ece06007b 100644 --- a/src/jeb/JebMcpConnection.ts +++ b/src/jeb/JebMcpConnection.ts @@ -4,7 +4,10 @@ import { } from "@modelcontextprotocol/client"; import { z } from "zod"; import { jsonValueSchema, type JsonValue } from "../domain/jsonValue.js"; -import { AnalysisProtocolError } from "../domain/analysisErrorCore.js"; +import { + AnalysisCancelledError, + AnalysisProtocolError, +} from "../domain/analysisErrorCore.js"; /** Injectable MCP boundary used by the provider and conformance tests. */ export interface JebMcpConnection { @@ -61,41 +64,165 @@ export type JebMcpConnectionFactory = (endpoint: URL) => JebMcpConnection; export const createStreamableHttpJebMcpConnection: JebMcpConnectionFactory = ( endpoint, ) => { - let client: Client | undefined; - let serverTools: readonly string[] | undefined; - return { - async connect(signal) { - client ??= new Client({ name: "rea-jeb", version: "0.0.0" }); - if (serverTools === undefined) { - await client.connect( - new StreamableHTTPClientTransport(endpoint, { - requestInit: { redirect: "error" }, - }), - signal === undefined ? {} : { signal }, - ); - const listing = await client.listTools( + interface ConnectionState { + readonly client: Client; + readonly transport: StreamableHTTPClientTransport; + readonly signal: AbortController; + tools: readonly string[] | undefined; + acquisition: Promise | undefined; + disposal: Promise | undefined; + } + let state: ConnectionState | undefined; + let closing: Promise | undefined; + let closed = false; + + const waitForCaller = ( + promise: Promise, + signal?: AbortSignal, + ): Promise => { + if (signal?.aborted === true) + return Promise.reject(new AnalysisCancelledError("jeb_mcp")); + if (signal === undefined) return promise; + return new Promise((resolve, reject) => { + const abort = () => { + signal.removeEventListener("abort", abort); + reject(new AnalysisCancelledError("jeb_mcp")); + }; + signal.addEventListener("abort", abort, { once: true }); + void promise.then( + (value) => { + signal.removeEventListener("abort", abort); + resolve(value); + }, + (cause: unknown) => { + signal.removeEventListener("abort", abort); + reject(cause); + }, + ); + }); + }; + + const dispose = async (owner: ConnectionState): Promise => { + if (owner.disposal !== undefined) return owner.disposal; + const cleanup = (async () => { + const settled = await Promise.allSettled([ + owner.client.close(), + owner.transport.close(), + ]); + const failures = settled.flatMap((result) => + result.status === "rejected" ? [result.reason] : [], + ); + if (failures.length > 0) + throw new AggregateError(failures, "JEB MCP connection cleanup failed"); + if (state === owner) state = undefined; + })(); + owner.disposal = cleanup; + try { + await cleanup; + } catch (cause: unknown) { + owner.disposal = undefined; + throw cause; + } + }; + + const connect = async (signal?: AbortSignal): Promise => { + if (signal?.aborted === true) throw new AnalysisCancelledError("jeb_mcp"); + if (closing !== undefined) await closing; + if (closed) { + if (state !== undefined) await dispose(state); + throw new AnalysisProtocolError("JEB MCP connection is closed"); + } + if (state?.tools !== undefined && state.disposal === undefined) + return state.tools; + if (state?.acquisition !== undefined) + return waitForCaller(state.acquisition, signal); + if (state !== undefined) { + await dispose(state); + if (state !== undefined || closing !== undefined) return connect(signal); + } + const owner: ConnectionState = { + client: new Client({ + name: "rea-jeb", + version: "0.0.0", + }), + transport: new StreamableHTTPClientTransport(endpoint, { + requestInit: { redirect: "error" }, + }), + signal: new AbortController(), + tools: undefined, + acquisition: undefined, + disposal: undefined, + }; + state = owner; + const acquisition = (async () => { + try { + await owner.client.connect(owner.transport, { + signal: owner.signal.signal, + }); + const listing = await owner.client.listTools( {}, - signal === undefined ? {} : { signal }, + { signal: owner.signal.signal }, ); - serverTools = listing.tools.map((tool) => tool.name); + const tools = listing.tools.map((tool) => tool.name); + owner.tools = tools; + return tools; + } catch (cause: unknown) { + owner.signal.abort(cause); + try { + await dispose(owner); + } catch (cleanupCause: unknown) { + throw new AggregateError( + [cause, cleanupCause], + "JEB MCP connection acquisition and cleanup failed", + { cause }, + ); + } + throw cause; } - return serverTools; - }, + })(); + owner.acquisition = acquisition; + void acquisition.then( + () => { + if (owner.acquisition === acquisition) owner.acquisition = undefined; + }, + () => { + if (owner.acquisition === acquisition) owner.acquisition = undefined; + }, + ); + return waitForCaller(acquisition, signal); + }; + + return { + connect, async call(name, args, signal) { - if (serverTools === undefined) await this.connect(signal); - if (client === undefined) + await connect(signal); + const owner = state; + if (owner === undefined || owner.tools === undefined) throw new AnalysisProtocolError("JEB MCP client failed to connect"); - const result = await client.callTool( + const result = await owner.client.callTool( { name, arguments: args }, signal === undefined ? {} : { signal }, ); return decodeJebToolResult(result); }, async close() { - const active = client; - if (active !== undefined) await active.close(); - client = undefined; - serverTools = undefined; + closed = true; + if (closing !== undefined) return closing; + const operation = (async () => { + const owner = state; + if (owner === undefined) return; + owner.tools = undefined; + owner.signal.abort(); + if (owner.acquisition !== undefined) + await owner.acquisition.catch(() => undefined); + if (state === owner) await dispose(owner); + })(); + closing = operation; + try { + await operation; + } finally { + if (closing === operation) closing = undefined; + } }, }; };