fix(daemon): retain Apple runners for leases kept on close

packages/daemon/src/core/__tests__/session-store-lifetime.test.ts

@@ -78,6 +78,25 @@ test('address reuse with the same record still starts a different lifetime', ()
   assert.equal(store.getRuntimeHints(ADDRESS), undefined);
 });
 
+test('retained execution-host ownership follows the current session lifetime', () => {
+  const store = makeSessionStore();
+  const session = makeSession('default', {
+    lease: { leaseId: 'lease-1', tenantId: 'tenant-a', runId: 'run-1' },
+  });
+  const old = store.publish(ADDRESS, session);
+  store.retire(old, { retainExecutionHost: true });
+  const [retained] = store.listRetainedExecutionHosts('lease-1');
+  assert.ok(retained);
+  assert.equal(store.ownsRetainedExecutionHost(retained), true);
+  const successor = store.publish(ADDRESS, session);
+  assert.equal(store.ownsRetainedExecutionHost(retained), false);
+  assert.equal(store.retire(old, { retainExecutionHost: true }), false);
+  assert.deepEqual(store.listRetainedExecutionHosts('lease-1'), []);
+  store.retire(successor, { retainExecutionHost: true });
+  assert.equal(store.releaseRetainedExecutionHost(retained), false);
+  assert.equal(store.listRetainedExecutionHosts('lease-1').length, 1);
+});
+
 test('retired updates cannot run their derivation or resurrect a record', () => {
   const store = makeSessionStore();
   const ref = store.publish(ADDRESS, makeSession('default'));

packages/daemon/src/core/platform-resource-cleanup.ts

@@ -16,6 +16,7 @@ export type PlatformResourceCleanup = Readonly<{
     device: DeviceInfo;
     shutdownRequested: boolean;
     hasScreenRecording: boolean;
-    hasLease: boolean;
+    /** True when the lease is not kept on close, or when its close decision could not be inspected. */
+    leaseWillEnd: boolean;
   }): boolean;
 }>;

packages/daemon/src/core/session-store.ts

@@ -49,9 +49,15 @@ import {
 const REPAIR_TOMBSTONE_TTL_MS = 60 * 60_000;
 type SessionEntry = { current: SessionState };
 type SessionPatch = Partial<SessionState> | ((current: SessionState) => Partial<SessionState>);
+export type RetainedExecutionHost = Readonly<{
+  address: string;
+  device: SessionState['device'];
+  lease: NonNullable<SessionState['lease']>;
+}>;
 
 export class SessionStore {
   private readonly sessions = new Map<string, SessionEntry>();
+  private readonly retainedExecutionHosts = new Map<string, RetainedExecutionHost>();
   private acceptingSessions = true;
   private readonly runtimeHints = new Map<string, SessionRuntimeHints>();
   private readonly sessionsDir: string;
@@ -105,6 +111,11 @@ export class SessionStore {
     this.assertPublishable(address);
     const entry = { current: session };
     this.sessions.set(address, entry);
+    for (const retained of this.retainedExecutionHosts.values()) {
+      if (devicesOverlap(retained.device.id, session.device.id)) {
+        this.releaseRetainedExecutionHost(retained);
+      }
+    }
     this.clearIdleExpiryTombstone(address);
     return this.captureRef(address, entry);
   }
@@ -141,12 +152,35 @@ export class SessionStore {
     return next;
   }
 
-  retire(ref: SessionRef): boolean {
-    if (!this.resolveCurrent(ref)) return false;
+  retire(ref: SessionRef, options: { retainExecutionHost?: boolean } = {}): boolean {
+    const session = this.resolveCurrent(ref);
+    if (!session) return false;
+    if (options.retainExecutionHost && session.lease) {
+      this.retainedExecutionHosts.set(session.device.id, {
+        address: ref.address,
+        device: session.device,
+        lease: { ...session.lease },
+      });
+    }
     this.runtimeHints.delete(ref.address);
     return this.sessions.delete(ref.address);
   }
 
+  listRetainedExecutionHosts(leaseId: string): RetainedExecutionHost[] {
+    return Array.from(this.retainedExecutionHosts.values()).filter(
+      (host) => host.lease.leaseId === leaseId,
+    );
+  }
+
+  ownsRetainedExecutionHost(host: RetainedExecutionHost): boolean {
+    return this.retainedExecutionHosts.get(host.device.id) === host;
+  }
+
+  releaseRetainedExecutionHost(host: RetainedExecutionHost): boolean {
+    if (!this.ownsRetainedExecutionHost(host)) return false;
+    return this.retainedExecutionHosts.delete(host.device.id);
+  }
+
   private captureRef(address: string, entry: SessionEntry): SessionRef {
     return Object.freeze({ address, session: entry.current, lifetime: entry });
   }

packages/daemon/src/server/__tests__/request-execution-scope-lease-session-end.test.ts

@@ -16,6 +16,7 @@ import { acquireDeviceClaim } from '@agent-device/daemon-contracts/device-claims
 import { existingSessionExecutionLockKeys } from '../request-binding.ts';
 import { withRequestExecutionLockKeys } from '../request-execution-locks.ts';
 import type { SessionStore } from '../../core/session-store.ts';
+import type { PlatformResourceCleanup } from '../../core/platform-resource-cleanup.ts';
 import type { SessionTeardown } from '../../sessions/lease-lifecycle.ts';
 import { createLeaseSessionEnd } from '../request-execution-scope.ts';
 
@@ -67,7 +68,13 @@ function holdSessionLocks(
   return { finish, held };
 }
 
-function leaseSessionEnd(sessionStore: SessionStore, teardownSession: SessionTeardown) {
+function leaseSessionEnd(
+  sessionStore: SessionStore,
+  teardownSession: SessionTeardown,
+  cleanupSessionlessExecutionHost: PlatformResourceCleanup['cleanupSessionlessExecutionHost'] = vi.fn(
+    async () => {},
+  ),
+) {
   const executionLocks = new Map<string, Promise<unknown>>();
   return {
     executionLocks,
@@ -76,12 +83,259 @@ function leaseSessionEnd(sessionStore: SessionStore, teardownSession: SessionTea
       sessionStore,
       executionLocks,
       teardownSession,
+      cleanupSessionlessExecutionHost,
     }),
   };
 }
 
 const settle = async (): Promise<void> => await new Promise((resolve) => setImmediate(resolve));
 
+function retainExecutionHost(sessionStore: SessionStore, name = 'leased', leaseId = 'lease-1') {
+  const ref = sessionStore.publish(
+    name,
+    makeIosSession(name, {
+      lease: { leaseId, tenantId: 'tenant-a', runId: 'run-1' },
+    }),
+  );
+  sessionStore.retire(ref, { retainExecutionHost: true });
+  return ref;
+}
+
+test.each(['LEASE_RELEASED', 'LEASE_EXPIRED'] as const)(
+  '%s stops a retained host after its session has retired',
+  async (reason) => {
+    const sessionStore = makeSessionStore();
+    const ref = retainExecutionHost(sessionStore);
+    const cleanup = vi.fn(async () => {});
+    const teardown = vi.fn(async () => {});
+    const { ender } = leaseSessionEnd(sessionStore, teardown, cleanup);
+
+    await ender.end('lease-1', reason);
+
+    expect(teardown).not.toHaveBeenCalled();
+    expect(cleanup).toHaveBeenCalledExactlyOnceWith(ref.session.device);
+    expect(sessionStore.listRetainedExecutionHosts('lease-1')).toEqual([]);
+  },
+);
+
+test('expiry waiting behind close cleans the host that close retains before retiring', async () => {
+  const sessionStore = makeSessionStore();
+  const ref = sessionStore.publish(
+    'leased',
+    makeIosSession('leased', {
+      lease: { leaseId: 'lease-1', tenantId: 'tenant-a', runId: 'run-1' },
+    }),
+  );
+  const cleanup = vi.fn(async () => {});
+  const teardown = vi.fn(async () => {});
+  const { executionLocks, ender } = leaseSessionEnd(sessionStore, teardown, cleanup);
+  const close = holdSessionLocks(executionLocks, sessionStore, 'leased');
+  const ending = ender.end('lease-1', 'LEASE_EXPIRED');
+  await settle();
+  expect(cleanup).not.toHaveBeenCalled();
+
+  sessionStore.retire(ref, { retainExecutionHost: true });
+  close.finish();
+  await close.held;
+  await ending;
+
+  expect(teardown).not.toHaveBeenCalled();
+  expect(cleanup).toHaveBeenCalledExactlyOnceWith(ref.session.device);
+  expect(sessionStore.listRetainedExecutionHosts('lease-1')).toEqual([]);
+});
+
+test.each([false, true])(
+  'lease end preserves a replacement owner with retained=%s',
+  async (retained) => {
+    const sessionStore = makeSessionStore();
+    const old = retainExecutionHost(sessionStore);
+    const cleanup = vi.fn(async () => {});
+    const { executionLocks, ender } = leaseSessionEnd(sessionStore, async () => {}, cleanup);
+    let unlock!: () => void;
+    const held = withRequestExecutionLockKeys(
+      executionLocks,
+      existingSessionExecutionLockKeys(old.address, old.session.device.id),
+      async () =>
+        await new Promise<void>((resolve) => {
+          unlock = resolve;
+        }),
+    );
+    await settle();
+    const ending = ender.end('lease-1', 'LEASE_RELEASED');
+    await settle();
+    const replacement = sessionStore.publish(
+      'replacement',
+      makeIosSession('replacement', {
+        lease: { leaseId: 'lease-2', tenantId: 'tenant-b', runId: 'run-2' },
+      }),
+    );
+    if (retained) sessionStore.retire(replacement, { retainExecutionHost: true });
+    unlock();
+    await held;
+    await ending;
+
+    expect(cleanup).not.toHaveBeenCalled();
+    expect(sessionStore.listRetainedExecutionHosts('lease-1')).toEqual([]);
+    if (retained) expect(sessionStore.listRetainedExecutionHosts('lease-2')).toHaveLength(1);
+    else expect(sessionStore.resolveCurrent(replacement)).toBeDefined();
+  },
+);
+
+test('release ends a same-lease reopen published at a new address while cleanup waits for its device', async () => {
+  const sessionStore = makeSessionStore();
+  const old = retainExecutionHost(sessionStore);
+  const cleanup = vi.fn(async () => {});
+  const teardown = vi.fn(async () => {});
+  const { executionLocks, ender } = leaseSessionEnd(sessionStore, teardown, cleanup);
+  let finishNativeOpen!: () => void;
+  const nativeOpen = new Promise<void>((resolve) => {
+    finishNativeOpen = resolve;
+  });
+  const reopening = withRequestExecutionLockKeys(
+    executionLocks,
+    existingSessionExecutionLockKeys('reopened', old.session.device.id),
+    async () => {
+      await nativeOpen;
+      sessionStore.publish('reopened', makeIosSession('reopened', { lease: old.session.lease }));
+    },
+  );
+  await settle();
+  const ending = ender.end('lease-1', 'LEASE_RELEASED');
+  await settle();
+  expect(teardown).not.toHaveBeenCalled();
+  finishNativeOpen();
+  await reopening;
+  await ending;
+
+  expect(cleanup).not.toHaveBeenCalled();
+  expect(teardown).toHaveBeenCalledExactlyOnceWith(
+    expect.objectContaining({ address: 'reopened' }),
+  );
+  expect(sessionStore.get('reopened')).toBeUndefined();
+  expect(sessionStore.listRetainedExecutionHosts('lease-1')).toEqual([]);
+});
+
+test('failed retained-host cleanup stays owned and throttles retries without another release request', async () => {
+  const now = vi.spyOn(Date, 'now').mockReturnValue(0);
+  try {
+    const sessionStore = makeSessionStore();
+    retainExecutionHost(sessionStore);
+    const cleanup = vi
+      .fn(async () => {})
+      .mockRejectedValueOnce(new Error('device unreachable'))
+      .mockRejectedValueOnce(new Error('device still unreachable'));
+    const { ender } = leaseSessionEnd(sessionStore, async () => {}, cleanup);
+
+    await ender.end('lease-1', 'LEASE_RELEASED');
+    expect(sessionStore.listRetainedExecutionHosts('lease-1')).toHaveLength(1);
+    for (let elapsed = 1_000; elapsed < 10_000; elapsed += 1_000) {
+      now.mockReturnValue(elapsed);
+      await ender.retryPending();
+      expect(cleanup).toHaveBeenCalledTimes(1);
+    }
+    now.mockReturnValue(10_000);
+    await ender.retryPending();
+    expect(cleanup).toHaveBeenCalledTimes(2);
+    expect(sessionStore.listRetainedExecutionHosts('lease-1')).toHaveLength(1);
+    now.mockReturnValue(19_999);
+    await ender.retryPending();
+    expect(cleanup).toHaveBeenCalledTimes(2);
+    now.mockReturnValue(20_000);
+    await ender.retryPending();
+    expect(cleanup).toHaveBeenCalledTimes(3);
+    expect(sessionStore.listRetainedExecutionHosts('lease-1')).toEqual([]);
+    now.mockReturnValue(30_000);
+    await ender.retryPending();
+    expect(cleanup).toHaveBeenCalledTimes(3);
+  } finally {
+    now.mockRestore();
+  }
+});
+
+test('cancel preserves a retained host queued for cleanup for daemon shutdown', async () => {
+  const sessionStore = makeSessionStore();
+  const ref = retainExecutionHost(sessionStore);
+  const cleanup = vi.fn(async () => {});
+  const { executionLocks, ender } = leaseSessionEnd(sessionStore, async () => {}, cleanup);
+  let unlock!: () => void;
+  const held = withRequestExecutionLockKeys(
+    executionLocks,
+    existingSessionExecutionLockKeys(ref.address, ref.session.device.id),
+    async () =>
+      await new Promise<void>((resolve) => {
+        unlock = resolve;
+      }),
+  );
+  await settle();
+  const ending = ender.end('lease-1', 'LEASE_EXPIRED');
+  await settle();
+  ender.cancel();
+  unlock();
+  await held;
+  await ending;
+  await ender.drain();
+  await ender.retryPending();
+
+  expect(cleanup).not.toHaveBeenCalled();
+  expect(sessionStore.listRetainedExecutionHosts('lease-1')).toHaveLength(1);
+});
+
+test('retry does not queue duplicate cleanup and shutdown drains a running retained-host stop', async () => {
+  const sessionStore = makeSessionStore();
+  retainExecutionHost(sessionStore);
+  let finish!: () => void;
+  const cleanup = vi.fn(
+    async () =>
+      await new Promise<void>((resolve) => {
+        finish = resolve;
+      }),
+  );
+  const { ender } = leaseSessionEnd(sessionStore, async () => {}, cleanup);
+  const ending = ender.end('lease-1', 'LEASE_EXPIRED');
+  await settle();
+  await ender.retryPending();
+  expect(cleanup).toHaveBeenCalledTimes(1);
+  ender.cancel();
+  let drained = false;
+  const drain = ender.drain().then(() => {
+    drained = true;
+  });
+  await settle();
+  expect(drained).toBe(false);
+  finish();
+  await ending;
+  await drain;
+  expect(sessionStore.listRetainedExecutionHosts('lease-1')).toEqual([]);
+});
+
+test('a request holding the device lock can end its lease ahead of a queued background end', async () => {
+  const sessionStore = makeSessionStore();
+  const ref = retainExecutionHost(sessionStore);
+  const cleanup = vi.fn(async () => {});
+  const { executionLocks, ender } = leaseSessionEnd(sessionStore, async () => {}, cleanup);
+  let release!: () => void;
+  const continueRequest = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  const request = withRequestExecutionLockKeys(
+    executionLocks,
+    existingSessionExecutionLockKeys(ref.address, ref.session.device.id),
+    async () => {
+      await continueRequest;
+      await ender.end('lease-1', 'LEASE_RELEASED');
+    },
+  );
+  await settle();
+  const expiry = ender.end('lease-1', 'LEASE_EXPIRED');
+  await settle();
+  release();
+  await request;
+  await expiry;
+
+  expect(cleanup).toHaveBeenCalledTimes(1);
+  expect(sessionStore.listRetainedExecutionHosts('lease-1')).toEqual([]);
+});
+
 test('end publishes the recorded script of the session it ends', async () => {
   const sessionStore = makeSessionStore('agent-device-lease-end-script-');
   await publishLeasedSession(sessionStore, 'leased', 'lease-1');