fix: make replay session handoff reliable

Author: thymikee

packages/contracts/src/replay.ts

@@ -7,10 +7,10 @@ export type ReplayCommandResult = {
   session: string;
   /**
    * True iff `session` still exists in the daemon's session store when the
-   * response is built — i.e. the replayed script had no terminal `close`
-   * (ADR 0016's consumption contract). The client uses this, not script
-   * parsing, to decide whether an owned one-shot daemon must stay alive so
-   * the caller can keep addressing this session.
+   * response is built. This remains true when replay suppresses an authored
+   * terminal `close` for an explicit live-session handoff. The client uses
+   * this, not script parsing, to decide whether an owned one-shot daemon must
+   * stay alive so the caller can keep addressing this session.
    */
   sessionActive: boolean;
   artifactPaths: string[];

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

@@ -74,7 +74,60 @@ test('a close-less replay reports sessionActive: true (real producer, session st
   expect(response.ok).toBe(true);
   if (!response.ok) return;
   expect(sessionStore.get(sessionName)).toBeDefined();
-  expect((response.data as { sessionActive: boolean }).sessionActive).toBe(true);
+  expect(response.data).toMatchObject({ sessionActive: true, replayed: 2 });
+});
+
+test('--keep-session suppresses a close that is terminal among executable actions', async () => {
+  const root = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-device-replay-keep-marker-tail-'));
+  const sessionStore = new SessionStore(path.join(root, 'sessions'));
+  const sessionName = 'default';
+  sessionStore.set(sessionName, makeIosSession(sessionName));
+  const filePath = writeReplayFile(root, ['open "Demo"', 'close', 'replay "./nested-flow.ad"']);
+  const commands: string[] = [];
+
+  const response = await runReplayScriptFile({
+    req: baseReq({ positionals: [filePath], flags: { replayKeepSession: true } }),
+    sessionName,
+    logPath: path.join(root, 'daemon.log'),
+    sessionStore,
+    invoke: async (req) => {
+      commands.push(req.command);
+      if (req.command === 'close') sessionStore.delete(sessionName);
+      return { ok: true, data: {} };
+    },
+  });
+
+  expect(response.ok).toBe(true);
+  expect(commands).toEqual(['open']);
+  if (!response.ok) return;
+  expect(response.data).toMatchObject({ sessionActive: true, replayed: 1 });
+});
+
+test('--keep-session fails explicitly when the completed replay has no live session', async () => {
+  const root = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-device-replay-keep-postcondition-'));
+  const sessionStore = new SessionStore(path.join(root, 'sessions'));
+  const sessionName = 'default';
+  sessionStore.set(sessionName, makeIosSession(sessionName));
+  const filePath = writeReplayFile(root, ['open "Demo"', 'click "Log out"']);
+
+  const response = await runReplayScriptFile({
+    req: baseReq({ positionals: [filePath], flags: { replayKeepSession: true } }),
+    sessionName,
+    logPath: path.join(root, 'daemon.log'),
+    sessionStore,
+    invoke: async (req) => {
+      if (req.command === 'click') sessionStore.delete(sessionName);
+      return { ok: true, data: {} };
+    },
+  });
+
+  expect(response).toMatchObject({
+    ok: false,
+    error: {
+      code: 'COMMAND_FAILED',
+      message: expect.stringContaining('--keep-session could not preserve session'),
+    },
+  });
 });
 
 test('a replay whose terminal close removes the session reports sessionActive: false', async () => {
@@ -128,7 +181,7 @@ test('--keep-session suppresses only the authored terminal close and reports the
   expect(response.ok).toBe(true);
   if (!response.ok) return;
   expect(sessionStore.get(sessionName)).toBeDefined();
-  expect((response.data as { sessionActive: boolean }).sessionActive).toBe(true);
+  expect(response.data).toMatchObject({ sessionActive: true, replayed: 2 });
 });
 
 test('--keep-session preserves an interior close instead of broad command filtering', async () => {
@@ -225,7 +278,7 @@ test('Maestro YAML uses the typed engine while .ad remains generic', async () =>
   const yamlResponse = await runReplayScriptFile({
     req: baseReq({
       positionals: [yamlPath],
-      flags: { replayBackend: 'maestro', platform: 'ios' },
+      flags: { replayBackend: 'maestro', platform: 'ios', replayKeepSession: false },
     }),
     sessionName,
     logPath: path.join(root, 'daemon.log'),

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

@@ -8,6 +8,8 @@ import { LeaseRegistry } from '../../lease-registry.ts';
 import type { DaemonRequest, DaemonResponse } from '../../types.ts';
 import { makeIosSession } from '../../../__tests__/test-utils/index.ts';
 import { buildNestedReplayFlags, handleSessionReplayCommands } from '../session-replay.ts';
+import { REPLAY_ONLY_TEST_FLAG_REJECTIONS } from '../session-replay-test-policy.ts';
+import { replayCommandFamily } from '../../../commands/replay/index.ts';
 
 const recordTraceMocks = vi.hoisted(() => ({
   handleRecordCommand: vi.fn(),
@@ -209,7 +211,7 @@ test('buildNestedReplayFlags strips test-only recordVideo before replay actions
   assert.deepEqual(result, { platform: 'ios' });
 });
 
-test('test --record-video records each replay attempt on the generated test session', async () => {
+test('test normalizes false replay-only booleans while recording each replay attempt', async () => {
   vi.useFakeTimers({ now: 1_000 });
   const { root, replayPath, sessionStore, nestedRequests, events } = createRecordVideoFixture();
   installMockRecordingHandler(sessionStore, { recordingPath: '', events });
@@ -220,7 +222,13 @@ test('test --record-video records each replay attempt on the generated test sess
       session: 'default',
       command: 'test',
       positionals: [replayPath],
-      flags: { recordVideo: true, artifactsDir: path.join(root, 'artifacts') },
+      flags: {
+        recordVideo: true,
+        replayKeepSession: false,
+        saveScript: false,
+        force: false,
+        artifactsDir: path.join(root, 'artifacts'),
+      },
       meta: { cwd: root, requestId: 'record-video-suite' },
     },
     sessionName: 'default',
@@ -272,6 +280,20 @@ test('test --record-video records each replay attempt on the generated test sess
 
 // --- ADR 0012 decision 4 / migration step 5: `--from` is replay-only ---
 
+test('raw test-request guards enumerate every daemon-visible replay-only CLI flag', () => {
+  const replayFlags = replayCommandFamily.cliSchemas.replay?.allowedFlags ?? [];
+  const testFlags = new Set(replayCommandFamily.cliSchemas.test?.allowedFlags ?? []);
+  const clientOnlyReplayFlags = new Set(['out']);
+  const expectedDaemonFlags = replayFlags
+    .filter((flag) => !testFlags.has(flag) && !clientOnlyReplayFlags.has(flag))
+    .sort();
+
+  const guardedDaemonFlags = REPLAY_ONLY_TEST_FLAG_REJECTIONS.flatMap(
+    (rejection) => rejection.keys,
+  ).sort();
+  assert.deepEqual(guardedDaemonFlags, expectedDaemonFlags);
+});
+
 test('test rejects raw --keep-session with INVALID_ARGS before running the suite', async () => {
   const root = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-device-test-keep-session-rejected-'));
   const replayPath = path.join(root, 'flow.ad');
@@ -396,3 +418,34 @@ test('test rejects --save-script with INVALID_ARGS before running the suite', as
   assert.equal(response.error.code, 'INVALID_ARGS');
   assert.match(response.error.message, /--save-script/);
 });
+
+test('test rejects raw --force without --save-script before running the suite', async () => {
+  const root = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-device-test-force-rejected-'));
+  const replayPath = path.join(root, 'flow.ad');
+  fs.writeFileSync(replayPath, 'open "Demo"\n');
+  const sessionStore = new SessionStore(path.join(root, 'sessions'));
+  const invoke = vi.fn(async () => ({ ok: true as const, data: {} }));
+
+  const response = await handleSessionReplayCommands({
+    req: {
+      token: 'token',
+      session: 'default',
+      command: 'test',
+      positionals: [replayPath],
+      flags: { force: true },
+      meta: { cwd: root },
+    },
+    sessionName: 'default',
+    logPath: path.join(root, 'daemon.log'),
+    sessionStore,
+    leaseRegistry: new LeaseRegistry(),
+    invoke,
+  });
+
+  if (!response) throw new Error('Expected response');
+  assert.equal(response.ok, false);
+  if (response.ok) return;
+  assert.equal(response.error.code, 'INVALID_ARGS');
+  assert.match(response.error.message, /--force/);
+  assert.equal(invoke.mock.calls.length, 0);
+});

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

@@ -220,6 +220,7 @@ export async function runReplayScriptFile(params: {
   }
 
   const startedAt = Date.now();
+  const keepSession = req.flags?.replayKeepSession === true;
   let resolved = '';
   const artifactPaths = new Set<string>();
   // #1478 P4b: the one locked coordinator this request reaches the repair
@@ -231,7 +232,7 @@ export async function runReplayScriptFile(params: {
       return errorResponse('INVALID_ARGS', maestroBackendRequiredMessage('replay', filePath));
     }
     if (resolveReplayFormat(resolved, req.flags?.replayBackend) === 'maestro') {
-      if (req.flags?.replayKeepSession === true) {
+      if (keepSession) {
         return errorResponse(
           'INVALID_ARGS',
           '--keep-session is supported only for native .ad replay; Maestro YAML owns its lifecycle.',
@@ -252,6 +253,7 @@ export async function runReplayScriptFile(params: {
       tracePath,
       resolved,
       coordinator,
+      keepSession,
     });
     if (!planPreparation.ok) return planPreparation.response;
     const {
@@ -264,6 +266,7 @@ export async function runReplayScriptFile(params: {
       scope,
       actionTracePath,
       snapshotDiagnosticSamples,
+      suppressedTerminalCloseIndex,
     } = planPreparation.value;
     const sessionPreparation = prepareReplaySession({
       req,
@@ -308,6 +311,7 @@ export async function runReplayScriptFile(params: {
       snapshotDiagnosticSamples,
       onStep,
       armSaveScript: sessionPreparation.armSaveScript,
+      suppressedTerminalCloseIndex,
     });
     if (failure) return failure;
     return completeReplayRun({
@@ -320,6 +324,8 @@ export async function runReplayScriptFile(params: {
       snapshotDiagnosticSamples,
       armSaveScript: sessionPreparation.armSaveScript,
       coordinator,
+      keepSession,
+      suppressedTerminalCloseIndex,
     });
   } catch (err) {
     const appErr = asAppError(err);
@@ -348,6 +354,7 @@ type ReplayActionExecution = {
   snapshotDiagnosticSamples: SnapshotTimingSample[];
   onStep: ReplayTestAttemptStepSink | undefined;
   armSaveScript: () => void;
+  suppressedTerminalCloseIndex: number | undefined;
 };
 
 async function executeReplayActions(
@@ -363,24 +370,15 @@ async function executeReplayActions(
     snapshotDiagnosticSamples,
     onStep,
     armSaveScript,
+    suppressedTerminalCloseIndex,
   } = params;
   for (let index = entryIndex; index < actions.length; index += 1) {
     const action = actions[index];
     if (!isExecutableReplayAction(action)) continue;
     // Arm before checking terminal close so `[open, close]` records the
     // session created by `open` before treating `close` as lifecycle.
     armSaveScript();
-    if (
-      shouldSkipTerminalClose({
-        action,
-        index,
-        totalActions: actions.length,
-        keepSession: params.req.flags?.replayKeepSession === true,
-        coordinator: stepContext.coordinator,
-      })
-    ) {
-      continue;
-    }
+    if (index === suppressedTerminalCloseIndex) continue;
     onStep?.(replayActionStep(index, actions.length, action));
     const sampleStart = readSessionSnapshotSampleCount(sessionStore, sessionName);
     const response = await resolveReplayStepResponse(stepContext, action, index, [
@@ -441,6 +439,8 @@ function completeReplayRun(params: {
   snapshotDiagnosticSamples: SnapshotTimingSample[];
   armSaveScript: () => void;
   coordinator: ReplayCoordinator;
+  keepSession: boolean;
+  suppressedTerminalCloseIndex: number | undefined;
 }): DaemonResponse {
   const {
     startedAt,
@@ -452,11 +452,24 @@ function completeReplayRun(params: {
     snapshotDiagnosticSamples,
     armSaveScript,
     coordinator,
+    keepSession,
+    suppressedTerminalCloseIndex,
   } = params;
   armSaveScript();
   coordinator.markCompleteIfArmed();
   const completedSession = sessionStore.get(sessionName);
-  const replayedCount = actions.length - entryIndex;
+  if (keepSession && !completedSession) {
+    return errorResponse(
+      'COMMAND_FAILED',
+      `Replay completed but --keep-session could not preserve session "${sessionName}". Run the script again after checking which action closed the session.`,
+      artifactPaths.size > 0 ? { artifactPaths: [...artifactPaths] } : undefined,
+    );
+  }
+  const replayedCount = countExecutedReplayActions({
+    actions,
+    entryIndex,
+    suppressedTerminalCloseIndex,
+  });
   const snapshotDiagnosticsSummary = summarizeSnapshotTimingSamples(snapshotDiagnosticSamples);
   return {
     ok: true,
@@ -521,6 +534,7 @@ type PreparedReplayPlan = {
   scope: ReplayVarScope;
   actionTracePath: string | undefined;
   snapshotDiagnosticSamples: SnapshotTimingSample[];
+  suppressedTerminalCloseIndex: number | undefined;
 };
 
 type ParsedReplayInput = ReturnType<typeof parseReplayInput>;
@@ -532,8 +546,9 @@ function prepareReplayPlan(params: {
   tracePath: string | undefined;
   resolved: string;
   coordinator: ReplayCoordinator;
+  keepSession: boolean;
 }): { ok: true; value: PreparedReplayPlan } | { ok: false; response: DaemonResponse } {
-  const { req, sessionName, sessionStore, tracePath, resolved, coordinator } = params;
+  const { req, sessionName, sessionStore, tracePath, resolved, coordinator, keepSession } = params;
   const parsedResult = parseReplayScript(resolved, req);
   if (!parsedResult.ok) return parsedResult;
   const parsed = parsedResult.value;
@@ -571,6 +586,12 @@ function prepareReplayPlan(params: {
       scope: buildPreparedReplayScope({ req, replayReq, sessionName, resolved, metadata }),
       actionTracePath: tracePath ?? preEntrySession?.trace?.outPath,
       snapshotDiagnosticSamples: [],
+      suppressedTerminalCloseIndex: resolveSuppressedTerminalCloseIndex({
+        actions,
+        keepSession,
+        saveScript: req.flags?.saveScript,
+        repairActive: coordinator.view()?.repairBoundary !== undefined,
+      }),
     },
   };
 }
@@ -794,24 +815,37 @@ function preflightSaveScriptTarget(params: {
 }
 
 /**
- * The one native replay lifecycle seam for an authored terminal `close`.
- * `--keep-session` suppresses it so callers can take over the live session;
- * ADR 0012 repair suppresses it so the agent can finalize through
- * `close --save-script`. Interior closes retain authored semantics. Repair is
- * checked against session state (not only this leg's flags), preserving R2
- * across separate `--from` continuations.
+ * Resolves the one native replay lifecycle seam once per plan. Terminal means
+ * the last executable action, because nested `replay` markers are plan
+ * metadata and never dispatch. The suppressed close is therefore neither
+ * divergence-checked nor included in the successful `replayed` count.
  */
-function shouldSkipTerminalClose(params: {
-  action: SessionAction;
-  index: number;
-  totalActions: number;
+function resolveSuppressedTerminalCloseIndex(params: {
+  actions: SessionAction[];
   keepSession: boolean;
-  coordinator: ReplayCoordinator;
-}): boolean {
-  const { action, index, totalActions, keepSession, coordinator } = params;
-  if (action.command !== 'close') return false;
-  if (index !== totalActions - 1) return false;
-  return keepSession || coordinator.view()?.repairBoundary !== undefined;
+  saveScript: boolean | string | undefined;
+  repairActive: boolean;
+}): number | undefined {
+  if (!params.keepSession && !params.saveScript && !params.repairActive) return undefined;
+  for (let index = params.actions.length - 1; index >= 0; index -= 1) {
+    const action = params.actions[index];
+    if (!isExecutableReplayAction(action)) continue;
+    return action.command === 'close' ? index : undefined;
+  }
+  return undefined;
+}
+
+function countExecutedReplayActions(params: {
+  actions: SessionAction[];
+  entryIndex: number;
+  suppressedTerminalCloseIndex: number | undefined;
+}): number {
+  let count = 0;
+  for (let index = params.entryIndex; index < params.actions.length; index += 1) {
+    if (index === params.suppressedTerminalCloseIndex) continue;
+    if (isExecutableReplayAction(params.actions[index])) count += 1;
+  }
+  return count;
 }
 
 /**