fix: complete request cancellation propagation (#1709)

* fix: route duration waits through cancellation-aware sleep

AI-assisted implementation. The code and validation evidence were reviewed before submission.

* refactor: expose shared cancellation-aware wait sleep

AI-assisted implementation. The code and validation evidence were reviewed before submission.

* fix: forward request cancellation through snapshot runtime

AI-assisted implementation. The code and validation evidence were reviewed before submission.

* test: cover duration wait cancellation authorities

AI-assisted implementation. The code and validation evidence were reviewed before submission.

* test: cover snapshot request cancellation propagation

AI-assisted implementation. The code and validation evidence were reviewed before submission.

* style: normalize wait cancellation test ending

AI-assisted cleanup. The remote content was compared byte-for-byte with the reviewed local test file.

* style: normalize snapshot cancellation test ending

AI-assisted cleanup. The remote content was compared byte-for-byte with the reviewed local test file.

Author: kinKingen

src/commands/interaction/runtime/selector-wait-cancellation.test.ts

@@ -0,0 +1,55 @@
+import assert from 'node:assert/strict';
+import { test } from 'vitest';
+import type { AgentDeviceBackend } from '../../../backend.ts';
+import { createLocalArtifactAdapter } from '../../../io.ts';
+import {
+  createAgentDevice,
+  createMemorySessionStore,
+  localCommandPolicy,
+} from '../../../runtime.ts';
+
+type Deferred = {
+  promise: Promise<void>;
+  resolve: () => void;
+};
+
+function deferred(): Deferred {
+  let resolve!: () => void;
+  const promise = new Promise<void>((done) => {
+    resolve = done;
+  });
+  return { promise, resolve };
+}
+
+for (const authority of ['runtime', 'command'] as const) {
+  test(`duration wait observes ${authority} cancellation`, async () => {
+    const controller = new AbortController();
+    const sleepStarted = deferred();
+    const releaseSleep = deferred();
+    const reason = new Error(`${authority} wait canceled`);
+    const device = createAgentDevice({
+      backend: { platform: 'ios' } satisfies AgentDeviceBackend,
+      artifacts: createLocalArtifactAdapter(),
+      sessions: createMemorySessionStore(),
+      policy: localCommandPolicy(),
+      clock: {
+        now: () => 0,
+        sleep: async () => {
+          sleepStarted.resolve();
+          await releaseSleep.promise;
+        },
+      },
+      ...(authority === 'runtime' ? { signal: controller.signal } : {}),
+    });
+
+    const waiting = device.selectors.wait({
+      target: { kind: 'sleep', durationMs: 1_000 },
+      ...(authority === 'command' ? { signal: controller.signal } : {}),
+    });
+    await sleepStarted.promise;
+    controller.abort(reason);
+    releaseSleep.resolve();
+
+    await assert.rejects(waiting, (error) => error === reason);
+  });
+}

src/commands/interaction/runtime/selector-wait.ts

@@ -25,6 +25,7 @@ import { findNodeByLabel, resolveRefLabel } from './selector-read-utils.ts';
 import {
   createWaitPolling,
   DEFAULT_WAIT_TIMEOUT_MS,
+  sleepWithWaitCancellation,
   type WaitPollDeadline,
   waitTimeoutError,
 } from './wait-polling.ts';
@@ -178,7 +179,7 @@ export function createSelectorWaitCommands<Runtime extends SelectorWaitRuntime>(
     options: WaitCommandOptions,
   ): Promise<WaitCommandResult> => {
     if (options.target.kind === 'sleep') {
-      await sleep(runtime, options.target.durationMs);
+      await sleepWithWaitCancellation(runtime, options, options.target.durationMs);
       return { kind: 'sleep', waitedMs: options.target.durationMs };
     }
     if (options.target.kind === 'ref') {
@@ -447,8 +448,3 @@ function backendContext(
     metadata: options.metadata,
   };
 }
-
-async function sleep(runtime: SelectorWaitRuntime, durationMs: number): Promise<void> {
-  if (runtime.clock) await runtime.clock.sleep(durationMs);
-  else await new Promise((resolve) => setTimeout(resolve, durationMs));
-}

src/commands/interaction/runtime/wait-polling.ts

@@ -87,7 +87,11 @@ export function createWaitPolling(
     }),
     rethrowIfNeverReadable: unreadable.rethrowIfNeverReadable,
     sleepUntilNextPoll: async () =>
-      await sleepWithinWait(runtime, options, Math.min(WAIT_POLL_INTERVAL_MS, remainingMs())),
+      await sleepWithWaitCancellation(
+        runtime,
+        options,
+        Math.min(WAIT_POLL_INTERVAL_MS, remainingMs()),
+      ),
     timeoutMs,
     waitedMs: () => now(runtime) - startedAtMs,
   };
@@ -162,7 +166,7 @@ function now(runtime: WaitPollingRuntime): number {
   return runtime.clock?.now() ?? Date.now();
 }
 
-async function sleepWithinWait(
+export async function sleepWithWaitCancellation(
   runtime: WaitPollingRuntime,
   options: WaitPollingOptions,
   durationMs: number,

src/daemon/__tests__/snapshot-runtime-cancellation.test.ts

@@ -0,0 +1,87 @@
+import { afterEach, expect, test, vi } from 'vitest';
+import { makeAndroidSession } from '../../__tests__/test-utils/session-factories.ts';
+import { makeSessionStore } from '../../__tests__/test-utils/store-factory.ts';
+import {
+  clearRequestAbortRegistration,
+  markRequestCanceled,
+  registerRequestAbort,
+} from '../../request/cancel.ts';
+import { dispatchSnapshotDiffViaRuntime, dispatchSnapshotViaRuntime } from '../snapshot-runtime.ts';
+
+const dispatchCommandMock = vi.hoisted(() => vi.fn());
+
+vi.mock('../../core/dispatch.ts', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('../../core/dispatch.ts')>();
+  return {
+    ...actual,
+    dispatchCommand: dispatchCommandMock,
+  };
+});
+
+type Deferred = {
+  promise: Promise<void>;
+  resolve: () => void;
+};
+
+function deferred(): Deferred {
+  let resolve!: () => void;
+  const promise = new Promise<void>((done) => {
+    resolve = done;
+  });
+  return { promise, resolve };
+}
+
+afterEach(() => {
+  dispatchCommandMock.mockReset();
+});
+
+for (const command of ['snapshot', 'diff snapshot'] as const) {
+  test(`${command} forwards request cancellation into snapshot dispatch`, async () => {
+    const sessionName = 'default';
+    const sessionStore = makeSessionStore('agent-device-snapshot-cancellation-');
+    sessionStore.set(sessionName, makeAndroidSession(sessionName));
+    const requestId = `snapshot-cancellation-${command.replace(' ', '-')}`;
+    const registration = registerRequestAbort(requestId);
+    if (!registration) throw new Error('expected request abort registration');
+    const dispatchEntered = deferred();
+    const releaseDispatch = deferred();
+    let observedSignal: AbortSignal | undefined;
+
+    dispatchCommandMock.mockImplementation(async (...args: unknown[]) => {
+      const context = args[4] as { signal?: AbortSignal } | undefined;
+      observedSignal = context?.signal;
+      dispatchEntered.resolve();
+      await releaseDispatch.promise;
+      context?.signal?.throwIfAborted();
+      return { nodes: [], truncated: false, backend: 'uiautomator' };
+    });
+
+    try {
+      const input = {
+        req: {
+          command: command === 'snapshot' ? 'snapshot' : 'diff',
+          positionals: command === 'snapshot' ? [] : ['snapshot'],
+          token: 't',
+          session: sessionName,
+          meta: { requestId },
+        },
+        sessionName,
+        logPath: '/tmp/agent-device-snapshot-cancellation.log',
+        sessionStore,
+      };
+      const running =
+        command === 'snapshot'
+          ? dispatchSnapshotViaRuntime(input)
+          : dispatchSnapshotDiffViaRuntime(input);
+      await dispatchEntered.promise;
+      markRequestCanceled(requestId);
+      releaseDispatch.resolve();
+
+      expect(observedSignal).toBe(registration.controller.signal);
+      await expect(running).rejects.toBe(registration.controller.signal.reason);
+    } finally {
+      releaseDispatch.resolve();
+      clearRequestAbortRegistration(registration);
+    }
+  });
+}

src/daemon/snapshot-runtime.ts

@@ -27,6 +27,7 @@ import {
 } from './snapshot-quality-latch.ts';
 import { createDaemonRuntimePolicy } from './runtime-policy.ts';
 import { createDaemonRuntimeSessionStore } from './runtime-session.ts';
+import { getRequestSignal } from '../request/cancel.ts';
 import { isInteractiveObservation } from './session-action-recorder.ts';
 import { setSnapshotLineage } from './session-snapshot.ts';
 import { SessionStore } from './session-store.ts';
@@ -284,6 +285,7 @@ function createSnapshotRuntime(params: {
       capturedQuality: params.capturedQuality,
     }),
     ...createDaemonRuntimePolicy('snapshot'),
+    signal: getRequestSignal(req.meta?.requestId),
     sessions: createDaemonRuntimeSessionStore({
       sessionName,
       getSession: () => sessionStore.get(sessionName),
@@ -394,14 +396,15 @@ function createDaemonSnapshotBackend(params: {
   const { req, logPath, session, device, snapshotScope } = params;
   return {
     platform: publicPlatformString(device),
-    captureSnapshot: async (_context, options): Promise<BackendSnapshotResult> => {
+    captureSnapshot: async (context, options): Promise<BackendSnapshotResult> => {
       const capture = await captureSnapshot({
         device,
         session,
         flags: req.flags,
         outPath: options?.outPath ?? req.flags?.out,
         logPath,
         snapshotScope,
+        signal: context.signal,
       });
       const annotations = snapshotCaptureAnnotationsFrom(capture);
       // Feed the latch seam the capture's own verdict: the stored session