diff --git a/docs/adr/0007-remote-device-leases.md b/docs/adr/0007-remote-device-leases.md index d64f08b6ae..d386648e08 100644 --- a/docs/adr/0007-remote-device-leases.md +++ b/docs/adr/0007-remote-device-leases.md @@ -67,6 +67,16 @@ one minute, while a cloud WebDriver connection profile asks for ten. A single co longer than its own lease is therefore ordinary on the default and only reachable through a profile on the longer one. +## Provider allocation + +Allocation is admitted work too. The registry records a lease before a hosted provider creates the +session behind it, and creating that session can take longer than the lease's inactivity TTL, so a +lease timed from its record was already expired, or nearly so, when its client first received it, and +the paid session it pointed at was orphaned. The allocation therefore holds a work pass while the +provider allocates: the lease cannot expire underneath it, and a successful allocation still wanted by +its requester starts the inactivity TTL from the moment allocation completed. A requester that hung up +preserves nothing, and its allocation is released as before. + ## Client-side work that precedes admission Protecting admitted work covers nothing that happens before a request is admitted. Installing an diff --git a/src/daemon/handlers/__tests__/lease-artifacts.test.ts b/src/daemon/handlers/__tests__/lease-artifacts.test.ts index 8067a4fdae..a6fd63a376 100644 --- a/src/daemon/handlers/__tests__/lease-artifacts.test.ts +++ b/src/daemon/handlers/__tests__/lease-artifacts.test.ts @@ -2,7 +2,8 @@ import assert from 'node:assert/strict'; import { test } from 'vitest'; import type { CloudArtifactsQuery } from '@agent-device/contracts/observability'; import type { DeviceLease } from '@agent-device/contracts/device'; -import { AppError } from '@agent-device/kernel/errors'; +import { AppError, isRequestCanceledError } from '@agent-device/kernel/errors'; +import { clearRequestCanceled, markRequestCanceled } from '@agent-device/host-kit/request'; import { makeSessionStore } from '../../../__tests__/test-utils/store-factory.ts'; import type { DaemonRequest, DaemonResponse } from '../../daemon-request.ts'; import { handleLeaseCommands } from '../lease.ts'; @@ -117,7 +118,7 @@ test('artifacts refuses an expired provider session after retention before lazy assert.deepEqual(world.providerCalls, []); }); -test('artifacts refuses a provider session returned after allocation expiry retention', async () => { +test('artifacts lists a provider session whose allocation outlasted the lease TTL', async () => { let now = 1_000; const world = createWorld({ now: () => now, @@ -128,16 +129,60 @@ test('artifacts refuses a provider session returned after allocation expiry rete }); world.lifecycle.allocate = async (lease) => { now = lease.expiresAt + 51; - return { providerSessionId: 'late-allocation-session' }; + return { providerSessionId: 'slow-allocation-session' }; }; await allocateLease(world, 'tenant-a', 'run-a'); - await assertProviderSessionNotOwned(world, { + const listed = await listArtifacts(world, { tenantId: 'tenant-a', runId: 'run-a', - providerSessionId: 'late-allocation-session', + providerSessionId: 'slow-allocation-session', }); - assert.deepEqual(world.providerCalls, []); + assert.equal(listed.ok, true); +}); + +test('artifacts refuses a provider session returned after a canceled allocation outlived retention', async () => { + let now = 1_000; + const world = createWorld({ + now: () => now, + defaultLeaseTtlMs: 100, + minLeaseTtlMs: 1, + maxLeaseTtlMs: 100, + providerSessionRetentionMs: 50, + }); + const requestId = 'late-canceled-allocation'; + world.lifecycle.allocate = async (lease) => { + markRequestCanceled(requestId); + now = lease.expiresAt + 51; + return { providerSessionId: 'late-allocation-session' }; + }; + + try { + await assert.rejects( + handleLeaseCommands({ + req: leaseRequest('lease_allocate', { + requestId, + tenantId: 'tenant-a', + runId: 'run-a', + leaseBackend: 'android-instance', + leaseProvider: CLOUD_PROVIDER, + }), + sessionName: 'artifact-test', + sessionStore: world.sessionStore, + leaseRegistry: world.leaseRegistry, + leaseLifecycleProvider: world.lifecycle, + }), + isRequestCanceledError, + ); + await assertProviderSessionNotOwned(world, { + tenantId: 'tenant-a', + runId: 'run-a', + providerSessionId: 'late-allocation-session', + }); + assert.deepEqual(world.providerCalls, []); + } finally { + clearRequestCanceled(requestId); + } }); test('artifacts refuses a provider session returned after release expiry retention', async () => { diff --git a/src/daemon/handlers/__tests__/lease.test.ts b/src/daemon/handlers/__tests__/lease.test.ts index 08fdb22c7f..28346814a9 100644 --- a/src/daemon/handlers/__tests__/lease.test.ts +++ b/src/daemon/handlers/__tests__/lease.test.ts @@ -2,7 +2,9 @@ import assert from 'node:assert/strict'; import { test } from 'vitest'; import { handleLeaseCommands } from '../lease.ts'; import { LeaseRegistry } from '../../lease-registry.ts'; +import type { DaemonRequest } from '../../daemon-request.ts'; import { makeSessionStore } from '../../../__tests__/test-utils/store-factory.ts'; +import type { DeviceLease } from '@agent-device/contracts/device'; import { AppError } from '@agent-device/kernel/errors'; import { clearRequestCanceled, markRequestCanceled } from '@agent-device/host-kit/request'; import { @@ -93,3 +95,83 @@ test('activation drains canceled provider allocation and its release cleanup', a clearRequestCanceled(requestId); } }); + +function allocateRequest(): DaemonRequest { + return { + token: 'test-token', + session: 'lease-ttl-test', + command: 'lease_allocate', + positionals: [], + flags: {}, + meta: { + tenantId: 'tenant-a', + runId: 'run-a', + clientId: 'client-a', + leaseBackend: 'android-instance', + leaseProvider: 'cloud', + }, + }; +} + +// A hosted provider can spend longer creating its session than the lease's inactivity +// TTL. Stamping the TTL when the registry record is created handed the client a lease +// that was already expired, and the paid session behind it was orphaned. +test('a lease whose provider allocation outlasts its TTL is active when allocation returns', async () => { + let now = 0; + const registry = new LeaseRegistry({ now: () => now, defaultLeaseTtlMs: 60_000 }); + const response = await handleLeaseCommands({ + req: allocateRequest(), + sessionName: 'lease-ttl-test', + sessionStore: makeSessionStore('agent-device-slow-provider-'), + leaseRegistry: registry, + leaseLifecycleProvider: { + allocate: async (lease) => { + now = 80_000; + assert.deepEqual( + registry.consumeExpiredLeases(), + [], + 'the sweeper must not reap a lease mid-allocation', + ); + now = 90_000; + return { providerSessionId: `session-${lease.leaseId}` }; + }, + }, + }); + + assert.equal(response?.ok, true); + const lease = (response?.ok ? response.data?.lease : undefined) as DeviceLease; + assert.equal(lease.expiresAt, 150_000); + assert.deepEqual( + registry.listActiveLeases().map((entry) => [entry.leaseId, entry.expiresAt]), + [[lease.leaseId, 150_000]], + ); + now = 149_999; + registry.assertLeaseAdmission({ + leaseId: lease.leaseId, + tenantId: lease.tenantId, + runId: lease.runId, + leaseBackend: lease.backend, + leaseProvider: lease.leaseProvider, + }); +}); + +test('a lease allocated without a provider keeps the TTL it was created with', async () => { + let now = 5_000; + const registry = new LeaseRegistry({ + now: () => (now += 1_000), + defaultLeaseTtlMs: 60_000, + }); + const response = await handleLeaseCommands({ + req: allocateRequest(), + sessionName: 'lease-ttl-test', + sessionStore: makeSessionStore('agent-device-no-provider-'), + leaseRegistry: registry, + }); + + assert.equal(response?.ok, true); + const lease = (response?.ok ? response.data?.lease : undefined) as DeviceLease; + assert.ok(now > lease.createdAt, 'the clock advanced while the lease was allocated'); + assert.equal(lease.heartbeatAt, lease.createdAt); + assert.equal(lease.expiresAt, lease.createdAt + 60_000); + assert.deepEqual(registry.listActiveLeases(), [lease]); +}); diff --git a/src/daemon/handlers/lease.ts b/src/daemon/handlers/lease.ts index fa597823e2..0aaac48f63 100644 --- a/src/daemon/handlers/lease.ts +++ b/src/daemon/handlers/lease.ts @@ -71,20 +71,28 @@ export async function handleLeaseCommands(args: LeaseHandlerArgs): Promise { let providerData: Record | undefined; + // A hosted provider can take longer than the lease TTL to create its session; the work + // pass keeps the lease alive until it does, and ending the pass restarts the TTL then. + const work = leaseLifecycleProvider?.allocate + ? leaseRegistry.retainLeaseWork(lease, () => !isRequestCanceled(requestId)) + : undefined; try { providerData = await leaseLifecycleProvider?.allocate?.(lease, { ...leaseLifecycleContext(req), - signal: getRequestSignal(req.meta?.requestId), + signal: getRequestSignal(requestId), deadline: Date.now() + LEASE_ALLOCATION_BUDGET_MS, }); recordProviderSession(leaseRegistry, lease, providerData); } catch (error) { leaseRegistry.releaseLease(leaseReleaseRequestFor(lease)); throw error; + } finally { + work?.release(); } - if (isRequestCanceled(req.meta?.requestId)) { + if (isRequestCanceled(requestId)) { // The requester left while the provider was allocating; the lease it // produced is real (and billed) and nobody will ever release it. throw await releaseAllocationForGoneRequester( @@ -93,9 +101,10 @@ export async function handleLeaseCommands(args: LeaseHandlerArgs): Promise` instead of exporting the environment variable also works, but only authenticates the single command it is passed to; subsequent commands need the token again through the env var, a `daemonAuthToken` entry in your remote config profile, or a repeated `--daemon-auth-token` flag. -`connect proxy` stores the proxy profile and client identity. Device leases are automatic on `open` and expire after five minutes without commands. That five minutes is the window `open` asks for; a lease allocated directly over the RPC without `ttlMs` keeps the daemon's one-minute inactivity default instead. `close` releases the active session and device lease; `disconnect` clears local connection state. +`connect proxy` stores the proxy profile and client identity. Device leases are automatic on `open` and expire after five minutes without commands. That five minutes is the window `open` asks for; a lease allocated directly over the RPC without `ttlMs` keeps the daemon's one-minute inactivity default instead. Either window starts when allocation completes, not when it was requested. `close` releases the active session and device lease; `disconnect` clears local connection state. Multiple agents can share one proxy when each uses the normal `connect proxy`, `open`, commands, `close`, and `disconnect` flow. A busy device error means another agent owns the device until it closes or its inactivity lease expires.