diff --git a/src/server/environment/StateResponseCleanup.ts b/src/server/environment/StateResponseCleanup.ts new file mode 100644 index 0000000000..72f0abd211 --- /dev/null +++ b/src/server/environment/StateResponseCleanup.ts @@ -0,0 +1,53 @@ +import { + ArtifactResourceScope, + type ArtifactResourceOwner, +} from "../../artifacts/ArtifactResourceScope.js"; + +/** Retain a cancelled response branch until its underlying cancellation settles. */ +export class StateResponseCleanup { + readonly #scope = new ArtifactResourceScope(); + + constructor(private readonly changed: () => void = () => {}) {} + + get cleanupResources(): readonly string[] { + return this.#scope.cleanupResources; + } + + cancel(cancel: () => Promise): void { + let state: + | { status: "pending" } + | { status: "done" } + | { status: "failed"; cause: unknown } = { status: "pending" }; + const owner: ArtifactResourceOwner = { + kind: "lease", + resource: "state response body", + lease: { + release: async () => { + if (state.status === "pending") + throw new Error("State response cancellation is still pending"); + if (state.status === "failed") throw state.cause; + }, + }, + }; + this.#scope.retain(owner); + this.changed(); + void Promise.resolve() + .then(cancel) + .then( + async () => { + state = { status: "done" }; + await this.#scope.release(owner); + this.changed(); + }, + (cause: unknown) => { + state = { status: "failed", cause }; + this.changed(); + }, + ) + .catch(() => undefined); + } + + close(): Promise { + return this.#scope.close(); + } +} diff --git a/src/server/environment/environmentDataPlane.ts b/src/server/environment/environmentDataPlane.ts index 10623c61a8..94002164e5 100644 --- a/src/server/environment/environmentDataPlane.ts +++ b/src/server/environment/environmentDataPlane.ts @@ -1,3 +1,4 @@ +import { StateResponseCleanup } from "./StateResponseCleanup.js"; import type { EnvironmentStateStore } from "../../application/environment/environmentStateStore.js"; import { ArtifactResourceScope } from "../../artifacts/ArtifactResourceScope.js"; import type { AnalysisInputError } from "../../domain/analysisErrorCore.js"; @@ -46,17 +47,24 @@ interface CallContext { readonly intake: StateIntake; readonly cancel: AbortSignal; readonly files: ArtifactResourceScope; + readonly responses: StateResponseCleanup; } /** `data/add`: exactly one part, whose document replaces the whole state only when it is valid and the call was not cancelled. */ const addState = async ( - { store, intake, cancel, files }: CallContext, + { store, intake, cancel, files, responses }: CallContext, params: unknown, ): Promise => { const parts = isRecord(params) ? params.parts : undefined; if (!Array.isArray(parts) || parts.length !== 1) return { error: invalidParams("data/add takes exactly one part") }; - const document = await partDocument(parts[0], intake, cancel, files); + const document = await partDocument( + parts[0], + intake, + cancel, + files, + responses, + ); if (!document.ok) return { error: document.error }; // A file read is not interruptible, so a cancellation during it is honoured here, before the state changes. if (cancel.aborted) return { error: CANCELLED }; @@ -101,6 +109,7 @@ export const stateDataPlane = ( intake: StateIntake = DEFAULT_STATE_INTAKE, stopped: AbortSignal = new AbortController().signal, files: ArtifactResourceScope = new ArtifactResourceScope(), + responses: StateResponseCleanup = new StateResponseCleanup(), ): RpcMethods => { let previous: Promise = Promise.resolve(); return (method, params, signal) => { @@ -111,7 +120,7 @@ export const stateDataPlane = ( const next = previous.then(() => cancel.aborted ? { error: CANCELLED } - : run({ store, intake, cancel, files }, params), + : run({ store, intake, cancel, files, responses }, params), ); previous = next.catch(() => undefined); return next; diff --git a/src/server/environment/serveEnvironment.ts b/src/server/environment/serveEnvironment.ts index 91bb183ea8..7fa0c13022 100644 --- a/src/server/environment/serveEnvironment.ts +++ b/src/server/environment/serveEnvironment.ts @@ -1,3 +1,4 @@ +import { StateResponseCleanup } from "./StateResponseCleanup.js"; import type { HttpHostname } from "../../domain/httpHostname.js"; import { STATE_INTAKE_LIMIT_BYTES } from "../../application/environment/environmentStateDocument.js"; import { @@ -477,6 +478,8 @@ 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 served = await serveStreamableHttp({ host, port, @@ -499,6 +502,7 @@ export const serveEnvironment = async ( }, stopping.signal, stateFiles, + stateResponses, ), ), ...(logger === undefined ? {} : { logger }), @@ -512,11 +516,14 @@ export const serveEnvironment = async ( }); let httpClosePending = false; let stateClosePending = false; + let responseClosePending = false; + let settledSuccessful = false; let closing: Promise | undefined; const shutdownFailures = new Set(); const owner: EnvironmentShutdownOwner = { get cleanupResources() { const stateResources = stateFiles.cleanupResources; + const responseResources = stateResponses.cleanupResources; return [ ...(httpClosePending ? [served.url] : []), ...(stateClosePending @@ -524,6 +531,11 @@ export const serveEnvironment = async ( ? stateResources : ["environment state-file scope"] : []), + ...(responseClosePending + ? responseResources.length > 0 + ? responseResources + : ["state response cleanup"] + : []), ]; }, close: () => { @@ -531,7 +543,7 @@ export const serveEnvironment = async ( const attempt = (async () => { stop(); const failures: { - readonly stage: "http-server" | "state-files"; + readonly stage: "http-server" | "state-files" | "state-responses"; readonly cause: unknown; }[] = []; try { @@ -548,6 +560,15 @@ export const serveEnvironment = async ( stateClosePending = true; failures.push({ stage: "state-files", cause }); } + try { + await stateResponses.close(); + if (stateResponses.cleanupResources.length > 0) + throw new Error("State response cleanup remains pending"); + responseClosePending = false; + } catch (cause: unknown) { + responseClosePending = true; + failures.push({ stage: "state-responses", cause }); + } if (failures.length > 0) { retainedEnvironmentShutdowns.add(owner); const failure = new AnalysisOutputError( @@ -583,6 +604,7 @@ export const serveEnvironment = async ( retainedEnvironmentShutdowns.delete(owner); for (const failure of shutdownFailures) failure.resolveCleanup(); shutdownFailures.clear(); + settledSuccessful = true; })(); closing = attempt; void attempt.catch(() => { @@ -591,6 +613,26 @@ export const serveEnvironment = async ( return attempt; }, }; + // A fetcher may hand back a response after an earlier stop completed. Keep + // that newly acquired body owned and make a later close observe its outcome. + responseStateChanged = () => { + if (!stopping.signal.aborted) return; + if (stateResponses.cleanupResources.length > 0) { + responseClosePending = true; + retainedEnvironmentShutdowns.add(owner); + if (settledSuccessful) { + settledSuccessful = false; + closing = undefined; + } + } else if ( + shutdownFailures.size === 0 && + !httpClosePending && + !stateClosePending + ) { + responseClosePending = false; + retainedEnvironmentShutdowns.delete(owner); + } + }; return ok({ ...served, card, diff --git a/src/server/environment/serveEnvironmentState.test.ts b/src/server/environment/serveEnvironmentState.test.ts index 6bfb839cc4..6937fd182b 100644 --- a/src/server/environment/serveEnvironmentState.test.ts +++ b/src/server/environment/serveEnvironmentState.test.ts @@ -1037,7 +1037,122 @@ describe("serveEnvironment with a state", () => { expect(missing.isError).toBe(true); }); }); +}); +describe("state response cleanup", () => { + it("reports pending teed response cleanup and retries after its sibling closes", async () => { + let sourceCancelled = false; + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array([1])); + }, + cancel() { + sourceCancelled = true; + }, + }); + const [body, sibling] = stream.tee(); + const environment = await serveState( + { source, state }, + { + fetchState: async () => + new Response(body, { + headers: { "content-length": String(STATE_INTAKE_LIMIT_BYTES + 1) }, + }), + }, + ); + let timer: ReturnType | undefined; + try { + const response = await fetch( + new URL(ENVIRONMENT_RPC_PATH, environment.url), + { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(addByUri(1)), + signal: AbortSignal.timeout(10_000), + }, + ); + await expect(response.json()).resolves.toMatchObject({ + error: { code: -32602 }, + }); + const closing = Promise.race([ + environment.close(), + new Promise((_, reject) => { + timer = setTimeout( + () => reject(new Error("server close waited for tee sibling")), + 10_000, + ); + }), + ]); + await expect(closing).rejects.toMatchObject({ + _tag: "AnalysisOutputError", + cleanupIncomplete: true, + diagnostics: { cleanup_stages: ["state-responses"] }, + }); + expect(sourceCancelled).toBe(false); + await sibling.cancel(); + expect(sourceCancelled).toBe(true); + await expect(environment.close()).resolves.toBeUndefined(); + } finally { + if (timer !== undefined) clearTimeout(timer); + await Promise.allSettled([body.cancel(), sibling.cancel()]); + await environment.close(); + } + }); + + it("retains a late response cleanup after an earlier successful stop", async () => { + let giveResponse!: (response: Response) => void; + const gate = new Promise((resolve) => { + giveResponse = resolve; + }); + const held = heldFetcher(() => gate); + let release!: () => void; + let acknowledge!: () => void; + const cleanup = new Promise((resolve) => { + release = resolve; + }); + const entered = new Promise((resolve) => { + acknowledge = resolve; + }); + const environment = await serveState( + { source, state }, + { fetchState: held.fetchState }, + ); + const answered = fetch(new URL(ENVIRONMENT_RPC_PATH, environment.url), { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(addByUri(1)), + signal: AbortSignal.timeout(10_000), + }).catch(() => undefined); + try { + await held.fetching; + await environment.close(); + giveResponse( + new Response( + new ReadableStream({ + cancel() { + acknowledge(); + return cleanup; + }, + }), + ), + ); + await entered; + await expect(environment.close()).rejects.toMatchObject({ + _tag: "AnalysisOutputError", + cleanupIncomplete: true, + diagnostics: { cleanup_stages: ["state-responses"] }, + }); + release(); + await expect(environment.close()).resolves.toBeUndefined(); + } finally { + release(); + await answered; + await environment.close(); + } + }); +}); + +describe("serveEnvironment state file cleanup", () => { it("stops without waiting for a stalled download", async () => { const held = heldFetcher(() => new Promise(() => undefined)); const environment = await serveState( diff --git a/src/server/environment/stateIntake.test.ts b/src/server/environment/stateIntake.test.ts index b5aaf98665..d2de84083b 100644 --- a/src/server/environment/stateIntake.test.ts +++ b/src/server/environment/stateIntake.test.ts @@ -1,3 +1,4 @@ +import { StateResponseCleanup } from "./StateResponseCleanup.js"; import { expect, it, vi } from "vitest"; import { ArtifactResourceScope } from "../../artifacts/ArtifactResourceScope.js"; @@ -11,14 +12,17 @@ import { const inlineDocument = async (bytes: string) => { const files = new ArtifactResourceScope(); + const responses = new StateResponseCleanup(); try { return await partDocument( { kind: "file", file: { bytes } }, DEFAULT_STATE_INTAKE, new AbortController().signal, files, + responses, ); } finally { + await responses.close(); await files.close(); } }; @@ -79,3 +83,109 @@ it("downloads with redirects refused and the 30 s timeout by default", async () fetched.mockRestore(); } }); + +it.each(["http-status", "declared-overflow", "stream-overflow"] as const)( + "returns %s while a teed response's other branch remains open", + async (scenario) => { + const source = new ReadableStream({ + start(controller) { + controller.enqueue( + new Uint8Array( + scenario === "stream-overflow" ? STATE_INTAKE_LIMIT_BYTES + 1 : 1, + ), + ); + }, + }); + const [body, sibling] = source.tee(); + const response = new Response(body, { + status: scenario === "http-status" ? 503 : 200, + ...(scenario === "declared-overflow" + ? { + headers: { "content-length": String(STATE_INTAKE_LIMIT_BYTES + 1) }, + } + : {}), + }); + const files = new ArtifactResourceScope(); + const responses = new StateResponseCleanup(); + const operation = partDocument( + { kind: "file", file: { uri: "https://state.example.test/state.json" } }, + { fetchState: async () => response, timeoutMs: 1_000 }, + new AbortController().signal, + files, + responses, + ); + let timer: ReturnType | undefined; + try { + const outcome = await Promise.race([ + operation.then((result) => ({ kind: "result" as const, result })), + new Promise<{ kind: "deadline" }>((resolve) => { + timer = setTimeout(() => resolve({ kind: "deadline" }), 3_000); + }), + ]); + expect(outcome.kind).toBe("result"); + expect(responses.cleanupResources).toEqual(["state response body"]); + await expect(responses.close()).rejects.toMatchObject({ + reason: "unavailable", + }); + if (outcome.kind === "result") { + expect(outcome.result).toMatchObject( + scenario === "http-status" + ? { ok: false, error: { data: { reason: "http-status" } } } + : { + ok: false, + error: { + message: expect.stringContaining( + String(STATE_INTAKE_LIMIT_BYTES), + ), + }, + }, + ); + } + } finally { + if (timer !== undefined) clearTimeout(timer); + await sibling.cancel(); + await operation; + await responses.close(); + await files.close(); + } + }, +); + +it("retains the exact failure when response cancellation rejects", async () => { + const failure = new Error("controlled cancellation failure"); + const body = new ReadableStream({ + cancel() { + throw failure; + }, + }); + const files = new ArtifactResourceScope(); + const responses = new StateResponseCleanup(); + const result = await partDocument( + { kind: "file", file: { uri: "https://state.example.test/state.json" } }, + { + fetchState: async () => new Response(body, { status: 503 }), + timeoutMs: 1_000, + }, + new AbortController().signal, + files, + responses, + ); + expect(result).toMatchObject({ + ok: false, + error: { data: { reason: "http-status" } }, + }); + await new Promise((resolve) => setImmediate(resolve)); + const rejected: unknown = await responses + .close() + .catch((cause: unknown) => cause); + expect(rejected).toMatchObject({ + reason: "unavailable", + cleanup: { resources: ["state response body"] }, + }); + expect((rejected as Error).cause).toBeInstanceOf(AggregateError); + expect(((rejected as Error).cause as AggregateError).errors).toContain( + failure, + ); + expect(responses.cleanupResources).toEqual(["state response body"]); + await files.close(); +}); diff --git a/src/server/environment/stateIntake.ts b/src/server/environment/stateIntake.ts index 5ce12e714b..153905ed86 100644 --- a/src/server/environment/stateIntake.ts +++ b/src/server/environment/stateIntake.ts @@ -1,3 +1,4 @@ +import { StateResponseCleanup } from "./StateResponseCleanup.js"; import { fileURLToPath } from "node:url"; import { mediaTypeEssence } from "../../domain/analysisView/recordedMediaType.js"; @@ -228,6 +229,7 @@ const downloadBounds = ( const readCapped = async ( body: ReadableStream | null, bounds: DownloadBounds, + responses: StateResponseCleanup, ): Promise> => { if (body === null) return ok(Buffer.alloc(0)); const reader = body.getReader(); @@ -240,13 +242,15 @@ const readCapped = async ( if (done) return ok(Buffer.concat(chunks, total)); total += value.byteLength; if (total > STATE_INTAKE_LIMIT_BYTES) { - await reader.cancel().catch(() => undefined); + // A tee's sibling can be caller-owned; cancelling this branch must not + // delay the refusal until that sibling finishes. + responses.cancel(() => reader.cancel()); return err(TOO_LARGE); } chunks.push(value); } } catch { - void reader.cancel().catch(() => undefined); + responses.cancel(() => reader.cancel()); return err(fetchFailed(bounds.failure())); } }; @@ -256,6 +260,7 @@ const downloadState = async ( url: URL, intake: StateIntake, cancel: AbortSignal, + responses: StateResponseCleanup, ): Promise> => { const bounds = downloadBounds(intake.timeoutMs, cancel); const fetching = Promise.resolve().then(() => @@ -268,7 +273,9 @@ const downloadState = async ( // A response that arrives after the download ended is released unread. void fetching .then( - (late) => late.body?.cancel(), + (late) => { + if (late.body !== null) responses.cancel(() => late.body!.cancel()); + }, () => undefined, ) .catch(() => undefined); @@ -276,7 +283,8 @@ const downloadState = async ( } const declared = Number(response.headers.get("content-length") ?? "0"); if (!response.ok || declared > STATE_INTAKE_LIMIT_BYTES) { - await response.body?.cancel().catch(() => undefined); + // Cancel this branch without waiting for a potentially independent tee sibling. + if (response.body !== null) responses.cancel(() => response.body!.cancel()); return err( response.ok ? TOO_LARGE @@ -286,7 +294,7 @@ const downloadState = async ( }, ); } - return readCapped(response.body, bounds); + return readCapped(response.body, bounds, responses); }; const NETWORK_SHARE = invalidParams( @@ -315,6 +323,7 @@ const fileBytes = async ( intake: StateIntake, cancel: AbortSignal, files: ArtifactResourceScope, + responses: StateResponseCleanup, ): Promise> => { if (typeof file.bytes === "string" && file.uri === undefined) { if (file.bytes.length > MAX_INLINE_BASE64_CHARACTERS) return err(TOO_LARGE); @@ -333,7 +342,8 @@ const fileBytes = async ( invalidParams("a file part carries exactly one of bytes or a uri"), ); const url = new URL(file.uri); - if (url.protocol === "https:") return downloadState(url, intake, cancel); + if (url.protocol === "https:") + return downloadState(url, intake, cancel, responses); if (url.protocol !== "file:") return err(invalidParams("a file uri is a file:// or https:// URL")); const path = localFilePath(url); @@ -346,6 +356,7 @@ export const partDocument = async ( intake: StateIntake, cancel: AbortSignal, files: ArtifactResourceScope, + responses: StateResponseCleanup, ): Promise> => { if (isRecord(part) && part.kind === "data" && isRecord(part.data)) return ok(part.data); @@ -362,6 +373,6 @@ export const partDocument = async ( mediaTypeEssence(mimeType) !== "application/json") ) return err(invalidParams("a file part's mimeType is application/json")); - const bytes = await fileBytes(part.file, intake, cancel, files); + const bytes = await fileBytes(part.file, intake, cancel, files, responses); return bytes.ok ? parseDocumentBytes(bytes.value) : bytes; };