diff --git a/src/cli/commands/snapshot.ts b/src/cli/commands/snapshot.ts index 42e00af1..7e00beba 100644 --- a/src/cli/commands/snapshot.ts +++ b/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 { - 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); } diff --git a/src/host/README.md b/src/host/README.md index 56b5c74e..ad4895bc 100644 --- a/src/host/README.md +++ b/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] diff --git a/src/host/rpcClient.ts b/src/host/rpcClient.ts index 6d78dc9e..4322cb90 100644 --- a/src/host/rpcClient.ts +++ b/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; diff --git a/src/host/rpcServer.ts b/src/host/rpcServer.ts index 461c8d23..788413e9 100644 --- a/src/host/rpcServer.ts +++ b/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`); } } diff --git a/src/protocol/README.md b/src/protocol/README.md index b8404970..38d8becf 100644 --- a/src/protocol/README.md +++ b/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`, so TypeScript requires one entry per method. [messages.ts:RpcMethodSchemas] - `RunParamsSchema` defaults `noWait` to `false`. [messages.ts:RunParamsSchema] diff --git a/src/protocol/messages.ts b/src/protocol/messages.ts index 1a1657e4..fed77655 100644 --- a/src/protocol/messages.ts +++ b/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), diff --git a/test/e2e/helpers.ts b/test/e2e/helpers.ts index a6936833..d2d1d63a 100644 --- a/test/e2e/helpers.ts +++ b/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`]; } diff --git a/test/e2e/large-snapshot.test.ts b/test/e2e/large-snapshot.test.ts new file mode 100644 index 00000000..92defa34 --- /dev/null +++ b/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 { + 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>( + [ + 'create', + '--cols', + String(COLS), + '--rows', + String(ROWS), + '--', + ...fixtureCommand('styled-fill'), + ], + env, + ); + const { sessionId } = createEnvelope.result; + createdSessionIds.push(sessionId); + + const waitEnvelope = runCliJson>( + ['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 & { + 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(); + }); +}); diff --git a/test/fixtures/apps/styled-fill/main.ts b/test/fixtures/apps/styled-fill/main.ts new file mode 100644 index 00000000..1e16927a --- /dev/null +++ b/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); diff --git a/test/unit/commands/snapshot.test.ts b/test/unit/commands/snapshot.test.ts index 44d28a67..06b3ead1 100644 --- a/test/unit/commands/snapshot.test.ts +++ b/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', diff --git a/test/unit/host/rpcServer.test.ts b/test/unit/host/rpcServer.test.ts index 22736e8a..6c670f89 100644 --- a/test/unit/host/rpcServer.test.ts +++ b/test/unit/host/rpcServer.test.ts @@ -274,3 +274,99 @@ describe('RpcServer CliError responses', () => { } }); }); + +describe('RpcServer large success responses', () => { + beforeEach(async () => { + tempDir = await mkdtemp(join(tmpdir(), 'agent-tty-rpc-server-')); + }); + + afterEach(async () => { + await rm(tempDir, { recursive: true, force: true }); + tempDir = ''; + }); + + async function textSnapshotFrom(text: string): Promise { + const socketFile = join(tempDir, 'rpc.sock'); + const server = new RpcServer(socketFile, { + snapshot: () => + Promise.resolve({ + format: 'text', + sessionId: 'session-1', + capturedAtSeq: 1, + cols: 80, + rows: 24, + cursorRow: 0, + cursorCol: 0, + text, + }), + }); + await server.listen(); + try { + return await sendRpc( + socketFile, + 'snapshot', + { format: 'text' }, + 20_000, + ).catch((error: unknown) => error); + } finally { + await server.close(); + } + } + + it('returns a success response larger than 1 MiB', async () => { + const result = await textSnapshotFrom('x'.repeat(2_000_000)); + + expect(result).toMatchObject({ format: 'text' }); + expect((result as { text: string }).text).toHaveLength(2_000_000); + }); + + it('scans each chunk of a large response for the newline only once', async () => { + // A socket read gives at most 64 KiB, so this response comes in 60 or more + // chunks. Scanning the whole buffer after each chunk would read the + // response about 30 times; scanning each new chunk reads it once. + // oxlint-disable-next-line typescript/unbound-method -- the spy below calls it with nativeIndexOf.call(this, ...). + const nativeIndexOf = String.prototype.indexOf; + let newlineScanChars = 0; + const indexOf = vi + .spyOn(String.prototype, 'indexOf') + .mockImplementation(function ( + this: string, + searchString: string, + position?: number, + ): number { + if (searchString === '\n') { + newlineScanChars += this.length; + } + return nativeIndexOf.call(this, searchString, position); + }); + + let result: unknown; + try { + result = await textSnapshotFrom('x'.repeat(4_000_000)); + } finally { + indexOf.mockRestore(); + } + + expect((result as { text: string }).text).toHaveLength(4_000_000); + expect(newlineScanChars).toBeLessThan(2 * 4_000_000); + }); + + it( + 'refuses a success response over 64 MiB with PROTOCOL_ERROR', + { timeout: 30_000 }, + async () => { + const error = await textSnapshotFrom('x'.repeat(67_108_864)); + + expect(error).toMatchObject({ + code: 'PROTOCOL_ERROR', + details: { maxResponseChars: 67_108_864 }, + }); + const { details, message } = error as CliError; + const responseChars = details.responseChars as number; + expect(responseChars).toBeGreaterThan(67_108_864); + expect(message).toBe( + `RPC response takes ${String(responseChars)} characters, over the 67108864-character limit.`, + ); + }, + ); +});