From ae563d4a57e21d9c12d411e2a5aa484c1d937598 Mon Sep 17 00:00:00 2001 From: ScriptedAlchemy Date: Tue, 1 Sep 2026 18:53:37 +0000 Subject: [PATCH] =?UTF-8?q?fix(runtime):=20review=20follow-ups=20=E2=80=94?= =?UTF-8?q?=20traversal=20bounds,=20demand-bounded=20progress,=20failure?= =?UTF-8?q?=20conversion?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enforce document limits during JSON and Flight decode walks, bound live progress by consumer demand, close the progress queue on setup failure, and convert synchronous host throws into stream failures. --- .changeset/runtime-review-followups.md | 7 + packages/rsc-runtime/src/agent-document.ts | 87 ++++-- packages/rsc-runtime/src/decode-document.ts | 31 +- packages/rsc-runtime/src/dispatcher.ts | 18 +- packages/rsc-runtime/src/lower-mcp.ts | 58 +++- packages/rsc-runtime/src/reconciler.ts | 266 ++++++++++++------ .../rsc-runtime/tests/agent-document.test.ts | 52 ++++ packages/rsc-runtime/tests/dispatcher.test.ts | 145 ++++++++++ 8 files changed, 550 insertions(+), 114 deletions(-) create mode 100644 .changeset/runtime-review-followups.md diff --git a/.changeset/runtime-review-followups.md b/.changeset/runtime-review-followups.md new file mode 100644 index 000000000..c2cc8715c --- /dev/null +++ b/.changeset/runtime-review-followups.md @@ -0,0 +1,7 @@ +--- +"@agent-bundle/runtime": patch +--- + +Enforce Agent Document bounds during JSON and Flight decode walks, bound live +progress by downstream demand, close the progress queue on setup failure, and +convert synchronous host throws into stream failures. diff --git a/packages/rsc-runtime/src/agent-document.ts b/packages/rsc-runtime/src/agent-document.ts index 7b778518d..3284e3b07 100644 --- a/packages/rsc-runtime/src/agent-document.ts +++ b/packages/rsc-runtime/src/agent-document.ts @@ -1,6 +1,6 @@ import { Buffer } from 'node:buffer'; -import { snapshotJsonValue, type JsonValue } from './lower-mcp.js'; +import { snapshotJsonValue, type JsonSnapshotBudget, type JsonValue } from './lower-mcp.js'; export const AGENT_DOCUMENT_VERSION = 1 as const; @@ -163,7 +163,7 @@ export class AgentContractError extends Error { } } -const resolveLimits = (overrides: Partial): AgentRenderLimits => { +export const resolveAgentRenderLimits = (overrides: Partial = {}): AgentRenderLimits => { const limits = { ...DEFAULT_AGENT_RENDER_LIMITS, ...overrides }; for (const [name, value] of Object.entries(limits)) { if (!Number.isSafeInteger(value) || value <= 0) { @@ -190,10 +190,51 @@ const text = (value: unknown, field: string): string => { const optionalString = (value: unknown, field: string): string | undefined => value === undefined ? undefined : requiredString(value, field); -const snapshotJson = (value: unknown, message: string): JsonValue => { +export const elapsedTimeExceeded = (maxElapsedMs: number): AgentContractError => + new AgentContractError( + 'elapsed-time-exceeded', + `Agent render elapsed time exceeds ${String(maxElapsedMs)}ms`, + ); + +const jsonBudget = (state: NodeSnapshotState): JsonSnapshotBudget => ({ + addBytes(n) { + state.bytes += n; + if (state.bytes > state.limits.maxDocumentBytes) { + throw new AgentContractError( + 'document-bytes-exceeded', + `Agent Document bytes exceed ${String(state.limits.maxDocumentBytes)}`, + ); + } + }, + addNode() { + state.nodes += 1; + if (state.nodes > state.limits.maxDocumentNodes) { + throw new AgentContractError( + 'document-node-count-exceeded', + `Agent Document node count exceeds ${String(state.limits.maxDocumentNodes)}`, + ); + } + }, + checkDepth(depth) { + if (depth > state.limits.maxDocumentDepth) { + throw new AgentContractError( + 'document-depth-exceeded', + `Agent Document depth exceeds ${String(state.limits.maxDocumentDepth)}`, + ); + } + }, +}); + +const snapshotJson = ( + value: unknown, + message: string, + depth: number, + state: NodeSnapshotState, +): JsonValue => { try { - return snapshotJsonValue(value, message); + return snapshotJsonValue(value, message, { depth, limits: jsonBudget(state) }); } catch (error) { + if (error instanceof AgentContractError) throw error; throw new AgentContractError('invalid-document', error instanceof Error ? error.message : message, { cause: error }); } }; @@ -208,6 +249,7 @@ const progressNumber = (value: unknown, field: string): number => { interface NodeSnapshotState { readonly ancestors: Set; readonly limits: AgentRenderLimits; + bytes: number; nodes: number; } @@ -242,7 +284,7 @@ const snapshotNode = (node: AgentDocumentNode, depth: number, state: NodeSnapsho const children = Object.freeze(node.children.map((child) => snapshotNode(child, depth + 1, state))); const metadata = node.metadata === undefined ? undefined - : snapshotJson(node.metadata, 'Agent result metadata must be JSON-serializable'); + : snapshotJson(node.metadata, 'Agent result metadata must be JSON-serializable', depth, state); return Object.freeze({ children, kind: 'result', @@ -258,7 +300,7 @@ const snapshotNode = (node: AgentDocumentNode, depth: number, state: NodeSnapsho case 'json': return Object.freeze({ kind: 'json', - value: snapshotJson(node.value, 'Agent JSON node value must be JSON-serializable'), + value: snapshotJson(node.value, 'Agent JSON node value must be JSON-serializable', depth, state), }); case 'progress': { const completed = progressNumber(node.completed, 'Agent progress completed'); @@ -337,11 +379,12 @@ export const createAgentDocument = ( `Unsupported Agent Document version: ${String(input.version)}`, ); } - const limits = resolveLimits(limitOverrides); - const root = snapshotNode(input.root, 1, { ancestors: new Set(), limits, nodes: 0 }); + const limits = resolveAgentRenderLimits(limitOverrides); + const state: NodeSnapshotState = { ancestors: new Set(), bytes: 0, limits, nodes: 0 }; + const root = snapshotNode(input.root, 1, state); const value = input.value === undefined ? undefined - : snapshotJson(input.value, 'Agent Document value must be JSON-serializable'); + : snapshotJson(input.value, 'Agent Document value must be JSON-serializable', 1, state); const document: AgentDocument = Object.freeze({ root, status: documentStatus(input.status), @@ -358,10 +401,15 @@ export const createAgentDocument = ( return document; }; -const snapshotRenderError = (error: AgentRenderError): AgentRenderError => { +const snapshotRenderError = (error: AgentRenderError, limits: AgentRenderLimits): AgentRenderError => { const data = error.data === undefined ? undefined - : snapshotJson(error.data, 'Agent render error data must be JSON-serializable'); + : snapshotJson( + error.data, + 'Agent render error data must be JSON-serializable', + 1, + { ancestors: new Set(), bytes: 0, limits, nodes: 0 }, + ); return Object.freeze({ code: requiredString(error.code, 'Agent render error code'), ...(data === undefined ? {} : { data }), @@ -403,7 +451,7 @@ const snapshotEvent = ( const boundaryId = optionalString(input.boundaryId, 'Agent render boundaryId'); return Object.freeze({ ...(boundaryId === undefined ? {} : { boundaryId }), - error: snapshotRenderError(input.error), + error: snapshotRenderError(input.error, limits), sequence, type: 'error', }); @@ -422,14 +470,16 @@ const snapshotEvent = ( export interface AgentRenderEventSequence { readonly completed: boolean; + readonly maxElapsedMs: number; readonly nextSequence: number; + readonly remainingMs: number; readonly emit: (input: AgentRenderEventInput) => AgentRenderEvent; } export const createAgentRenderEventSequence = ( limitOverrides: Partial = {}, ): AgentRenderEventSequence => { - const limits = resolveLimits(limitOverrides); + const limits = resolveAgentRenderLimits(limitOverrides); const startedAt = Date.now(); const recentTimes: number[] = []; let completed = false; @@ -438,6 +488,9 @@ export const createAgentRenderEventSequence = ( get completed() { return completed; }, + get maxElapsedMs() { + return limits.maxElapsedMs; + }, emit(input: AgentRenderEventInput): AgentRenderEvent { if (completed) { throw new AgentContractError( @@ -447,10 +500,7 @@ export const createAgentRenderEventSequence = ( } const now = Date.now(); if (now - startedAt > limits.maxElapsedMs) { - throw new AgentContractError( - 'elapsed-time-exceeded', - `Agent render elapsed time exceeds ${String(limits.maxElapsedMs)}ms`, - ); + throw elapsedTimeExceeded(limits.maxElapsedMs); } recentTimes.push(now); const windowStart = now - 1000; @@ -484,5 +534,8 @@ export const createAgentRenderEventSequence = ( get nextSequence() { return nextSequence; }, + get remainingMs() { + return limits.maxElapsedMs - (Date.now() - startedAt); + }, }); }; diff --git a/packages/rsc-runtime/src/decode-document.ts b/packages/rsc-runtime/src/decode-document.ts index 363642f58..2deffded6 100644 --- a/packages/rsc-runtime/src/decode-document.ts +++ b/packages/rsc-runtime/src/decode-document.ts @@ -3,6 +3,7 @@ import { Children, isValidElement, type ReactNode } from 'react'; import { AgentContractError, createAgentDocument, + resolveAgentRenderLimits, type AgentDocument, type AgentDocumentNode, type AgentRenderLimits, @@ -55,16 +56,35 @@ const textChild = (children: unknown, type: AgentElementType): string => { }; interface DecodeState { + readonly limits: AgentRenderLimits; + nodes: number; representedError: boolean; } -const decodeNode = (node: ReactNode, state: DecodeState): AgentDocumentNode => { +const enterDecodeNode = (depth: number, state: DecodeState): void => { + if (depth > state.limits.maxDocumentDepth) { + throw new AgentContractError( + 'document-depth-exceeded', + `Agent Document depth exceeds ${String(state.limits.maxDocumentDepth)}`, + ); + } + state.nodes += 1; + if (state.nodes > state.limits.maxDocumentNodes) { + throw new AgentContractError( + 'document-node-count-exceeded', + `Agent Document node count exceeds ${String(state.limits.maxDocumentNodes)}`, + ); + } +}; + +const decodeNode = (node: ReactNode, depth: number, state: DecodeState): AgentDocumentNode => { + enterDecodeNode(depth, state); const element = protocolElement(node); const { props } = element; switch (element.type) { case 'agent-result': return { - children: Children.toArray(props.children as ReactNode).map((child) => decodeNode(child, state)), + children: Children.toArray(props.children as ReactNode).map((child) => decodeNode(child, depth + 1, state)), kind: 'result', ...(props.metadata === undefined ? {} : { metadata: props.metadata as JsonValue }), }; @@ -112,17 +132,18 @@ export const decodeAgentDocument = ( node: ReactNode, limits: Partial = {}, ): AgentDocument => { + const resolved = resolveAgentRenderLimits(limits); const root = protocolElement(node); if (root.type !== 'agent-result') { throw new AgentContractError('invalid-document', 'Flight output must have Agent.Result as its root'); } - const state: DecodeState = { representedError: false }; - const documentRoot = decodeNode(node, state); + const state: DecodeState = { limits: resolved, nodes: 0, representedError: false }; + const documentRoot = decodeNode(node, 1, state); return createAgentDocument({ root: documentRoot, status: state.representedError ? 'represented-error' : 'success', ...(root.props.value === undefined ? {} : { value: root.props.value as JsonValue }), version: 1, - }, limits); + }, resolved); }; diff --git a/packages/rsc-runtime/src/dispatcher.ts b/packages/rsc-runtime/src/dispatcher.ts index 747f6e5f4..f41e75314 100644 --- a/packages/rsc-runtime/src/dispatcher.ts +++ b/packages/rsc-runtime/src/dispatcher.ts @@ -92,11 +92,19 @@ export const createAgentRenderDispatcher = ( limits: options.limits, signal: request.signal, }); - pendingFlight.current = host.execute({ - invocation: request.invocation, - progress: session.progress, - signal: request.signal, - }); + const rememberFlight = (flight: Promise>): Promise> => { + void flight.catch(() => undefined); + return flight; + }; + try { + pendingFlight.current = rememberFlight(host.execute({ + invocation: request.invocation, + progress: session.progress, + signal: request.signal, + })); + } catch (error) { + pendingFlight.current = rememberFlight(Promise.reject(request.signal.aborted ? abortError() : error)); + } return toPublicEventStream(session.events, demand, request.signal); }; diff --git a/packages/rsc-runtime/src/lower-mcp.ts b/packages/rsc-runtime/src/lower-mcp.ts index 30b7b9a22..1e948400a 100644 --- a/packages/rsc-runtime/src/lower-mcp.ts +++ b/packages/rsc-runtime/src/lower-mcp.ts @@ -1,3 +1,5 @@ +import { Buffer } from 'node:buffer'; + import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; import { Children, isValidElement, type ReactElement, type ReactNode } from 'react'; @@ -50,6 +52,16 @@ export interface JsonObject { export type JsonValue = null | boolean | number | string | readonly JsonValue[] | JsonObject; +/** Incremental depth / node / byte checks while cloning JSON (Agent Document bounds). */ +export interface JsonSnapshotBudget { + readonly addBytes: (n: number) => void; + readonly addNode: () => void; + readonly checkDepth: (depth: number) => void; +} + +const jsonLeafBytes = (value: null | boolean | number | string): number => + Buffer.byteLength(JSON.stringify(value), 'utf8'); + const isArrayIndex = (key: string, length: number): boolean => { if (key === '0') return length > 0; if (!/^[1-9]\d*$/.test(key)) return false; @@ -60,10 +72,22 @@ const isArrayIndex = (key: string, length: number): boolean => { const jsonPathError = (reason: string, path: string): Error => new Error(path === '' ? reason : `${reason} at ${path}`); -const cloneJsonValue = (value: unknown, ancestors: Set, path: string): JsonValue => { - if (value === null || typeof value === 'boolean' || typeof value === 'string') return value; +const cloneJsonValue = ( + value: unknown, + ancestors: Set, + path: string, + depth = 0, + budget?: JsonSnapshotBudget, +): JsonValue => { + budget?.checkDepth(depth); + budget?.addNode(); + if (value === null || typeof value === 'boolean' || typeof value === 'string') { + budget?.addBytes(jsonLeafBytes(value)); + return value; + } if (typeof value === 'number') { if (!Number.isFinite(value)) throw jsonPathError('non-finite number', path); + budget?.addBytes(jsonLeafBytes(value)); return value; } if (typeof value !== 'object') throw jsonPathError('non-JSON value', path); @@ -80,21 +104,29 @@ const cloneJsonValue = (value: unknown, ancestors: Set, path: string): J throw jsonPathError('sparse or decorated array', path); } + budget?.addBytes(2); const clone: JsonValue[] = []; for (let index = 0; index < value.length; index += 1) { + if (index > 0) budget?.addBytes(1); const elementPath = `${path}[${index}]`; if (!Object.hasOwn(value, index)) throw jsonPathError('sparse array', elementPath); const descriptor = Object.getOwnPropertyDescriptor(value, String(index)); if (descriptor === undefined || !('value' in descriptor)) throw jsonPathError('array accessor', elementPath); // JSON.stringify serializes undefined array elements as null; match the SDK wire shape. - clone.push(descriptor.value === undefined ? null : cloneJsonValue(descriptor.value, ancestors, elementPath)); + clone.push( + descriptor.value === undefined + ? cloneJsonValue(null, ancestors, elementPath, depth + 1, budget) + : cloneJsonValue(descriptor.value, ancestors, elementPath, depth + 1, budget), + ); } return clone; } const prototype = Object.getPrototypeOf(value); if (prototype !== Object.prototype && prototype !== null) throw jsonPathError('non-plain object', path); + budget?.addBytes(2); const clone: { [key: string]: JsonValue } = Object.create(null) as { [key: string]: JsonValue }; + let properties = 0; for (const key of Reflect.ownKeys(value)) { if (typeof key !== 'string') throw jsonPathError('symbol key', path); const propertyPath = path === '' ? key : `${path}.${key}`; @@ -104,7 +136,10 @@ const cloneJsonValue = (value: unknown, ancestors: Set, path: string): J } // JSON.stringify drops undefined-valued properties; match the SDK wire shape. if (descriptor.value === undefined) continue; - clone[key] = cloneJsonValue(descriptor.value, ancestors, propertyPath); + if (properties > 0) budget?.addBytes(1); + budget?.addBytes(jsonLeafBytes(key) + 1); + properties += 1; + clone[key] = cloneJsonValue(descriptor.value, ancestors, propertyPath, depth + 1, budget); } return clone; } finally { @@ -135,10 +170,21 @@ const deepFreezeJson = (value: JsonValue): JsonValue => { return value; }; -export const snapshotJsonValue = (value: unknown, message: string): JsonValue => { +export const snapshotJsonValue = ( + value: unknown, + message: string, + budget?: { readonly depth: number; readonly limits: JsonSnapshotBudget }, +): JsonValue => { try { - return deepFreezeJson(cloneJsonValue(value, new Set(), '')); + return deepFreezeJson(cloneJsonValue( + value, + new Set(), + '', + budget?.depth ?? 0, + budget?.limits, + )); } catch (error) { + if (error instanceof Error && error.name === 'AgentContractError') throw error; throw new Error(`${message} (${error instanceof Error ? error.message : String(error)})`, { cause: error }); } }; diff --git a/packages/rsc-runtime/src/reconciler.ts b/packages/rsc-runtime/src/reconciler.ts index 6b9f0c40e..0d331a4b1 100644 --- a/packages/rsc-runtime/src/reconciler.ts +++ b/packages/rsc-runtime/src/reconciler.ts @@ -1,13 +1,16 @@ -import { Effect, Option, Queue, Stream, type Scope } from 'effect'; +import { Duration, Effect, Option, Queue, Stream, type Scope } from 'effect'; import { createElement, isValidElement, type ReactElement, type ReactNode } from 'react'; import { createFromReadableStream } from 'react-server-dom-rspack/client.node'; import { AgentContractError, createAgentRenderEventSequence, + elapsedTimeExceeded, + resolveAgentRenderLimits, type AgentRenderError, type AgentRenderEvent, type AgentRenderEventInput, + type AgentRenderEventSequence, type AgentRenderLimits, } from './agent-document.js'; import type { AgentProgressReporter, AgentProgressUpdate } from './agent-request.js'; @@ -290,71 +293,121 @@ const waitSettledBoundary = ( ), ); -type ReconcileState = - | { readonly kind: 'shell'; readonly snapshot: TreeSnapshot } - | { readonly kind: 'loop'; readonly snapshot: TreeSnapshot }; - -const reconcileInputStream = (root: ReactNode): Stream.Stream => { - const ids = new Map(); - const initial = snapshotTree(root, ids); - return Stream.paginate( - { kind: 'shell', snapshot: initial } satisfies ReconcileState, - (state): Effect.Effect< - readonly [readonly AgentRenderEventInput[], Option.Option], +type SettledBoundary = { + readonly boundary: PendingBoundary; + readonly error?: unknown; + readonly ok: boolean; +}; + +const waitPendingOrDeadline = ( + pending: readonly PendingBoundary[], + sequence: AgentRenderEventSequence, +): Effect.Effect => { + const remaining = sequence.remainingMs; + if (remaining <= 0) return Effect.fail(elapsedTimeExceeded(sequence.maxElapsedMs)); + return Effect.raceFirst( + waitSettledBoundary(pending), + Effect.sleep(Duration.millis(remaining)).pipe( + Effect.flatMap(() => Effect.fail(elapsedTimeExceeded(sequence.maxElapsedMs))), + ), + ); +}; + +const settledBoundaryInputs = ( + previous: TreeSnapshot, + next: TreeSnapshot, + winner: SettledBoundary, + limits: Partial | undefined, +): readonly AgentRenderEventInput[] => { + const stillPending = new Set(next.pending.map((boundary) => boundary.id)); + const rejectedById = new Map(next.rejected.map((entry) => [entry.id, entry.error] as const)); + const document = decodeAgentDocument(next.tree, limits); + const inputs: AgentRenderEventInput[] = []; + const emitFor = (id: string, fallback?: SettledBoundary): void => { + const rejected = rejectedById.get(id); + if (rejected !== undefined) { + inputs.push({ boundaryId: id, error: rejected, type: 'error' }); + return; + } + if (fallback !== undefined && !fallback.ok) { + inputs.push({ boundaryId: id, error: renderErrorFrom(fallback.error), type: 'error' }); + return; + } + if (stillPending.has(id) && id !== winner.boundary.id) return; + inputs.push({ boundaryId: id, document, type: 'replace' }); + }; + emitFor(winner.boundary.id, winner); + for (const boundary of previous.pending) { + if (boundary.id === winner.boundary.id) continue; + emitFor(boundary.id); + } + return inputs; +}; + +type LoopWait = + | { readonly kind: 'boundary'; readonly winner: SettledBoundary } + | { readonly kind: 'progress'; readonly input: AgentRenderEventInput }; + +const reconcileLoopStream = ( + root: ReactNode, + ids: Map, + initial: TreeSnapshot, + limits: Partial | undefined, + progressInputs: Queue.Queue, + sequence: AgentRenderEventSequence, +): Stream.Stream => + Stream.paginate( + initial, + (snapshot): Effect.Effect< + readonly [readonly AgentRenderEventInput[], Option.Option], Error > => { - switch (state.kind) { - case 'shell': - return Effect.try({ - catch: (error) => toRuntimeError(error), - try: () => - [ - [{ document: decodeAgentDocument(state.snapshot.tree), type: 'shell' as const }], - Option.some({ kind: 'loop' as const, snapshot: state.snapshot }), - ] as const, - }); - case 'loop': { - if (state.snapshot.pending.length === 0) { - return Effect.try({ + if (snapshot.pending.length === 0) { + return Queue.clear(progressInputs).pipe( + Effect.flatMap((queued) => + Effect.try({ catch: (error) => toRuntimeError(error), try: () => [ - [{ document: decodeAgentDocument(state.snapshot.tree), type: 'complete' as const }], + [ + ...queued, + { document: decodeAgentDocument(snapshot.tree, limits), type: 'complete' as const }, + ], Option.none(), ] as const, - }); - } - return waitSettledBoundary(state.snapshot.pending).pipe( - Effect.flatMap((settled) => - Effect.try({ + }), + ), + ); + } + return Effect.raceFirst( + waitPendingOrDeadline(snapshot.pending, sequence).pipe( + Effect.map((winner): LoopWait => ({ kind: 'boundary', winner })), + ), + Queue.take(progressInputs).pipe( + Effect.map((input): LoopWait => ({ kind: 'progress', input })), + ), + ).pipe( + Effect.flatMap((event) => { + switch (event.kind) { + case 'progress': + return Effect.succeed([[event.input], Option.some(snapshot)] as const); + case 'boundary': + return Effect.try({ catch: (error) => toRuntimeError(error), try: () => { - const snapshot = snapshotTree(root, ids); - const input: AgentRenderEventInput = settled.ok - ? { - boundaryId: settled.boundary.id, - document: decodeAgentDocument(snapshot.tree), - type: 'replace', - } - : { - boundaryId: settled.boundary.id, - error: renderErrorFrom(settled.error), - type: 'error', - }; - return [[input], Option.some({ kind: 'loop' as const, snapshot })] as const; + const next = snapshotTree(root, ids); + return [settledBoundaryInputs(snapshot, next, event.winner, limits), Option.some(next)] as const; }, - }), - ), - ); - } - default: { - const exhaustive: never = state; - return exhaustive; - } - } + }); + default: { + const exhaustive: never = event; + return exhaustive; + } + } + }), + ); }, ); -}; const gatedFlightStream = ( flight: ReadableStream, @@ -424,28 +477,39 @@ export interface AgentRenderEventSession { readonly progress: AgentProgressReporter; } +const handoffRequired = (): AgentContractError => + new AgentContractError( + 'handoff-required', + 'The render is complete; later work requires a new invocation handoff', + ); + /** * Invocation-local render pipeline: Flight bytes as a pull-gated Stream, * pending boundaries as `Stream.paginate`, contract bounds as the emit * stage. `progress` is created synchronously so the host can execute in the - * same turn as `stream()`. + * same turn as `stream()`. Pre-shell reports buffer (capped by maxEvents); + * live reports wait on a demand-bounded queue (capacity 0). */ export const createAgentRenderEventSession = ( options: AgentRenderEventStreamOptions, ): AgentRenderEventSession => { const sequence = createAgentRenderEventSequence(options.limits); + const maxBufferedProgress = resolveAgentRenderLimits(options.limits).maxEvents; let offerProgress: ((input: AgentRenderEventInput) => Effect.Effect) | undefined; + let progressFailure: Error | undefined; const bufferedProgress: AgentRenderEventInput[] = []; const progress: AgentProgressReporter = Object.freeze({ report: async (update: AgentProgressUpdate) => { - if (sequence.completed) { - throw new AgentContractError( - 'handoff-required', - 'The render is complete; later work requires a new invocation handoff', - ); - } + if (progressFailure !== undefined) throw progressFailure; + if (sequence.completed) throw handoffRequired(); const input = progressInput(update); if (offerProgress === undefined) { + if (bufferedProgress.length >= maxBufferedProgress) { + throw new AgentContractError( + 'event-count-exceeded', + `Agent render event count exceeds ${String(maxBufferedProgress)}`, + ); + } bufferedProgress.push(input); return; } @@ -454,27 +518,67 @@ export const createAgentRenderEventSession = ( }); const events = Stream.unwrap( Effect.gen(function*() { - const progressInputs = yield* Queue.unbounded(); - offerProgress = (input) => Queue.offer(progressInputs, input).pipe(Effect.asVoid); - for (const input of bufferedProgress) { - yield* Queue.offer(progressInputs, input); - } - const flight = yield* Effect.tryPromise({ - catch: (error) => hostError(options.signal, error), - try: () => options.flight, + const progressInputs = yield* Queue.bounded(0); + const bindProgress = (): void => { + offerProgress = (input) => + Queue.offer(progressInputs, input).pipe( + Effect.flatMap((accepted) => { + if (progressFailure !== undefined) return Effect.fail(progressFailure); + if (!accepted) return Effect.fail(handoffRequired()); + return Effect.void; + }), + ); + }; + const finalizeProgress = (error: Error): Effect.Effect => + Effect.sync(() => { + progressFailure = error; + }).pipe(Effect.andThen(Queue.shutdown(progressInputs))); + const setup = Effect.gen(function*() { + const flight = yield* Effect.tryPromise({ + catch: (error) => hostError(options.signal, error), + try: () => options.flight, + }); + if (options.signal.aborted) return yield* Effect.fail(abortError()); + const root = yield* decodeFlightRoot(flight, options.demand, options.signal); + if (options.signal.aborted) return yield* Effect.fail(abortError()); + const prepared = yield* Effect.try({ + catch: (error) => toRuntimeError(error), + try: () => { + const ids = new Map(); + const initial = snapshotTree(root, ids); + return { + ids, + initial, + shellInput: { + document: decodeAgentDocument(initial.tree, options.limits), + type: 'shell' as const, + } satisfies AgentRenderEventInput, + }; + }, + }); + bindProgress(); + return Stream.concat( + Stream.fromArray([prepared.shellInput, ...bufferedProgress]), + reconcileLoopStream( + root, + prepared.ids, + prepared.initial, + options.limits, + progressInputs, + sequence, + ), + ).pipe( + Stream.mapEffect((input) => emitBoundRenderEvent(sequence, input)), + Stream.tap((event) => (event.type === 'shell' ? options.demand.markShell : Effect.void)), + Stream.tapError((error) => Effect.sync(() => { + progressFailure = error; + })), + Stream.takeUntil((event) => event.type === 'complete'), + Stream.ensuring(Effect.suspend(() => finalizeProgress(progressFailure ?? handoffRequired()))), + ); }); - if (options.signal.aborted) return yield* Effect.fail(abortError()); - const root = yield* decodeFlightRoot(flight, options.demand, options.signal); - if (options.signal.aborted) return yield* Effect.fail(abortError()); - return Stream.merge( - reconcileInputStream(root), - Stream.fromQueue(progressInputs), - { haltStrategy: 'left' }, - ).pipe( - Stream.mapEffect((input) => emitBoundRenderEvent(sequence, input)), - Stream.tap((event) => (event.type === 'shell' ? options.demand.markShell : Effect.void)), - Stream.takeUntil((event) => event.type === 'complete'), - Stream.ensuring(Queue.shutdown(progressInputs)), + return yield* setup.pipe( + Effect.catch((error) => finalizeProgress(error).pipe(Effect.andThen(Effect.fail(error)))), ); }), ); diff --git a/packages/rsc-runtime/tests/agent-document.test.ts b/packages/rsc-runtime/tests/agent-document.test.ts index 47e816b91..7383f51d1 100644 --- a/packages/rsc-runtime/tests/agent-document.test.ts +++ b/packages/rsc-runtime/tests/agent-document.test.ts @@ -120,6 +120,58 @@ describe('AgentDocument', () => { version: 1, })).toThrow('Unsupported Agent Document node kind: div'); }); + + it('enforces depth, node-count, and byte bounds while cloning JSON payloads', () => { + const nest = (depth: number): { readonly child: unknown } | string => + depth === 0 ? 'leaf' : { child: nest(depth - 1) }; + try { + createAgentDocument({ + root: { kind: 'json', value: nest(70) }, + status: 'success', + version: 1, + }); + throw new Error('expected nested JSON value to exceed depth'); + } catch (error) { + expect(error).toMatchObject({ code: 'document-depth-exceeded' }); + } + try { + createAgentDocument({ + root: { kind: 'text', text: 'ok' }, + status: 'success', + value: nest(70), + version: 1, + }); + throw new Error('expected document value to exceed depth'); + } catch (error) { + expect(error).toMatchObject({ code: 'document-depth-exceeded' }); + } + try { + createAgentDocument( + { + root: { kind: 'json', value: 'x'.repeat(8_000) }, + status: 'success', + version: 1, + }, + { maxDocumentBytes: 100 }, + ); + throw new Error('expected oversized JSON string to exceed bytes'); + } catch (error) { + expect(error).toMatchObject({ code: 'document-bytes-exceeded' }); + } + try { + createAgentDocument( + { + root: { kind: 'json', value: Array.from({ length: 20 }, (_, index) => index) }, + status: 'success', + version: 1, + }, + { maxDocumentNodes: 4 }, + ); + throw new Error('expected wide JSON array to exceed node count'); + } catch (error) { + expect(error).toMatchObject({ code: 'document-node-count-exceeded' }); + } + }); }); describe('Agent render events', () => { diff --git a/packages/rsc-runtime/tests/dispatcher.test.ts b/packages/rsc-runtime/tests/dispatcher.test.ts index 38efd6ecc..d7893e8e2 100644 --- a/packages/rsc-runtime/tests/dispatcher.test.ts +++ b/packages/rsc-runtime/tests/dispatcher.test.ts @@ -53,6 +53,32 @@ describe('decodeAgentDocument', () => { expect(invoked).toBe(false); expect(() => decodeAgentDocument(createElement('div'))).toThrow('protocol element'); }); + + it('enforces configured document limits during the decode walk', () => { + const nest = (depth: number): ReturnType => + depth <= 1 + ? createElement('agent-result', null, createElement('agent-text', null, 'leaf')) + : createElement('agent-result', null, nest(depth - 1)); + try { + decodeAgentDocument(nest(5), { maxDocumentDepth: 2 }); + throw new Error('expected deep Result tree to exceed depth during decode'); + } catch (error) { + expect(error).toMatchObject({ code: 'document-depth-exceeded' }); + } + expect(decodeAgentDocument(nest(5), { maxDocumentDepth: 10 }).root).toMatchObject({ kind: 'result' }); + + const wide = createElement( + 'agent-result', + null, + ...Array.from({ length: 20 }, (_, index) => createElement('agent-text', null, `n${String(index)}`)), + ); + try { + decodeAgentDocument(wide, { maxDocumentNodes: 3 }); + throw new Error('expected broad Result tree to exceed node count during decode'); + } catch (error) { + expect(error).toMatchObject({ code: 'document-node-count-exceeded' }); + } + }); }); describe('AgentRenderDispatcher', () => { @@ -423,4 +449,123 @@ describe('AgentRenderDispatcher streaming', () => { await expect(reader.read()).rejects.toBeInstanceOf(AgentContractError); await expect(reader.read()).rejects.toMatchObject({ code: 'event-count-exceeded' }); }); + + it('holds pre-shell progress until the shell event is emitted', { retry: 2 }, async () => { + const inner = createWorkerHost('ready'); + const host: AgentFlightExecutionHost = { + execute: async (request) => { + if (request.progress === undefined) throw new Error('expected a progress reporter'); + await request.progress.report({ completed: 1, message: 'pre-shell' }); + return inner.execute(request); + }, + }; + const dispatcher = createAgentRenderDispatcher(host); + const events = await collectEvents(dispatcher.stream({ invocation, signal: new AbortController().signal })); + expect(eventTypes(events)).toEqual(['shell', 'progress', 'complete']); + expect(events[1]).toMatchObject({ completed: 1, message: 'pre-shell', sequence: 1, type: 'progress' }); + }); + + it('emits a replace or error for every boundary that settled before resnapshot', { retry: 2 }, async () => { + const host = createWorkerHost('dual'); + const dispatcher = createAgentRenderDispatcher(host); + const reader = dispatcher.stream({ invocation, signal: new AbortController().signal }).getReader(); + const shell = await reader.read(); + if (shell.value?.type !== 'shell') throw new Error('expected a shell event'); + host.resolve('a'); + host.resolve('b'); + const rest: AgentRenderEvent[] = []; + while (true) { + const next = await reader.read(); + if (next.done) break; + if (next.value === undefined) throw new Error('expected a render event'); + rest.push(next.value); + } + const replacements = rest.filter((event) => event.type === 'replace'); + expect(replacements.map((event) => event.boundaryId).sort()).toEqual(['b:1', 'b:2']); + expect(rest.at(-1)?.type).toBe('complete'); + }); + + it('fails a permanently pending boundary when the elapsed deadline expires', { retry: 2 }, async () => { + const host = createWorkerHost('single'); + const dispatcher = createAgentRenderDispatcher(host, { limits: { maxElapsedMs: 150 } }); + const reader = dispatcher.stream({ invocation, signal: new AbortController().signal }).getReader(); + const shell = await reader.read(); + if (shell.value?.type !== 'shell') throw new Error('expected a shell event'); + await expect(reader.read()).rejects.toBeInstanceOf(AgentContractError); + await expect(reader.read()).rejects.toMatchObject({ code: 'elapsed-time-exceeded' }); + }); + + it('converts a synchronous host throw into a stream failure', async () => { + const host: AgentFlightExecutionHost = { + execute: () => { + throw new Error('sync host setup'); + }, + }; + const dispatcher = createAgentRenderDispatcher(host); + const stream = dispatcher.stream({ invocation, signal: new AbortController().signal }); + await expect(stream.getReader().read()).rejects.toThrow('sync host setup'); + await expect(dispatcher.dispatch({ invocation, signal: new AbortController().signal })).rejects.toThrow( + 'sync host setup', + ); + }); + + it('holds later progress reports until the consumer accepts the prior update', { retry: 2 }, async () => { + let progress: AgentProgressReporter | undefined; + const inner = createWorkerHost('single'); + const host: AgentFlightExecutionHost = { + execute: async (request) => { + progress = request.progress; + return inner.execute(request); + }, + }; + const dispatcher = createAgentRenderDispatcher(host); + const reader = dispatcher.stream({ invocation, signal: new AbortController().signal }).getReader(); + const shell = await reader.read(); + if (shell.value?.type !== 'shell') throw new Error('expected a shell event'); + if (progress === undefined) throw new Error('expected the dispatcher to install a progress reporter'); + + const reports = [1, 2, 3].map((completed) => { + let settled = false; + const done = progress.report({ completed, message: `n${String(completed)}` }).then(() => { + settled = true; + }); + return { done, get settled() { return settled; } }; + }); + await new Promise((resolve) => { + setTimeout(resolve, 30); + }); + expect(reports[2]?.settled).toBe(false); + + const firstEvent = await reader.read(); + if (firstEvent.value?.type !== 'progress') throw new Error('expected the first progress event'); + expect(firstEvent.value).toMatchObject({ completed: 1, message: 'n1' }); + await reports[0]?.done; + const secondEvent = await reader.read(); + if (secondEvent.value?.type !== 'progress') throw new Error('expected the second progress event'); + expect(secondEvent.value).toMatchObject({ completed: 2, message: 'n2' }); + await reports[1]?.done; + const thirdEvent = await reader.read(); + if (thirdEvent.value?.type !== 'progress') throw new Error('expected the third progress event'); + expect(thirdEvent.value).toMatchObject({ completed: 3, message: 'n3' }); + await reports[2]?.done; + + inner.resolve('a'); + expect((await reader.read()).value?.type).toBe('replace'); + expect((await reader.read()).value?.type).toBe('complete'); + }); + + it('rejects progress after Flight setup fails and shuts down the queue', async () => { + let progress: AgentProgressReporter | undefined; + const host: AgentFlightExecutionHost = { + execute: async (request) => { + progress = request.progress; + throw new Error('flight setup failed'); + }, + }; + const dispatcher = createAgentRenderDispatcher(host); + const reader = dispatcher.stream({ invocation, signal: new AbortController().signal }).getReader(); + await expect(reader.read()).rejects.toThrow('flight setup failed'); + if (progress === undefined) throw new Error('expected the dispatcher to install a progress reporter'); + await expect(progress.report({ completed: 1, message: 'after-fail' })).rejects.toThrow('flight setup failed'); + }); });