fix(environment): retain response cleanup through shutdown

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>): 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<void> {
+    return this.#scope.close();
+  }
+}

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<RpcOutcome> => {
   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<unknown> = 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;

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<void> | undefined;
   const shutdownFailures = new Set<AnalysisError>();
   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,

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<Uint8Array>({
+      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<typeof setTimeout> | 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<never>((_, 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<Response>((resolve) => {
+      giveResponse = resolve;
+    });
+    const held = heldFetcher(() => gate);
+    let release!: () => void;
+    let acknowledge!: () => void;
+    const cleanup = new Promise<void>((resolve) => {
+      release = resolve;
+    });
+    const entered = new Promise<void>((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<Uint8Array>({
+            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<Response>(() => undefined));
     const environment = await serveState(

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<Uint8Array>({
+      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<typeof setTimeout> | 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<Uint8Array>({
+    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<void>((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();
+});

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<Uint8Array> | null,
   bounds: DownloadBounds,
+  responses: StateResponseCleanup,
 ): Promise<Result<Buffer, RpcError>> => {
   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<Result<Buffer, RpcError>> => {
   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<Result<Uint8Array, RpcError>> => {
   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<Result<unknown, RpcError>> => {
   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;
 };