fix: keep request preparation in its session lifetime

src/daemon/__tests__/request-recording-health.test.ts

@@ -1,5 +1,11 @@
 import { test, expect, vi, beforeEach } from 'vitest';
 import type { SessionState } from '../session-state.ts';
+import { makeSessionStore } from '../../__tests__/test-utils/store-factory.ts';
+import {
+  createRequestExecutionScope,
+  prepareLockedRequestScope,
+} from '../request-execution-scope.ts';
+import { LeaseRegistry } from '../lease-registry.ts';
 import { makeTestScreenRecordingResource } from '../../__tests__/test-utils/screen-recording-live-handle.ts';
 
 vi.mock('../../platform-runtime-apple-resources.ts', async (importOriginal) => ({
@@ -54,7 +60,9 @@ test('runner-backed iOS recordings still invalidate on runner restarts', async (
     sessionId: 'runner-after',
   });
 
-  await refreshRecordingHealth(session);
+  const store = makeSessionStore();
+  const ref = store.publish(session.name, session);
+  await refreshRecordingHealth(store, ref);
 
   expect(mockObserveRunnerSession).toHaveBeenCalledWith('sim-1');
   expect(session.screenRecording?.handle.inspect().invalidatedReason).toBe(
@@ -78,7 +86,9 @@ test.each([
   });
   mockObserveRunnerSession.mockResolvedValue(snapshot);
 
-  await refreshRecordingHealth(session);
+  const store = makeSessionStore();
+  const ref = store.publish(session.name, session);
+  await refreshRecordingHealth(store, ref);
 
   expect(session.screenRecording.handle.inspect().invalidatedReason).toBe(reason);
 });
@@ -91,9 +101,138 @@ test('a recording without a runner identity adopts the first live observation',
   });
   mockObserveRunnerSession.mockResolvedValue({ alive: true, sessionId: 'runner-first' });
 
-  await refreshRecordingHealth(session);
+  const store = makeSessionStore();
+  const ref = store.publish(session.name, session);
+  await refreshRecordingHealth(store, ref);
 
   const recording = session.screenRecording.handle.inspect();
   expect(recording.runnerSessionId).toBe('runner-first');
   expect(recording.invalidatedReason).toBeUndefined();
 });
+
+test.each(['rebuild', 'retire', 'handle', 'token', 'generation'] as const)(
+  'a held health observation respects the current lifetime and resource: %s',
+  async (change) => {
+    const store = makeSessionStore();
+    const session = makeIosSimulatorSession(true);
+    const active = makeTestScreenRecordingResource(session, {
+      backend: 'runner AVAssetWriter',
+      showTouches: true,
+      runnerSessionId: 'runner-before',
+    });
+    session.screenRecording = active;
+    const ref = store.publish('default', session);
+    let finish!: (value: { alive: boolean; sessionId: string }) => void;
+    mockObserveRunnerSession.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          finish = resolve;
+        }),
+    );
+    const observation = refreshRecordingHealth(store, ref);
+    expect(mockObserveRunnerSession).toHaveBeenCalledWith('sim-1');
+    if (change === 'rebuild') {
+      store.update(ref, { appName: 'Intervening app' });
+    } else if (change === 'retire') {
+      store.retire(ref);
+      store.publish('default', session);
+    } else {
+      const replacement = makeTestScreenRecordingResource(session, {
+        backend: 'runner AVAssetWriter',
+        showTouches: true,
+        runnerSessionId: 'successor-runner',
+      });
+      store.update(ref, {
+        screenRecording: {
+          ...active,
+          handle: change === 'handle' ? replacement.handle : active.handle,
+          envelope: {
+            ...active.envelope,
+            fence: {
+              ...active.envelope.fence,
+              token: change === 'token' ? 'new-token' : active.envelope.fence.token,
+              generation: active.envelope.fence.generation + (change === 'generation' ? 1 : 0),
+            },
+          },
+        },
+      });
+    }
+    const current = store.get('default')!;
+    finish({ alive: true, sessionId: 'runner-after' });
+    await observation;
+    expect(store.get('default')).toBe(current);
+    expect(active.handle.inspect().invalidatedReason).toBe(
+      change === 'rebuild' ? 'iOS runner session restarted during recording' : undefined,
+    );
+    if (change === 'rebuild') expect(current.appName).toBe('Intervening app');
+    else expect(current.screenRecording?.handle.inspect().invalidatedReason).toBeUndefined();
+  },
+);
+
+test.each(['rebuild', 'retire'] as const)(
+  'locked request preparation keeps its captured lifetime after runner observation: %s',
+  async (change) => {
+    const store = makeSessionStore();
+    const session = makeIosSimulatorSession(true);
+    session.screenRecording = makeTestScreenRecordingResource(session, {
+      backend: 'runner AVAssetWriter',
+      showTouches: true,
+      runnerSessionId: 'runner-before',
+    });
+    const ref = store.publish('default', session);
+    let observed!: () => void;
+    const started = new Promise<void>((resolve) => {
+      observed = resolve;
+    });
+    let finish!: (value: { alive: boolean; sessionId: string }) => void;
+    mockObserveRunnerSession.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          finish = resolve;
+          observed();
+        }),
+    );
+    await using scope = await createRequestExecutionScope({
+      req: { token: 'token', session: 'default', command: 'snapshot', positionals: [] },
+      sessionStore: store,
+      leaseRegistry: new LeaseRegistry(),
+    });
+    const prepared = scope.runLocked(() =>
+      prepareLockedRequestScope({
+        scope,
+        sessionStore: store,
+        trackDownloadableArtifact: () => 'artifact',
+      }),
+    );
+    const outcome = prepared.then(
+      (value) => ({ value }),
+      (error) => ({ error }),
+    );
+    await started;
+    let current;
+    if (change === 'rebuild') current = store.update(ref, { appName: 'Latest app' });
+    else {
+      store.retire(ref);
+      current = makeIosSimulatorSession(false);
+      store.publish('default', current);
+      store.setRuntimeHints('default', { metroPort: 8083 });
+    }
+    finish({ alive: true, sessionId: 'runner-before' });
+    const result = await outcome;
+    expect(store.get('default')).toBe(current);
+    if (change === 'rebuild') {
+      expect(result).toMatchObject({
+        value: { type: 'scope', scope: { existingSession: current } },
+      });
+      expect(current.appName).toBe('Latest app');
+      if ('value' in result && result.value.type === 'scope') {
+        store.retire(ref);
+        store.publish('default', { ...makeIosSimulatorSession(false), surface: 'app' });
+        expect(result.value.scope.handlerContextFromFlags(undefined).surface).toBeUndefined();
+      }
+    } else {
+      expect(result).toMatchObject({ error: { details: { reason: 'session_lifetime_ended' } } });
+      expect(store.getRuntimeHints('default')).toEqual({ metroPort: 8083 });
+    }
+  },
+);

src/daemon/request-binding.ts

@@ -64,10 +64,9 @@ export async function resolveRequestExecutionLockPlan(params: {
 
 export function prepareLockedRequestBinding(params: {
   req: DaemonRequest;
-  sessionName: string;
-  sessionStore: SessionStore;
+  existingRef: SessionRef | undefined;
 }): LockedRequestBinding {
-  const existingRef = params.sessionStore.lookup(params.sessionName);
+  const { existingRef } = params;
   return {
     req: applyRequestLockPolicy(params.req, existingRef),
     existingRef,

src/daemon/request-execution-scope.ts

@@ -486,16 +486,15 @@ export async function prepareLockedRequestScope(params: {
   const { scope, sessionStore, trackDownloadableArtifact } = params;
   const logPath = scope.runnerLogPath;
   scope.throwIfCanceled();
-  const seededSession = sessionStore.get(scope.sessionName);
-  if (seededSession) {
-    // Called under runLocked: refreshRecordingHealth may mutate session recording state.
-    await refreshRecordingHealth(seededSession);
-    sessionStore.set(scope.sessionName, seededSession);
+  const seededRef = sessionStore.lookup(scope.sessionName);
+  if (seededRef) {
+    await refreshRecordingHealth(sessionStore, seededRef);
+    scope.throwIfCanceled();
+    sessionStore.requireCurrent(seededRef);
   }
   const binding = prepareLockedRequestBinding({
     req: scope.req,
-    sessionName: scope.sessionName,
-    sessionStore,
+    existingRef: seededRef ? sessionStore.refresh(seededRef) : undefined,
   });
   const lockedReq = binding.req;
   // `scope.sessionName` is the resolved store key, so `existingRef` carries the address every
@@ -563,7 +562,10 @@ export async function prepareLockedRequestScope(params: {
         ({
           ...contextFromFlags(flags, appBundleId, traceLogPath),
           // Handlers may update surface during the request, so read the current session state.
-          surface: sessionStore.get(scope.sessionName)?.surface,
+          surface: (seededRef
+            ? sessionStore.resolveCurrent(seededRef)
+            : sessionStore.get(scope.sessionName)
+          )?.surface,
         }) satisfies DaemonCommandContext,
     },
   };

src/daemon/request-recording-health.ts

@@ -1,15 +1,19 @@
 import { isIosFamily } from '@agent-device/kernel/device';
 import { appleSessionObservation } from '../platform-runtime-apple-resources.ts';
-import type { SessionState } from './session-state.ts';
+import type { SessionRef, SessionState } from './session-state.ts';
+import type { SessionStore } from './session-store.ts';
 
-export async function refreshRecordingHealth(session: SessionState): Promise<void> {
+export async function refreshRecordingHealth(store: SessionStore, ref: SessionRef): Promise<void> {
+  const session = store.requireCurrent(ref);
   if (!recordingRequiresRunnerHealth(session)) {
     return;
   }
-  const recording = session.screenRecording!.handle;
-  const state = recording.inspect();
+  const resource = session.screenRecording!;
+  const recording = resource.handle;
 
   const snapshot = await appleSessionObservation.observeRunnerSession(session.device.id);
+  if (store.resolveCurrent(ref)?.screenRecording !== resource) return;
+  const state = recording.inspect();
   if (!state.runnerSessionId) {
     if (snapshot?.alive) {
       recording.setRunnerSessionId(snapshot.sessionId);