diff --git a/src/application/environment/environmentState.test.ts b/src/application/environment/environmentState.test.ts index bdb66a88b9..655414555b 100644 --- a/src/application/environment/environmentState.test.ts +++ b/src/application/environment/environmentState.test.ts @@ -418,6 +418,30 @@ describe("EnvironmentStateStore", () => { }); }); +describe("caller-owned state replacement", () => { + it("does not retain mutable data from a caller-owned replacement", () => { + const store = new EnvironmentStateStore(state); + const document = documentWith((value) => { + const entity = value.rea_state.entities.find((item) => item.id === "2"); + if (!entity) throw new Error("Expected the user entity"); + entity.fields = { ...(entity.fields as object), login: "renamed" }; + }); + expect(store.replace(document).ok).toBe(true); + const callerEntity = document.rea_state.entities.find( + (entity) => entity.id === "2", + ); + if (!callerEntity) throw new Error("Expected the caller's user entity"); + (callerEntity.fields as Record).login = "later edit"; + + expect(store.snapshot().rea_state.entities).toContainEqual( + expect.objectContaining({ + id: "2", + fields: expect.objectContaining({ login: "renamed" }), + }), + ); + }); +}); + describe("rea-state/v1 validation", () => { it.each([ [ @@ -570,3 +594,36 @@ describe("rea-state/v1 validation", () => { ).toMatchObject({ path: "rea_state/version" }); }); }); + +describe("wide state value validation", () => { + it("preserves a wide caller value through the state replacement boundary", () => { + const wide = Array.from({ length: 200_000 }, (_, index) => index % 10); + const wideObject = Object.fromEntries( + Array.from({ length: 20_000 }, (_, index) => [ + `field_${String(index)}`, + index, + ]), + ); + const document = documentWith((value) => { + const fields = value.rea_state.entities[0]!.fields as Record< + string, + unknown + >; + value.rea_state.entities[0]!.fields = { + ...fields, + caller_wide: wide, + caller_wide_object: wideObject, + }; + }); + const store = new EnvironmentStateStore(state); + + expect(store.replace(document).ok).toBe(true); + const retained = store.snapshot().rea_state.entities[0]?.fields.caller_wide; + if (!Array.isArray(retained)) throw new Error("Expected the wide value"); + expect(retained).toHaveLength(wide.length); + expect(retained.at(-1)).toBe(wide.at(-1)); + const retainedObject = + store.snapshot().rea_state.entities[0]?.fields.caller_wide_object; + expect(retainedObject).toEqual(wideObject); + }); +}); diff --git a/src/application/environment/environmentStateDocument.ts b/src/application/environment/environmentStateDocument.ts index d4d21603a3..a0c11e93a6 100644 --- a/src/application/environment/environmentStateDocument.ts +++ b/src/application/environment/environmentStateDocument.ts @@ -1,3 +1,4 @@ +import { types } from "node:util"; import { z } from "zod"; import { @@ -16,7 +17,11 @@ import { } from "../../domain/analysisErrorCore.js"; import { canonicalJson } from "../../domain/comparisonSemantics.js"; import { projectInputIssues } from "../../domain/inputIssueProjection.js"; -import { jsonValueSchema } from "../../domain/jsonValue.js"; +import { + jsonValueSchema, + ownedJsonValueSchema, + type JsonValue, +} from "../../domain/jsonValue.js"; import { err, ok, type Result } from "../../domain/result.js"; import { compositeKey } from "../../domain/unicodeCodePointOrder.js"; import type { EnvironmentState } from "./environmentState.js"; @@ -37,37 +42,42 @@ export interface StateContent { readonly documents: readonly HeldDocument[]; } -const fieldsSchema = z.record(z.string(), jsonValueSchema); -const ordinalSchema = z.number().int().nonnegative(); +const stateDocumentSchemaFor = (jsonValue: z.ZodType) => { + const fieldsSchema = z.record(z.string(), jsonValue); + const ordinalSchema = z.number().int().nonnegative(); -const stateDocumentSchema = z.strictObject({ - rea_state: z.strictObject({ - version: z.literal(1), - // `data/get` reports these; accepting them back lets its output be loaded unchanged. - state_evidence_id: z.string().optional(), - types: z.array(jsonValueSchema).optional(), - entities: z.array( - z.strictObject({ - type: z.string(), - collection: z.string(), - id: z.string(), - fields: fieldsSchema, - view_fields: z - .array(z.strictObject({ tool: z.string(), fields: fieldsSchema })) - .optional(), - record_ordinals: z.array(ordinalSchema).optional(), - }), - ), - documents: z.array( - z.strictObject({ - tool: z.string(), - path: z.string(), - value: jsonValueSchema, - record_ordinal: ordinalSchema.optional(), - }), - ), - }), -}); + return z.strictObject({ + rea_state: z.strictObject({ + version: z.literal(1), + // `data/get` reports these; accepting them back lets its output be loaded unchanged. + state_evidence_id: z.string().optional(), + types: z.array(jsonValue).optional(), + entities: z.array( + z.strictObject({ + type: z.string(), + collection: z.string(), + id: z.string(), + fields: fieldsSchema, + view_fields: z + .array(z.strictObject({ tool: z.string(), fields: fieldsSchema })) + .optional(), + record_ordinals: z.array(ordinalSchema).optional(), + }), + ), + documents: z.array( + z.strictObject({ + tool: z.string(), + path: z.string(), + value: jsonValue, + record_ordinal: ordinalSchema.optional(), + }), + ), + }), + }); +}; + +const stateDocumentSchema = stateDocumentSchemaFor(jsonValueSchema); +const ownedStateDocumentSchema = stateDocumentSchemaFor(ownedJsonValueSchema); type ParsedDocument = z.output["rea_state"]; type ParsedEntity = ParsedDocument["entities"][number]; @@ -105,31 +115,82 @@ const memberProblem = (path: Path): AnalysisInputIssue => { /** The path of the first `__proto__` member in document order; schema parsing would drop such a member rather than refuse it. */ const prototypeMemberPath = (root: unknown): Path | undefined => { - interface Node { + interface Cursor { readonly value: unknown; - readonly parent: Node | undefined; + readonly parent: Cursor | undefined; readonly key: string | number; + readonly entries?: Generator; + index: number; } - const pathOf = (node: Node): (string | number)[] => { + const pathOf = (node: Cursor): (string | number)[] => { const path: (string | number)[] = []; - for (let at: Node = node; at.parent !== undefined; at = at.parent) + for (let at: Cursor = node; at.parent !== undefined; at = at.parent) path.push(at.key); return path.reverse(); }; - const pending: Node[] = [{ value: root, parent: undefined, key: "" }]; - for (let node = pending.pop(); node !== undefined; node = pending.pop()) { + function* ownKeys(value: Record): Generator { + // `for...in` is lazy. The ownership check keeps inherited enumerable + // properties out, matching Object.entries without allocating all entries. + for (const key in value) if (Object.hasOwn(value, key)) yield key; + } + const cursor = ( + value: unknown, + parent: Cursor | undefined, + key: string | number, + ): Cursor => ({ + value, + parent, + key, + index: 0, + ...(!Array.isArray(value) && typeof value === "object" && value !== null + ? { entries: ownKeys(value as Record) } + : {}), + }); + const pending: Cursor[] = [cursor(root, undefined, "")]; + while (pending.length > 0) { + const node = pending[pending.length - 1]!; const { value } = node; - if (typeof value !== "object" || value === null) continue; + if (typeof value !== "object" || value === null) { + pending.pop(); + continue; + } + // The canonical JSON validator reports proxies without invoking their traps. + if (types.isProxy(value)) { + pending.pop(); + continue; + } const array = Array.isArray(value); if (!array && Object.hasOwn(value, "__proto__")) return [...pathOf(node), "__proto__"]; - const entries = Object.entries(value); - for (let index = entries.length - 1; index >= 0; index -= 1) - pending.push({ - value: entries[index]![1], - parent: node, - key: array ? index : entries[index]![0], - }); + if (array) { + if (node.index >= value.length) { + pending.pop(); + continue; + } + const index = node.index++; + const descriptor = Object.getOwnPropertyDescriptor(value, index); + if ( + descriptor?.enumerable !== true || + !Object.hasOwn(descriptor, "value") + ) + continue; + const child: unknown = descriptor.value; + if (typeof child === "object" && child !== null) + pending.push(cursor(child, node, index)); + continue; + } + const next = node.entries?.next(); + if (next === undefined || next.done) { + pending.pop(); + continue; + } + const key = next.value; + const descriptor = Object.getOwnPropertyDescriptor(value, key); + if (descriptor?.enumerable !== true || !Object.hasOwn(descriptor, "value")) + continue; + const child: unknown = descriptor.value; + if (typeof child === "object" && child !== null) + pending.push(cursor(child, node, key)); } return undefined; }; @@ -365,9 +426,12 @@ const heldContent = ( export const parseStateDocument = ( state: EnvironmentState, value: unknown, + owned = false, ): Result => { const prototype = prototypeMemberPath(value); - const parsed = stateDocumentSchema.safeParse(value); + const parsed = ( + owned ? ownedStateDocumentSchema : stateDocumentSchema + ).safeParse(value); const initial = new Map( state.entities.map((entity) => [entityKey(entity), entity]), ); diff --git a/src/application/environment/environmentStateStore.ts b/src/application/environment/environmentStateStore.ts index 66d13eaa70..c3b34ac28b 100644 --- a/src/application/environment/environmentStateStore.ts +++ b/src/application/environment/environmentStateStore.ts @@ -58,9 +58,21 @@ export class EnvironmentStateStore { this.#install(this.state); } - /** Validate a `rea-state/v1` document and, only when it is valid, make it the whole state. */ + /** Validate a caller-owned document and copy its JSON values before retaining them. */ replace(document: unknown): Result { - const content = parseStateDocument(this.state, document); + return this.#replace(document, false); + } + + /** Validate and retain a JSON document created exclusively by the serving boundary. */ + replaceOwnedDocument(document: unknown): Result { + return this.#replace(document, true); + } + + #replace( + document: unknown, + owned: boolean, + ): Result { + const content = parseStateDocument(this.state, document, owned); if (!content.ok) return content; // data/get writes members a document may omit, so the state it reports must still be one data/add loads. if (jsonBytes(this.#snapshotOf(content.value)) > STATE_INTAKE_LIMIT_BYTES) diff --git a/src/domain/jsonValue.test.ts b/src/domain/jsonValue.test.ts index faca71f535..727ff85cfe 100644 --- a/src/domain/jsonValue.test.ts +++ b/src/domain/jsonValue.test.ts @@ -64,7 +64,9 @@ describe("jsonValueSchema depth bound", () => { const result = jsonValueSchema.safeParse(deepObject(MAX_JSON_DEPTH + 1)); expect(result.success).toBe(false); if (result.success) return; - expect(result.error.issues[0]?.message).toContain("maximum nesting depth"); + expect(result.error.issues[0]?.message).toBe( + `JSON value exceeds maximum nesting depth of ${MAX_JSON_DEPTH}. This release has no runtime override.`, + ); }); it.each(["object", "array"])( diff --git a/src/domain/jsonValue.ts b/src/domain/jsonValue.ts index 250ce8ea47..ed464cea01 100644 --- a/src/domain/jsonValue.ts +++ b/src/domain/jsonValue.ts @@ -27,6 +27,10 @@ export const jsonValueValidationIssue = (root: unknown): string | undefined => { } }; +function* ownEnumerableKeys(value: Record): Generator { + for (const key in value) if (Object.hasOwn(value, key)) yield key; +} + /** Validate JSON in bounded steps; cache only authenticated immutable roots. */ export function* jsonValueValidationSteps( root: unknown, @@ -41,57 +45,104 @@ export function* jsonValueValidationSteps( if (immutable !== undefined) immutableValidation.set(immutable, issue); return issue; }; - const pending: { value: unknown; depth: number }[] = [ - { value: root, depth: 0 }, - ]; + type Frame = + | { + readonly kind: "array"; + readonly value: unknown[]; + readonly depth: number; + index: number; + } + | { + readonly kind: "object"; + readonly value: Record; + readonly depth: number; + readonly keys: Generator; + }; + const pending: Frame[] = []; + let value = root; + let depth = 0; + let inspect = true; let examined = 0; - while (pending.length > 0) { - if (examined++ === 4096) { - examined = 0; - yield; - } - const item = pending.pop(); - if (item === undefined) break; - if (item.depth > MAX_JSON_DEPTH) - return finish( - `JSON value exceeds maximum nesting depth of ${MAX_JSON_DEPTH}. This release has no runtime override.`, - ); - const value = item.value; - if (typeof value === "object" && value !== null && types.isProxy(value)) - return finish("JSON values cannot contain proxies"); - if ( - value === null || - typeof value === "string" || - typeof value === "boolean" - ) + for (;;) { + if (inspect) { + if (examined++ === 4096) { + examined = 0; + yield; + } + if (depth > MAX_JSON_DEPTH) + return finish( + `JSON value exceeds maximum nesting depth of ${MAX_JSON_DEPTH}. This release has no runtime override.`, + ); + if (typeof value === "object" && value !== null && types.isProxy(value)) + return finish("JSON values cannot contain proxies"); + if ( + value === null || + typeof value === "string" || + typeof value === "boolean" + ) { + inspect = false; + } else if (typeof value === "number") { + if (!Number.isFinite(value)) + return finish("JSON numbers must be finite"); + inspect = false; + } else if (Array.isArray(value)) { + pending.push({ + kind: "array", + value, + depth, + index: value.length - 1, + }); + inspect = false; + } else if (isRecord(value)) { + const prototype: unknown = Object.getPrototypeOf(value); + if (prototype !== Object.prototype && prototype !== null) + return finish("JSON objects require ordinary data properties"); + for (const key of Object.getOwnPropertySymbols(value)) + if (Object.getOwnPropertyDescriptor(value, key)?.enumerable === true) + return finish("JSON object keys must be strings"); + pending.push({ + kind: "object", + value, + depth, + keys: ownEnumerableKeys(value), + }); + inspect = false; + } else return finish("Invalid JSON value"); continue; - if (typeof value === "number") { - if (!Number.isFinite(value)) return finish("JSON numbers must be finite"); - } else if (Array.isArray(value)) { - for (let index = 0; index < value.length; index += 1) { - const descriptor = Object.getOwnPropertyDescriptor(value, index); - if (descriptor === undefined || !Object.hasOwn(descriptor, "value")) - return finish( - "JSON arrays require dense data properties without accessors", - ); - pending.push({ value: descriptor.value, depth: item.depth + 1 }); + } + + const frame = pending[pending.length - 1]; + if (frame === undefined) return finish(undefined); + if (frame.kind === "array") { + if (frame.index < 0) { + pending.pop(); + continue; } - } else if (isRecord(value)) { - const prototype: unknown = Object.getPrototypeOf(value); - if (prototype !== Object.prototype && prototype !== null) - return finish("JSON objects require ordinary data properties"); - for (const key of Reflect.ownKeys(value)) { - const descriptor = Object.getOwnPropertyDescriptor(value, key); - if (descriptor?.enumerable !== true) continue; - if (typeof key !== "string") - return finish("JSON object keys must be strings"); - if (!Object.hasOwn(descriptor, "value")) - return finish("JSON objects cannot contain accessors"); - pending.push({ value: descriptor.value, depth: item.depth + 1 }); - } - } else return finish("Invalid JSON value"); + const index = frame.index--; + const descriptor = Object.getOwnPropertyDescriptor(frame.value, index); + if (descriptor === undefined || !Object.hasOwn(descriptor, "value")) + return finish( + "JSON arrays require dense data properties without accessors", + ); + value = descriptor.value; + depth = frame.depth + 1; + inspect = true; + continue; + } + + const next = frame.keys.next(); + if (next.done) { + pending.pop(); + continue; + } + const descriptor = Object.getOwnPropertyDescriptor(frame.value, next.value); + if (descriptor?.enumerable !== true) continue; + if (!Object.hasOwn(descriptor, "value")) + return finish("JSON objects cannot contain accessors"); + value = descriptor.value; + depth = frame.depth + 1; + inspect = true; } - return finish(undefined); } const preserveJsonValue = (value: JsonValue): JsonValue => { @@ -103,13 +154,17 @@ const preserveJsonValue = (value: JsonValue): JsonValue => { copy.push(preserveJsonValue(value[index]!)); return copy; } - if (value !== null && typeof value === "object") - return Object.fromEntries( - Object.entries(value).map(([key, child]) => [ - key, - preserveJsonValue(child), - ]), - ); + if (value !== null && typeof value === "object") { + const copy: Record = {}; + for (const key of ownEnumerableKeys(value)) + Object.defineProperty(copy, key, { + value: preserveJsonValue(value[key]!), + enumerable: true, + configurable: true, + writable: true, + }); + return copy; + } return value; }; @@ -132,6 +187,13 @@ export const jsonValueSchema = z .transform(preserveJsonValue); jsonValueSchema._zod.toJSONSchema = () => ({}); +/** Validate JSON that the current boundary exclusively owns, without a second full-tree copy. */ +export const ownedJsonValueSchema = z + .unknown() + .nonoptional() + .pipe(validatedJsonValue); +ownedJsonValueSchema._zod.toJSONSchema = () => ({}); + const objectValidation = z.custom>( isRecord, "Expected JSON object", diff --git a/src/server/environment/environmentDataPlane.ts b/src/server/environment/environmentDataPlane.ts index b1363cb361..10623c61a8 100644 --- a/src/server/environment/environmentDataPlane.ts +++ b/src/server/environment/environmentDataPlane.ts @@ -60,7 +60,7 @@ const addState = async ( 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 }; - const replaced = store.replace(document.value); + const replaced = store.replaceOwnedDocument(document.value); return replaced.ok ? { result: {} } : { error: documentError(replaced.error) }; diff --git a/src/server/environment/serveEnvironment.ts b/src/server/environment/serveEnvironment.ts index 2415c8deb0..aa9a0220a0 100644 --- a/src/server/environment/serveEnvironment.ts +++ b/src/server/environment/serveEnvironment.ts @@ -193,6 +193,65 @@ const startRpc = ( ); }; +interface CachedStateReplySize { + readonly types: unknown; + readonly documents: unknown; + readonly bytes: number; +} + +// Repeated `data/get` snapshots have fresh wrappers but share the installed +// state's member arrays. The request-scoped cache avoids retaining stale sizes. +const cachedStateReplyResultBytes = ( + stateReplySizes: WeakMap, + reply: RpcReply, +): number | undefined => { + if (!("result" in reply) || !isRecord(reply.result)) return undefined; + const parts = reply.result.parts; + if (!Array.isArray(parts) || parts.length !== 1 || !isRecord(parts[0])) + return undefined; + const data = parts[0].data; + if (!isRecord(data) || !isRecord(data.rea_state)) return undefined; + const { types, entities, documents } = data.rea_state; + if (!Array.isArray(entities) || !Array.isArray(documents)) return undefined; + const cached = stateReplySizes.get(entities); + return cached !== undefined && + cached.types === types && + cached.documents === documents + ? cached.bytes + : undefined; +}; + +const rememberStateReplyResultBytes = ( + stateReplySizes: WeakMap, + reply: RpcReply, + replyBytes: number, +): void => { + if (!("result" in reply) || !isRecord(reply.result)) return; + const parts = reply.result.parts; + if (!Array.isArray(parts) || parts.length !== 1 || !isRecord(parts[0])) + return; + const data = parts[0].data; + if (!isRecord(data) || !isRecord(data.rea_state)) return; + const { types, entities, documents } = data.rea_state; + if (!Array.isArray(entities) || !Array.isArray(documents)) return; + // A one-character string marks the result slot. Removing its compact JSON + // envelope yields the exact result size, independent of the RPC id. + const markerBytes = Buffer.byteLength( + JSON.stringify({ jsonrpc: "2.0", id: reply.id, result: "x" }), + "utf8", + ); + const bytes = replyBytes - markerBytes + 3; + stateReplySizes.set(entities, { types, documents, bytes }); +}; + +const replyBytesFor = (reply: RpcReply, resultBytes: number): number => + resultBytes + + Buffer.byteLength( + JSON.stringify({ jsonrpc: "2.0", id: reply.id, result: "x" }), + "utf8", + ) - + 3; + const overBudget = (replyBytes: number): RpcError => ({ code: SERVER_ERROR, message: `the reply would take this response past its ${String(MAX_RPC_RESPONSE_BYTES)}-byte budget; the call ran, so repeat a read in another request`, @@ -224,11 +283,30 @@ const replyEncoder = ( const refusals = messages.filter(expectsReply).map(refusalBytes); let reserved = refusals.reduce((sum, bytes) => sum + bytes, 0); let used = framed ? 2 + Math.max(0, refusals.length - 1) : 0; + const stateReplySizes = new WeakMap(); let index = 0; return (reply) => { reserved -= refusals[index] ?? 0; index += 1; const room = MAX_RPC_RESPONSE_BYTES - used - reserved; + const cachedResultBytes = cachedStateReplyResultBytes( + stateReplySizes, + reply, + ); + const knownBytes = + cachedResultBytes === undefined + ? undefined + : replyBytesFor(reply, cachedResultBytes); + if (knownBytes !== undefined && knownBytes > room) { + const refused = JSON.stringify({ + jsonrpc: "2.0", + id: reply.id, + error: overBudget(knownBytes), + }); + const refusedBytes = Buffer.byteLength(refused, "utf8"); + used += refusedBytes; + return refused; + } const parts: string[] = []; let bytes = 0; for (const part of bufferedJsonParts(jsonParts(reply))) { @@ -236,9 +314,13 @@ const replyEncoder = ( if (bytes <= room) parts.push(part); } if (bytes <= room) { + if (cachedResultBytes === undefined) + rememberStateReplyResultBytes(stateReplySizes, reply, bytes); used += bytes; return parts.join(""); } + if (cachedResultBytes === undefined) + rememberStateReplyResultBytes(stateReplySizes, reply, bytes); const refused = JSON.stringify({ jsonrpc: "2.0", id: reply.id, @@ -418,8 +500,23 @@ export const serveEnvironment = async ( }, close: async () => { stop(); - await served.close(); - await stateFiles.close(); + const failures: unknown[] = []; + try { + await served.close(); + } catch (cause: unknown) { + failures.push(cause); + } + try { + await stateFiles.close(); + } catch (cause: unknown) { + failures.push(cause); + } + if (failures.length === 1) throw failures[0]; + if (failures.length > 1) + throw new AggregateError( + failures, + "Environment server and state-file cleanup both failed", + ); }, }); }; diff --git a/src/server/environment/serveEnvironmentState.test.ts b/src/server/environment/serveEnvironmentState.test.ts index ddc9773fc5..b6693abc81 100644 --- a/src/server/environment/serveEnvironmentState.test.ts +++ b/src/server/environment/serveEnvironmentState.test.ts @@ -7,6 +7,32 @@ import { join } from "node:path"; import { pathToFileURL } from "node:url"; import { afterEach, describe, expect, it, onTestFinished, vi } from "vitest"; +const closeFailureControl = vi.hoisted(() => ({ + failure: null as Error | null, +})); + +vi.mock("../http/serveStreamableHttp.js", async (importOriginal) => { + const original = + await importOriginal(); + return { + ...original, + serveStreamableHttp: async ( + ...args: Parameters + ) => { + const served = await original.serveStreamableHttp(...args); + return { + ...served, + close: async () => { + await served.close(); + const failure = closeFailureControl.failure; + closeFailureControl.failure = null; + if (failure !== null) throw failure; + }, + }; + }, + }; +}); + import { timedCaptureEvidence } from "../../../tests/fixtures/analysisView.js"; import { environmentStateServeFixture, @@ -672,6 +698,7 @@ describe("state data plane response budget", () => { typeof value === "object" && value !== null && "result" in value && + typeof (value as { result?: unknown }).result === "object" && (value as { id?: unknown }).id === 3, ), ).toBe(false); @@ -681,6 +708,56 @@ describe("state data plane response budget", () => { }); }); +describe("repeated state reads", () => { + it("reuses the encoded size of a state for refused repeated data/get calls", async () => { + const plane = dataPlane(); + const document = renamedDocument(); + const large = { + rea_state: { + ...document.rea_state, + entities: document.rea_state.entities.map((entity) => + entity.type === "/api/users" && entity.id === "2" + ? { + ...entity, + fields: { ...entity.fields, large: "x".repeat(1024 * 1024) }, + } + : entity, + ), + }, + }; + await plane.add({ kind: "data", data: large }); + const fields = plane.store + .snapshot() + .rea_state.entities.find((entity) => entity.id === "2")?.fields; + if (fields === undefined) throw new Error("Expected the loaded entity"); + + const originalKeys = Object.keys; + let fieldEnumerations = 0; + const keys = vi.spyOn(Object, "keys").mockImplementation((value) => { + if (value === fields) fieldEnumerations += 1; + return originalKeys(value); + }); + try { + const batch = await plane.post( + Array.from({ length: 100 }, (_, index) => ({ + jsonrpc: "2.0", + id: index + 1, + method: "data/get", + })), + ); + expect( + (batch.json as { error?: unknown }[]).filter((reply) => reply.error), + ).not.toHaveLength(0); + expect(fieldEnumerations).toBeLessThan(65); + expect(Buffer.byteLength(JSON.stringify(batch.json))).toBeLessThanOrEqual( + MAX_RPC_RESPONSE_BYTES, + ); + } finally { + keys.mockRestore(); + } + }); +}); + describe("state data plane response budget reserve", () => { it("keeps room for every later refusal and the batch's framing, so the body never passes the budget", async () => { const document = renamedDocument(); @@ -891,7 +968,47 @@ describe("serveEnvironment with a state", () => { await expect(environment.close()).resolves.toBeUndefined(); expect(owners[0]?.handle.fd).toBe(-1); }); +}); +describe("serveEnvironment cleanup retries", () => { + it("attempts state-file cleanup after listener failure and keeps both failures retryable", async () => { + const environment = await serveState({ source, state }); + const uri = await renamedStateUri(); + let failing = true; + const owners = failFirstFileClose(() => failing); + const response = await fetch( + new URL(ENVIRONMENT_RPC_PATH, environment.url), + { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(addByUri(1, uri)), + }, + ); + await expect(response.json()).resolves.toMatchObject({ + error: { + code: -32000, + message: expect.stringContaining("could not be closed"), + }, + }); + + const listenerFailure = new Error("injected listener cleanup failure"); + closeFailureControl.failure = listenerFailure; + await expect(environment.close()).rejects.toSatisfy((cause: unknown) => { + if (!(cause instanceof AggregateError)) return false; + return ( + cause.errors[0] === listenerFailure && + cause.errors[1] instanceof ArtifactResourceRetryFailure + ); + }); + expect(owners[0]?.handle.fd).toBeGreaterThanOrEqual(0); + + failing = false; + await expect(environment.close()).resolves.toBeUndefined(); + expect(owners[0]?.handle.fd).toBe(-1); + }); +}); + +describe("serveEnvironment with a state", () => { it("serves query POSTs from their recordings and leaves state-changing writes out", async () => { const environment = await serveState( environmentStateServeInputs(