fix(runtime): review follow-ups — traversal bounds, demand-bounded progress, failure conversion (#172)

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.

Author: ScriptedAlchemy

.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.

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>): AgentRenderLimits => {
+export const resolveAgentRenderLimits = (overrides: Partial<AgentRenderLimits> = {}): 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<object>;
   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<AgentRenderLimits> = {},
 ): 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);
+    },
   });
 };

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<AgentRenderLimits> = {},
 ): 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);
 };
 

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<ReadableStream<Uint8Array>>): Promise<ReadableStream<Uint8Array>> => {
+      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);
   };