fix(webdriver): give cloud session creation its own budget and stop leaking billed sessions

packages/contracts/src/device-provider.ts

@@ -25,6 +25,14 @@ export type DeviceLease = {
 export type LeaseLifecycleContext = {
   flags?: Readonly<Record<string, unknown>>;
   cwd?: string;
+  /** Request-bound cancellation (explicit cancel or client disconnect). */
+  signal?: AbortSignal;
+  /**
+   * Epoch-ms deadline by which `allocate` must have settled; derived from the
+   * same budget as the client's `lease_allocate` envelope, so a provider that
+   * fits its remote phases within it is never abandoned by a client first.
+   */
+  deadline?: number;
 };
 
 export type LeaseLifecycleProvider = {

packages/kernel/src/errors.ts

@@ -128,6 +128,45 @@ export function throwDaemonError(error: DaemonError): never {
   });
 }
 
+/**
+ * `details.reason` of a request its requester abandoned — an explicit cancel or
+ * a client disconnect. One definition, so every layer that must let a
+ * cancellation through untouched (retry loops, provider adapters, runner
+ * transports) dispatches on the same typed reason.
+ */
+const REQUEST_CANCELED_REASON = 'request_canceled';
+const REQUEST_CANCELED_MESSAGE = 'request canceled';
+const REQUEST_CANCELED_HINT =
+  'The request was canceled intentionally (explicit cancel or client disconnect) — no retry is needed unless the cancellation was unintended.';
+
+/**
+ * The canceled-request error. `details` may add evidence (what was released,
+ * which command was interrupted) or override the hint; the reason itself is
+ * not overridable, so a caller cannot build one this predicate misses.
+ */
+export function createRequestCanceledError(details?: AppErrorDetails, cause?: unknown): AppError {
+  return new AppError(
+    'COMMAND_FAILED',
+    REQUEST_CANCELED_MESSAGE,
+    { hint: REQUEST_CANCELED_HINT, ...details, reason: REQUEST_CANCELED_REASON },
+    cause,
+  );
+}
+
+export function isRequestCanceledError(error: unknown): boolean {
+  if (!(error instanceof AppError)) return false;
+  if (error.code !== 'COMMAND_FAILED') return false;
+  if (error.details?.reason === REQUEST_CANCELED_REASON) return true;
+  // Owned debt: canceled errors that crossed a wire without their details keep
+  // the message; do not add new message sniffs beside it.
+  return error.message === REQUEST_CANCELED_MESSAGE;
+}
+
+/** The message of whatever was thrown, for diagnostics that must not themselves throw. */
+export function errorMessage(error: unknown): string {
+  return error instanceof Error ? error.message : String(error);
+}
+
 export function asAppError(err: unknown, fallbackCode: AppErrorCode = 'UNKNOWN'): AppError {
   if (err instanceof AppError) return err;
   if (err instanceof Error) {

packages/maestro/src/internal/engine-flow.ts

@@ -1,6 +1,5 @@
 import path from 'node:path';
-import { AppError } from '@agent-device/kernel/errors';
-import { createRequestCanceledError } from './shared.ts';
+import { AppError, createRequestCanceledError } from '@agent-device/kernel/errors';
 import {
   MAESTRO_NUMERIC_FIELD_CONSTRAINTS,
   numericDescription,

packages/maestro/src/internal/program-loader.ts

@@ -1,6 +1,6 @@
 import fs from 'node:fs';
 import path from 'node:path';
-import { createRequestCanceledError } from './shared.ts';
+import { createRequestCanceledError } from '@agent-device/kernel/errors';
 import type { MaestroProgram } from './program-ir.ts';
 import { parseMaestroProgram } from './program-ir-parser.ts';
 

packages/maestro/src/internal/shared.ts

@@ -1,4 +1,3 @@
-import { AppError } from '@agent-device/kernel/errors';
 import type { Point, Rect, SnapshotNode } from '@agent-device/kernel/snapshot';
 
 export function stripUndefined<T extends Record<string, unknown>>(value: T): T {
@@ -45,10 +44,3 @@ export function extractNodeText(node: SnapshotNode): string {
       ?.trim() ?? ''
   );
 }
-
-export function createRequestCanceledError(): AppError {
-  return new AppError('COMMAND_FAILED', 'request canceled', {
-    reason: 'request_canceled',
-    hint: 'The request was canceled intentionally (explicit cancel or client disconnect) — no retry is needed unless the cancellation was unintended.',
-  });
-}

packages/provider-webdriver/src/aws-device-farm.test.ts

@@ -0,0 +1,181 @@
+import assert from 'node:assert/strict';
+import { afterEach, test, vi } from 'vitest';
+import type { DeviceLease } from '@agent-device/contracts/device';
+import {
+  AppError,
+  createRequestCanceledError,
+  isRequestCanceledError,
+} from '@agent-device/kernel/errors';
+import {
+  createAwsDeviceFarmPrepareSession,
+  type AwsDeviceFarmClient,
+  type AwsDeviceFarmRemoteAccessSession,
+} from './aws-device-farm.ts';
+import { buildCloudWebDriverBaseCapabilities } from './runtime.ts';
+
+const ARN = 'arn:aws:devicefarm:us-west-2:1:session/pending';
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+// Once `create-remote-access-session` answers, the ARN is a billed session that
+// nothing else will stop. Every way the startup wait can end short of RUNNING
+// must stop it before the failure surfaces (#1774 ownership rule, one phase
+// earlier than the WebDriver session).
+test('a startup timeout stops the remote-access session it created', async () => {
+  const client = fakeClient({ status: 'PENDING' });
+  const prepare = createAwsDeviceFarmPrepareSession({
+    ...baseOptions(client),
+    startupTimeoutMs: 30,
+    pollIntervalMs: 5,
+  });
+
+  await assert.rejects(
+    () => prepare({ lease: makeLease(), base: baseSession() }),
+    (error: unknown) => {
+      assert.ok(error instanceof AppError);
+      assert.match(error.message, /Timed out waiting for AWS Device Farm/);
+      return true;
+    },
+  );
+  assert.deepEqual(client.stopped, [ARN]);
+});
+
+test('a canceled request stops the remote-access session mid-startup', async () => {
+  const controller = new AbortController();
+  const client = fakeClient({ status: 'PENDING' }, () =>
+    controller.abort(createRequestCanceledError()),
+  );
+  const prepare = createAwsDeviceFarmPrepareSession({
+    ...baseOptions(client),
+    startupTimeoutMs: 5_000,
+    pollIntervalMs: 5,
+  });
+
+  await assert.rejects(
+    () => prepare({ lease: makeLease(), req: { signal: controller.signal }, base: baseSession() }),
+    (error: unknown) => isRequestCanceledError(error),
+  );
+  assert.deepEqual(client.stopped, [ARN]);
+});
+
+// The daemon's allocation deadline caps the startup wait so a client that
+// stops waiting at that deadline never abandons a still-polling daemon.
+test('the allocation deadline caps the startup wait below its own default', async () => {
+  const client = fakeClient({ status: 'PENDING' });
+  const prepare = createAwsDeviceFarmPrepareSession({
+    ...baseOptions(client),
+    startupTimeoutMs: 60_000,
+    pollIntervalMs: 5,
+  });
+  const startedAt = Date.now();
+
+  await assert.rejects(
+    () => prepare({ lease: makeLease(), req: { deadline: Date.now() + 40 }, base: baseSession() }),
+    /Timed out waiting for AWS Device Farm/,
+  );
+  assert.ok(Date.now() - startedAt < 5_000, 'the wait must end at the deadline, not the default');
+  assert.deepEqual(client.stopped, [ARN]);
+});
+
+// Live iOS real devices needed ~128s to reach RUNNING while the daemon's 300s
+// allocation budget still had room, and the standalone 120s default cut them
+// off. When the daemon supplies a deadline it is THE bound; the default only
+// applies without one. Time is a virtual clock advanced 10s per poll, so this is
+// deterministic and fails on the old `min(default, deadline)` logic (which
+// throws at 120s, before the 150s RUNNING).
+test('the allocation deadline lets startup run past the standalone 120s default', async () => {
+  const startedAt = 1_700_000_000_000;
+  let virtualNow = startedAt;
+  vi.spyOn(Date, 'now').mockImplementation(() => virtualNow);
+  const client = fakeClient({ status: 'PENDING' }, () => {
+    virtualNow += 10_000;
+    if (virtualNow - startedAt >= 150_000) {
+      client.session.status = 'RUNNING';
+      client.session.endpoints = { appium: 'https://appium.example/wd/hub' };
+    }
+  });
+  const prepare = createAwsDeviceFarmPrepareSession({
+    ...baseOptions(client),
+    // The standalone default; a daemon-supplied deadline must override it.
+    startupTimeoutMs: 120_000,
+    pollIntervalMs: 1,
+  });
+
+  const prepared = await prepare({
+    lease: makeLease(),
+    req: { deadline: startedAt + 300_000 },
+    base: baseSession(),
+  });
+  assert.equal(prepared.providerSessionId, ARN);
+  assert.ok(virtualNow - startedAt >= 150_000, 'RUNNING must have been observed after 120s');
+  assert.deepEqual(client.stopped, []);
+});
+
+test('a session that reaches RUNNING is handed on and not stopped', async () => {
+  const client = fakeClient({
+    status: 'RUNNING',
+    endpoints: { appium: 'https://appium.example/wd/hub' },
+  });
+  const prepare = createAwsDeviceFarmPrepareSession(baseOptions(client));
+
+  const prepared = await prepare({ lease: makeLease(), base: baseSession() });
+  assert.equal(prepared.providerSessionId, ARN);
+  assert.equal(prepared.endpoint, 'https://appium.example/wd/hub');
+  assert.deepEqual(client.stopped, []);
+});
+
+function fakeClient(
+  session: Partial<AwsDeviceFarmRemoteAccessSession>,
+  onPoll?: () => void,
+): AwsDeviceFarmClient & { stopped: string[]; session: Partial<AwsDeviceFarmRemoteAccessSession> } {
+  const stopped: string[] = [];
+  return {
+    stopped,
+    session,
+    createRemoteAccessSession: async () => ({ arn: ARN, status: 'PENDING' }),
+    getRemoteAccessSession: async (arn) => {
+      onPoll?.();
+      return { arn, ...session };
+    },
+    stopRemoteAccessSession: async (arn) => {
+      stopped.push(arn);
+      return { arn, status: 'STOPPING' };
+    },
+    listArtifacts: async () => [],
+  };
+}
+
+function baseOptions(client: AwsDeviceFarmClient) {
+  return {
+    client,
+    platform: 'android' as const,
+    deviceName: 'Pixel',
+    projectArn: 'arn:aws:devicefarm:us-west-2:1:project/p',
+    deviceArn: 'arn:aws:devicefarm:us-west-2::device/d',
+  };
+}
+
+function baseSession() {
+  return {
+    provider: 'aws-device-farm',
+    endpoint: 'http://127.0.0.1/',
+    platform: 'android' as const,
+    deviceName: 'Pixel',
+    webdriverCapabilities: buildCloudWebDriverBaseCapabilities('android', 'Pixel'),
+  };
+}
+
+function makeLease(): DeviceLease {
+  return {
+    leaseId: 'lease-aws',
+    tenantId: 'team-a',
+    runId: 'run-a',
+    leaseProvider: 'aws-device-farm',
+    backend: 'android-instance',
+    createdAt: 1,
+    expiresAt: 2,
+    heartbeatAt: 1,
+  };
+}

packages/provider-webdriver/src/aws-device-farm.ts

@@ -15,12 +15,16 @@ import type {
   CloudWebDriverRuntimeOptions,
   CloudWebDriverPrepareSession,
 } from './runtime.ts';
-import type { DeviceLease, ProviderDeviceRuntime } from '@agent-device/contracts/device';
+import type {
+  DeviceLease,
+  LeaseLifecycleContext,
+  ProviderDeviceRuntime,
+} from '@agent-device/contracts/device';
 import { setTimeout as sleep } from 'node:timers/promises';
 import { AppError } from '@agent-device/kernel/errors';
 import type { RunHostCommand } from './dependencies.ts';
 import { CLOUD_WEBDRIVER_PROVIDERS } from './providers.ts';
-import { resolveLeaseValue, type LeaseValue } from './webdriver-utils.ts';
+import { releaseOnFailure, resolveLeaseValue, type LeaseValue } from './webdriver-utils.ts';
 
 const AWS_DEVICE_FARM_PROVIDER = CLOUD_WEBDRIVER_PROVIDERS.awsDeviceFarm;
 export const AWS_DEVICE_FARM_CAPABILITY_OVERRIDES = {
@@ -203,7 +207,7 @@ export function createAwsDeviceFarmPrepareSession(
       'client' | 'platform' | 'deviceName' | 'clientVersion'
     >,
 ): CloudWebDriverPrepareSession {
-  return async ({ lease, base }) => {
+  return async ({ lease, req, base }) => {
     const remoteAccess = await options.client.createRemoteAccessSession({
       projectArn: options.projectArn,
       deviceArn: options.deviceArn,
@@ -212,13 +216,23 @@ export function createAwsDeviceFarmPrepareSession(
       interactionMode: options.interactionMode,
       configuration: options.configuration,
     });
-    const running = await waitForRunningRemoteAccessSession(remoteAccess.arn, options);
-    const endpoint = selectAwsDeviceFarmWebDriverEndpoint(running);
-    if (!endpoint) {
-      throw new AppError('COMMAND_FAILED', 'AWS Device Farm did not expose a WebDriver endpoint.', {
-        sessionArn: running.arn,
-        status: running.status,
-      });
+    // The ARN is a billed session from here on; any failure short of RUNNING
+    // must stop it before surfacing, or it bills until AWS reaps it.
+    let running: AwsDeviceFarmRemoteAccessSession;
+    let endpoint: string | undefined;
+    try {
+      running = await waitForRunningRemoteAccessSession(remoteAccess.arn, options, req);
+      endpoint = selectAwsDeviceFarmWebDriverEndpoint(running);
+      if (!endpoint) {
+        throw new AppError(
+          'COMMAND_FAILED',
+          'AWS Device Farm did not expose a WebDriver endpoint.',
+          { sessionArn: running.arn, status: running.status },
+        );
+      }
+    } catch (error) {
+      await releaseOnFailure(error, () => options.client.stopRemoteAccessSession(remoteAccess.arn));
+      throw error;
     }
     const deviceName = running.device?.name ?? options.deviceName;
     const configured =
@@ -276,28 +290,40 @@ async function waitForRunningRemoteAccessSession(
     pollIntervalMs?: number;
     startupTimeoutMs?: number;
   },
+  req: LeaseLifecycleContext | undefined,
 ): Promise<AwsDeviceFarmRemoteAccessSession> {
-  const timeoutMs = options.startupTimeoutMs ?? 120_000;
   const pollIntervalMs = options.pollIntervalMs ?? 5_000;
   const startedAt = Date.now();
+  // The daemon's allocation deadline is the bound when present (real-device
+  // startup routinely needs the whole ~2 min); the standalone default only
+  // applies when no allocation budget was supplied.
+  const deadline = req?.deadline ?? startedAt + (options.startupTimeoutMs ?? 120_000);
+  const signal = req?.signal;
   let last = await options.client.getRemoteAccessSession(arn);
-  while (Date.now() - startedAt < timeoutMs) {
+  while (Date.now() < deadline) {
+    signal?.throwIfAborted();
     if (last.status === 'RUNNING') return last;
-    if (last.status === 'ERRORED' || last.status === 'STOPPED' || last.status === 'COMPLETED') {
-      throw new AppError('COMMAND_FAILED', 'AWS Device Farm remote access session did not start.', {
-        sessionArn: arn,
-        status: last.status,
-        result: last.result,
-      });
-    }
-    await sleep(pollIntervalMs);
+    throwIfRemoteAccessSessionEnded(last);
+    // Wake early on cancellation; the typed reason is rethrown at the loop top.
+    await sleep(pollIntervalMs, undefined, { signal }).catch(() => signal?.throwIfAborted());
     last = await options.client.getRemoteAccessSession(arn);
   }
   throw new AppError('COMMAND_FAILED', 'Timed out waiting for AWS Device Farm remote access.', {
     sessionArn: arn,
     status: last.status,
     result: last.result,
-    timeoutMs,
+    timeoutMs: deadline - startedAt,
+  });
+}
+
+const ENDED_REMOTE_ACCESS_STATUSES = new Set(['ERRORED', 'STOPPED', 'COMPLETED']);
+
+function throwIfRemoteAccessSessionEnded(session: AwsDeviceFarmRemoteAccessSession): void {
+  if (!session.status || !ENDED_REMOTE_ACCESS_STATUSES.has(session.status)) return;
+  throw new AppError('COMMAND_FAILED', 'AWS Device Farm remote access session did not start.', {
+    sessionArn: session.arn,
+    status: session.status,
+    result: session.result,
   });
 }