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/server/environment/StateResponseCleanup.ts

@@ -1,53 +0,0 @@
-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,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<unknown> = Promise.resolve();
   return (method, params, signal) => {

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,

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" } },
     {

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<Uint8Array> | null,
   bounds: DownloadBounds,
-  responses: StateResponseCleanup,
+  responses: ResponseBodyCleanup,
 ): Promise<Result<Buffer, RpcError>> => {
   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<Result<Buffer, RpcError>> => {
   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<Result<Uint8Array, RpcError>> => {
   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<Result<unknown, RpcError>> => {
   if (isRecord(part) && part.kind === "data" && isRecord(part.data))
     return ok(part.data);