From 46d79ecd78b306bb1cfb9250d238d05cd846fb7f Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Wed, 11 Nov 2026 17:23:44 +0800 Subject: [PATCH] fix(ida): retain HTTP response cleanup through close (#9975) * fix(http): retain response cleanup through connection close * fix(ida): cancel acquired diagnostics during close --- src/ida/IdaMcpConnection.test.ts | 356 ++++++++++++++++++ src/ida/IdaMcpConnection.ts | 201 +++++++--- src/ida/IdaSessionClient.test.ts | 70 ++++ .../ResponseBodyCleanup.ts} | 13 +- .../environment/environmentDataPlane.ts | 8 +- src/server/environment/serveEnvironment.ts | 6 +- src/server/environment/stateIntake.test.ts | 8 +- src/server/environment/stateIntake.ts | 10 +- 8 files changed, 609 insertions(+), 63 deletions(-) rename src/{server/environment/StateResponseCleanup.ts => providers/ResponseBodyCleanup.ts} (79%) diff --git a/src/ida/IdaMcpConnection.test.ts b/src/ida/IdaMcpConnection.test.ts index 1d15e1b886..5cd98628f1 100644 --- a/src/ida/IdaMcpConnection.test.ts +++ b/src/ida/IdaMcpConnection.test.ts @@ -3,6 +3,7 @@ import { Client, SdkErrorCode, SdkHttpError, + StreamableHTTPClientTransport, } from "@modelcontextprotocol/client"; import { afterEach, describe, expect, it, vi } from "vitest"; import { @@ -770,3 +771,358 @@ it("requires complete native IDA JSON text despite usable structured content", ( ), ).toThrow(); }); + +const turn = () => new Promise((resolve) => setImmediate(resolve)); +const cleanupConnection = () => + createIdaMcpConnection( + { + url: "http://127.0.0.1:12345/mcp", + headers: {}, + mode: "attached", + timeoutMs: 1000, + }, + {}, + ); + +it("refuses an oversized HTTP error promptly and retains its exact pending cancellation for retry", async () => { + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const cancel = vi.fn(() => gate); + vi.stubGlobal( + "fetch", + vi.fn( + async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(65537)); + }, + cancel, + }), + { status: 500, headers: { "content-type": "text/plain" } }, + ), + ), + ); + const connection = cleanupConnection(); + try { + const failure = await connection.connect().catch((cause: unknown) => cause); + expect(failure).not.toBeInstanceOf(AnalysisTimeoutError); + expect(failure).toMatchObject({ + message: expect.stringContaining("exceeds 65536 UTF-8 bytes"), + }); + const closeFailure = await connection + .close() + .catch((cause: unknown) => cause); + expect(closeFailure).toMatchObject({ + cleanupIncomplete: true, + cleanupResources: ["IDA MCP HTTP response body"], + }); + await expect(connection.close()).rejects.toMatchObject({ + cleanupIncomplete: true, + }); + expect(cancel).toHaveBeenCalledOnce(); + release(); + await turn(); + await expect(connection.close()).resolves.toBeUndefined(); + expect(closeFailure).toMatchObject({ cleanupIncomplete: false }); + } finally { + release(); + await turn(); + await connection.close(); + } +}); + +it("preserves independent transport-close and exact rejected response-cancellation failures", async () => { + const cancellationFailure = new Error( + "controlled response cancellation rejected", + ); + const cancel = vi.fn(async () => { + throw cancellationFailure; + }); + vi.stubGlobal( + "fetch", + vi.fn( + async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(65537)); + }, + cancel, + }), + { status: 500, headers: { "content-type": "text/plain" } }, + ), + ), + ); + const connection = cleanupConnection(); + await connection.connect().catch(() => undefined); + await turn(); + const transportClose = vi + .spyOn(Client.prototype, "close") + .mockRejectedValueOnce(new Error("controlled transport close rejected")); + try { + const mixed = await connection.close().catch((cause: unknown) => cause); + expect(mixed).toBeInstanceOf(AggregateError); + if (!(mixed instanceof AggregateError)) + throw new Error("Expected both cleanup failures"); + expect(mixed.errors[0]).toMatchObject({ + message: "controlled transport close rejected", + }); + expect(mixed.errors[1]).toMatchObject({ + cleanupIncomplete: true, + cleanupResources: ["IDA MCP HTTP response body"], + }); + const responseFailure = mixed.errors[1] as Error; + expect(responseFailure.cause).toMatchObject({ + cause: { errors: [cancellationFailure] }, + }); + await expect(connection.close()).rejects.toMatchObject({ + cleanupIncomplete: true, + }); + expect(cancel).toHaveBeenCalledOnce(); + } finally { + transportClose.mockRestore(); + // A rejected Web Stream cancel cannot be retried; keep its exact failure. + await connection.close().catch(() => undefined); + } +}); + +it.each([200, 500])( + "retains an ignored-abort late HTTP %s acquisition and its later body cancellation", + async (status) => { + let deliver!: (response: Response) => void; + let entered!: () => void; + const enteredFetch = new Promise((resolve) => { + entered = resolve; + }); + const pendingFetch = new Promise((resolve) => { + deliver = resolve; + }); + let release!: () => void; + const cleanup = new Promise((resolve) => { + release = resolve; + }); + const cancel = vi.fn(() => cleanup); + vi.stubGlobal( + "fetch", + vi.fn(() => { + entered(); + return pendingFetch; + }), + ); + const connection = cleanupConnection(); + const connecting = connection.connect().catch((cause: unknown) => cause); + try { + await enteredFetch; + await expect(connection.close()).rejects.toMatchObject({ + cleanupIncomplete: true, + cleanupResources: ["IDA MCP HTTP response acquisition"], + }); + deliver( + new Response(new ReadableStream({ cancel }), { status }), + ); + await connecting; + await turn(); + await expect(connection.close()).rejects.toMatchObject({ + cleanupIncomplete: true, + cleanupResources: ["IDA MCP HTTP response body"], + }); + expect(cancel).toHaveBeenCalledOnce(); + release(); + await turn(); + await expect(connection.close()).resolves.toBeUndefined(); + } finally { + deliver(new Response(null, { status: 500 })); + release(); + await connecting; + await turn(); + await connection.close(); + } + }, +); + +it("allows SDK transport-close session DELETE traffic to finish its own response", async () => { + const methods: string[] = []; + vi.stubGlobal( + "fetch", + vi.fn(async (_input: unknown, init?: RequestInit) => { + methods.push(init?.method ?? "GET"); + if (init?.method === "GET") return new Response(null, { status: 405 }); + if (init?.method === "DELETE") return new Response(null, { status: 204 }); + const request = JSON.parse(String(init?.body)) as MockedIdaHttpRequest; + if (request.id === undefined) return new Response(null, { status: 202 }); + const response = jsonRpcResponse({ + jsonrpc: "2.0", + id: request.id, + result: + request.method === "initialize" + ? { + protocolVersion: request.params?.protocolVersion, + capabilities: { tools: {} }, + serverInfo: { name: "close-traffic-fixture", version: "1" }, + } + : { tools: [] }, + }); + response.headers.set("mcp-session-id", "owned-fixture-session"); + return response; + }), + ); + const connection = cleanupConnection(); + await connection.connect(); + const originalClose = StreamableHTTPClientTransport.prototype.close; + const close = vi + .spyOn(StreamableHTTPClientTransport.prototype, "close") + .mockImplementationOnce(async function ( + this: StreamableHTTPClientTransport, + ) { + await this.terminateSession(); + await originalClose.call(this); + }); + try { + await expect(connection.close()).resolves.toBeUndefined(); + expect(methods).toContain("DELETE"); + } finally { + close.mockRestore(); + await connection.close(); + } +}); + +it.each(["completed", "pending", "rejected"] as const)( + "cancels an acquired parked HTTP error reader at close after deadline with %s cleanup", + async (outcome) => { + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const rejection = new Error( + "controlled parked-reader cancellation failure", + ); + const cancel = vi.fn(async () => { + if (outcome === "pending") await gate; + if (outcome === "rejected") throw rejection; + }); + vi.stubGlobal( + "fetch", + vi.fn( + async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue( + new TextEncoder().encode("partial diagnostic"), + ); + }, + cancel, + }), + { status: 500, headers: { "content-type": "text/plain" } }, + ), + ), + ); + const connection = createIdaMcpConnection( + { + url: "http://127.0.0.1:12345/mcp", + headers: {}, + mode: "attached", + timeoutMs: 30, + }, + {}, + ); + try { + const failure = await connection + .connect() + .catch((cause: unknown) => cause); + expect(failure).toBeInstanceOf(AnalysisTimeoutError); + await connection.close().catch(() => undefined); + await turn(); + expect(cancel).toHaveBeenCalledOnce(); + if (outcome === "completed") + await expect(connection.close()).resolves.toBeUndefined(); + else { + const retained = await connection + .close() + .catch((cause: unknown) => cause); + expect(retained).toMatchObject({ + cleanupIncomplete: true, + cleanupResources: ["IDA MCP HTTP response body"], + }); + if (outcome === "rejected") { + expect(retained).toMatchObject({ + cause: { cause: { errors: [rejection] } }, + }); + await expect(connection.close()).rejects.toMatchObject({ + cleanupIncomplete: true, + }); + } else { + release(); + await turn(); + await expect(connection.close()).resolves.toBeUndefined(); + } + } + expect(cancel).toHaveBeenCalledOnce(); + } finally { + release(); + await turn(); + await connection.close().catch(() => undefined); + } + }, +); + +it("closes an acquired parked error reader after caller abort and preserves the exact reason", async () => { + let entered!: () => void; + let release!: () => void; + const parked = new Promise((resolve) => { + entered = resolve; + }); + const gate = new Promise((resolve) => { + release = resolve; + }); + const cancel = vi.fn(() => gate); + mockedHttpBoundary({ + serverName: "parked-error-abort-fixture", + tools: [], + call: () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("partial diagnostic")); + }, + pull() { + entered(); + }, + cancel, + }), + { status: 500, headers: { "content-type": "text/plain" } }, + ), + }); + const connection = cleanupConnection(); + const controller = new AbortController(); + const reason = { caller: "parked diagnostic abort" }; + try { + await connection.connect(); + const calling = connection + .call("decompile", {}, { signal: controller.signal }) + .catch((cause: unknown) => cause); + await parked; + controller.abort(reason); + const failure = await calling; + expect(failure).toBeInstanceOf(AnalysisCancelledError); + expect(failure).toMatchObject({ cause: reason }); + await connection.close().catch(() => undefined); + await turn(); + expect(cancel).toHaveBeenCalledOnce(); + await expect(connection.close()).rejects.toMatchObject({ + cleanupIncomplete: true, + cleanupResources: ["IDA MCP HTTP response body"], + }); + release(); + await turn(); + await expect(connection.close()).resolves.toBeUndefined(); + expect(cancel).toHaveBeenCalledOnce(); + } finally { + release(); + await turn(); + await connection.close().catch(() => undefined); + } +}); diff --git a/src/ida/IdaMcpConnection.ts b/src/ida/IdaMcpConnection.ts index f8032c40a8..eb9ca880af 100644 --- a/src/ida/IdaMcpConnection.ts +++ b/src/ida/IdaMcpConnection.ts @@ -1,3 +1,10 @@ +import { + ArtifactResourceScope, + type ArtifactResourceOwner, +} from "../artifacts/ArtifactResourceScope.js"; +import { ResponseBodyCleanup } from "../providers/ResponseBodyCleanup.js"; +import { bindArtifactFileCleanup } from "../artifacts/ArtifactFileOperation.js"; +import { ProviderCleanupError } from "../domain/providerCleanupError.js"; import { mcpToolResultSchema, mcpTextBlocks, @@ -272,71 +279,125 @@ const httpErrorBody = ( /** Keep error-body materialization within the diagnostic source-byte budget. */ const readBoundedResponseText = async ( response: Response, + responses: ResponseBodyCleanup, + signal: AbortSignal | null | undefined, ): Promise => { if (response.body === null) return ""; const reader = response.body.getReader(); const chunks: Buffer[] = []; let total = 0; + let cancellationOwnsReader = false; + const cancel = (): void => { + if (cancellationOwnsReader) return; + cancellationOwnsReader = true; + responses.cancel(async () => { + try { + await reader.cancel(signal?.reason); + } finally { + reader.releaseLock(); + } + }); + }; + signal?.addEventListener("abort", cancel, { once: true }); try { for (;;) { + if (signal?.aborted) { + cancel(); + throw signal.reason; + } const { done, value } = await reader.read(); + if (signal?.aborted) throw signal.reason; if (done) return Buffer.concat(chunks, total).toString("utf8"); total += value.byteLength; if (total > IDA_DIAGNOSTIC_SOURCE_BYTES) { - await reader.cancel().catch(() => undefined); + cancel(); return undefined; } chunks.push(Buffer.from(value)); } } finally { - reader.releaseLock(); + signal?.removeEventListener("abort", cancel); + if (!cancellationOwnsReader) reader.releaseLock(); } }; const diagnosticFetch = ( redactor: IdaDiagnosticRedactor, + responses: ResponseBodyCleanup, + acquisitions: ArtifactResourceScope, + stopped: () => boolean, rejectTransport: (failure: McpHttpUtf8Error) => void, ): typeof fetch => async (input, init) => { - const response = validateMcpHttpResponse( - await fetch(input, init), - rejectTransport, - ); - const readText = response.text.bind(response); - let safeBody: string | undefined; - if (!response.ok) { - // Normalize before the SDK can swallow a body-consumer failure or drop - // its format metadata. Both consumers use this one owned diagnostic. - const sourceBody = await readBoundedResponseText(response); - safeBody = - sourceBody === undefined - ? redactor.text( - `IDA MCP response body omitted: exceeds ${IDA_DIAGNOSTIC_SOURCE_BYTES} UTF-8 bytes.`, - ) - : httpErrorBody(response, sourceBody, redactor); - const body = safeBody; - Object.defineProperty(response, "text", { - value: async (): Promise => body, - }); - } - // Successful JSON remains lazy, and SSE keeps its original body stream. - Object.defineProperty(response, "json", { - value: async (): Promise => { - const text = safeBody ?? (await readText()); - try { - const value: unknown = JSON.parse(text); - return value; - } catch { - // Parsing failure owns diagnostics even when the HTTP status succeeded. - throw new IdaHttpJsonError( - safeBody ?? httpErrorBody(response, text, redactor), - redactor, - ); - } + const beganBeforeStop = !stopped(); + let pending = true; + const owner: ArtifactResourceOwner = { + kind: "lease", + resource: "IDA MCP HTTP response acquisition", + lease: { + release: async () => { + if (pending) + throw new Error( + "IDA MCP HTTP response acquisition is still pending", + ); + }, }, - }); - return response; + }; + acquisitions.retain(owner); + try { + const acquired = await fetch(input, init); + if (init?.signal?.aborted || (beganBeforeStop && stopped())) { + if (acquired.body !== null) + responses.cancel(() => acquired.body!.cancel()); + throw ( + init?.signal?.reason ?? + new Error("IDA MCP connection closed before response acquisition") + ); + } + const response = validateMcpHttpResponse(acquired, rejectTransport); + const readText = response.text.bind(response); + let safeBody: string | undefined; + if (!response.ok) { + // Normalize before the SDK can swallow a body-consumer failure or drop + // its format metadata. Both consumers use this one owned diagnostic. + const sourceBody = await readBoundedResponseText( + response, + responses, + init?.signal, + ); + safeBody = + sourceBody === undefined + ? redactor.text( + `IDA MCP response body omitted: exceeds ${IDA_DIAGNOSTIC_SOURCE_BYTES} UTF-8 bytes.`, + ) + : httpErrorBody(response, sourceBody, redactor); + const body = safeBody; + Object.defineProperty(response, "text", { + value: async (): Promise => body, + }); + } + // Successful JSON remains lazy, and SSE keeps its original body stream. + Object.defineProperty(response, "json", { + value: async (): Promise => { + const text = safeBody ?? (await readText()); + try { + const value: unknown = JSON.parse(text); + return value; + } catch { + // Parsing failure owns diagnostics even when the HTTP status succeeded. + throw new IdaHttpJsonError( + safeBody ?? httpErrorBody(response, text, redactor), + redactor, + ); + } + }, + }); + return response; + } finally { + pending = false; + await acquisitions.release(owner); + } }; const retainedTransportCause = (cause: unknown, message: string): Error => { @@ -363,6 +424,11 @@ export const createIdaMcpConnection = ( const redactor = ( dependencies.createDiagnosticRedactor ?? createIdaDiagnosticRedactor )(config); + const responses = new ResponseBodyCleanup("IDA MCP HTTP response body"); + const acquisitions = new ArtifactResourceScope(); + let stopping = false; + let transportClosed = false; + let closing: Promise | undefined; let stderr = Buffer.alloc(0); let stderrBytes = 0; let httpUtf8Failure: McpHttpUtf8Error | undefined; @@ -500,10 +566,16 @@ export const createIdaMcpConnection = ( }) : new StreamableHTTPClientTransport(new URL(config.url), { requestInit: { headers: config.headers, redirect: "error" }, - fetch: diagnosticFetch(redactor, (failure) => { - httpUtf8Failure ??= failure; - void transport.close().catch(() => undefined); - }), + fetch: diagnosticFetch( + redactor, + responses, + acquisitions, + () => stopping, + (failure) => { + httpUtf8Failure ??= failure; + void transport.close().catch(() => undefined); + }, + ), }); // Keep a bounded diagnostic tail for genuine failures; never persist ambient variables. if (transport instanceof StdioClientTransport) @@ -567,7 +639,48 @@ export const createIdaMcpConnection = ( return decodeIdaJsonResource(uri, response); }, serverInfo: () => jsonValueSchema.parse(client.getServerVersion() ?? null), - close: () => client.close().catch(transportFailure), + close() { + if (closing !== undefined) return closing; + stopping = true; + const attempt = (async () => { + const failures: unknown[] = []; + if (!transportClosed) { + try { + await client.close().catch(transportFailure); + transportClosed = true; + } catch (cause: unknown) { + failures.push(cause); + } + } + for (const resources of [acquisitions, responses]) { + try { + await resources.close(); + } catch (cause: unknown) { + const failure = new ProviderCleanupError( + "ida", + resources.cleanupResources, + { reason: "IDA MCP HTTP response cleanup remains incomplete." }, + { cause, operation: "ida_mcp_close" }, + ); + bindArtifactFileCleanup(cause, failure); + failures.push(failure); + } + } + if (failures.length === 1) throw failures[0]; + if (failures.length > 1) + throw new AggregateError( + failures, + "IDA MCP transport and response cleanup failed", + ); + })(); + closing = attempt; + void attempt + .finally(() => { + if (closing === attempt) closing = undefined; + }) + .catch(() => undefined); + return attempt; + }, detach() { if (!(transport instanceof StdioClientTransport)) return; // The SDK exposes only the child's pid, not the handles it holds. diff --git a/src/ida/IdaSessionClient.test.ts b/src/ida/IdaSessionClient.test.ts index 91d0f3ad29..7192124e57 100644 --- a/src/ida/IdaSessionClient.test.ts +++ b/src/ida/IdaSessionClient.test.ts @@ -1,3 +1,4 @@ +import { createIdaMcpConnection } from "./IdaMcpConnection.js"; import { mkdir, readFile, @@ -2283,3 +2284,72 @@ it("attributes native IDA metadata timeout to the requested analysis operation", expect(result.error.cause).toBe(timeout); await expectCleanClose(producer, client); }); + +it("keeps the session cleanup owner until an ignored-abort HTTP acquisition and body settle", async () => { + const { root, target } = await createIdaTarget(); + roots.push(root); + let deliver!: (response: Response) => void; + let entered!: () => void; + const started = new Promise((resolve) => { + entered = resolve; + }); + const response = new Promise((resolve) => { + deliver = resolve; + }); + let release!: () => void; + const bodyCleanup = new Promise((resolve) => { + release = resolve; + }); + const cancel = vi.fn(() => bodyCleanup); + vi.stubGlobal( + "fetch", + vi.fn(() => { + entered(); + return response; + }), + ); + const config = { + url: "http://127.0.0.1:12345/mcp", + headers: {}, + mode: "attached" as const, + timeoutMs: 30, + }; + const client = new IdaSessionClient( + config, + target, + createIdaMcpConnection(config, {}), + ); + const executing = client.execute("health", {}); + try { + await started; + const result = await executing; + expect(result).toMatchObject({ + ok: false, + error: { _tag: "AnalysisTimeoutError" }, + }); + const firstClose = await client.close(); + expect(firstClose).toMatchObject({ + ok: false, + error: { cleanupIncomplete: true }, + }); + deliver(new Response(new ReadableStream({ cancel }))); + await new Promise((resolve) => setImmediate(resolve)); + expect(await client.close()).toMatchObject({ + ok: false, + error: { cleanupIncomplete: true }, + }); + expect(cancel).toHaveBeenCalledOnce(); + release(); + await new Promise((resolve) => setImmediate(resolve)); + expect(await client.close()).toMatchObject({ ok: true }); + if (firstClose.ok) throw new Error("Expected retained outer cleanup"); + expect(firstClose.error.cleanupIncomplete).toBe(false); + } finally { + deliver(new Response(null, { status: 500 })); + release(); + await executing; + await new Promise((resolve) => setImmediate(resolve)); + await client.close(); + vi.unstubAllGlobals(); + } +}); diff --git a/src/server/environment/StateResponseCleanup.ts b/src/providers/ResponseBodyCleanup.ts similarity index 79% rename from src/server/environment/StateResponseCleanup.ts rename to src/providers/ResponseBodyCleanup.ts index 72f0abd211..fa70f9b4af 100644 --- a/src/server/environment/StateResponseCleanup.ts +++ b/src/providers/ResponseBodyCleanup.ts @@ -1,13 +1,16 @@ import { ArtifactResourceScope, type ArtifactResourceOwner, -} from "../../artifacts/ArtifactResourceScope.js"; +} from "../artifacts/ArtifactResourceScope.js"; /** Retain a cancelled response branch until its underlying cancellation settles. */ -export class StateResponseCleanup { +export class ResponseBodyCleanup { readonly #scope = new ArtifactResourceScope(); - constructor(private readonly changed: () => void = () => {}) {} + constructor( + private readonly resource: string, + private readonly changed: () => void = () => {}, + ) {} get cleanupResources(): readonly string[] { return this.#scope.cleanupResources; @@ -20,11 +23,11 @@ export class StateResponseCleanup { | { status: "failed"; cause: unknown } = { status: "pending" }; const owner: ArtifactResourceOwner = { kind: "lease", - resource: "state response body", + resource: this.resource, lease: { release: async () => { if (state.status === "pending") - throw new Error("State response cancellation is still pending"); + throw new Error(`${this.resource} cancellation is still pending`); if (state.status === "failed") throw state.cause; }, }, diff --git a/src/server/environment/environmentDataPlane.ts b/src/server/environment/environmentDataPlane.ts index 94002164e5..0129b8a6aa 100644 --- a/src/server/environment/environmentDataPlane.ts +++ b/src/server/environment/environmentDataPlane.ts @@ -1,4 +1,4 @@ -import { StateResponseCleanup } from "./StateResponseCleanup.js"; +import { ResponseBodyCleanup } from "../../providers/ResponseBodyCleanup.js"; import type { EnvironmentStateStore } from "../../application/environment/environmentStateStore.js"; import { ArtifactResourceScope } from "../../artifacts/ArtifactResourceScope.js"; import type { AnalysisInputError } from "../../domain/analysisErrorCore.js"; @@ -47,7 +47,7 @@ interface CallContext { readonly intake: StateIntake; readonly cancel: AbortSignal; readonly files: ArtifactResourceScope; - readonly responses: StateResponseCleanup; + readonly responses: ResponseBodyCleanup; } /** `data/add`: exactly one part, whose document replaces the whole state only when it is valid and the call was not cancelled. */ @@ -109,7 +109,9 @@ export const stateDataPlane = ( intake: StateIntake = DEFAULT_STATE_INTAKE, stopped: AbortSignal = new AbortController().signal, files: ArtifactResourceScope = new ArtifactResourceScope(), - responses: StateResponseCleanup = new StateResponseCleanup(), + responses: ResponseBodyCleanup = new ResponseBodyCleanup( + "state response body", + ), ): RpcMethods => { let previous: Promise = Promise.resolve(); return (method, params, signal) => { diff --git a/src/server/environment/serveEnvironment.ts b/src/server/environment/serveEnvironment.ts index 7fa0c13022..99b9a22329 100644 --- a/src/server/environment/serveEnvironment.ts +++ b/src/server/environment/serveEnvironment.ts @@ -1,4 +1,4 @@ -import { StateResponseCleanup } from "./StateResponseCleanup.js"; +import { ResponseBodyCleanup } from "../../providers/ResponseBodyCleanup.js"; import type { HttpHostname } from "../../domain/httpHostname.js"; import { STATE_INTAKE_LIMIT_BYTES } from "../../application/environment/environmentStateDocument.js"; import { @@ -479,7 +479,9 @@ export const serveEnvironment = async ( // A state file whose close failed stays owned here: later file reads retry it, and close reports it if it is still open. const stateFiles = new ArtifactResourceScope(); let responseStateChanged = (): void => undefined; - const stateResponses = new StateResponseCleanup(() => responseStateChanged()); + const stateResponses = new ResponseBodyCleanup("state response body", () => + responseStateChanged(), + ); const served = await serveStreamableHttp({ host, port, diff --git a/src/server/environment/stateIntake.test.ts b/src/server/environment/stateIntake.test.ts index d2de84083b..9bab225321 100644 --- a/src/server/environment/stateIntake.test.ts +++ b/src/server/environment/stateIntake.test.ts @@ -1,4 +1,4 @@ -import { StateResponseCleanup } from "./StateResponseCleanup.js"; +import { ResponseBodyCleanup } from "../../providers/ResponseBodyCleanup.js"; import { expect, it, vi } from "vitest"; import { ArtifactResourceScope } from "../../artifacts/ArtifactResourceScope.js"; @@ -12,7 +12,7 @@ import { const inlineDocument = async (bytes: string) => { const files = new ArtifactResourceScope(); - const responses = new StateResponseCleanup(); + const responses = new ResponseBodyCleanup("state response body"); try { return await partDocument( { kind: "file", file: { bytes } }, @@ -106,7 +106,7 @@ it.each(["http-status", "declared-overflow", "stream-overflow"] as const)( : {}), }); const files = new ArtifactResourceScope(); - const responses = new StateResponseCleanup(); + const responses = new ResponseBodyCleanup("state response body"); const operation = partDocument( { kind: "file", file: { uri: "https://state.example.test/state.json" } }, { fetchState: async () => response, timeoutMs: 1_000 }, @@ -159,7 +159,7 @@ it("retains the exact failure when response cancellation rejects", async () => { }, }); const files = new ArtifactResourceScope(); - const responses = new StateResponseCleanup(); + const responses = new ResponseBodyCleanup("state response body"); const result = await partDocument( { kind: "file", file: { uri: "https://state.example.test/state.json" } }, { diff --git a/src/server/environment/stateIntake.ts b/src/server/environment/stateIntake.ts index 153905ed86..d4de76c47a 100644 --- a/src/server/environment/stateIntake.ts +++ b/src/server/environment/stateIntake.ts @@ -1,4 +1,4 @@ -import { StateResponseCleanup } from "./StateResponseCleanup.js"; +import { ResponseBodyCleanup } from "../../providers/ResponseBodyCleanup.js"; import { fileURLToPath } from "node:url"; import { mediaTypeEssence } from "../../domain/analysisView/recordedMediaType.js"; @@ -229,7 +229,7 @@ const downloadBounds = ( const readCapped = async ( body: ReadableStream | null, bounds: DownloadBounds, - responses: StateResponseCleanup, + responses: ResponseBodyCleanup, ): Promise> => { if (body === null) return ok(Buffer.alloc(0)); const reader = body.getReader(); @@ -260,7 +260,7 @@ const downloadState = async ( url: URL, intake: StateIntake, cancel: AbortSignal, - responses: StateResponseCleanup, + responses: ResponseBodyCleanup, ): Promise> => { const bounds = downloadBounds(intake.timeoutMs, cancel); const fetching = Promise.resolve().then(() => @@ -323,7 +323,7 @@ const fileBytes = async ( intake: StateIntake, cancel: AbortSignal, files: ArtifactResourceScope, - responses: StateResponseCleanup, + responses: ResponseBodyCleanup, ): Promise> => { if (typeof file.bytes === "string" && file.uri === undefined) { if (file.bytes.length > MAX_INLINE_BASE64_CHARACTERS) return err(TOO_LARGE); @@ -356,7 +356,7 @@ export const partDocument = async ( intake: StateIntake, cancel: AbortSignal, files: ArtifactResourceScope, - responses: StateResponseCleanup, + responses: ResponseBodyCleanup, ): Promise> => { if (isRecord(part) && part.kind === "data" && isRecord(part.data)) return ok(part.data);