fix: report snapshot timing diagnostics (#798)

* fix: report snapshot timing diagnostics

* fix: refine snapshot diagnostics reporting

* fix: preserve replay snapshot diagnostics on failure

* refactor: simplify snapshot diagnostics plumbing

---------

Co-authored-by: Michał Pierzchała <thymikee@gmail.com>

Author: pupuking723

src/__tests__/client-shared.test.ts

@@ -151,3 +151,31 @@ test('serializeSnapshotResult maps capture quality annotation to public snapshot
     snapshotQuality,
   });
 });
+
+test('serializeSnapshotResult includes snapshot diagnostics', () => {
+  const snapshotDiagnostics = {
+    stats: {
+      count: 3,
+      p50Ms: 450,
+      p95Ms: 1_800,
+      maxMs: 1_800,
+      slowThresholdMs: 1_500,
+      platform: 'android',
+    },
+    warning: 'Warning: android snapshots are slow in this run: p95 1800ms over 3 captures.',
+  } as const;
+  const data = serializeSnapshotResult({
+    nodes: [],
+    truncated: false,
+    snapshotDiagnostics,
+    identifiers: {
+      session: 'qa',
+    },
+  });
+
+  assert.deepEqual(data, {
+    nodes: [],
+    truncated: false,
+    snapshotDiagnostics,
+  });
+});

src/__tests__/snapshot-diagnostics.test.ts

@@ -0,0 +1,60 @@
+import { expect, test } from 'vitest';
+import {
+  mergeSnapshotDiagnostics,
+  recordSnapshotTiming,
+  summarizeSnapshotDiagnostics,
+} from '../snapshot-diagnostics.ts';
+
+test('records session snapshot timing stats', () => {
+  const session = {};
+
+  recordSnapshotTiming(session, { durationMs: 400, backend: 'android', platform: 'android' });
+  recordSnapshotTiming(session, { durationMs: 2_100, backend: 'android', platform: 'android' });
+
+  expect(summarizeSnapshotDiagnostics(session)).toEqual({
+    stats: {
+      count: 2,
+      p50Ms: 400,
+      p95Ms: 2_100,
+      maxMs: 2_100,
+      slowThresholdMs: 1_500,
+      platform: 'android',
+      backends: { android: 2 },
+    },
+    warning: expect.stringContaining('p95 2100ms over 2 captures'),
+  });
+});
+
+test('merges snapshot diagnostics without inflating capture count', () => {
+  const merged = mergeSnapshotDiagnostics([
+    {
+      stats: {
+        count: 1,
+        p50Ms: 300,
+        p95Ms: 300,
+        maxMs: 300,
+        slowThresholdMs: 1_500,
+        platform: 'android',
+      },
+    },
+    {
+      stats: {
+        count: 2,
+        p50Ms: 500,
+        p95Ms: 1_900,
+        maxMs: 1_900,
+        slowThresholdMs: 1_500,
+        platform: 'android',
+      },
+    },
+  ]);
+
+  expect(merged?.stats).toMatchObject({
+    count: 3,
+    p50Ms: 500,
+    p95Ms: 1_900,
+    maxMs: 1_900,
+    platform: 'android',
+  });
+  expect(merged?.warning).toContain('p95 1900ms over 3 captures');
+});

src/backend.ts

@@ -15,6 +15,7 @@ import type { ClickButton } from './core/click-button.ts';
 import type { DeviceRotation } from './core/device-rotation.ts';
 import type { ScrollDirection } from './core/scroll-gesture.ts';
 import type { SessionSurface } from './core/session-surface.ts';
+import type { SnapshotDiagnosticsSummary } from './snapshot-diagnostics.ts';
 import type {
   SnapshotCaptureAnalysis,
   SnapshotCaptureAnnotations,
@@ -49,6 +50,7 @@ export type BackendSnapshotResult = {
   snapshot?: SnapshotState;
   appName?: string;
   appBundleId?: string;
+  snapshotDiagnostics?: SnapshotDiagnosticsSummary;
 } & SnapshotCaptureAnnotations;
 
 export type BackendSnapshotOptions = SnapshotOptions & {

src/client-shared.ts

@@ -178,6 +178,7 @@ export function serializeSnapshotResult(result: CaptureSnapshotResult): Record<s
     ...(result.visibility ? { visibility: result.visibility } : {}),
     ...publicSnapshotCaptureAnnotations(snapshotResultAnnotations(result)),
     ...(result.unchanged ? { unchanged: result.unchanged } : {}),
+    ...(result.snapshotDiagnostics ? { snapshotDiagnostics: result.snapshotDiagnostics } : {}),
   };
 }
 

src/client-types.ts

@@ -1,4 +1,5 @@
 import type { PublicSnapshotCaptureAnnotations } from './snapshot-capture-annotations.ts';
+import type { SnapshotDiagnosticsSummary } from './snapshot-diagnostics.ts';
 import type {
   DaemonResponseData,
   DaemonInstallSource,
@@ -336,6 +337,7 @@ export type CaptureSnapshotResult = {
   appBundleId?: string;
   visibility?: SnapshotVisibility;
   unchanged?: SnapshotUnchanged;
+  snapshotDiagnostics?: SnapshotDiagnosticsSummary;
   identifiers: AgentDeviceIdentifiers;
 } & PublicSnapshotCaptureAnnotations;
 

src/client.ts

@@ -43,6 +43,7 @@ import type {
   MetroPrepareOptions,
 } from './client-types.ts';
 import { readSerializedSnapshotCaptureAnnotations } from './snapshot-capture-annotations.ts';
+import { readSnapshotDiagnosticsSummary } from './snapshot-diagnostics.ts';
 
 export function createAgentDeviceClient(
   config: AgentDeviceClientConfig = {},
@@ -343,15 +344,22 @@ function optionalSnapshotResponseFields(
 ): Partial<
   Pick<
     CaptureSnapshotResult,
-    'androidSnapshot' | 'unchanged' | 'visibility' | 'warnings' | 'snapshotQuality'
+    | 'androidSnapshot'
+    | 'unchanged'
+    | 'visibility'
+    | 'warnings'
+    | 'snapshotQuality'
+    | 'snapshotDiagnostics'
   >
 > {
   const visibility = readObject(data.visibility);
   const unchanged = readObject(data.unchanged);
+  const snapshotDiagnostics = readSnapshotDiagnosticsSummary(data.snapshotDiagnostics);
   return {
     ...(visibility ? { visibility: visibility as CaptureSnapshotResult['visibility'] } : {}),
     ...readSerializedSnapshotCaptureAnnotations(data),
     ...(unchanged ? { unchanged: unchanged as CaptureSnapshotResult['unchanged'] } : {}),
+    ...(snapshotDiagnostics ? { snapshotDiagnostics } : {}),
   };
 }
 

src/commands/capture/index.test.ts

@@ -12,6 +12,7 @@ import {
   waitCliReader,
   waitDaemonWriter,
 } from './index.ts';
+import { snapshotCliOutput } from './output.ts';
 
 function flags(overrides: Partial<CliFlags> = {}): CliFlags {
   return overrides as CliFlags;
@@ -50,6 +51,32 @@ describe('capture command interface', () => {
     });
   });
 
+  test('routes snapshot diagnostics warning to stderr output', () => {
+    const output = snapshotCliOutput({
+      result: {
+        nodes: [],
+        truncated: false,
+        identifiers: {},
+        snapshotDiagnostics: {
+          stats: {
+            count: 2,
+            p50Ms: 400,
+            p95Ms: 1_900,
+            maxMs: 1_900,
+            slowThresholdMs: 1_500,
+            platform: 'ios',
+          },
+          warning: 'Warning: ios snapshots are slow in this run: p95 1900ms over 2 captures.',
+        },
+      },
+    });
+
+    expect(output.stderr).toBe(
+      'Warning: ios snapshots are slow in this run: p95 1900ms over 2 captures.\n',
+    );
+    expect(output.text).not.toContain('snapshots are slow');
+  });
+
   test('reads screenshot path and writes screenshot flags', () => {
     const input = screenshotCliReader(
       ['page.png'],

src/commands/capture/output.ts

@@ -16,6 +16,9 @@ export function snapshotCliOutput(params: {
     data,
     // Programmatic SDK callers can see `unchanged`; CLI --json hides it for schema compatibility.
     jsonData: withoutUnchanged(data),
+    stderr: params.result.snapshotDiagnostics?.warning
+      ? `${params.result.snapshotDiagnostics.warning}\n`
+      : undefined,
     text: formatSnapshotText(data, {
       raw: params.raw,
       flatten: params.interactiveOnly,

src/commands/capture/runtime/snapshot.ts

@@ -1,4 +1,5 @@
 import type { BackendSnapshotResult } from '../../../backend.ts';
+import type { SnapshotDiagnosticsSummary } from '../../../snapshot-diagnostics.ts';
 import type { AgentDeviceRuntime, CommandSessionRecord } from '../../../runtime-contract.ts';
 import {
   publicSnapshotCaptureAnnotations,
@@ -38,6 +39,7 @@ export type SnapshotCommandResult = {
   appBundleId?: string;
   visibility?: SnapshotVisibility;
   unchanged?: SnapshotUnchanged;
+  snapshotDiagnostics?: SnapshotDiagnosticsSummary;
 } & PublicSnapshotCaptureAnnotations;
 
 export type DiffSnapshotCommandResult = {
@@ -84,6 +86,9 @@ export const snapshotCommand: RuntimeCommand<
       warnings: capture.warnings,
     }),
     ...(unchanged ? { unchanged } : {}),
+    ...(capture.result.snapshotDiagnostics
+      ? { snapshotDiagnostics: capture.result.snapshotDiagnostics }
+      : {}),
     ...snapshotAppFields(capture),
   };
 };

src/daemon/handlers/__tests__/session-replay-vars.test.ts

@@ -8,6 +8,7 @@ import { runCmdBackground, type ExecBackgroundResult } from '../../../utils/exec
 import type { DaemonInvokeFn, DaemonRequest, DaemonResponse, SessionAction } from '../../types.ts';
 import type { CommandFlags } from '../../../core/dispatch.ts';
 import { SessionStore } from '../../session-store.ts';
+import { makeIosSession } from '../../../__tests__/test-utils/index.ts';
 import {
   buildReplayVarScope,
   collectReplayShellEnv,
@@ -475,6 +476,110 @@ test('runReplayScriptFile dispatches resolved literals with file env overridden
   }
 });
 
+test('runReplayScriptFile reports snapshot diagnostics from per-action session samples', async () => {
+  const root = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-device-replay-snapshot-samples-'));
+  const scriptPath = path.join(root, 'flow.ad');
+  fs.writeFileSync(scriptPath, ['snapshot', 'snapshot', ''].join('\n'));
+  const sessionStore = new SessionStore(path.join(root, 'state'));
+  sessionStore.set(
+    's',
+    makeIosSession('s', {
+      snapshotDiagnostics: { samples: [] },
+    }),
+  );
+  let captures = 0;
+
+  const response = await runReplayScriptFile({
+    req: {
+      token: 't',
+      session: 's',
+      command: 'replay',
+      positionals: [scriptPath],
+      meta: { cwd: root },
+    },
+    sessionName: 's',
+    logPath: path.join(root, 'log'),
+    sessionStore,
+    invoke: async (): Promise<DaemonResponse> => {
+      captures += 1;
+      const session = sessionStore.get('s');
+      session?.snapshotDiagnostics?.samples.push({
+        durationMs: captures === 1 ? 400 : 1_900,
+        backend: 'xctest',
+        platform: 'ios',
+      });
+      return {
+        ok: true,
+        data: {
+          snapshotDiagnostics: {
+            stats: {
+              count: captures,
+              p50Ms: captures === 1 ? 400 : 1_900,
+              p95Ms: captures === 1 ? 400 : 1_900,
+              maxMs: captures === 1 ? 400 : 1_900,
+              slowThresholdMs: 1_500,
+              platform: 'ios',
+            },
+          },
+        },
+      };
+    },
+  });
+
+  assert.equal(response.ok, true);
+  const diagnostics = response.data?.snapshotDiagnostics as
+    | { stats?: { count?: number }; warning?: string }
+    | undefined;
+  assert.equal(diagnostics?.stats?.count, 2);
+  assert.match(String(diagnostics?.warning), /p95 1900ms over 2 captures/);
+});
+
+test('runReplayScriptFile reports snapshot diagnostics on replay failure', async () => {
+  const root = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-device-replay-snapshot-failure-'));
+  const scriptPath = path.join(root, 'flow.ad');
+  fs.writeFileSync(scriptPath, ['snapshot', 'click "Missing"', ''].join('\n'));
+  const sessionStore = new SessionStore(path.join(root, 'state'));
+  sessionStore.set(
+    's',
+    makeIosSession('s', {
+      snapshotDiagnostics: { samples: [] },
+    }),
+  );
+  let captures = 0;
+
+  const response = await runReplayScriptFile({
+    req: {
+      token: 't',
+      session: 's',
+      command: 'replay',
+      positionals: [scriptPath],
+      meta: { cwd: root },
+    },
+    sessionName: 's',
+    logPath: path.join(root, 'log'),
+    sessionStore,
+    invoke: async (): Promise<DaemonResponse> => {
+      captures += 1;
+      const session = sessionStore.get('s');
+      session?.snapshotDiagnostics?.samples.push({
+        durationMs: captures === 1 ? 450 : 2_100,
+        backend: 'xctest',
+        platform: 'ios',
+      });
+      if (captures === 1) return { ok: true, data: {} };
+      return { ok: false, error: { code: 'COMMAND_FAILED', message: 'button missing' } };
+    },
+  });
+
+  assert.equal(response.ok, false);
+  const diagnostics = response.error.details?.snapshotDiagnostics as
+    | { stats?: { count?: number; p95Ms?: number }; warning?: string }
+    | undefined;
+  assert.equal(diagnostics?.stats?.count, 2);
+  assert.equal(diagnostics?.stats?.p95Ms, 2_100);
+  assert.match(String(diagnostics?.warning), /p95 2100ms over 2 captures/);
+});
+
 test('runReplayScriptFile applies CLI env overrides before Maestro compat mapping', async () => {
   const { response, calls } = await runReplayFixture({
     label: 'maestro-env',