diff --git a/src/daemon/__tests__/lease-lifecycle.test.ts b/src/daemon/__tests__/lease-lifecycle.test.ts index 8fce339a62..0ef074372c 100644 --- a/src/daemon/__tests__/lease-lifecycle.test.ts +++ b/src/daemon/__tests__/lease-lifecycle.test.ts @@ -83,7 +83,9 @@ test('cleanupExpiredLeasedSession consumes expired lease and deletes the session }); expect(cleaned).toBe(true); - expect(teardownSession).toHaveBeenCalledWith(session, 'default'); + expect(teardownSession).toHaveBeenCalledWith( + expect.objectContaining({ address: 'default', session }), + ); expect(sessionStore.get('default')).toBeUndefined(); expect(leaseRegistry.listActiveLeases()).toHaveLength(0); }); diff --git a/src/daemon/__tests__/request-execution-scope-lease-expiry.test.ts b/src/daemon/__tests__/request-execution-scope-lease-expiry.test.ts new file mode 100644 index 0000000000..16d48f801d --- /dev/null +++ b/src/daemon/__tests__/request-execution-scope-lease-expiry.test.ts @@ -0,0 +1,132 @@ +import { expect, test, vi } from 'vitest'; +import { 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 { createRequestExecutionScope } from '../request-execution-scope.ts'; +import type { DaemonRequest } from '../daemon-request.ts'; + +function makeRequest(overrides: Partial): DaemonRequest { + return { + token: 't', + session: 'default', + command: 'snapshot', + positionals: [], + flags: {}, + ...overrides, + }; +} + +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, + leaseBackend: lease.backend, + leaseProvider: 'proxy', + deviceKey: 'ios:SIM-001', + expiresAt: lease.expiresAt, + }, + }), + ); + now = 1_011; + + const scope = await createRequestExecutionScope({ + req: makeRequest({ command: 'snapshot' }), + sessionStore, + leaseRegistry, + }); + await scope.runLocked(async () => 'ran'); + + expect(sessionStore.get('default')).toBeUndefined(); + const nextLease = leaseRegistry.allocateLease({ tenantId: 'tenant-b', runId: 'run-2' }); + expect(nextLease.tenantId).toBe('tenant-b'); +}); + +test.each(['rebuild', 'retire'] as const)( + 'scoped lease expiry preserves a public-name session and a %s during held teardown', + async (change) => { + let now = 1_000; + const store = makeSessionStore('request-scope-lease-lifetime-'); + const leases = new LeaseRegistry({ defaultLeaseTtlMs: 10, minLeaseTtlMs: 1, now: () => now }); + const lease = leases.allocateLease({ tenantId: 'tenant-a', runId: 'run-1' }); + const address = 'cwd:ownership:default'; + const original = store.publish( + address, + makeSession('default', { + device: LINUX_DEVICE, + sessionScope: { kind: 'cwd', id: 'ownership' }, + lease: { + leaseId: lease.leaseId, + tenantId: lease.tenantId, + runId: lease.runId, + leaseBackend: lease.backend, + }, + }), + ); + const decoy = store.publish( + 'default', + makeSession('default', { + device: { ...LINUX_DEVICE, id: 'public-name-decoy' }, + appName: 'decoy', + }), + ); + let enter!: () => void; + let resume!: () => void; + const entered = new Promise((resolve) => { + enter = resolve; + }); + const release = new Promise((resolve) => { + resume = resolve; + }); + const stopSnapshotHelper = vi.fn(async () => { + enter(); + await release; + }); + now = 1_011; + const scope = await createRequestExecutionScope({ + req: makeRequest({ session: address, flags: { session: address } }), + sessionStore: store, + leaseRegistry: leases, + platformResourceCleanup: { + stopSnapshotHelper, + closeManagedBrowser: async () => {}, + cleanupSessionlessExecutionHost: async () => {}, + retainExecutionHostAfterClose: () => false, + }, + }); + expect(scope.sessionName).toBe(address); + const running = scope.runLocked(async () => 'ran'); + await entered; + let successor: ReturnType | undefined; + if (change === 'retire') { + store.retire(original); + successor = store.publish( + address, + makeSession('default', { + device: { ...LINUX_DEVICE, id: 'successor' }, + appName: 'successor', + }), + ); + } else store.update(original, { appName: 'latest' }); + resume(); + await expect(running).resolves.toBe('ran'); + expect(stopSnapshotHelper).toHaveBeenCalledExactlyOnceWith(original.session.device); + expect(store.requireCurrent(decoy)).toBe(decoy.session); + if (successor) expect(store.requireCurrent(successor)).toBe(successor.session); + else expect(store.lookup(address)).toBeUndefined(); + }, +); diff --git a/src/daemon/__tests__/request-execution-scope.test.ts b/src/daemon/__tests__/request-execution-scope.test.ts index 5da8da19d9..a4f0ce5a32 100644 --- a/src/daemon/__tests__/request-execution-scope.test.ts +++ b/src/daemon/__tests__/request-execution-scope.test.ts @@ -9,9 +9,7 @@ import { 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 '@agent-device/host-kit/request'; @@ -497,45 +495,6 @@ test('provider lease admission succeeds without a device key', async () => { 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, - leaseBackend: lease.backend, - leaseProvider: 'proxy', - deviceKey: 'ios:SIM-001', - expiresAt: lease.expiresAt, - }, - }), - ); - now = 1_011; - - const scope = await createRequestExecutionScope({ - req: makeRequest({ command: 'snapshot' }), - sessionStore, - leaseRegistry, - }); - await scope.runLocked(async () => 'ran'); - - expect(sessionStore.get('default')).toBeUndefined(); - const nextLease = leaseRegistry.allocateLease({ tenantId: 'tenant-b', runId: 'run-2' }); - expect(nextLease.tenantId).toBe('tenant-b'); -}); - // A lease renewed only at admission lets one command slower than its inactivity TTL // expire the lease paying for the device it is using, and expiry then tears the // provider session down under the client still waiting for that same command. Found diff --git a/src/daemon/lease-lifecycle.ts b/src/daemon/lease-lifecycle.ts index cc70befb02..04866c889b 100644 --- a/src/daemon/lease-lifecycle.ts +++ b/src/daemon/lease-lifecycle.ts @@ -14,12 +14,12 @@ import { } from './request-admission.ts'; import type { SessionStore } from './session-store.ts'; import type { DaemonRequest } from './daemon-request.ts'; -import type { SessionState } from './session-state.ts'; +import type { SessionRef, SessionState } from './session-state.ts'; import { providerSessionIdFromData } from './provider-session-ownership.ts'; export type ExpiredProviderLeaseRecovery = (lease: DeviceLease) => Promise; -export type SessionTeardown = (session: SessionState, sessionName: string) => Promise; +export type SessionTeardown = (ref: SessionRef) => Promise; export async function releaseExpiredProviderLease( recoverExpiredLease: ExpiredProviderLeaseRecovery | undefined, @@ -62,9 +62,10 @@ export async function cleanupExpiredLeasedSession(params: { leaseRegistry: LeaseRegistry; teardownSession: SessionTeardown; }): Promise { - const session = params.sessionStore.get(params.sessionName); + const ref = params.sessionStore.lookup(params.sessionName); + const session = ref?.session; const lease = session?.lease; - if (!session || !lease) return false; + if (!ref || !session || !lease) return false; const expiredLease = params.leaseRegistry.consumeExpiredLease(lease.leaseId); if (!expiredLease) return false; emitDiagnostic({ @@ -77,7 +78,7 @@ export async function cleanupExpiredLeasedSession(params: { deviceKey: lease.deviceKey, }, }); - await params.teardownSession(session, session.name).catch((error) => { + await params.teardownSession(ref).catch((error) => { emitDiagnostic({ level: 'debug', phase: 'leased_session_expiry_cleanup_failed', @@ -100,7 +101,7 @@ export async function cleanupExpiredLeasedSession(params: { }, }); }); - params.sessionStore.delete(session.name); + params.sessionStore.retire(ref); return true; } @@ -112,7 +113,8 @@ export function admitRequestLeaseForLockedScope(params: { providerAppCatalog?: ProviderAppCatalog; }): DaemonRequest { const { sessionName, sessionStore, leaseRegistry } = params; - const existingSession = sessionStore.get(sessionName); + const ref = sessionStore.lookup(sessionName); + const existingSession = ref?.session; const activeLease = assertRequestLeaseAdmission(params.req, leaseRegistry, existingSession, { providerAppCatalog: params.providerAppCatalog, }); @@ -125,15 +127,14 @@ export function admitRequestLeaseForLockedScope(params: { admittedLease: activeLease, }, }; - if (existingSession?.lease) { - sessionStore.set(sessionName, { - ...existingSession, + if (ref && existingSession?.lease) { + sessionStore.update(ref, (current) => ({ lease: { - ...existingSession.lease, + ...current.lease!, leaseBackend: activeLease.backend, expiresAt: activeLease.expiresAt, }, - }); + })); } return nextReq; } diff --git a/src/daemon/request-execution-scope.ts b/src/daemon/request-execution-scope.ts index b19eaa3bcd..4430b38d86 100644 --- a/src/daemon/request-execution-scope.ts +++ b/src/daemon/request-execution-scope.ts @@ -45,7 +45,7 @@ import type { LeaseRegistry } from './lease-registry.ts'; import { type SessionStore } from './session-store.ts'; import { resolveSessionRequestLog, resolveSessionRunnerLogPath } from './session-artifact-paths.ts'; import type { DaemonRequest, DaemonResponse } from './daemon-request.ts'; -import type { SessionState } from './session-state.ts'; +import type { SessionRef, SessionState } from './session-state.ts'; import { teardownSessionResources } from './session-teardown.ts'; import { finalizeBoundSessionApplicationLifecycle } from './application-lifecycle-recovery.ts'; import { runtimeHintValues } from './session-runtime.ts'; @@ -273,10 +273,9 @@ export async function createRequestExecutionScope(params: { sessionName, sessionStore, leaseRegistry, - teardownSession: async (session, expiredSessionName) => + teardownSession: async (ref) => await teardownExpiredSession({ - session, - sessionName: expiredSessionName, + ref, sessionStore, inspectFacts: scope.inspectFacts, bindDevice: scope.bindDevice, @@ -410,16 +409,16 @@ function createRequestDeviceAccess(params: { } async function teardownExpiredSession(params: { - session: SessionState; - sessionName: string; + ref: SessionRef; sessionStore: SessionStore; inspectFacts: InspectDeviceRuntimeFacts; bindDevice: BindDeviceRuntime; platformCleanup: PlatformResourceCleanup; }): Promise { - const { session, sessionName, sessionStore, inspectFacts, bindDevice, platformCleanup } = params; - const ref = sessionStore.lookup(sessionName); - if (!ref) return; + const { ref, sessionStore, inspectFacts, bindDevice, platformCleanup } = params; + const session = sessionStore.resolveCurrent(ref) ?? ref.session; + const sessionName = ref.address; + const runtimeHints = runtimeHintValues(sessionStore.getRuntimeHints(ref.address)); let primaryError: unknown; try { await teardownSessionResources({ @@ -437,7 +436,7 @@ async function teardownExpiredSession(params: { bindDevice, session, stateDir: sessionStore.resolveDaemonStateDir(), - runtimeHints: runtimeHintValues(sessionStore.getRuntimeHints(sessionName)), + runtimeHints, }); } catch (cleanupError) { if (primaryError !== undefined) {