fix: enforce device leases for daemon sessions

src/daemon/__tests__/request-execution-scope.test.ts

@@ -6,7 +6,9 @@ import { flushDiagnosticsToSessionFile, withDiagnosticsScope } from '../../utils
 import {
   makeAndroidSession,
   makeIosSession,
+  makeSession,
 } from '../../__tests__/test-utils/session-factories.ts';
+import { LINUX_DEVICE } from '../../__tests__/test-utils/device-fixtures.ts';
 import { makeSessionStore } from '../../__tests__/test-utils/store-factory.ts';
 import { LeaseRegistry } from '../lease-registry.ts';
 import { clearRequestCanceled, markRequestCanceled } from '../request-cancel.ts';
@@ -114,6 +116,179 @@ test('createRequestExecutionScope rejects tenant requests without an active leas
   ).rejects.toThrow(/Lease is not active/);
 });
 
+test('leased session admission uses stored lease metadata and heartbeats', async () => {
+  let now = 1_000;
+  const sessionStore = makeSessionStore('agent-device-request-scope-');
+  const leaseRegistry = new LeaseRegistry({ now: () => now });
+  const lease = leaseRegistry.allocateLease({ tenantId: 'tenant-a', runId: 'run-1' });
+  sessionStore.set(
+    'default',
+    makeIosSession('default', {
+      lease: {
+        leaseId: lease.leaseId,
+        tenantId: lease.tenantId,
+        runId: lease.runId,
+        backend: lease.backend,
+        leaseProvider: 'proxy',
+        clientId: 'client-a',
+        deviceKey: 'ios:sim-1',
+        expiresAt: lease.expiresAt,
+      },
+    }),
+  );
+  now = 2_000;
+
+  const scope = await createRequestExecutionScope({
+    req: makeRequest({ command: 'snapshot' }),
+    sessionStore,
+    leaseRegistry,
+  });
+
+  expect(scope.sessionName).toBe('default');
+  const activeLease = leaseRegistry.listActiveLeases()[0];
+  expect(activeLease?.heartbeatAt).toBe(2_000);
+  expect(activeLease?.expiresAt).toBe(302_000);
+  expect(sessionStore.get('default')?.lease?.expiresAt).toBe(302_000);
+});
+
+test('leased session rejects mismatched lease id before dispatch', async () => {
+  const sessionStore = makeSessionStore('agent-device-request-scope-');
+  const leaseRegistry = new LeaseRegistry();
+  const lease = leaseRegistry.allocateLease({ tenantId: 'tenant-a', runId: 'run-1' });
+  sessionStore.set(
+    'default',
+    makeIosSession('default', {
+      lease: {
+        leaseId: lease.leaseId,
+        tenantId: lease.tenantId,
+        runId: lease.runId,
+        backend: lease.backend,
+      },
+    }),
+  );
+
+  await expect(
+    createRequestExecutionScope({
+      req: makeRequest({ command: 'snapshot', meta: { leaseId: '1'.repeat(32) } }),
+      sessionStore,
+      leaseRegistry,
+    }),
+  ).rejects.toThrow(/Lease does not match session owner \(leaseId\)/);
+});
+
+test.each([
+  ['leaseProvider', { leaseProvider: 'cloud' }],
+  ['clientId', { clientId: 'client-b' }],
+  ['deviceKey', { deviceKey: 'ios:SIM-002' }],
+] as const)('leased session rejects mismatched %s before dispatch', async (_field, meta) => {
+  const sessionStore = makeSessionStore('agent-device-request-scope-');
+  const leaseRegistry = new LeaseRegistry();
+  const lease = leaseRegistry.allocateLease({ tenantId: 'tenant-a', runId: 'run-1' });
+  sessionStore.set(
+    'default',
+    makeIosSession('default', {
+      lease: {
+        leaseId: lease.leaseId,
+        tenantId: lease.tenantId,
+        runId: lease.runId,
+        backend: lease.backend,
+        leaseProvider: 'proxy',
+        clientId: 'client-a',
+        deviceKey: 'ios:SIM-001',
+      },
+    }),
+  );
+
+  await expect(
+    createRequestExecutionScope({
+      req: makeRequest({ command: 'snapshot', meta }),
+      sessionStore,
+      leaseRegistry,
+    }),
+  ).rejects.toThrow(/Lease does not match session owner/);
+});
+
+test('local unleased session admission still succeeds', async () => {
+  const sessionStore = makeSessionStore('agent-device-request-scope-');
+  sessionStore.set('default', makeIosSession('default'));
+
+  const scope = await createRequestExecutionScope({
+    req: makeRequest({ command: 'snapshot' }),
+    sessionStore,
+    leaseRegistry: new LeaseRegistry(),
+  });
+
+  expect(scope.sessionName).toBe('default');
+});
+
+test('provider lease admission succeeds without a device key', async () => {
+  const sessionStore = makeSessionStore('agent-device-request-scope-');
+  const leaseRegistry = new LeaseRegistry();
+  const lease = leaseRegistry.allocateLease({
+    tenantId: 'tenant-a',
+    runId: 'run-1',
+    backend: 'android-instance',
+  });
+  sessionStore.set(
+    'default',
+    makeAndroidSession('default', {
+      lease: {
+        leaseId: lease.leaseId,
+        tenantId: lease.tenantId,
+        runId: lease.runId,
+        backend: lease.backend,
+        leaseProvider: 'limrun',
+      },
+    }),
+  );
+
+  const scope = await createRequestExecutionScope({
+    req: makeRequest({ command: 'snapshot' }),
+    sessionStore,
+    leaseRegistry,
+  });
+
+  expect(scope.sessionName).toBe('default');
+});
+
+test('expired leases remove owned sessions before the next command and free capacity', async () => {
+  let now = 1_000;
+  const sessionStore = makeSessionStore('agent-device-request-scope-');
+  const leaseRegistry = new LeaseRegistry({
+    maxActiveSimulatorLeases: 1,
+    defaultLeaseTtlMs: 10,
+    minLeaseTtlMs: 1,
+    now: () => now,
+  });
+  const lease = leaseRegistry.allocateLease({ tenantId: 'tenant-a', runId: 'run-1' });
+  sessionStore.set(
+    'default',
+    makeSession('default', {
+      device: LINUX_DEVICE,
+      lease: {
+        leaseId: lease.leaseId,
+        tenantId: lease.tenantId,
+        runId: lease.runId,
+        backend: lease.backend,
+        leaseProvider: 'proxy',
+        deviceKey: 'ios:SIM-001',
+        expiresAt: lease.expiresAt,
+      },
+    }),
+  );
+  now = 1_011;
+
+  await createRequestExecutionScope({
+    req: makeRequest({ command: 'snapshot' }),
+    sessionStore,
+    leaseRegistry,
+  });
+
+  expect(sessionStore.get('default')).toBeUndefined();
+  const nextLease = leaseRegistry.allocateLease({ tenantId: 'tenant-b', runId: 'run-2' });
+  expect(nextLease.tenantId).toBe('tenant-b');
+});
+
 test('tenant lease rejection flushes diagnostics into the effective session request log', async () => {
   const sessionStore = makeSessionStore('agent-device-request-scope-');
   const requestId = 'tenant-lease-rejection';

src/daemon/__tests__/request-router-open.test.ts

@@ -29,24 +29,32 @@ function makeIosDevice(id: string): DeviceInfo {
   };
 }
 
-function createOpenHandler(sessionStore: ReturnType<typeof makeSessionStore>) {
+function createOpenHandler(
+  sessionStore: ReturnType<typeof makeSessionStore>,
+  leaseRegistry = new LeaseRegistry(),
+) {
   return createRequestHandler({
     logPath: path.join(os.tmpdir(), 'daemon.log'),
     token: 'test-token',
     sessionStore,
-    leaseRegistry: new LeaseRegistry(),
+    leaseRegistry,
     trackDownloadableArtifact: () => 'artifact-id',
   });
 }
 
-function openRequest(session: string, flags: Record<string, unknown>, requestId: string) {
+function openRequest(
+  session: string,
+  flags: Record<string, unknown>,
+  requestId: string,
+  meta: Record<string, unknown> = {},
+) {
   return {
     token: 'test-token',
     session,
     command: 'open',
     positionals: [],
     flags,
-    meta: { requestId },
+    meta: { requestId, ...meta },
   };
 }
 
@@ -82,6 +90,128 @@ test('open returns and creates the session state directory', async () => {
   }
 });
 
+test('open stores admitted lease metadata on the session', async () => {
+  const sessionStore = makeSessionStore('agent-device-router-open-');
+  const leaseRegistry = new LeaseRegistry();
+  const lease = leaseRegistry.allocateLease({ tenantId: 'tenant-a', runId: 'run-1' });
+  const device = makeIosDevice('SIM-LEASED');
+  mockResolveTargetDevice.mockResolvedValue(device);
+
+  const handler = createOpenHandler(sessionStore, leaseRegistry);
+
+  const response = await handler(
+    openRequest('default', { platform: 'ios' }, 'req-open-lease', {
+      tenantId: 'tenant-a',
+      runId: 'run-1',
+      leaseId: lease.leaseId,
+      sessionIsolation: 'tenant',
+      leaseProvider: 'proxy',
+      clientId: 'client-a',
+      deviceKey: 'ios:SIM-LEASED',
+      leaseBackend: 'ios-simulator',
+    }),
+  );
+
+  expect(response.ok).toBe(true);
+  expect(sessionStore.get('tenant-a:default')?.lease).toEqual({
+    leaseId: lease.leaseId,
+    tenantId: 'tenant-a',
+    runId: 'run-1',
+    backend: 'ios-simulator',
+    leaseProvider: 'proxy',
+    clientId: 'client-a',
+    deviceKey: 'ios:SIM-LEASED',
+    expiresAt: undefined,
+  });
+});
+
+test('proxy open without required lease metadata fails before device resolution', async () => {
+  const sessionStore = makeSessionStore('agent-device-router-open-');
+  const handler = createOpenHandler(sessionStore, new LeaseRegistry());
+
+  const response = await handler(
+    openRequest('default', { platform: 'ios' }, 'req-open-proxy-missing', {
+      tenantId: 'tenant-a',
+      runId: 'run-1',
+      leaseProvider: 'proxy',
+      sessionIsolation: 'tenant',
+    }),
+  );
+
+  expect(response.ok).toBe(false);
+  if (!response.ok) {
+    expect(response.error.code).toBe('INVALID_ARGS');
+    expect(response.error.message).toMatch(/Proxy open requires leaseId/);
+  }
+  expect(mockResolveTargetDevice).not.toHaveBeenCalled();
+  expect(mockDispatch).not.toHaveBeenCalled();
+});
+
+test('close releases the session lease', async () => {
+  const sessionStore = makeSessionStore('agent-device-router-open-');
+  const leaseRegistry = new LeaseRegistry();
+  const lease = leaseRegistry.allocateLease({ tenantId: 'tenant-a', runId: 'run-1' });
+  sessionStore.set('default', {
+    name: 'default',
+    device: makeIosDevice('SIM-CLOSE'),
+    createdAt: Date.now(),
+    actions: [],
+    lease: {
+      leaseId: lease.leaseId,
+      tenantId: lease.tenantId,
+      runId: lease.runId,
+      backend: lease.backend,
+      clientId: 'client-a',
+    },
+  });
+  const handler = createOpenHandler(sessionStore, leaseRegistry);
+
+  const response = await handler({
+    token: 'test-token',
+    session: 'default',
+    command: 'close',
+    positionals: [],
+    meta: { requestId: 'req-close-lease' },
+  });
+
+  expect(response.ok).toBe(true);
+  expect(sessionStore.get('default')).toBeUndefined();
+  expect(leaseRegistry.listActiveLeases()).toHaveLength(0);
+});
+
+test('close rejects a different client before cleanup', async () => {
+  const sessionStore = makeSessionStore('agent-device-router-open-');
+  const leaseRegistry = new LeaseRegistry();
+  const lease = leaseRegistry.allocateLease({ tenantId: 'tenant-a', runId: 'run-1' });
+  sessionStore.set('default', {
+    name: 'default',
+    device: makeIosDevice('SIM-CLOSE-CLIENT'),
+    createdAt: Date.now(),
+    actions: [],
+    lease: {
+      leaseId: lease.leaseId,
+      tenantId: lease.tenantId,
+      runId: lease.runId,
+      backend: lease.backend,
+      clientId: 'client-a',
+    },
+  });
+  const handler = createOpenHandler(sessionStore, leaseRegistry);
+
+  const response = await handler({
+    token: 'test-token',
+    session: 'default',
+    command: 'close',
+    positionals: [],
+    meta: { requestId: 'req-close-wrong-client', clientId: 'client-b' },
+  });
+
+  expect(response.ok).toBe(false);
+  expect(sessionStore.get('default')).toBeDefined();
+  expect(leaseRegistry.listActiveLeases()).toHaveLength(1);
+  expect(mockDispatch).not.toHaveBeenCalled();
+});
+
 test('router serializes same-device open requests before first session creation finishes', async () => {
   const sessionStore = makeSessionStore('agent-device-router-open-');
   const sameDevice = makeIosDevice('SIM-001');

src/daemon/__tests__/session-store.test.ts

@@ -108,6 +108,32 @@ test('defaultTracePath sanitizes session name', () => {
   assert.match(tracePath, /\.trace\.log$/);
 });
 
+test('session lease metadata round-trips through the store', () => {
+  const { store, session } = makeFixture('agent-device-session-lease-');
+  session.lease = {
+    leaseId: 'f'.repeat(32),
+    tenantId: 'tenant-a',
+    runId: 'run-1',
+    clientId: 'client-a',
+    backend: 'ios-simulator',
+    leaseProvider: 'proxy',
+    deviceKey: 'ios:SIM-001',
+    expiresAt: 123_456,
+  };
+
+  store.set(session.name, session);
+
+  assert.deepEqual(store.get(session.name)?.lease, session.lease);
+});
+
+test('sessions without lease metadata remain valid', () => {
+  const { store, session } = makeFixture('agent-device-session-unleased-');
+
+  store.set(session.name, session);
+
+  assert.equal(store.get(session.name)?.lease, undefined);
+});
+
 test('saveScript flag enables .ad session log writing', () => {
   const { root, store, session } = makeFixture('agent-device-session-log-enabled-');
   recordOpen(store, session);