From 76dbc32f3990164f5b8e30dc3bf4870480eec50e Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Thu, 22 Oct 2026 06:06:36 +0800 Subject: [PATCH] fix(environment): bound state reads and retain cleanup ownership (#6158) * feat(environment): extract the initial state a capture's reads observed Add the environment-state analysis view. From the parents the environment manifest accepts, it extracts the data the capture's reads observed before the first related write: entity types with an identity field (proved by a path parameter or inferred from an identity name), entities with their recorded values and provenance, documents for written non-entity paths, and a render plan per read tool that is kept only when it reproduces the tool's latest initial recordings from the state; every other read tool is a recorded view with the reason. The Evidence is always inferred. Response-body decoding moves to a domain owner shared with replay, exact JSON parsing moves to the owner of the exact-number policy, the manifest and state views share parent admission and record grouping, and projectAnalysisView takes its accumulators as one object. Co-Authored-By: Claude Opus 5.5 * fix(environment): bound state at redacted writes; infer identity only from items A scenario write whose URL was redacted has no request line, so it was dropped and reads after it counted as initial state. Such a successful write now joins the timeline with an empty path, whose root relates to every type, so it ends every type's state; a limitation counts them. A list whose recordings were all empty passed the "every listed item carries it" check vacuously and published an identity field nothing showed; it now needs at least one item, else the view stays recorded as no-identity. A value that differed between two request lines of one item tool, such as a query variant, was called volatile, and the gate then substituted each recording's own value; only repeats of one request line now show volatility, so differing lines go through fields-differ. List membership is compared only while a list can still be the whole collection, and merged entities are indexed by type and collection once, so paged captures no longer cost pages times entities (8,000 pages of 40,000 entities: 5.7 s to 0.23 s). Docs and the tool contract say scenario state keeps the capture's redaction markers and how redacted writes bound it. Co-Authored-By: Claude Opus 5.5 * fix(environment): keep first sightings, one list pointer and linear field unions Entities took the order of the first plan observed, and item plans run before list plans, so a complete list read before its detail reads came out reordered and lost "all" membership; each entity now keeps its earliest recorded sighting. A list plan reused the first recording's pointer for later recordings, so a later response that was itself an entity wrapping a list was read as that list; every recording must now locate the same list. The field union scanned its ordered list for every key, quadratic in an object's width; a set now answers membership. Co-Authored-By: Claude Opus 5.5 * fix(environment): keep post-write values and partial lists out of the initial state A value that changed between repeats of one request line was named volatile whenever no write under the type's own root lay between, so a write under another root (for example POST /api/dashboards/db before a second GET /api/search) left its effect in the initial state. Any recorded write between the two observations now leaves every tool of the type recorded with reason fields-differ instead. A list was published with membership "all" even when a total or count envelope member exceeded the number it listed. Such a counter now makes the membership "recorded". Also drop a stale doc comment on ObservedEntity.order. Co-Authored-By: Claude Opus 5.5 * docs(environment): scope item-shaped and deleted written paths out of documents State documents are held only for list-shaped paths a POST, PUT or PATCH tool writes. Reads of item-shaped written paths and of paths that are only deleted stay recorded views; the guide now says so and the spec defers them to a later increment of #4933. Co-Authored-By: Claude Opus 5.5 * fix(environment): count writes in flight when a value changes between reads The volatility check compared write start positions with the two reads, so a write already in flight when the earlier read started was missed and its effect was named volatile. A write now counts when its interval overlaps the span from the earlier read's start to the later read's end. The check also ran once per changed field and scanned every write. It now runs once per pair of observations, answered by a binary search over the writes in timeline order and the latest end among each prefix. Co-Authored-By: Claude Opus 5.5 * fix(environment): catch cross-tool changes across writes; read only item totals A value that changed across a write was caught only between repeats of one tool's request line, so a list read before a write and an item read after it put the post-write value into the initial state. Changes are now classified per entity in timeline order after observation: across any recorded write, a change between repeats of a line or between two lines of the type rewrites the type (every tool recorded as fields-differ); with no write between, a repeat's change is volatile. The classification lives in environmentStateMerge.ts. Any numeric member whose name contained "total" or "count" counted as an item total, so total_pages, account_id or page_count made complete lists partial. Item totals are now members whose name tokens include total or count and none of page, pages or per; a string of digits counts as a number; and a set has_more, next or next-page token member also marks the list partial. Co-Authored-By: Claude Opus 5.5 * fix(environment): keep post-write observations out of initial state * fix(environment): retain field provenance and bound state inference work * feat(environment): serve tool calls from the environment state `environment-serve --state ` serves an environment from the initial state an environment-state view extracted. Read tools whose views are state-backed answer from a live state store: item views render the held entity into the recorded envelope, list views render the collection or the recorded members, and document views return the held value; a missing entity or one lacking a required field is refused, and requests the state cannot answer fall back to their recordings with the reason. The /agentenv data plane implements data/reset, data/add (a rea-state/v1 document as a data part, base64 bytes, a file:// path or an https URL, validated and replaced atomically) and data/get, and the card advertises those operations and an intake declaration. Stage 1 is read-only: write tools are not served in state mode. Without --state nothing changes. Co-Authored-By: Claude Opus 5.5 * docs(environment): record verification after merging the state view's review fixes Co-Authored-By: Claude Opus 5.5 * feat(environment)!: serve the environment state by default environment-serve derives the environment state at startup by running the environment-state view in-process on the capture with the manifest's template rules, so an environment comes with its state and a user supplies nothing extra. The derived Evidence is the same one inspect-analysis-view produces by hand. --state is removed; --replay-only serves exactly the stateless replay of main. When the state cannot be derived, or its catalog does not fit the discovery budget, the server starts in replay mode, logs the reason and reports it in the stop summary's `state`. Exported bundles serve the state by default too: export derives it the same way, so the verifier's tool count and checked answer match what is served, and the README describes the served mode. BREAKING CHANGE: environment-serve answers read tools from the derived state and does not serve write tools unless --replay-only is given; the stop summary reports `state` instead of `state_evidence_id`. Co-Authored-By: Claude Opus 5.5 * fix(environment): bound data-plane replies, keep batch order, cancel intake on stop - One /agentenv response is held to 64 MiB: replies are measured as they are produced and one past the budget becomes a -32000 error stating it. - Every message of a batch enters the arrival-order queue on arrival, so another request's reset never lands between them. - Stopping the server, or a caller leaving, cancels the active download and answers queued calls `cancelled` without running them. - file:// URIs name only this host: another host or a network share is refused before anything is opened. - A loaded view value may not give an entity another identity. - List answers refuse a member lacking a field every recording carried, as item answers do. Co-Authored-By: Claude Opus 5.5 * fix(environment): keep bodies too deep for the state's nesting out of the state A recorded body at the JSON depth limit was accepted by extraction, then nested a few levels deeper as a document or entity value, so sealing the state Evidence threw and environment-serve and environment-export stopped instead of falling back. stateBodyJson, the one rule extraction and serving share, now treats a body nesting deeper than the limit minus 16 levels as unusable, so its tool stays a recorded view. Co-Authored-By: Claude Opus 5.5 * fix(environment): measure data-plane replies before building them 83fcc739 bounded one /agentenv response to 64 MiB, but it still serialized each reply whole before measuring it, so a batch of data/get calls on a large state built every refused reply in full only to drop it. Each reply is now measured in bounded JSON parts (the shared jsonParts encoder), its parts kept only while it fits; a reply past the budget is counted for its exact reply_bytes but never assembled. A test checks that a refused read is never handed to JSON.stringify whole and that its reply_bytes is exact. Co-Authored-By: Claude Opus 5.5 * fix(environment): release batch replies as they are encoded; judge exact numbers by value answerBatch kept every settled reply in its pending array until the batch ended, so a batch alternating data/add and data/get held each loaded state alive; each reply now leaves the queue as it is encoded. The exact-number policy judged a number by its recorded spelling, so a state loaded with 1e20 came back from data/get as 100000000000000000000, which data/add refused, breaking the promise that data/get output loads back unchanged. A number is exact only when both its recorded spelling and the one serialization writes are safe, so integers from 2^53 below 1e21 are refused in any spelling; replay keeps such bodies as recorded text. A domain test checks that every accepted number is accepted again once written back. Co-Authored-By: Claude Opus 5.5 * fix(environment): keep room for refusals and framing in the reply budget The response budget counted accepted replies only: a budget error that replaced a reply, and a batch's brackets and commas, were not charged, so near 64 MiB later refusals pushed the body past the advertised limit. Every reply still to come now keeps room for its own budget error, and a batch keeps room for its framing, so a refusal always fits and the body never passes the budget. Reply ids and which messages are answered come from one place shared by the replies and the reserve. Co-Authored-By: Claude Opus 5.5 * fix(environment): refuse a loaded state that data/get would write past the intake limit data/add measured only the document it was given, while data/get writes the state back with the members a document may omit (types, the state Evidence id, view fields and record citations). A document just under 32 MiB could therefore load and be reported larger than any intake accepts, breaking the data/get then data/add round trip. replace now measures the state as data/get writes it, part by part, and refuses it past the limit, leaving the state unchanged. The limit moves to the state document module, so the application store can read it; the intake re-exports it. Co-Authored-By: Claude Opus 5.5 * fix(environment): keep merged export checks on the shared test graph Main's export tests installed module mocks for the whole services project, which reuses one module graph. A tarball read failure then missed its importer, and a one-shot heap mock could be taken by another file. Read those through bindings the tests replace for one case. Co-authored-by: morluto * fix(environment): keep a state file whose close failed owned until it closes A file:// data/add that could not close its descriptor dropped the handle, and an admission failure whose cleanup close failed dropped its owner, so nothing could retry the close or report it. The read now runs through withOwnedRegularFile, and a failed close is retained in a server-lifetime ArtifactResourceScope. The next file read retries it first and answers -32000 while it is still open; inline and https documents still load. Closing the served environment retries it again and rejects with the unfinished close. Co-Authored-By: Claude Opus 5.5 * fix(environment): attempt all state cleanup on close * fix(environment): bound wide state JSON traversal * fix(environment): scope cached state reply sizes to requests * fix(domain): preserve JSON depth diagnostic guidance --------- Co-authored-by: Edgar Arakelyan Co-authored-by: Claude Opus 5.5 Co-authored-by: Kaoru0822 <208281322+kaoru0822-kitauji@users.noreply.github.com> Co-authored-by: gensokyo-dev Co-authored-by: Cursor Agent Co-authored-by: morluto --- .../environment/environmentState.test.ts | 57 ++++++ .../environment/environmentStateDocument.ts | 156 +++++++++++----- .../environment/environmentStateStore.ts | 16 +- src/domain/jsonValue.test.ts | 4 +- src/domain/jsonValue.ts | 170 ++++++++++++------ .../environment/environmentDataPlane.ts | 2 +- src/server/environment/serveEnvironment.ts | 101 ++++++++++- .../environment/serveEnvironmentState.test.ts | 117 ++++++++++++ 8 files changed, 517 insertions(+), 106 deletions(-) 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(