fix: preserve scoped Android snapshot refs

Author: thymikee

src/daemon/handlers/__tests__/snapshot-handler.test.ts

@@ -153,6 +153,145 @@ test('snapshot surfaces filtered-to-zero Android guidance for interactive snapsh
   }
 });
 
+test('snapshot resolves @ref scope with the stored source after scoped output replaces refs', async () => {
+  const sessionStore = makeSessionStore();
+  const sessionName = 'android-ref-scope-repeat';
+  const session = makeSession(sessionName, androidDevice);
+  session.snapshot = {
+    nodes: [
+      {
+        ref: 'e1',
+        index: 0,
+        depth: 0,
+        type: 'android.widget.FrameLayout',
+        rect: { x: 0, y: 0, width: 390, height: 844 },
+      },
+      {
+        ref: 'e2',
+        index: 1,
+        depth: 1,
+        parentIndex: 0,
+        type: 'androidx.recyclerview.widget.RecyclerView',
+        rect: { x: 0, y: 80, width: 390, height: 600 },
+      },
+      {
+        ref: 'e3',
+        index: 2,
+        depth: 2,
+        parentIndex: 1,
+        type: 'android.widget.TextView',
+        label: 'Unable to load script. Make sure you are running Metro.',
+        value: 'Unable to load script. Make sure you are running Metro.',
+        rect: { x: 16, y: 120, width: 358, height: 200 },
+      },
+      {
+        ref: 'e4',
+        index: 3,
+        depth: 3,
+        parentIndex: 2,
+        type: 'android.widget.TextView',
+        label: 'loadJSBundleFromAssets',
+        rect: { x: 16, y: 140, width: 358, height: 40 },
+      },
+    ],
+    createdAt: Date.now(),
+    backend: 'android',
+  };
+  sessionStore.set(sessionName, session);
+
+  mockDispatch.mockResolvedValue({
+    nodes: [
+      {
+        index: 0,
+        depth: 0,
+        type: 'android.widget.TextView',
+        label: 'Unable to load script. Make sure you are running Metro.',
+        value: 'Unable to load script. Make sure you are running Metro.',
+        rect: { x: 16, y: 120, width: 358, height: 200 },
+      },
+      {
+        index: 1,
+        depth: 1,
+        parentIndex: 0,
+        type: 'android.widget.TextView',
+        label: 'loadJSBundleFromAssets',
+        rect: { x: 16, y: 140, width: 358, height: 40 },
+      },
+    ],
+    truncated: false,
+    backend: 'android',
+  });
+
+  for (let attempt = 0; attempt < 2; attempt += 1) {
+    const response = await handleSnapshotCommands({
+      req: {
+        token: 't',
+        session: sessionName,
+        command: 'snapshot',
+        positionals: [],
+        flags: { snapshotScope: '@e3' },
+      },
+      sessionName,
+      logPath: '/tmp/daemon.log',
+      sessionStore,
+    });
+
+    expect(response?.ok).toBe(true);
+    if (response?.ok) {
+      expect(response.data?.nodes).toHaveLength(2);
+    }
+  }
+
+  expect(mockDispatch).toHaveBeenCalledTimes(2);
+  for (const call of mockDispatch.mock.calls) {
+    expect((call[4] as { snapshotScope?: string })?.snapshotScope).toBe(
+      'Unable to load script. Make sure you are running Metro.',
+    );
+  }
+  const updated = sessionStore.get(sessionName);
+  expect(updated?.snapshot?.nodes).toHaveLength(2);
+  expect(updated?.snapshotScopeSource?.nodes[2]?.ref).toBe('e3');
+});
+
+test('empty @ref-scoped snapshot output does not replace the stored session snapshot', async () => {
+  const sessionStore = makeSessionStore();
+  const sessionName = 'android-empty-scope-preserve';
+  const session = makeSession(sessionName, androidDevice);
+  session.snapshot = {
+    nodes: [
+      { ref: 'e1', index: 0, depth: 0, type: 'android.widget.TextView', label: 'Current screen' },
+    ],
+    createdAt: Date.now(),
+    backend: 'android',
+  };
+  sessionStore.set(sessionName, session);
+
+  mockDispatch.mockResolvedValue({
+    nodes: [],
+    truncated: false,
+    backend: 'android',
+  });
+
+  const response = await handleSnapshotCommands({
+    req: {
+      token: 't',
+      session: sessionName,
+      command: 'snapshot',
+      positionals: [],
+      flags: { snapshotScope: '@e1' },
+    },
+    sessionName,
+    logPath: '/tmp/daemon.log',
+    sessionStore,
+  });
+
+  expect(response?.ok).toBe(true);
+  if (response?.ok) {
+    expect(response.data?.nodes).toEqual([]);
+  }
+  expect(sessionStore.get(sessionName)?.snapshot?.nodes[0]?.label).toBe('Current screen');
+});
+
 test('snapshot warns when recent snapshot node count collapses sharply', async () => {
   const sessionStore = makeSessionStore();
   const sessionName = 'android-stale-collapse';

src/daemon/handlers/find.ts

@@ -9,6 +9,7 @@ import { ensureDeviceReady } from '../device-ready.ts';
 import { extractNodeText, findNearestHittableAncestor } from '../snapshot-processing.ts';
 import { readTextForNode } from './interaction-read.ts';
 import { captureSnapshot } from './snapshot-capture.ts';
+import { setSessionSnapshot } from '../session-snapshot.ts';
 import { errorResponse } from './response.ts';
 import { getActiveAndroidSnapshotFreshness } from '../android-snapshot-freshness.ts';
 import { dispatchFindReadOnlyViaRuntime } from '../selector-runtime.ts';
@@ -109,7 +110,7 @@ export async function handleFindCommands(params: {
     lastSnapshotAt = now;
     lastNodes = nodes;
     if (session) {
-      session.snapshot = snapshot;
+      setSessionSnapshot(session, snapshot);
       sessionStore.set(sessionName, session);
     }
     return { nodes, truncated: snapshot.truncated, backend: snapshot.backend };

src/daemon/handlers/interaction-runtime.ts

@@ -8,6 +8,7 @@ import { createAgentDevice, localCommandPolicy } from '../../runtime.ts';
 import { AppError } from '../../utils/errors.ts';
 import type { SessionState } from '../types.ts';
 import { createUnsupportedArtifactAdapter } from '../runtime-artifacts.ts';
+import { setSessionSnapshot } from '../session-snapshot.ts';
 import type { InteractionHandlerParams } from './interaction-common.ts';
 import type { CaptureSnapshotForSession } from './interaction-snapshot.ts';
 
@@ -34,7 +35,7 @@ export function createInteractionRuntime(
           : undefined,
       set: (record) => {
         if (!record.snapshot) return;
-        session.snapshot = record.snapshot;
+        setSessionSnapshot(session, record.snapshot);
         params.sessionStore.set(params.sessionName, session);
       },
     },

src/daemon/handlers/interaction-snapshot.ts

@@ -4,6 +4,7 @@ import type { SessionState } from '../types.ts';
 import type { SnapshotState } from '../../utils/snapshot.ts';
 import type { ContextFromFlags } from './interaction-common.ts';
 import { captureSnapshot } from './snapshot-capture.ts';
+import { setSessionSnapshot } from '../session-snapshot.ts';
 
 export type CaptureSnapshotForSession = (
   session: SessionState,
@@ -38,7 +39,7 @@ export async function captureSnapshotForSession(
     logPath: dispatchContext.logPath ?? '',
     androidFreshnessMode: options.androidFreshnessMode,
   });
-  session.snapshot = snapshot;
+  setSessionSnapshot(session, snapshot);
   sessionStore.set(session.name, session);
-  return session.snapshot;
+  return snapshot;
 }

src/daemon/handlers/session-replay-heal.ts

@@ -1,4 +1,5 @@
 import { dispatchCommand } from '../../core/dispatch.ts';
+import { setSessionSnapshot } from '../session-snapshot.ts';
 import {
   attachRefs,
   type RawSnapshotNode,
@@ -198,7 +199,7 @@ async function captureSnapshotForReplay(
     createdAt: Date.now(),
     backend: data?.backend,
   };
-  session.snapshot = snapshot;
+  setSessionSnapshot(session, snapshot);
   sessionStore.set(session.name, session);
   return snapshot;
 }

src/daemon/handlers/snapshot-capture.ts

@@ -60,7 +60,7 @@ export async function captureSnapshot(params: CaptureSnapshotParams): Promise<{
   const data = await captureSnapshotData(params);
   clearAndroidSnapshotFreshness(params.session);
   return {
-    snapshot: buildSnapshotState(data, params.flags),
+    snapshot: buildSnapshotState(data, resolveSnapshotStateFlags(params)),
     analysis: data.analysis,
   };
 }
@@ -146,7 +146,19 @@ async function captureSnapshotAttempt(
   const data = await captureSnapshotData(params);
   return {
     data,
-    snapshot: buildSnapshotState(data, params.flags),
+    snapshot: buildSnapshotState(data, resolveSnapshotStateFlags(params)),
+  };
+}
+
+function resolveSnapshotStateFlags(
+  params: Pick<CaptureSnapshotParams, 'flags' | 'snapshotScope'>,
+): CaptureSnapshotParams['flags'] {
+  if (params.snapshotScope === undefined) {
+    return params.flags;
+  }
+  return {
+    ...params.flags,
+    snapshotScope: params.snapshotScope,
   };
 }
 
@@ -350,8 +362,16 @@ export function resolveSnapshotScope(
   if (!ref) {
     return errorResponse('INVALID_ARGS', `Invalid ref scope: ${snapshotScope}`);
   }
-  const node = findNodeByRef(session.snapshot.nodes, ref);
-  const resolved = node ? resolveRefLabel(node, session.snapshot.nodes) : undefined;
+  const candidates = [
+    session.snapshot,
+    ...(session.snapshotScopeSource ? [session.snapshotScopeSource] : []),
+  ];
+  let resolved: string | undefined;
+  for (const snapshot of candidates) {
+    const node = findNodeByRef(snapshot.nodes, ref);
+    resolved = node ? resolveRefLabel(node, snapshot.nodes) : undefined;
+    if (resolved) break;
+  }
   if (!resolved) {
     return errorResponse('COMMAND_FAILED', `Ref ${snapshotScope} not found or has no label`);
   }

src/daemon/request-router.ts

@@ -23,6 +23,7 @@ import { handleInteractionCommands } from './handlers/interaction.ts';
 import { handleLeaseCommands } from './handlers/lease.ts';
 import { buildSnapshotState, captureSnapshotData } from './handlers/snapshot-capture.ts';
 import { assertSessionSelectorMatches } from './session-selector.ts';
+import { setSessionSnapshot } from './session-snapshot.ts';
 import { applyRequestLockPolicy } from './request-lock-policy.ts';
 import { resolveEffectiveSessionName } from './session-routing.ts';
 import { normalizeTenantId, resolveSessionIsolationMode } from './config.ts';
@@ -473,7 +474,7 @@ async function applyScreenshotOverlay(
     snapshotScope: undefined,
   });
   const overlaySnapshot = buildSnapshotState(overlaySnapshotData, undefined);
-  session.snapshot = overlaySnapshot;
+  setSessionSnapshot(session, overlaySnapshot);
   const overlayRefs = await annotateScreenshotWithRefs({
     screenshotPath: data.path as string,
     snapshot: overlaySnapshot,

src/daemon/selector-runtime.ts

@@ -31,6 +31,7 @@ import type { IsCommandOptions } from '../commands/selector-read.ts';
 import { isSupportedPredicate } from './is-predicates.ts';
 import type { ContextFromFlags } from './handlers/interaction-common.ts';
 import { createUnsupportedArtifactAdapter } from './runtime-artifacts.ts';
+import { setSessionSnapshot } from './session-snapshot.ts';
 import { getActiveAndroidSnapshotFreshness } from './android-snapshot-freshness.ts';
 import {
   describeAndroidEscapeSurface,
@@ -232,7 +233,7 @@ function createSelectorRuntimeForDevice(params: {
       get: (name) => (name === params.sessionName ? toCommandSession(params.session) : undefined),
       set: (record) => {
         if (!params.session || !record.snapshot) return;
-        params.session.snapshot = record.snapshot;
+        setSessionSnapshot(params.session, record.snapshot);
         params.sessionStore.set(params.sessionName, params.session);
       },
     },
@@ -312,7 +313,7 @@ function createSelectorBackend(params: {
         snapshotScope,
       });
       if (session) {
-        session.snapshot = capture.snapshot;
+        setSessionSnapshot(session, capture.snapshot);
         sessionStore.set(sessionName, session);
       }
       lastSnapshotAt = timestamp;
@@ -405,7 +406,7 @@ async function captureWaitSnapshot(params: {
     logPath: params.logPath ?? '',
   });
   if (params.session) {
-    params.session.snapshot = capture.snapshot;
+    setSessionSnapshot(params.session, capture.snapshot);
     params.sessionStore.set(params.sessionName, params.session);
   }
   return capture.snapshot;

src/daemon/session-snapshot.ts

@@ -0,0 +1,7 @@
+import type { SnapshotState } from '../utils/snapshot.ts';
+import type { SessionState } from './types.ts';
+
+export function setSessionSnapshot(session: SessionState, snapshot: SnapshotState): void {
+  session.snapshot = snapshot;
+  session.snapshotScopeSource = undefined;
+}

src/daemon/snapshot-runtime.ts

@@ -138,13 +138,24 @@ function createSnapshotRuntime(params: {
           throw new AppError('UNKNOWN', 'snapshot runtime did not produce session state');
         }
         const current = sessionStore.get(sessionName);
+        const refScopedSnapshot = req.flags?.snapshotScope?.trim().startsWith('@') === true;
+        const keepCurrentSnapshot =
+          refScopedSnapshot && record.snapshot.nodes.length === 0 && current?.snapshot;
+        const snapshot = keepCurrentSnapshot ? current.snapshot : record.snapshot;
         const nextSession = buildSnapshotSession({
           session: current,
           sessionName,
           device,
-          snapshot: record.snapshot,
+          snapshot,
           appBundleId: record.appBundleId,
         });
+        if (!refScopedSnapshot) {
+          nextSession.snapshotScopeSource = undefined;
+        } else if (keepCurrentSnapshot) {
+          nextSession.snapshotScopeSource = current?.snapshotScopeSource;
+        } else {
+          nextSession.snapshotScopeSource = current?.snapshotScopeSource ?? current?.snapshot;
+        }
         if (record.appName) nextSession.appName = record.appName;
         sessionStore.set(sessionName, nextSession);
       },