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'); + }); });