fix(lease): return a retried allocation's own lease and refuse an unbound renewed lease

packages/capture-kit/src/durable-resource-envelope.ts

@@ -1,5 +1,6 @@
 import {
   isAppleOs,
+  isDeviceTarget,
   isPlatform,
   type AppleOS,
   type DeviceIdentity,
@@ -233,10 +234,6 @@ function isDeviceKind(value: unknown): value is DeviceKind {
   return value === 'simulator' || value === 'emulator' || value === 'device';
 }
 
-function isDeviceTarget(value: unknown): value is DeviceTarget {
-  return value === 'mobile' || value === 'tv' || value === 'desktop';
-}
-
 function isNonEmptyString(value: unknown): value is string {
   return typeof value === 'string' && value.trim().length > 0;
 }

packages/daemon/src/resources/__tests__/lease-registry.test.ts

@@ -75,6 +75,87 @@ test('allocateLease keeps a caller-chosen lease id and refuses any id it already
   }
 });
 
+test('an allocation repeating its own live lease id gets that lease back, on any device it names or none', () => {
+  const registry = new LeaseRegistry();
+  const leaseId = 'abcdef0123456789abcdef0123456789';
+  const owner = {
+    tenantId: 'tenant-a',
+    runId: 'run-1',
+    leaseBackend: 'ios-instance',
+    leaseProvider: 'proxy',
+    clientId: 'client-1',
+  } as const;
+  const lease = registry.allocateLease({ ...owner, leaseId, deviceKey: 'ios:mobile:sim-1' });
+
+  assert.equal(registry.renewRequestedLease({ ...owner, leaseId })?.leaseId, leaseId);
+  assert.equal(registry.allocateLease({ ...owner, leaseId }).deviceKey, lease.deviceKey);
+  assert.equal(
+    registry.renewRequestedLease({ ...owner, leaseId, deviceKey: 'ios:mobile:sim-2' }),
+    undefined,
+  );
+  assert.equal(
+    registry.renewRequestedLease({ ...owner, leaseId, clientId: 'client-2' }),
+    undefined,
+  );
+  assert.equal(registry.renewRequestedLease({ ...owner, leaseId, runId: 'run-2' }), undefined);
+  assert.equal(registry.renewRequestedLease(owner), undefined);
+  assert.throws(
+    () => registry.allocateLease({ ...owner, leaseId, deviceKey: 'ios:mobile:sim-2' }),
+    isLeaseIdInUse,
+  );
+  assert.equal(registry.listActiveLeases().length, 1);
+});
+
+test('a retried allocation renews its own lease for the window that lease carries', () => {
+  let now = 1_000;
+  const registry = new LeaseRegistry({ now: () => now, defaultLeaseTtlMs: 10_000 });
+  const leaseId = 'abcdef0123456789abcdef0123456789';
+  const owner = { tenantId: 'tenant-a', runId: 'run-1', clientId: 'client-1' } as const;
+  const lease = registry.allocateLease({ ...owner, leaseId, ttlMs: 60_000 });
+
+  now = 31_000;
+  const retried = registry.renewRequestedLease({ ...owner, leaseId });
+
+  assert.equal(lease.expiresAt, 61_000);
+  assert.equal(retried?.heartbeatAt, 31_000);
+  assert.equal(retried?.expiresAt, 91_000);
+});
+
+test('an allocation naming an issued lease id gets exactly that lease or LEASE_ID_IN_USE, never another lease of its run', () => {
+  const registry = new LeaseRegistry();
+  const owner = { tenantId: 'tenant-a', runId: 'run-1', clientId: 'client-1' } as const;
+  const first = registry.allocateLease({
+    ...owner,
+    leaseId: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa',
+    deviceKey: 'ios:mobile:sim-1',
+  });
+  const second = registry.allocateLease({
+    ...owner,
+    leaseId: 'bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb',
+    deviceKey: 'ios:mobile:sim-2',
+  });
+
+  assert.notEqual(first.leaseId, second.leaseId);
+  assert.throws(
+    () => registry.allocateLease({ ...owner, leaseId: first.leaseId, deviceKey: second.deviceKey }),
+    isLeaseIdInUse,
+  );
+});
+
+test('a retried allocation of a host-pinned lease never outlives the host window', () => {
+  let now = 1_000;
+  const registry = new LeaseRegistry({ now: () => now });
+  const leaseId = 'cccccccccccccccccccccccccccccccc';
+  const owner = { tenantId: 'tenant-a', runId: 'run-1' } as const;
+  registry.putHostLease(leaseId, { ...owner, ttlMs: 10_000 });
+
+  now = 5_000;
+  const retried = registry.allocateLease({ ...owner, leaseId, ttlMs: 600_000 });
+
+  assert.equal(retried.leaseId, leaseId);
+  assert.equal(retried.expiresAt, 15_000);
+});
+
 test('an expired provider lease keeps its id and provider session from a caller-chosen reuse', () => {
   let now = 1_000;
   const registry = new LeaseRegistry({ now: () => now, defaultLeaseTtlMs: 10_000 });

packages/daemon/src/resources/lease-registry.ts

@@ -79,19 +79,17 @@ export class LeaseRegistry {
   }
 
   allocateLease(request: AllocateLeaseRequest): DeviceLease {
-    const chosenLeaseId =
-      request.leaseId === undefined ? undefined : normalizeRequiredLeaseId(request.leaseId);
-    const normalized = normalizeAllocateLeaseRequest(request);
-    this.cleanupExpiredLeases();
-    const leaseTtlMs = this.resolveLeaseTtlMs(normalized.ttlMs);
-    this.assertHumanControlAdmission(normalized);
-    const existingLease = this.refreshExistingRunBinding(normalized, leaseTtlMs);
-    if (existingLease) return existingLease;
+    const { chosenLeaseId, normalized, leaseTtlMs } = this.prepareAllocation(request);
+    const requestedLease = this.refreshRequestedLease(chosenLeaseId, normalized, leaseTtlMs);
+    if (requestedLease) return requestedLease;
     if (chosenLeaseId && this.issuedLeaseIds.has(chosenLeaseId)) {
       throw new AppError('INVALID_ARGS', 'Lease id already names a lease this daemon issued.', {
         reason: 'LEASE_ID_IN_USE',
       });
     }
+    this.assertHumanControlAdmission(normalized);
+    const existingLease = this.refreshExistingRunBinding(normalized, leaseTtlMs);
+    if (existingLease) return existingLease;
     this.assertDeviceAvailable(normalized);
     this.enforceCapacity(normalized.backend);
     const lease = createDeviceLease(normalized, leaseTtlMs, this.now(), chosenLeaseId);
@@ -134,6 +132,45 @@ export class LeaseRegistry {
     return { ...lease };
   }
 
+  /**
+   * The live lease an allocation gets back when it repeats the `leaseId` it named before, for the
+   * same owner: a retry after a lost reply receives the lease it was issued, not a second one. A
+   * request naming no device matches whichever device the lease binds, and one naming no `ttlMs`
+   * keeps the window the lease already carries.
+   */
+  renewRequestedLease(request: AllocateLeaseRequest): DeviceLease | undefined {
+    const { chosenLeaseId, normalized, leaseTtlMs } = this.prepareAllocation(request);
+    return this.refreshRequestedLease(chosenLeaseId, normalized, leaseTtlMs);
+  }
+
+  private prepareAllocation(request: AllocateLeaseRequest): {
+    chosenLeaseId?: string;
+    normalized: NormalizedAllocateLeaseRequest;
+    leaseTtlMs: number;
+  } {
+    const chosenLeaseId =
+      request.leaseId === undefined ? undefined : normalizeRequiredLeaseId(request.leaseId);
+    const normalized = normalizeAllocateLeaseRequest(request);
+    this.cleanupExpiredLeases();
+    return { chosenLeaseId, normalized, leaseTtlMs: this.resolveLeaseTtlMs(normalized.ttlMs) };
+  }
+
+  private refreshRequestedLease(
+    chosenLeaseId: string | undefined,
+    request: NormalizedAllocateLeaseRequest,
+    leaseTtlMs: number,
+  ): DeviceLease | undefined {
+    const lease = chosenLeaseId === undefined ? undefined : this.leases.get(chosenLeaseId);
+    if (!lease) return undefined;
+    const scope = { ...request, deviceKey: request.deviceKey ?? lease.deviceKey };
+    if (!hasAllocatedScope(lease, scope)) return undefined;
+    this.assertHumanControlAdmission(lease);
+    return this.refreshLease(
+      request.retainOnClose ? { ...lease, retainOnClose: true } : lease,
+      request.ttlMs === undefined ? leaseOwnTtlMs(lease) : leaseTtlMs,
+    );
+  }
+
   private issueLease(lease: DeviceLease): void {
     this.issuedLeaseIds.add(lease.leaseId);
     this.leases.set(lease.leaseId, lease);
@@ -180,11 +217,7 @@ export class LeaseRegistry {
     assertLeaseScopeMatch(lease, request);
     const requestedTtlMs =
       request.ttlMs === undefined ? leaseOwnTtlMs(lease) : this.resolveLeaseTtlMs(request.ttlMs);
-    const hostTtlMs = this.hostTtlMsByLeaseId.get(leaseId);
-    return this.refreshLease(
-      lease,
-      hostTtlMs === undefined ? requestedTtlMs : Math.min(requestedTtlMs, hostTtlMs),
-    );
+    return this.refreshLease(lease, requestedTtlMs);
   }
 
   /**
@@ -584,11 +617,13 @@ export class LeaseRegistry {
     });
   }
 
+  /** Every renewal route ends here, so a host-pinned lease never outlives the host's window. */
   private refreshLease(lease: DeviceLease, ttlMs: number, now = this.now()): DeviceLease {
+    const hostTtlMs = this.hostTtlMsByLeaseId.get(lease.leaseId);
     const updated: DeviceLease = {
       ...lease,
       heartbeatAt: now,
-      expiresAt: now + ttlMs,
+      expiresAt: now + (hostTtlMs === undefined ? ttlMs : Math.min(ttlMs, hostTtlMs)),
     };
     this.leases.set(updated.leaseId, updated);
     this.bindLease(updated);

packages/daemon/src/sessions/handlers/__tests__/lease.test.ts

@@ -483,6 +483,7 @@ function proxyAllocateRequest(
 async function allocateAgainstInventory(
   request: DaemonRequest,
   devices: readonly DeviceInfo[],
+  leaseRegistry = new LeaseRegistry(),
 ): Promise<{ lease: DeviceLease; inventoryReads: number }> {
   let inventoryReads = 0;
   const response = await withTestDeviceInventory(
@@ -497,11 +498,13 @@ async function allocateAgainstInventory(
         req: request,
         sessionName: request.session,
         sessionStore: makeSessionStore('agent-device-proxy-allocate-by-spec-'),
-        leaseRegistry: new LeaseRegistry(),
+        leaseRegistry,
       }),
   );
   assert.ok(response?.ok);
-  return { lease: response.data?.lease as DeviceLease, inventoryReads };
+  const lease = response.data?.lease as DeviceLease | undefined;
+  assert.ok(lease);
+  return { lease, inventoryReads };
 }
 
 test('a proxy allocation that names no device picks one by its selectors and binds its deviceKey', async () => {
@@ -548,3 +551,43 @@ test('an allocation that names its device, names no spec, or is not a proxy leas
   assert.equal(unspecified.lease.deviceKey, undefined);
   assert.equal(unspecified.inventoryReads, 0);
 });
+
+test('a proxy allocation retried with its lease id gets its lease back without picking again', async () => {
+  const leaseId = 'abcdef0123456789abcdef0123456789';
+  const registry = new LeaseRegistry();
+  const request = proxyAllocateRequest({ platform: 'ios' }, { leaseId });
+  const first = await allocateAgainstInventory(request, [IOS_SIMULATOR], registry);
+  assert.equal(first.lease.deviceKey, 'ios:mobile:sim-1');
+
+  const deviceGone = await allocateAgainstInventory(request, [], registry);
+  assert.equal(deviceGone.lease.leaseId, leaseId);
+  assert.equal(deviceGone.lease.deviceKey, 'ios:mobile:sim-1');
+  assert.equal(deviceGone.inventoryReads, 0);
+
+  const otherCandidate = await allocateAgainstInventory(request, [SECOND_IOS_SIMULATOR], registry);
+  assert.equal(otherCandidate.lease.leaseId, leaseId);
+  assert.equal(otherCandidate.inventoryReads, 0);
+  assert.deepEqual(
+    registry.listActiveLeases().map((lease) => lease.leaseId),
+    [leaseId],
+  );
+});
+
+test("a proxy allocation naming another owner's lease id still picks and is refused", async () => {
+  const leaseId = 'abcdef0123456789abcdef0123456789';
+  const registry = new LeaseRegistry();
+  await allocateAgainstInventory(
+    proxyAllocateRequest({ platform: 'ios' }, { leaseId }),
+    [IOS_SIMULATOR],
+    registry,
+  );
+  await assert.rejects(
+    allocateAgainstInventory(
+      proxyAllocateRequest({ platform: 'ios' }, { leaseId, clientId: 'client-2' }),
+      [SECOND_IOS_SIMULATOR],
+      registry,
+    ),
+    (error: AppError) =>
+      error.code === 'INVALID_ARGS' && error.details?.reason === 'LEASE_ID_IN_USE',
+  );
+});

packages/daemon/src/sessions/handlers/lease.ts

@@ -90,16 +90,7 @@ export async function handleLeaseCommands(args: LeaseHandlerArgs): Promise<Daemo
         req.meta?.providerCredentialFingerprint,
         providerCredentials,
       );
-      const pickedDeviceKey = await pickProxyLeaseDeviceKey(req, leaseScope);
-      const activeLeaseIds = new Set(
-        leaseRegistry.listActiveLeases().map((entry) => entry.leaseId),
-      );
-      const lease = leaseRegistry.allocateLease(
-        leaseScopeToAllocateRequest(
-          pickedDeviceKey ? { ...leaseScope, deviceKey: pickedDeviceKey } : leaseScope,
-        ),
-      );
-      const reused = activeLeaseIds.has(lease.leaseId);
+      const { lease, reused } = await allocateOrRenewLease(req, leaseScope, leaseRegistry);
       const requestId = req.meta?.requestId;
       return await leaseRegistry.runDeviceMutation(lease, async () => {
         let providerData: Record<string, unknown> | undefined;
@@ -186,6 +177,26 @@ export async function handleLeaseCommands(args: LeaseHandlerArgs): Promise<Daemo
   }
 }
 
+/**
+ * A request repeating the `leaseId` of its own live lease gets that lease back before any device is
+ * picked, so a retry neither fails on a device that left the inventory nor binds another one.
+ */
+async function allocateOrRenewLease(
+  req: DaemonRequest,
+  scope: LeaseScope,
+  leaseRegistry: LeaseRegistry,
+): Promise<{ lease: DeviceLease; reused: boolean }> {
+  const request = leaseScopeToAllocateRequest(scope);
+  const renewed = leaseRegistry.renewRequestedLease(request);
+  if (renewed) return { lease: renewed, reused: true };
+  const pickedDeviceKey = await pickProxyLeaseDeviceKey(req, scope);
+  const activeLeaseIds = new Set(leaseRegistry.listActiveLeases().map((entry) => entry.leaseId));
+  const lease = leaseRegistry.allocateLease(
+    pickedDeviceKey ? { ...request, deviceKey: pickedDeviceKey } : request,
+  );
+  return { lease, reused: activeLeaseIds.has(lease.leaseId) };
+}
+
 /**
  * A proxy allocation that names no `deviceKey` but does name a device spec leaves the device to this
  * daemon: it resolves the request's selectors as a command would, and the lease binds that device.

src/__tests__/remote-connection-proxy-allocate-by-spec.test.ts

@@ -263,6 +263,40 @@ test('a daemon whose lease names no device gets that lease released and the devi
   fs.rmSync(tempRoot, { recursive: true, force: true });
 });
 
+test('a renewed lease that names no device is refused and left to its owner', async () => {
+  const { tempRoot, stateDir, flags } = proxyWorkspace('agent-device-proxy-by-spec-keyless-', {
+    leaseId: 'lease-1',
+  });
+
+  await assert.rejects(
+    materializeRemoteConnectionForCommand({
+      command: 'open',
+      flags: { ...flags, device: 'iPhone 17' },
+      client: createTestClient({
+        listDevices: async () => [IPHONE_SIMULATOR],
+        allocate: unexpectedCommandCall,
+        release: unexpectedCommandCall,
+        heartbeat: async (request) => ({
+          leaseId: request.leaseId,
+          tenantId: 'proxy',
+          runId: 'proxy-client-1',
+          backend: 'ios-instance',
+          leaseProvider: 'proxy',
+          clientId: 'client-1',
+        }),
+      }),
+    }),
+    (error: AppError) =>
+      error.code === 'UNAUTHORIZED' &&
+      error.details?.reason === 'LEASE_SCOPE_MISMATCH' &&
+      error.details?.leaseId === 'lease-1',
+  );
+  const state = readRemoteConnectionState({ stateDir, session: 'adc-proxy' });
+  assert.equal(state?.leaseId, 'lease-1');
+  assert.equal(state?.deviceKey, undefined);
+  fs.rmSync(tempRoot, { recursive: true, force: true });
+});
+
 test('proxy open with an exact udid still picks locally and allocates that device once', async () => {
   const { tempRoot, flags } = proxyWorkspace('agent-device-proxy-by-spec-exact-');
   const requests: LeaseAllocateRequest[] = [];

src/__tests__/remote-connection.fixtures.ts

@@ -104,6 +104,8 @@ export function createTestClient(
           tenantId: request.tenant ?? 'acme',
           runId: request.runId ?? 'run-123',
           backend: request.leaseBackend ?? 'android-instance',
+          // A daemon renews only a lease whose device matches the key the heartbeat names.
+          ...(request.deviceKey ? { deviceKey: request.deviceKey } : {}),
         })),
       release: options.release ?? (async () => ({ released: true })),
     }),