fix: refine prepare and snapshot timeouts

Author: thymikee

src/client-types.ts

@@ -306,6 +306,7 @@ export type CaptureSnapshotOptions = AgentDeviceRequestOverrides &
     scope?: string;
     raw?: boolean;
     forceFull?: boolean;
+    timeoutMs?: number;
   };
 
 export type CaptureSnapshotResult = {

src/command-catalog.ts

@@ -130,6 +130,7 @@ const MCP_UNEXPOSED_CLI_COMMANDS = commandSet(
   LOCAL_CLI_COMMANDS.disconnect,
   LOCAL_CLI_COMMANDS.mcp,
   LOCAL_CLI_COMMANDS.reactDevtools,
+  PUBLIC_COMMANDS.prepare,
 );
 
 const CAPABILITY_EXEMPT_CLI_COMMANDS = commandSet(
@@ -158,7 +159,6 @@ export const DAEMON_COMMAND_GROUPS = {
     PUBLIC_COMMANDS.apps,
   ),
   state: commandSet(PUBLIC_COMMANDS.boot, PUBLIC_COMMANDS.appState),
-  prepare: commandSet(PUBLIC_COMMANDS.prepare),
   observability: commandSet(PUBLIC_COMMANDS.perf, PUBLIC_COMMANDS.logs, PUBLIC_COMMANDS.network),
   replay: commandSet(PUBLIC_COMMANDS.replay, PUBLIC_COMMANDS.test),
   snapshot: commandSet(

src/commands/__tests__/command-surface-metadata.test.ts

@@ -19,3 +19,7 @@ test('MCP exposed command names have metadata and executable command definitions
     assert.ok(executableNames.has(name), `${name} must have an executable command definition`);
   }
 });
+
+test('CI-only prepare command stays out of MCP tool surface', () => {
+  assert.equal(listMcpExposedCommandNames().includes('prepare'), false);
+});

src/commands/cli-grammar/capture.ts

@@ -39,6 +39,7 @@ export const captureCliReaders = {
     scope: flags.snapshotScope,
     raw: flags.snapshotRaw,
     forceFull: flags.snapshotForceFull,
+    timeoutMs: flags.timeoutMs,
   }),
   screenshot: (positionals, flags) => ({
     ...commonInputFromFlags(flags),

src/commands/client-command-metadata.ts

@@ -105,6 +105,7 @@ export const clientCommandMetadata = [
     scope: stringField(),
     raw: booleanField(),
     forceFull: booleanField(),
+    timeoutMs: integerField('Maximum wall-clock time for the snapshot command.'),
   }),
   defineClientCommandMetadata('screenshot', {
     path: stringField('Output path.'),

src/daemon-client.ts

@@ -117,6 +117,7 @@ type ResolvedDaemonTransport = 'socket' | 'http';
 
 const REQUEST_TIMEOUT_MS = 90_000;
 const SNAPSHOT_REQUEST_TIMEOUT_MS = 30_000;
+const PREPARE_REQUEST_TIMEOUT_MS = 240_000;
 const DAEMON_STARTUP_TIMEOUT_MS = 15_000;
 const DAEMON_STARTUP_ATTEMPTS = 2;
 const DAEMON_TAKEOVER_TERM_TIMEOUT_MS = 3000;
@@ -197,16 +198,22 @@ export function resolveDaemonRequestTimeoutMs(
   req: Omit<DaemonRequest, 'token'>,
 ): number | undefined {
   if (req.command === PUBLIC_COMMANDS.test) return undefined;
-  if (
-    (req.command === PUBLIC_COMMANDS.replay || req.command === PUBLIC_COMMANDS.prepare) &&
-    typeof req.flags?.timeoutMs === 'number'
-  ) {
+  if (typeof req.flags?.timeoutMs === 'number' && isExplicitTimeoutCommand(req.command)) {
     return req.flags.timeoutMs;
   }
+  if (req.command === PUBLIC_COMMANDS.prepare) return PREPARE_REQUEST_TIMEOUT_MS;
   if (req.command === PUBLIC_COMMANDS.snapshot) return SNAPSHOT_REQUEST_TIMEOUT_MS;
   return REQUEST_TIMEOUT_MS;
 }
 
+function isExplicitTimeoutCommand(command: string | undefined): boolean {
+  return (
+    command === PUBLIC_COMMANDS.prepare ||
+    command === PUBLIC_COMMANDS.replay ||
+    command === PUBLIC_COMMANDS.snapshot
+  );
+}
+
 export async function openApp(options: OpenAppOptions = {}): Promise<DaemonResponse> {
   const {
     session = 'default',

src/utils/__tests__/args.test.ts

@@ -1320,14 +1320,18 @@ test('strict mode rejects click-only button flag on press', () => {
 });
 
 test('snapshot command accepts command-specific flags', () => {
-  const parsed = parseArgs(['snapshot', '-i', '-c', '--depth', '3', '-s', 'Login'], {
-    strictFlags: true,
-  });
+  const parsed = parseArgs(
+    ['snapshot', '-i', '-c', '--depth', '3', '-s', 'Login', '--timeout', '120000'],
+    {
+      strictFlags: true,
+    },
+  );
   assert.equal(parsed.command, 'snapshot');
   assert.equal(parsed.flags.snapshotInteractiveOnly, true);
   assert.equal(parsed.flags.snapshotCompact, true);
   assert.equal(parsed.flags.snapshotDepth, 3);
   assert.equal(parsed.flags.snapshotScope, 'Login');
+  assert.equal(parsed.flags.timeoutMs, 120000);
 });
 
 test('snapshot command accepts diff alias flag', () => {
@@ -1483,6 +1487,7 @@ test('snapshot command usage documents diff alias', () => {
   const help = usageForCommand('snapshot');
   if (help === null) throw new Error('Expected command help text');
   assert.match(help, /agent-device snapshot \[--diff\]/);
+  assert.match(help, /--timeout <ms>/);
   assert.match(help, /Capture accessibility tree or diff against the previous session baseline/);
 });
 

src/utils/__tests__/daemon-client.test.ts

@@ -166,7 +166,7 @@ test('snapshot request timeout preserves daemon metadata for follow-up evidence
   assert.equal(shouldResetDaemonAfterRequestTimeout(undefined), true);
 });
 
-test('snapshot uses a shorter daemon request timeout than runner prepare', () => {
+test('snapshot uses a shorter daemon request timeout with an explicit override', () => {
   const base = {
     session: 'default',
     positionals: [],
@@ -175,7 +175,23 @@ test('snapshot uses a shorter daemon request timeout than runner prepare', () =>
   };
 
   assert.equal(resolveDaemonRequestTimeoutMs({ ...base, command: 'snapshot' }), 30_000);
+  assert.equal(
+    resolveDaemonRequestTimeoutMs({
+      ...base,
+      command: 'snapshot',
+      flags: { timeoutMs: 120_000 },
+    }),
+    120_000,
+  );
   assert.equal(resolveDaemonRequestTimeoutMs({ ...base, command: 'screenshot' }), 90_000);
+  assert.equal(
+    resolveDaemonRequestTimeoutMs({
+      ...base,
+      command: 'prepare',
+      positionals: ['ios-runner'],
+    }),
+    240_000,
+  );
   assert.equal(
     resolveDaemonRequestTimeoutMs({
       ...base,

src/utils/cli-command-overrides.ts

@@ -116,9 +116,10 @@ const CLI_COMMAND_OVERRIDES = {
     positionalArgs: ['bundleOrPackage', 'payloadOrJson'],
   },
   snapshot: {
-    usageOverride: 'snapshot [--diff] [-i] [-c] [-d <depth>] [-s <scope>] [--raw] [--force-full]',
+    usageOverride:
+      'snapshot [--diff] [-i] [-c] [-d <depth>] [-s <scope>] [--raw] [--force-full] [--timeout <ms>]',
     helpDescription: 'Capture accessibility tree or diff against the previous session baseline',
-    allowedFlags: ['snapshotDiff', ...SNAPSHOT_FLAGS, 'snapshotForceFull'],
+    allowedFlags: ['snapshotDiff', ...SNAPSHOT_FLAGS, 'snapshotForceFull', 'timeoutMs'],
   },
   diff: {
     usageOverride:

src/utils/cli-flags.ts

@@ -800,7 +800,7 @@ const FLAG_DEFINITIONS: readonly FlagDefinition[] = [
     type: 'int',
     min: 1,
     usageLabel: '--timeout <ms>',
-    usageDescription: 'Prepare/Replay/Test: maximum wall-clock time for the command or attempt',
+    usageDescription: 'Prepare/Replay/Snapshot/Test: maximum wall-clock time for the command or attempt',
   },
   {
     key: 'retries',