fix: accept RPC responses up to 64 MiB and refuse a larger success with PROTOCOL_ERROR

src/cli/commands/snapshot.ts

@@ -20,6 +20,12 @@ import { manifestPath, socketPath } from '../../storage/sessionPaths.js';
 import { checkedSessionDir } from '../sessionGuards.js';
 
 const DEFAULT_SNAPSHOT_FORMAT = 'structured';
+// The host sends no bytes until it has built every cell. A fully styled
+// 500x200 snapshot took 1.5 to 1.8 s to its first byte at load average 100
+// (about 17 microseconds per cell), so a grid near MAX_RPC_RESPONSE_CHARS
+// (about 1.2 million styled cells) takes about 20 s. The 5,000 ms default
+// would end it with HOST_TIMEOUT.
+const SNAPSHOT_WITH_CELLS_RPC_TIMEOUT_MS = 60_000;
 
 interface CommandOptions {
   context: CommandContext;
@@ -85,16 +91,15 @@ async function runRpcSnapshot(
   includeScrollback: boolean,
   includeCells: boolean,
 ): Promise<SnapshotResult> {
-  const rawResult: unknown = await sendRpc(
-    socketPath(sessionDirectory),
-    'snapshot',
-    {
-      format,
-      includeScrollback,
-      includeCells,
-      rendererName,
-    },
-  );
+  const params = { format, includeScrollback, includeCells, rendererName };
+  const rawResult: unknown = includeCells
+    ? await sendRpc(
+        socketPath(sessionDirectory),
+        'snapshot',
+        params,
+        SNAPSHOT_WITH_CELLS_RPC_TIMEOUT_MS,
+      )
+    : await sendRpc(socketPath(sessionDirectory), 'snapshot', params);
 
   return parseSnapshotResult(rawResult);
 }

src/host/README.md

@@ -62,8 +62,8 @@
 - `allocateSession` throws `INVALID_DIMENSIONS` for non-positive `cols` or `rows`, `STORAGE_READ_ERROR` when `cwd` is not a directory and `INVALID_INPUT` when the shell path is not a file. An empty `command` defaults to the shell path. It creates the session directory and calls `launch` only after these checks pass. [lifecycle.ts:allocateSession]
 - `buildReplayInput` validates the manifest with `SessionRecordSchema` and requires its `sessionId` to match. It validates the events with `validateEventRecords` unless `trustValidated` is set. `targetSeq` is -1 for no events, else from 0 to the last `seq`. [replay.ts:buildReplayInput]
 - `RpcServer.listen` removes a stale socket file (no answer within 1,000 ms, `ECONNREFUSED` or `ENOENT`) and throws when another server is listening. It asserts a socket path of at most 104 characters. [rpcServer.ts:RpcServer]
-- `RpcServer` and `sendRpc` cap a message at 1,048,576 characters. The server validates the request with `RpcRequestSchema`, the params and result with `RpcMethodSchemas`, and the response with `RpcResponseSchema` (all from `src/protocol/messages.ts`). A handler that throws a `CliError` gets an error response with its code, message and, when it has any, its `details`. It leaves `details` out when they do not serialize to a JSON object (a `BigInt`, a circular object, a `toJSON` that returns something else) or take more than 65,536 characters. It cuts an error message to 65,536 characters, without splitting a surrogate pair, and ends it with `[message truncated]`. A request id over 128 characters fails `RpcRequestSchema`, and the error response uses the id `unknown`. So an error response stays under the 1,048,576-character cap and keeps the error's code. `sendRpc` rebuilds that `CliError` with the same details. Any other thrown error becomes `RPC_ERROR` without details. If building or sending the response throws, the server logs the error and answers `RPC_ERROR` "RPC request handling failed unexpectedly.", so the host keeps running. Unknown methods get an error response. [rpcServer.ts:RpcServer] [rpcClient.ts:sendRpc]
-- `sendRpc` waits 5,000 ms by default. It requires the response `id` to match the request. It maps `ECONNREFUSED`, `ENOENT` and `ECONNRESET` to `HOST_UNREACHABLE`, a socket timeout to `HOST_TIMEOUT` and other transport errors to `RPC_ERROR`. [rpcClient.ts:sendRpc]
+- `RpcServer` caps a request at 1,048,576 characters. `sendRpc` caps a response at `MAX_RPC_RESPONSE_CHARS` (67,108,864 characters, from `src/protocol/messages.ts`), so a structured snapshot with cells of a 500x200 screen (about 5.6 million characters when every cell is styled) fits. The server answers a success response above that cap with `PROTOCOL_ERROR` "RPC response takes N characters, over the 67108864-character limit.", with `responseChars` and `maxResponseChars` in `details`. The server validates the request with `RpcRequestSchema`, the params and result with `RpcMethodSchemas`, and the response with `RpcResponseSchema` (all from `src/protocol/messages.ts`). A handler that throws a `CliError` gets an error response with its code, message and, when it has any, its `details`. It leaves `details` out when they do not serialize to a JSON object (a `BigInt`, a circular object, a `toJSON` that returns something else) or take more than 65,536 characters. It cuts an error message to 65,536 characters, without splitting a surrogate pair, and ends it with `[message truncated]`. A request id over 128 characters fails `RpcRequestSchema`, and the error response uses the id `unknown`. So an error response stays under 1,048,576 characters, which a CLI from before the larger response cap also reads, and keeps the error's code. `sendRpc` rebuilds that `CliError` with the same details. Any other thrown error becomes `RPC_ERROR` without details. If building or sending the response throws, the server logs the error and answers `RPC_ERROR` "RPC request handling failed unexpectedly.", so the host keeps running. Unknown methods get an error response. [rpcServer.ts:RpcServer] [rpcClient.ts:sendRpc]
+- `sendRpc` waits 5,000 ms by default, as a socket inactivity timeout. The `snapshot` command passes 60,000 ms when it asks for cells, because the host sends nothing until it has built every cell. `sendRpc` scans each response chunk once for the newline. It requires the response `id` to match the request. It maps `ECONNREFUSED`, `ENOENT` and `ECONNRESET` to `HOST_UNREACHABLE`, a socket timeout to `HOST_TIMEOUT` and other transport errors to `RPC_ERROR`. [rpcClient.ts:sendRpc]
 - `HostRendererManager` runs backend operations one at a time. `withBackend` requires the replay input `sessionId` to match, and recreates the backend when the initial size changes. [renderer.ts:HostRendererManager]
 - A waited `run` calls `EventLog.startAppend`, `registerWaitedRun` and `pty.write` with no `await` in between, then awaits `written`. When `registerWaitedRun` or `pty.write` throws, it waits for `written` to settle and rethrows that error, so a failed write never becomes an unhandled rejection. So no other input reaches the PTY or the log between the command's `seq` and its write, and the marker is registered before the command can print. [hostMain.ts:runHost]
 - `RUN_MARKER_PATTERN` is `/^__AT_MARKER_([0-9a-f]{32})__$/u`. `registerWaitedRun` requires a unique marker and a non-negative integer `inputRunSeq`. `buildRunCompletePostamble` limits the postamble to 64 characters before its newline, and the postamble must not contain the marker. [runCompletionSentinel.ts:RUN_MARKER_PATTERN] [runCompletionCoordinator.ts:RunCompletionCoordinator] [runCompletionCoordinator.ts:buildRunCompletePostamble]

src/host/rpcClient.ts

@@ -8,6 +8,7 @@ import {
   type ProtocolErrorCode,
 } from '../protocol/errors.js';
 import {
+  MAX_RPC_RESPONSE_CHARS,
   RpcMethodSchemas,
   RpcRequestSchema,
   RpcResponseSchema,
@@ -25,7 +26,6 @@ import { ResourceScope } from '../util/resourceScope.js';
 import { isOwnLegacySocket, socketRootExists } from './socketRoot.js';
 
 const DEFAULT_TIMEOUT_MS = 5_000;
-const MAX_RPC_BUFFER_BYTES = 1_048_576;
 const HOST_UNREACHABLE_SOCKET_CODES = new Set([
   'ECONNREFUSED',
   'ENOENT',
@@ -233,7 +233,7 @@ async function sendRpcAt(
         return;
       }
 
-      if (buffer.length + chunk.length > MAX_RPC_BUFFER_BYTES) {
+      if (buffer.length + chunk.length > MAX_RPC_RESPONSE_CHARS) {
         settlers.reject(
           makeCliError(ERROR_CODES.RPC_ERROR, {
             message: 'RPC response exceeds maximum buffer size.',
@@ -243,15 +243,17 @@ async function sendRpcAt(
         return;
       }
 
-      buffer += chunk;
-      const newlineIndex = buffer.indexOf('\n');
+      // Scan only the new chunk: the buffer before it has no newline, and a
+      // rescan of the whole buffer after every chunk is quadratic.
+      const newlineIndex = chunk.indexOf('\n');
 
       if (newlineIndex < 0) {
+        buffer += chunk;
         return;
       }
 
       responseHandled = true;
-      const line = buffer.slice(0, newlineIndex);
+      const line = buffer + chunk.slice(0, newlineIndex);
 
       try {
         const rawResponse = JSON.parse(line) as unknown;

src/host/rpcServer.ts

@@ -2,9 +2,10 @@ import { chmod, stat, unlink } from 'node:fs/promises';
 import net from 'node:net';
 
 import { CliError } from '../cli/errors.js';
-import { ERROR_CODES } from '../protocol/errors.js';
+import { ERROR_CODES, makeCliError } from '../protocol/errors.js';
 import {
   MAX_RPC_REQUEST_ID_CHARS,
+  MAX_RPC_RESPONSE_CHARS,
   RpcMethodSchemas,
   RpcRequestSchema,
   RpcResponseSchema,
@@ -19,18 +20,19 @@ import { invariant } from '../util/assert.js';
 import { ResourceScope } from '../util/resourceScope.js';
 
 const MAX_UNIX_SOCKET_PATH = 104;
-const MAX_RPC_BUFFER_BYTES = 1_048_576;
+const MAX_RPC_REQUEST_CHARS = 1_048_576;
 const MAX_RPC_ERROR_DETAILS_CHARS = 65_536;
 const MAX_RPC_ERROR_MESSAGE_CHARS = 65_536;
 const TRUNCATED_MESSAGE_SUFFIX = ' [message truncated]';
 // JSON writes a control character as 6 characters (\u0001), so an error
-// response with an id, a message and details at their caps still fits in one
-// buffer.
+// response with an id, a message and details at their caps still fits in
+// 1,048,576 characters, the response cap of a CLI from before
+// MAX_RPC_RESPONSE_CHARS.
 invariant(
   6 * (MAX_RPC_REQUEST_ID_CHARS + MAX_RPC_ERROR_MESSAGE_CHARS) +
     MAX_RPC_ERROR_DETAILS_CHARS +
     1_024 <
-    MAX_RPC_BUFFER_BYTES,
+    1_048_576,
   'RPC error message and details caps must fit in one RPC response.',
 );
 
@@ -361,7 +363,7 @@ export class RpcServer {
         return;
       }
 
-      if (buffer.length + chunk.length > MAX_RPC_BUFFER_BYTES) {
+      if (buffer.length + chunk.length > MAX_RPC_REQUEST_CHARS) {
         handled = true;
         this.sendResponse(
           socket,
@@ -539,6 +541,28 @@ export class RpcServer {
   }
 
   private sendResponse(socket: net.Socket, response: RpcResponse): void {
-    socket.end(`${JSON.stringify(response)}\n`);
+    const message = `${JSON.stringify(response)}\n`;
+    if (message.length <= MAX_RPC_RESPONSE_CHARS) {
+      socket.end(message);
+      return;
+    }
+
+    // An error response stays under 1,048,576 characters (see the invariant
+    // at the top of this file), so only a success response gets here.
+    invariant(
+      response.ok,
+      'Only an RPC success response can exceed the response cap.',
+    );
+    const refusal = buildCliErrorResponse(
+      response.id,
+      makeCliError(ERROR_CODES.PROTOCOL_ERROR, {
+        message: `RPC response takes ${String(message.length)} characters, over the ${String(MAX_RPC_RESPONSE_CHARS)}-character limit.`,
+        details: {
+          responseChars: message.length,
+          maxResponseChars: MAX_RPC_RESPONSE_CHARS,
+        },
+      }),
+    );
+    socket.end(`${JSON.stringify(refusal)}\n`);
   }
 }

src/protocol/README.md

@@ -34,7 +34,7 @@
 - `InputRunEventPayloadSchema` requires `marker` when `noWait` is false and rejects it when `noWait` is true. [schemas.ts:InputRunEventPayloadSchema]
 - `WaitForRenderParamsSchema` requires at least one of `text`, `regex`, `screenStableMs`, `cursorRow` or `cursorCol`. `text` and `regex` are mutually exclusive, and scope `cursor-line` needs `text` or `regex`. [schemas.ts:WaitForRenderParamsSchema]
 - `RecordDiffResultSchema` checks that `identical` matches the equality of the two screen hashes. An identical result must have an empty diff and equal side row counts. A side with `capturedAtSeq` -1 must hash to a blank screen. For a non-identical result, the rows rebuilt from the diff must hash to each side's `screenHash`. [schemas.ts:RecordDiffResultSchema]
-- `RpcRequestSchema` requires non-empty `id` and `method`. The `id` takes at most `MAX_RPC_REQUEST_ID_CHARS` (128) characters, so the id that the server echoes in a response stays small. It does not restrict `method` to `RpcMethodSchema`. [messages.ts:RpcRequestSchema]
+- `RpcRequestSchema` requires non-empty `id` and `method`. The `id` takes at most `MAX_RPC_REQUEST_ID_CHARS` (128) characters, so the id that the server echoes in a response stays small. `MAX_RPC_RESPONSE_CHARS` (67,108,864) is the largest response line, newline included, that `sendRpc` reads and `RpcServer` sends. It does not restrict `method` to `RpcMethodSchema`. [messages.ts:RpcRequestSchema]
 - `RpcErrorSchema` requires non-empty `code` and `message`, and allows an optional `details` record. A CLI from before `details` was added rejects an error response that has it, so that CLI reports `RPC_ERROR` (`retryable: true`) for such an error from a newer host, and retrying does not help. A newer CLI reads an older host's error responses. [messages.ts:RpcErrorSchema]
 - `RpcMethodSchemas` uses `satisfies Record<RpcMethod, ...>`, so TypeScript requires one entry per method. [messages.ts:RpcMethodSchemas]
 - `RunParamsSchema` defaults `noWait` to `false`. [messages.ts:RunParamsSchema]

src/protocol/messages.ts

@@ -73,6 +73,12 @@ const DurationSchema = z.number().int().positive();
 // server echoes in every response.
 export const MAX_RPC_REQUEST_ID_CHARS = 128;
 
+// A fully styled cell takes about 56 characters in a structured snapshot with
+// cells, so a 500x200 snapshot takes about 5.6 million. This cap fits about
+// 1.2 million styled cells (a 1,000x1,000 screen) and still bounds the memory
+// that one response can take.
+export const MAX_RPC_RESPONSE_CHARS = 67_108_864;
+
 export const RpcRequestSchema = z
   .object({
     id: z.string().min(1).max(MAX_RPC_REQUEST_ID_CHARS),

test/e2e/helpers.ts

@@ -82,7 +82,8 @@ export function fixtureCommand(
     | 'alt-screen-demo'
     | 'crash-demo'
     | 'unicode-grid'
-    | 'scrollback-demo',
+    | 'scrollback-demo'
+    | 'styled-fill',
 ): string[] {
   return ['node', '--import', 'tsx', `test/fixtures/apps/${appName}/main.ts`];
 }

test/e2e/large-snapshot.test.ts

@@ -0,0 +1,113 @@
+import { spawnSync } from 'node:child_process';
+
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import type { SnapshotResult } from '../../src/protocol/messages.js';
+import {
+  cleanupHome,
+  createIsolatedHome,
+  DEFAULT_CLI_TIMEOUT_MS,
+  fixtureCommand,
+  runCli,
+  runCliJson,
+  type SuccessEnvelope,
+  type WaitResult,
+} from './helpers.js';
+
+interface CreateResult {
+  sessionId: string;
+}
+
+const COLS = 300;
+const ROWS = 100;
+
+function testEnv(home: string): Record<string, string> {
+  return { AGENT_TTY_HOME: home };
+}
+
+describe('large snapshot e2e', { timeout: 90_000 }, () => {
+  let testHome = '';
+  let createdSessionIds: string[] = [];
+
+  beforeEach(async () => {
+    testHome = await createIsolatedHome();
+    createdSessionIds = [];
+  });
+
+  afterEach(async () => {
+    const env = testEnv(testHome);
+
+    for (const sessionId of createdSessionIds) {
+      runCli(['destroy', sessionId, '--force', '--json'], env);
+    }
+
+    await cleanupHome(testHome);
+  });
+
+  it('returns a styled 300x100 snapshot with cells, larger than 1 MiB', () => {
+    const env = testEnv(testHome);
+    const createEnvelope = runCliJson<SuccessEnvelope<CreateResult>>(
+      [
+        'create',
+        '--cols',
+        String(COLS),
+        '--rows',
+        String(ROWS),
+        '--',
+        ...fixtureCommand('styled-fill'),
+      ],
+      env,
+    );
+    const { sessionId } = createEnvelope.result;
+    createdSessionIds.push(sessionId);
+
+    const waitEnvelope = runCliJson<SuccessEnvelope<WaitResult>>(
+      ['wait', sessionId, '--text', 'FILL-DONE', '--timeout', '30000'],
+      env,
+    );
+    expect(waitEnvelope.ok).toBe(true);
+
+    // spawnSync keeps at most 1 MiB of stdout unless maxBuffer is raised.
+    const snapshot = spawnSync(
+      process.execPath,
+      [
+        '--import',
+        'tsx',
+        './src/cli/main.ts',
+        'snapshot',
+        sessionId,
+        '--include-cells',
+        '--json',
+      ],
+      {
+        cwd: process.cwd(),
+        encoding: 'utf8',
+        env: { ...process.env, ...env },
+        maxBuffer: 64 * 1_048_576,
+        timeout: DEFAULT_CLI_TIMEOUT_MS,
+      },
+    );
+
+    const envelope = JSON.parse(
+      snapshot.stdout,
+    ) as SuccessEnvelope<SnapshotResult> & {
+      error?: unknown;
+    };
+    expect(envelope.error).toBeUndefined();
+    expect(envelope.ok).toBe(true);
+    expect(snapshot.status).toBe(0);
+    expect(snapshot.stdout.length).toBeGreaterThan(1_048_576);
+    const { result } = envelope;
+    if (result.format !== 'structured') {
+      throw new Error('expected a structured snapshot');
+    }
+    expect(result.cells).toHaveLength(ROWS);
+    expect(result.cells?.[0]?.cells).toHaveLength(COLS);
+    expect(result.cells?.[0]?.cells[0]).toMatchObject({
+      char: 'a',
+      bold: true,
+    });
+    expect(result.cells?.[0]?.cells[0]?.fg).toBeDefined();
+    expect(result.cells?.[0]?.cells[0]?.bg).toBeDefined();
+  });
+});

test/fixtures/apps/styled-fill/main.ts

@@ -0,0 +1,36 @@
+import assert from 'node:assert/strict';
+import process from 'node:process';
+
+// Fills the whole screen with cells that each have their own truecolor
+// foreground, background and bold, the largest per-cell snapshot data.
+const MARKER = 'FILL-DONE';
+const HOLD_OPEN_MS = 60_000;
+
+const cols = process.stdout.columns;
+const rows = process.stdout.rows;
+assert(
+  Number.isInteger(cols) && Number.isInteger(rows) && cols > MARKER.length,
+  'styled-fill needs a TTY wider than its marker',
+);
+
+let output = '\u001b[2J';
+for (let row = 0; row < rows; row += 1) {
+  output += `\u001b[${String(row + 1)};1H`;
+  const isLastRow = row === rows - 1;
+  // Leave the last cell of the screen empty so the screen does not scroll.
+  const width = isLastRow ? cols - 1 - MARKER.length : cols;
+  for (let col = 0; col < width; col += 1) {
+    const fg = `${String(col % 256)};${String(row % 256)};${String((col * 7) % 256)}`;
+    const bg = `${String((row * 3) % 256)};${String((col * 5) % 256)};${String(255 - (col % 256))}`;
+    const char = String.fromCharCode(97 + ((row * cols + col) % 26));
+    output += `\u001b[1;38;2;${fg}m\u001b[48;2;${bg}m${char}`;
+  }
+  if (isLastRow) {
+    output += `\u001b[0m${MARKER}`;
+  }
+}
+process.stdout.write(output);
+
+setTimeout(() => {
+  process.exit(0);
+}, HOLD_OPEN_MS);

test/unit/commands/snapshot.test.ts

@@ -358,6 +358,9 @@ describe('snapshot command', () => {
         includeCells: true,
         rendererName: 'ghostty-web',
       },
+      // The host sends nothing until it has built every cell, so a large
+      // grid needs more than the 5,000 ms default.
+      60_000,
     );
     expect(mocks.emitSuccess).toHaveBeenCalledWith({
       command: 'snapshot',