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

.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);
   };