diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0c71247..3952d6c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -37,7 +37,7 @@ jobs: run: git diff --exit-code -- specs/testing-contract-fixtures.json test: - name: Test + name: test runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 diff --git a/control-plane/src/domain/errors.ts b/control-plane/src/domain/errors.ts index bdb5995..7214292 100644 --- a/control-plane/src/domain/errors.ts +++ b/control-plane/src/domain/errors.ts @@ -74,6 +74,8 @@ export const unauthorized = (message: string): TalosError => export const forbidden = (message: string): TalosError => new TalosError('forbidden', message, 403); export const modeForbidden = (message: string): TalosError => new TalosError('mode_forbidden', message, 403); export const conflict = (message: string): TalosError => new TalosError('conflict', message, 409); +export const concurrentUpdate = (message = 'task state changed concurrently'): TalosError => + new TalosError('concurrent_update', message, 409); export const actionAlreadyCompleted = (): TalosError => new TalosError('action_already_completed', 'session action result was already stored', 409); export const taskCancelled = (message = 'task was cancelled'): TalosError => new TalosError('task_cancelled', message, 409); diff --git a/control-plane/src/domain/types.ts b/control-plane/src/domain/types.ts index bb59c66..bda4bb5 100644 --- a/control-plane/src/domain/types.ts +++ b/control-plane/src/domain/types.ts @@ -14,6 +14,29 @@ export type TaskStatus = | 'failed' | 'cancelled'; +export type TaskClaimRecoveryReason = + | 'partial_claim_identity' + | 'invalid_claim_generation' + | 'invalid_claim_credentials' + | 'invalid_lease_expiry' + | 'active_claim_missing_credentials' + | 'active_claim_marked_released' + | 'legacy_profile_identity_conflict'; + +export interface TaskClaimRecovery { + schemaVersion: 'talos.task-claim-recovery/v1'; + recoveryId: string; + kind: 'legacy' | 'malformed'; + phase: 'draining' | 'finalizing' | 'quarantined'; + sourceStatus: TaskStatus; + sourceMachineId?: string; + sourceProfileId?: string; + restoredQueuePriority: number; + reasonCode?: TaskClaimRecoveryReason; + startedAt: string; + updatedAt: string; +} + export type CapabilityTag = | 'os' | 'region' @@ -65,6 +88,12 @@ interface TaskBase { claimedAt?: string; leaseExpiresAt?: string; leaseToken?: string; + claimId?: string; + claimGeneration?: number; + taskVersion?: number; + claimCommitted?: boolean; + claimReleased?: boolean; + claimQueuePriority?: number; workerId?: string; machineId?: string; findings: readonly TaskFinding[]; @@ -74,6 +103,7 @@ interface TaskBase { handoff?: { url: string; expiresAt: string }; pendingActionId?: string; lastActionId?: string; + claimRecovery?: TaskClaimRecovery; } export interface BrowserTask extends TaskBase { @@ -142,6 +172,7 @@ export interface Machine { tags: Readonly>; capacity: number; activeLeases: number; + leaseReservations?: readonly MachineLeaseReservation[]; online: boolean; workerTokenHash: string; } @@ -151,9 +182,39 @@ export interface Profile { userId: string; machineId?: string; lockedByTaskId?: string; + lockedByClaimId?: string; + lockedByClaimGeneration?: number; lockExpiresAt?: string; } +export interface TaskClaimGuard { + claimId: string; + claimGeneration: number; + taskVersion: number; + status: TaskStatus; +} + +export interface TaskActiveClaimGuard extends TaskClaimGuard { + leaseExpiresAt: string; +} + +export interface TaskRecoveryGuard { + status: TaskStatus; + taskVersion: number; + updatedAt: string; + claimId?: string; + claimGeneration?: number; + recoveryId?: string; + recoveryPhase?: TaskClaimRecovery['phase']; +} + +export interface MachineLeaseReservation { + claimId: string; + claimGeneration: number; + taskId: string; + expiresAt: string; +} + export interface HandoffLink { id: string; taskId: string; diff --git a/control-plane/src/http/server.test.ts b/control-plane/src/http/server.test.ts index 244f89b..4a4df40 100644 --- a/control-plane/src/http/server.test.ts +++ b/control-plane/src/http/server.test.ts @@ -5,6 +5,7 @@ import { Scheduler } from '../services/scheduler.js'; import { TaskService } from '../services/task-service.js'; import { WebhookSigner } from '../services/webhook-signer.js'; import { MemoryRepository } from '../storage/memory-repository.js'; +import type { Repository } from '../storage/repository.js'; import { createApiServer } from './server.js'; import { loadOpenApiDocument } from '../openapi.js'; @@ -58,6 +59,114 @@ describe('control-plane HTTP API', () => { server.close(); }); + it('maps repository failures to an opaque public error', async () => { + const repository = new MemoryRepository(); + repository.getTask = async () => { throw new Error('claim-secret-sentinel lease-token-sentinel'); }; + const service = new TaskService(repository, new Scheduler(repository), new ProfileLockService(repository), new WebhookSigner('webhook-secret-1234')); + const server = createApiServer(service, repository); + await new Promise((resolve) => server.listen(0, resolve)); + const address = server.address(); + if (address === null || typeof address === 'string') throw new Error('server did not bind'); + + const response = await fetch(`http://127.0.0.1:${address.port}/v1/tasks/task`, { + headers: { 'x-nyxid-identity-token': 'user:user-a' } + }); + expect(response.status).toBe(500); + const body = JSON.stringify(await response.json()); + expect(body).toContain('internal_error'); + expect(body).not.toContain('claim-secret-sentinel'); + expect(body).not.toContain('lease-token-sentinel'); + server.close(); + }); + + it('returns a retryable public conflict after authorized task CAS exhaustion', async () => { + const storage = new MemoryRepository(); + let attempts = 0; + const repository = new Proxy(storage, { + get(target, property) { + if (property === 'replaceSubmittedTask') { + return async (...args: Parameters): Promise => { + if (args[0].status === 'cancelled') { + attempts += 1; + return false; + } + return target.replaceSubmittedTask(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('webhook-secret-1234') + ); + const server = createApiServer(service, repository); + await new Promise((resolve) => server.listen(0, resolve)); + const address = server.address(); + if (address === null || typeof address === 'string') throw new Error('server did not bind'); + const base = `http://127.0.0.1:${address.port}`; + const headers = { 'content-type': 'application/json', 'x-nyxid-identity-token': 'user:user-a' }; + const created = await fetch(`${base}/v1/tasks`, { + method: 'POST', + headers, + body: JSON.stringify({ kind: 'browse', goal: 'public CAS exhaustion' }) + }); + const task = await created.json() as { id: string }; + + const response = await fetch(`${base}/v1/tasks/${task.id}/cancel`, { method: 'POST', headers }); + expect(response.status).toBe(409); + expect(await response.json()).toEqual({ + error: { + code: 'concurrent_update', + message: 'task state changed concurrently', + retryable: true + } + }); + expect(attempts).toBe(3); + server.close(); + }); + + it('reports an expired handoff claim as a non-retryable public conflict', async () => { + const clock = { value: Date.now() }; + const repository = new MemoryRepository(() => clock.value); + await repository.savePool({ id: 'expired-handoff-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'expired-handoff-machine', poolId: 'expired-handoff-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('webhook-secret-1234'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + const task = await service.createTask('user-a', { kind: 'browse', goal: 'expired handoff' }); + const claim = await service.claim('worker-a', 'expired-handoff-machine', clock.value); + clock.value = Date.parse(claim.task.leaseExpiresAt!); + const server = createApiServer(service, repository, { clock: () => clock.value }); + await new Promise((resolve) => server.listen(0, resolve)); + const address = server.address(); + if (address === null || typeof address === 'string') throw new Error('server did not bind'); + const base = `http://127.0.0.1:${address.port}`; + const response = await fetch(`${base}/v1/tasks/${task.id}/handoff`, { + method: 'POST', + headers: { 'content-type': 'application/json', 'x-nyxid-identity-token': 'user:user-a' }, + body: '{}' + }); + + expect(response.status).toBe(409); + expect(await response.json()).toEqual({ + error: { + code: 'conflict', + message: 'task claim is not active', + retryable: false + } + }); + expect(await repository.getTask(task.id)).not.toHaveProperty('handoff'); + server.close(); + }); + it('enforces NyxID and worker authentication across lifecycle routes', async () => { const repository = new MemoryRepository(); await repository.savePool({ id: 'pool', visibility: 'platform', tags: {} }); @@ -77,8 +186,30 @@ describe('control-plane HTTP API', () => { expect(badWorker.status).toBe(401); const claimResponse = await fetch(`${base}/v1/worker/claim`, { method: 'POST', headers: { authorization: 'Bearer worker-token-123456', 'x-talos-worker-id': 'w', 'x-talos-machine-id': 'machine', 'content-type': 'application/json' }, body: JSON.stringify({ worker_id: 'w', machine_id: 'machine' }) }); expect(claimResponse.status).toBe(200); - const claim = await claimResponse.json() as { task: { id: string }; leaseToken: string }; + const claim = await claimResponse.json() as { task: Record & { id: string }; leaseToken: string }; expect(claim.task.id).toBe(created.id); + const internalAuthorityFields = [ + 'claimId', + 'claimGeneration', + 'taskVersion', + 'claimCommitted', + 'claimReleased', + 'claimQueuePriority', + 'queuePriority', + 'workerId', + 'machineId', + 'leaseExpiresAt', + 'leaseToken', + 'claimRecovery' + ]; + for (const field of internalAuthorityFields) expect(claim.task).not.toHaveProperty(field); + const publicTaskResponse = await fetch(`${base}/v1/tasks/${created.id}`, { + headers: { 'x-nyxid-identity-token': 'user:user-a' } + }); + expect(publicTaskResponse.status).toBe(200); + const publicTask = await publicTaskResponse.json() as Record; + for (const field of internalAuthorityFields) expect(publicTask).not.toHaveProperty(field); + expect(JSON.stringify(publicTask)).not.toContain(claim.leaseToken); const heartbeat = await fetch(`${base}/v1/worker/tasks/${created.id}/heartbeat`, { method: 'POST', headers: { authorization: 'Bearer worker-token-123456', 'x-talos-worker-id': 'w', 'x-talos-machine-id': 'machine', 'content-type': 'application/json' }, body: JSON.stringify({ lease_token: claim.leaseToken }) }); expect(heartbeat.status).toBe(200); server.close(); @@ -218,6 +349,7 @@ describe('control-plane HTTP API', () => { expect((await fetch(`${base}/v1/pools`, { method: 'POST', headers: user('bob'), body: JSON.stringify({ id: 'bob-pool' }) })).status).toBe(201); expect((await fetch(`${base}/v1/pools/bob-pool/machines`, { method: 'POST', headers: user('bob'), body: JSON.stringify({ id: 'bob-machine' }) })).status).toBe(201); expect((await fetch(`${base}/v1/profiles`, { method: 'POST', headers: user('bob'), body: JSON.stringify({ id: 'bob-profile', machine_id: 'bob-machine' }) })).status).toBe(201); + expect((await fetch(`${base}/v1/profiles`, { method: 'POST', headers: user('bob'), body: JSON.stringify({ id: 'bob-profile' }) })).status).toBe(409); const bobPools = await fetch(`${base}/v1/pools`, { headers: user('bob') }); expect(bobPools.status).toBe(200); expect((await bobPools.json() as Array<{ id: string }>).map((pool) => pool.id)).toEqual(['bob-pool']); @@ -290,7 +422,7 @@ describe('control-plane HTTP API', () => { const repository = new MemoryRepository(); await repository.savePool({ id: 'pool', visibility: 'platform', tags: {} }); await repository.saveMachine({ id: 'machine', poolId: 'pool', tags: {}, capacity: 2, activeLeases: 0, online: true, workerTokenHash: hashWorkerToken('worker-token-123456') }); - await repository.saveProfile({ id: 'p', userId: 'u' }); + await repository.createProfile({ id: 'p', userId: 'u' }); const service = new TaskService(repository, new Scheduler(repository), new ProfileLockService(repository), new WebhookSigner('webhook-secret-1234')); const server = createApiServer(service, repository); await new Promise((resolve) => server.listen(0, resolve)); diff --git a/control-plane/src/http/server.ts b/control-plane/src/http/server.ts index 0877981..35b7cbb 100644 --- a/control-plane/src/http/server.ts +++ b/control-plane/src/http/server.ts @@ -217,7 +217,9 @@ const route = async ( await assertPoolOwner(repository, machine.poolId, userId); selfRotateMachineSchema.parse(await readBody(request, options.maxBodyBytes)); const workerToken = issueWorkerToken(); - await repository.saveMachine({ ...machine, workerTokenHash: hashWorkerToken(workerToken) }); + if (!await repository.rotateMachineToken(machine.id, machine.workerTokenHash, hashWorkerToken(workerToken))) { + throw conflict('machine token changed concurrently'); + } return send(response, 200, { id: machine.id, rotated: true, worker_token: workerToken }); } if (parts[1] === 'profiles' && parts.length === 2 && method === 'POST') { @@ -228,8 +230,9 @@ const route = async ( await assertPoolOwner(repository, machine.poolId, userId); } const id = input.id ?? newId('profile'); - if (await repository.getProfile(id) !== undefined) throw conflict('profile already exists'); - await repository.saveProfile({ id, userId, ...(input.machine_id === undefined ? {} : { machineId: input.machine_id }) }); + if (!await repository.createProfile({ id, userId, ...(input.machine_id === undefined ? {} : { machineId: input.machine_id }) })) { + throw conflict('profile already exists'); + } return send(response, 201, { id, userId, @@ -385,7 +388,9 @@ const adminRoute = async ( const machine = await repository.getMachine(parts[3]); if (machine === undefined) throw notFound('machine not found'); const workerToken = input.worker_token ?? issueWorkerToken(); - await repository.saveMachine({ ...machine, workerTokenHash: hashWorkerToken(workerToken) }); + if (!await repository.rotateMachineToken(machine.id, machine.workerTokenHash, hashWorkerToken(workerToken))) { + throw conflict('machine token changed concurrently'); + } return send(response, 200, { id: machine.id, rotated: true, worker_token: workerToken }); } if (parts[2] === 'machines') { @@ -398,8 +403,9 @@ const adminRoute = async ( } if (parts[2] === 'profiles') { const input = adminProfileSchema.parse(await readBody(request, options.maxBodyBytes)); - if (await repository.getProfile(input.id) !== undefined) throw conflict('profile already exists'); - await repository.saveProfile({ id: input.id, userId: input.user_id, ...(input.machine_id === undefined ? {} : { machineId: input.machine_id }) }); + if (!await repository.createProfile({ id: input.id, userId: input.user_id, ...(input.machine_id === undefined ? {} : { machineId: input.machine_id }) })) { + throw conflict('profile already exists'); + } return send(response, 201, { id: input.id }); } return send(response, 404, publicErrorEnvelope('not_found', 'route not found', 404)); diff --git a/control-plane/src/http/session-routes.integration.test.ts b/control-plane/src/http/session-routes.integration.test.ts index bad8cbd..98dbd15 100644 --- a/control-plane/src/http/session-routes.integration.test.ts +++ b/control-plane/src/http/session-routes.integration.test.ts @@ -87,6 +87,22 @@ describe('interactive session HTTP API', () => { }) }); const claim = await claimResponse.json() as { leaseToken: string }; + const publicSession = await (await fetch(`${base}/v1/sessions/${session.id}`, { headers: alice })).json() as Record; + for (const field of [ + 'claimId', + 'claimGeneration', + 'taskVersion', + 'claimCommitted', + 'claimReleased', + 'claimQueuePriority', + 'queuePriority', + 'workerId', + 'machineId', + 'leaseExpiresAt', + 'leaseToken', + 'claimRecovery' + ]) expect(publicSession).not.toHaveProperty(field); + expect(JSON.stringify(publicSession)).not.toContain(claim.leaseToken); const forbiddenAction = await fetch(`${base}/v1/sessions/${session.id}/actions?wait_seconds=0`, { method: 'POST', headers: alice, diff --git a/control-plane/src/index.test.ts b/control-plane/src/index.test.ts index 477abf2..4c9c8ab 100644 --- a/control-plane/src/index.test.ts +++ b/control-plane/src/index.test.ts @@ -188,7 +188,7 @@ describe('control-plane factory', () => { const repository = new MemoryRepository(); let releaseFirst: ((tasks: readonly []) => void) | undefined; const blocked = new Promise((resolve) => { releaseFirst = resolve; }); - const listTasks = vi.spyOn(repository, 'listTasks') + const listExpirableTasks = vi.spyOn(repository, 'listExpirableTasks') .mockImplementationOnce(async () => blocked) .mockResolvedValue([]); const server = createControlPlane(repository, 'webhook-secret-1234', { @@ -197,13 +197,13 @@ describe('control-plane factory', () => { }); await vi.advanceTimersByTimeAsync(11); - expect(listTasks).toHaveBeenCalledTimes(1); + expect(listExpirableTasks).toHaveBeenCalledTimes(1); await vi.advanceTimersByTimeAsync(50); - expect(listTasks).toHaveBeenCalledTimes(1); + expect(listExpirableTasks).toHaveBeenCalledTimes(1); releaseFirst?.([]); await Promise.resolve(); await vi.advanceTimersByTimeAsync(11); - expect(listTasks).toHaveBeenCalledTimes(2); + expect(listExpirableTasks).toHaveBeenCalledTimes(2); server.stopSweep(); }); @@ -211,11 +211,11 @@ describe('control-plane factory', () => { vi.useFakeTimers(); const repository = new MemoryRepository(); const stderr = vi.spyOn(process.stderr, 'write').mockImplementation(() => true); - const listTasks = vi.spyOn(repository, 'listTasks') - .mockRejectedValueOnce(new Error('task sweep unavailable')) + const listTaskMaintenancePage = vi.spyOn(repository, 'listTaskMaintenancePage') + .mockRejectedValueOnce(new Error('claim-secret-sentinel lease-token-sentinel')) .mockResolvedValue([]); const listTestingRuns = vi.spyOn(repository, 'listTestingRuns') - .mockRejectedValueOnce(new Error('testing sweep unavailable')) + .mockRejectedValueOnce(new Error('testing-fence-secret-sentinel')) .mockResolvedValue([]); const server = createControlPlane(repository, 'webhook-secret-1234', { sweepIntervalMs: 10, @@ -224,10 +224,16 @@ describe('control-plane factory', () => { await vi.advanceTimersByTimeAsync(11); await vi.advanceTimersByTimeAsync(11); - expect(listTasks.mock.calls.length).toBeGreaterThanOrEqual(2); + expect(listTaskMaintenancePage.mock.calls.length).toBeGreaterThanOrEqual(2); expect(listTestingRuns.mock.calls.length).toBeGreaterThanOrEqual(2); - expect(stderr.mock.calls.flat().join('\n')).toContain('task lease sweep failed'); - expect(stderr.mock.calls.flat().join('\n')).toContain('testing attempt sweep failed'); + const output = stderr.mock.calls.flat().join('\n'); + expect(output).toContain('task lease sweep failed'); + expect(output).toContain('task_lease_sweep_failed'); + expect(output).toContain('testing attempt sweep failed'); + expect(output).toContain('testing_attempt_sweep_failed'); + expect(output).not.toContain('claim-secret-sentinel'); + expect(output).not.toContain('lease-token-sentinel'); + expect(output).not.toContain('testing-fence-secret-sentinel'); server.stopSweep(); }); }); diff --git a/control-plane/src/index.ts b/control-plane/src/index.ts index 92d8b3c..ed260ca 100644 --- a/control-plane/src/index.ts +++ b/control-plane/src/index.ts @@ -115,13 +115,13 @@ export const createControlPlane = ( try { try { await service.expireLeases(); - } catch (error) { - logger.error('task lease sweep failed', { error: errorMessage(error) }); + } catch { + logger.error('task lease sweep failed', { error: 'task_lease_sweep_failed' }); } try { await testingAttempts.sweep(); - } catch (error) { - logger.error('testing attempt sweep failed', { error: errorMessage(error) }); + } catch { + logger.error('testing attempt sweep failed', { error: 'testing_attempt_sweep_failed' }); } } finally { sweepInFlight = false; @@ -193,5 +193,3 @@ export * from './http/server.js'; export * from './http/testing-run-routes.js'; export * from './identity.js'; export * from './openapi.js'; - -const errorMessage = (error: unknown): string => error instanceof Error ? error.message : String(error); diff --git a/control-plane/src/openapi.test.ts b/control-plane/src/openapi.test.ts index 58f3ed2..8a46c1f 100644 --- a/control-plane/src/openapi.test.ts +++ b/control-plane/src/openapi.test.ts @@ -70,6 +70,39 @@ describe('OpenAPI loader', () => { } const schema = (name: string): Record => asObject(parsed.components.schemas[name]); const properties = (name: string): Record => asObject(schema(name).properties); + const internalTaskAuthorityFields = [ + 'claimId', + 'claimGeneration', + 'taskVersion', + 'claimCommitted', + 'claimReleased', + 'claimQueuePriority', + 'queuePriority', + 'workerId', + 'machineId', + 'leaseExpiresAt', + 'leaseToken', + 'claimRecovery' + ]; + for (const publicSchema of ['Task', 'Session']) { + const publicProperties = properties(publicSchema); + expect(schema(publicSchema).additionalProperties, publicSchema).toBe(false); + for (const field of internalTaskAuthorityFields) { + expect(publicProperties, `${publicSchema}.${field}`).not.toHaveProperty(field); + } + } + expect(schema('ClaimResponse').additionalProperties).toBe(false); + expect(schema('Lease').additionalProperties).toBe(false); + for (const closedSchema of [ + 'PublicErrorDetail', + 'WebhookEvent', + 'StateChangedPayload', + 'NeedsInputPayload', + 'HandoffRequestedPayload', + 'CompletedPayload', + 'WebhookDelivery', + 'TaskError' + ]) expect(schema(closedSchema).additionalProperties, closedSchema).toBe(false); for (const strictSchema of [ 'TestingToolRequest', 'TestingCapabilities', diff --git a/control-plane/src/services/profile-lock.ts b/control-plane/src/services/profile-lock.ts index b224923..960109f 100644 --- a/control-plane/src/services/profile-lock.ts +++ b/control-plane/src/services/profile-lock.ts @@ -1,9 +1,9 @@ import { conflict, forbidden, notFound } from '../domain/errors.js'; -import type { Profile } from '../domain/types.js'; +import type { MachineLeaseReservation, Profile } from '../domain/types.js'; import type { Repository } from '../storage/repository.js'; export class ProfileLockService { - public constructor(private readonly repository: Repository, private readonly leaseSeconds = 300) {} + public constructor(private readonly repository: Repository, _leaseSeconds = 300) {} public async assertOwner(profileId: string, userId: string): Promise { const profile = await this.repository.getProfile(profileId); @@ -12,25 +12,18 @@ export class ProfileLockService { return profile; } - public async acquire(profileId: string, userId: string, taskId: string, now = Date.now(), machineId?: string, leaseSeconds = this.leaseSeconds): Promise { - const profile = await this.assertOwner(profileId, userId); - if (profile.lockedByTaskId !== undefined && profile.lockExpiresAt !== undefined && Date.parse(profile.lockExpiresAt) > now && profile.lockedByTaskId !== taskId) { - throw conflict('profile already has an active session'); - } - const next: Profile = { ...profile, ...(machineId === undefined ? {} : { machineId }), lockedByTaskId: taskId, lockExpiresAt: new Date(now + leaseSeconds * 1000).toISOString() }; - await this.repository.saveProfile(next); - return next; - } - - public async release(profileId: string, taskId: string): Promise { - const profile = await this.repository.getProfile(profileId); - if (profile?.lockedByTaskId === taskId) await this.repository.saveProfile({ id: profile.id, userId: profile.userId, ...(profile.machineId === undefined ? {} : { machineId: profile.machineId }) }); + public async acquire( + profileId: string, + userId: string, + machineId: string, + reservation: MachineLeaseReservation + ): Promise { + const profile = await this.repository.acquireProfileLease(profileId, userId, machineId, reservation); + if (profile === undefined) throw conflict('profile already has an active session'); + return profile; } - public async renew(profileId: string, taskId: string, now = Date.now(), leaseSeconds = this.leaseSeconds): Promise { - const profile = await this.repository.getProfile(profileId); - if (profile?.lockedByTaskId === taskId) { - await this.repository.saveProfile({ ...profile, lockExpiresAt: new Date(now + leaseSeconds * 1000).toISOString() }); - } + public async release(profileId: string, reservation: Omit): Promise { + return this.repository.releaseProfileLease(profileId, reservation); } } diff --git a/control-plane/src/services/scheduler.test.ts b/control-plane/src/services/scheduler.test.ts index 6bb486a..01f355f 100644 --- a/control-plane/src/services/scheduler.test.ts +++ b/control-plane/src/services/scheduler.test.ts @@ -32,7 +32,7 @@ describe('Scheduler eligibility', () => { await repository.savePool({ id: 'platform', visibility: 'platform', tags: {} }); await repository.saveMachine({ id: 'a', poolId: 'private', tags: { os: 'macos' }, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'x' }); await repository.saveMachine({ id: 'b', poolId: 'platform', tags: { os: 'linux', computer_use: true }, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'x' }); - await repository.saveProfile({ id: 'p', userId: 'u', machineId: 'b' }); + await repository.createProfile({ id: 'p', userId: 'u', machineId: 'b' }); expect(await scheduler.isEligible(task({ os: 'linux' }), 'b', 'u')).toBeDefined(); expect(await scheduler.isEligible(task({ os: 'macos' }), 'b', 'u')).toBeUndefined(); expect(await scheduler.isEligible(task({ os: 'macos' }), 'a', 'u')).toBeUndefined(); diff --git a/control-plane/src/services/task-service.test.ts b/control-plane/src/services/task-service.test.ts index 57ce525..0c1c11a 100644 --- a/control-plane/src/services/task-service.test.ts +++ b/control-plane/src/services/task-service.test.ts @@ -6,16 +6,75 @@ import { Scheduler } from './scheduler.js'; import { TaskService } from './task-service.js'; import { WebhookSigner } from './webhook-signer.js'; import type { Task } from '../domain/types.js'; +import type { Repository } from '../storage/repository.js'; import { computeTestingTaskPayloadDigest, testingTaskSchema } from '@talos/testing-protocol'; const setup = (clock: { value: number } = { value: Date.now() }) => { - const repository = new MemoryRepository(); + const repository = new MemoryRepository(() => clock.value); const profiles = new ProfileLockService(repository); const scheduler = new Scheduler(repository); const service = new TaskService(repository, scheduler, profiles, new WebhookSigner('test-webhook-secret'), { clock: () => clock.value, leaseSeconds: 10 }); return { repository, profiles, service, clock }; }; +const advanceClaimGeneration = async ( + repository: Repository, + taskId: string, + status: 'running' | 'needs_input' +): Promise => { + const current = await repository.getTask(taskId); + if ( + current === undefined || + current.claimId === undefined || + current.claimGeneration === undefined + ) throw new Error('test task does not have an active claim'); + const requeued = { + ...current, + status: 'submitted' as const, + workerId: undefined, + leaseToken: undefined, + leaseExpiresAt: undefined, + claimCommitted: false, + claimReleased: true, + queuePriority: current.claimQueuePriority, + claimQueuePriority: undefined + }; + if (!await repository.replaceTaskForClaim(requeued, { + claimId: current.claimId, + claimGeneration: current.claimGeneration, + taskVersion: current.taskVersion ?? 0, + status: current.status + })) throw new Error('test could not requeue the current claim'); + const submitted = await repository.getTask(taskId); + if (submitted === undefined) throw new Error('test task disappeared while requeueing'); + const nextGeneration = current.claimGeneration + 1; + const reclaimed = await repository.claimTask({ + ...submitted, + status: 'claimed', + workerId: 'worker-next', + machineId: current.machineId ?? 'machine', + leaseToken: 'lease-next', + leaseExpiresAt: '2100-01-01T00:00:00.000Z', + claimId: 'claim-next', + claimGeneration: nextGeneration, + taskVersion: (submitted.taskVersion ?? 0) + 1, + claimCommitted: true, + claimReleased: false, + claimQueuePriority: submitted.queuePriority, + queuePriority: undefined + }, current.claimGeneration, submitted.taskVersion ?? 0); + if (reclaimed === undefined) throw new Error('test could not claim the next generation'); + if (!await repository.replaceTaskForClaim({ ...reclaimed, status }, { + claimId: reclaimed.claimId!, + claimGeneration: reclaimed.claimGeneration!, + taskVersion: reclaimed.taskVersion ?? 0, + status: reclaimed.status + })) throw new Error('test could not advance the next generation status'); + const persisted = await repository.getTask(taskId); + if (persisted === undefined) throw new Error('test task disappeared after reclaim'); + return persisted; +}; + describe('task service', () => { it('runs submit, claim, heartbeat, result and preserves identity', async () => { const { repository, service } = setup(); @@ -30,6 +89,101 @@ describe('task service', () => { await expect(service.getTask(task.id, 'user-b')).rejects.toMatchObject({ code: 'forbidden' }); }); + it('rejects heartbeat when its claim CAS resumes after lease expiry', async () => { + const clock = { value: 1_000 }; + const storage = new MemoryRepository(() => clock.value); + let reachClaimCas!: () => void; + let resumeClaimCas!: () => void; + const claimCasReached = new Promise((resolve) => { reachClaimCas = resolve; }); + const claimCasResume = new Promise((resolve) => { resumeClaimCas = resolve; }); + const repository = new Proxy(storage, { + get(target, property) { + if (property === 'replaceTaskForActiveClaim') { + return async (task: Parameters[0], guard: Parameters[1]) => { + reachClaimCas(); + await claimCasResume; + return target.replaceTaskForActiveClaim(task, guard); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + await repository.savePool({ id: 'expiry-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'expiry-machine', poolId: 'expiry-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'x' }); + const task = await service.createTask('user-a', { kind: 'browse', goal: 'expiry race' }); + const claim = await service.claim('worker-a', 'expiry-machine'); + + const heartbeat = service.heartbeat(task.id, 'worker-a', claim.leaseToken, 30); + await claimCasReached; + clock.value = Date.parse(claim.task.leaseExpiresAt!); + resumeClaimCas(); + + await expect(heartbeat).rejects.toMatchObject({ code: 'unauthorized' }); + expect(await repository.getTask(task.id)).toMatchObject({ + status: 'claimed', + leaseExpiresAt: claim.task.leaseExpiresAt, + taskVersion: claim.task.taskVersion + }); + }); + + it('does not take over a stale profile projection while its task renewal is authoritative', async () => { + const clock = { value: 1_000 }; + const storage = new MemoryRepository(() => clock.value); + let renewalCommitted!: () => void; + let resumeHeartbeat!: () => void; + const renewalCommit = new Promise((resolve) => { renewalCommitted = resolve; }); + const heartbeatResume = new Promise((resolve) => { resumeHeartbeat = resolve; }); + const repository = new Proxy(storage, { + get(target, property) { + if (property === 'replaceTaskForActiveClaim') { + return async (task: Parameters[0], guard: Parameters[1]) => { + const replaced = await target.replaceTaskForActiveClaim(task, guard); + renewalCommitted(); + await heartbeatResume; + return replaced; + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + await repository.createProfile({ id: 'renewal-profile', userId: 'user-a' }); + await repository.savePool({ id: 'renewal-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'renewal-machine', poolId: 'renewal-pool', tags: {}, capacity: 2, activeLeases: 0, online: true, workerTokenHash: 'x' }); + const first = await service.createTask('user-a', { kind: 'browse', goal: 'first', profile_id: 'renewal-profile' }); + await service.createTask('user-a', { kind: 'browse', goal: 'second', profile_id: 'renewal-profile' }); + const claim = await service.claim('worker-a', 'renewal-machine'); + + const heartbeat = service.heartbeat(first.id, 'worker-a', claim.leaseToken, 30); + await renewalCommit; + clock.value = Date.parse(claim.task.leaseExpiresAt!); + try { + await expect(service.claim('worker-b', 'renewal-machine')).rejects.toMatchObject({ code: 'not_found' }); + } finally { + resumeHeartbeat(); + } + await expect(heartbeat).resolves.toMatchObject({ status: 'running' }); + expect(await repository.getProfile('renewal-profile')).toMatchObject({ + lockedByTaskId: first.id, + lockedByClaimId: claim.task.claimId + }); + }); + it('redacts private input and scheduling fields from public tasks', async () => { const { repository, service } = setup(); await repository.savePool({ id: 'pool', visibility: 'platform', tags: {} }); @@ -41,10 +195,98 @@ describe('task service', () => { const publicTask = service.toPublicTask((await repository.getTask(task.id))!); expect(publicTask).not.toHaveProperty('input'); expect(publicTask).not.toHaveProperty('leaseToken'); + expect(publicTask).not.toHaveProperty('claimId'); + expect(publicTask).not.toHaveProperty('claimGeneration'); expect(publicTask).not.toHaveProperty('workerId'); expect(publicTask).not.toHaveProperty('machineId'); expect(publicTask).not.toHaveProperty('leaseExpiresAt'); expect(publicTask).not.toHaveProperty('queuePriority'); + const publicInteractiveTask = service.toPublicTask({ + ...(await repository.getTask(task.id))!, + pendingActionId: 'pending-action-secret', + lastActionId: 'last-action-secret', + claimRecovery: { + schemaVersion: 'talos.task-claim-recovery/v1', + recoveryId: 'recovery-secret', + kind: 'malformed', + phase: 'quarantined', + sourceStatus: 'running', + sourceMachineId: 'machine-secret', + restoredQueuePriority: 0, + reasonCode: 'partial_claim_identity', + startedAt: '2025-01-01T00:00:00.000Z', + updatedAt: '2025-01-01T00:00:00.000Z' + } + }); + expect(publicInteractiveTask).not.toHaveProperty('pendingActionId'); + expect(publicInteractiveTask).not.toHaveProperty('lastActionId'); + expect(publicInteractiveTask).not.toHaveProperty('claimRecovery'); + }); + + it('never includes internal claim authority in webhook payloads', async () => { + const { repository, service } = setup(); + await repository.savePool({ id: 'webhook-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ + id: 'webhook-machine', + poolId: 'webhook-pool', + tags: {}, + capacity: 1, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + const task = await service.createTask('user-a', { kind: 'browse', goal: 'webhook disclosure' }); + const claim = await service.claim('webhook-worker', 'webhook-machine'); + await service.complete(task.id, 'webhook-worker', claim.leaseToken, 'completed', []); + + const serialized = JSON.stringify(await repository.listWebhooks()); + const stored = await repository.getTask(task.id); + for (const field of [ + 'claimId', + 'claimGeneration', + 'taskVersion', + 'claimCommitted', + 'claimReleased', + 'claimQueuePriority', + 'queuePriority', + 'workerId', + 'machineId', + 'leaseExpiresAt', + 'leaseToken', + 'claimRecovery' + ]) expect(serialized).not.toContain(`\"${field}\"`); + expect(serialized).not.toContain(claim.leaseToken); + expect(serialized).not.toContain(stored?.claimId); + }); + + it('logs a stable code when webhook delivery rejects with internal authority', async () => { + const repository = new MemoryRepository(); + const warnings: Array<{ message: string; fields?: Record }> = []; + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('test-webhook-secret'), + { + onWebhook: async () => { throw new Error('claim-secret-sentinel lease-token-sentinel'); }, + logger: { warn: (message, fields) => warnings.push({ message, fields }) } + } + ); + + await service.createTask('user-a', { + kind: 'browse', + goal: 'webhook failure disclosure', + callback: 'https://example.invalid/webhook' + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(warnings).toContainEqual({ + message: 'webhook delivery failed', + fields: expect.objectContaining({ error: 'webhook_delivery_failed' }) + }); + const serialized = JSON.stringify(warnings); + expect(serialized).not.toContain('claim-secret-sentinel'); + expect(serialized).not.toContain('lease-token-sentinel'); }); it('uses a unique artifact id and injected clock', async () => { @@ -75,7 +317,7 @@ describe('task service', () => { it('enforces profile ownership and one concurrent lock', async () => { const { repository, service } = setup(); - await repository.saveProfile({ id: 'profile', userId: 'user-a' }); + await repository.createProfile({ id: 'profile', userId: 'user-a' }); await repository.savePool({ id: 'pool', visibility: 'platform', tags: {} }); await repository.saveMachine({ id: 'machine', poolId: 'pool', tags: {}, capacity: 2, activeLeases: 0, online: true, workerTokenHash: hashWorkerToken('worker-token-123456') }); await expect(service.createTask('user-b', { kind: 'browse', goal: 'x', profile_id: 'profile' })).rejects.toMatchObject({ code: 'forbidden' }); @@ -98,7 +340,7 @@ describe('task service', () => { online: true, workerTokenHash: 'x' }); - await repository.saveProfile({ id: 'profile', userId: 'user-a', machineId: 'remote-machine' }); + await repository.createProfile({ id: 'profile', userId: 'user-a', machineId: 'remote-machine' }); await expect(service.createTask('user-a', { kind: 'browse', @@ -113,7 +355,7 @@ describe('task service', () => { it('keeps a profile lock renewed and resumes late input without requeue', async () => { const { repository, service, clock } = setup({ value: 1000 }); - await repository.saveProfile({ id: 'profile', userId: 'user-a' }); + await repository.createProfile({ id: 'profile', userId: 'user-a' }); await repository.savePool({ id: 'pool', visibility: 'platform', tags: {} }); await repository.saveMachine({ id: 'machine', poolId: 'pool', tags: {}, capacity: 2, activeLeases: 0, online: true, workerTokenHash: hashWorkerToken('worker-token-123456') }); const task = await service.createTask('user-a', { kind: 'browse', goal: 'input', profile_id: 'profile' }); @@ -130,6 +372,287 @@ describe('task service', () => { expect((await service.getWorkerInput(task.id, 'worker-a', claim.leaseToken))?.value).toBe('123456'); }); + it('returns a retryable conflict after three cancellation CAS misses without side effects', async () => { + const clock = { value: 1_000 }; + const storage = new MemoryRepository(() => clock.value); + const authority = new TaskService( + storage, + new Scheduler(storage), + new ProfileLockService(storage), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + await storage.savePool({ id: 'cancel-pool', visibility: 'platform', tags: {} }); + await storage.saveMachine({ id: 'cancel-machine', poolId: 'cancel-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + const task = await authority.createTask('user-a', { kind: 'browse', goal: 'cancel CAS exhaustion' }); + await authority.claim('worker-a', 'cancel-machine'); + let attempts = 0; + let releases = 0; + const repository = new Proxy(storage, { + get(target, property) { + if (property === 'replaceTaskForClaim') { + return async (...args: Parameters): Promise => { + if (args[0].status === 'cancelled' && args[0].claimReleased !== true) { + attempts += 1; + return false; + } + return target.replaceTaskForClaim(...args); + }; + } + if (property === 'releaseMachineLease') { + return async (...args: Parameters): Promise => { + releases += 1; + return target.releaseMachineLease(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + + await expect(service.cancel(task.id, task.userId)).rejects.toMatchObject({ + code: 'concurrent_update', + status: 409 + }); + expect(attempts).toBe(3); + expect(releases).toBe(0); + expect((await storage.getTask(task.id))?.status).toBe('claimed'); + expect((await storage.listWebhooks()).filter((event) => event.payload.status === 'cancelled')).toHaveLength(0); + }); + + it('does not persist input or handoff side effects when their CAS retries are exhausted', async () => { + const clock = { value: 1_000 }; + const storage = new MemoryRepository(() => clock.value); + const authority = new TaskService( + storage, + new Scheduler(storage), + new ProfileLockService(storage), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + await storage.savePool({ id: 'side-effect-pool', visibility: 'platform', tags: {} }); + await storage.saveMachine({ id: 'side-effect-machine', poolId: 'side-effect-pool', tags: {}, capacity: 2, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + const inputTask = await authority.createTask('user-a', { kind: 'browse', goal: 'input CAS exhaustion' }); + const inputClaim = await authority.claim('worker-input', 'side-effect-machine'); + await authority.needsInput(inputTask.id, 'worker-input', inputClaim.leaseToken); + const handoffTask = await authority.createTask('user-a', { kind: 'browse', goal: 'handoff CAS exhaustion' }); + await authority.claim('worker-handoff', 'side-effect-machine'); + let inputAttempts = 0; + let handoffAttempts = 0; + let pendingInputWrites = 0; + let handoffWrites = 0; + const repository = new Proxy(storage, { + get(target, property) { + if (property === 'replaceTaskForClaim') { + return async (...args: Parameters): Promise => { + if (args[0].id === inputTask.id && args[0].status === 'running' && args[1].status === 'needs_input') { + inputAttempts += 1; + return false; + } + return target.replaceTaskForClaim(...args); + }; + } + if (property === 'replaceTaskForActiveClaim') { + return async (...args: Parameters): Promise => { + if (args[0].id === handoffTask.id && args[0].status === 'handoff') { + handoffAttempts += 1; + await authority.heartbeat( + handoffTask.id, + 'worker-handoff', + (await storage.getTask(handoffTask.id))?.leaseToken ?? '', + 30 + handoffAttempts + ); + return false; + } + return target.replaceTaskForActiveClaim(...args); + }; + } + if (property === 'savePendingInput') { + return async (...args: Parameters): Promise => { + pendingInputWrites += 1; + return target.savePendingInput(...args); + }; + } + if (property === 'saveHandoff') { + return async (...args: Parameters): Promise => { + handoffWrites += 1; + return target.saveHandoff(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + + await expect(service.provideInput(inputTask.id, inputTask.userId, { kind: 'text', value: 'secret' })) + .rejects.toMatchObject({ code: 'concurrent_update', status: 409 }); + await expect(service.requestHandoff(handoffTask.id, handoffTask.userId, 60)) + .rejects.toMatchObject({ code: 'concurrent_update', status: 409 }); + expect(inputAttempts).toBe(3); + expect(handoffAttempts).toBe(3); + expect(pendingInputWrites).toBe(0); + expect(handoffWrites).toBe(0); + expect(await storage.takePendingInput(inputTask.id)).toBeUndefined(); + expect(await storage.getTask(handoffTask.id)).not.toHaveProperty('handoff'); + }); + + it('rejects handoff while a claim is provisional and writes no handoff record', async () => { + const clock = { value: 1_000 }; + const storage = new MemoryRepository(() => clock.value); + let provisionalClaimReached!: () => void; + let resumeClaim!: () => void; + const claimReached = new Promise((resolve) => { provisionalClaimReached = resolve; }); + const claimResume = new Promise((resolve) => { resumeClaim = resolve; }); + let handoffWrites = 0; + const repository = new Proxy(storage, { + get(target, property) { + if (property === 'claimTask') { + return async (...args: Parameters): Promise>> => { + const claimed = await target.claimTask(...args); + if (claimed !== undefined) { + provisionalClaimReached(); + await claimResume; + } + return claimed; + }; + } + if (property === 'saveHandoff') { + return async (...args: Parameters): Promise => { + handoffWrites += 1; + return target.saveHandoff(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + await storage.savePool({ id: 'provisional-pool', visibility: 'platform', tags: {} }); + await storage.saveMachine({ id: 'provisional-machine', poolId: 'provisional-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + const task = await service.createTask('user-a', { kind: 'browse', goal: 'provisional handoff' }); + + const claiming = service.claim('worker-a', 'provisional-machine', clock.value); + await claimReached; + try { + await expect(service.requestHandoff(task.id, task.userId, 60)).rejects.toMatchObject({ + code: 'conflict', + status: 409, + message: 'task claim is not active' + }); + expect(handoffWrites).toBe(0); + expect(await storage.getTask(task.id)).toMatchObject({ status: 'claimed', claimCommitted: false }); + expect(await storage.getTask(task.id)).not.toHaveProperty('handoff'); + } finally { + resumeClaim(); + } + await expect(claiming).resolves.toMatchObject({ task: { id: task.id, claimCommitted: true } }); + }); + + it('fences input and handoff intent from a reclaimed generation', async () => { + const clock = { value: 1_000 }; + const storage = new MemoryRepository(() => clock.value); + const authority = new TaskService( + storage, + new Scheduler(storage), + new ProfileLockService(storage), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + await storage.savePool({ id: 'generation-pool', visibility: 'platform', tags: {} }); + await storage.saveMachine({ id: 'generation-machine', poolId: 'generation-pool', tags: {}, capacity: 2, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + const inputTask = await authority.createTask('user-a', { kind: 'browse', goal: 'generation-bound input' }); + const inputClaim = await authority.claim('worker-input', 'generation-machine'); + await authority.needsInput(inputTask.id, 'worker-input', inputClaim.leaseToken); + const handoffTask = await authority.createTask('user-a', { kind: 'browse', goal: 'generation-bound handoff' }); + const handoffClaim = await authority.claim('worker-handoff', 'generation-machine'); + await authority.heartbeat(handoffTask.id, 'worker-handoff', handoffClaim.leaseToken, 30); + let inputReclaimed = false; + let handoffReclaimed = false; + let pendingInputWrites = 0; + let handoffWrites = 0; + const repository = new Proxy(storage, { + get(target, property) { + if (property === 'replaceTaskForClaim') { + return async (...args: Parameters): Promise => { + if (!inputReclaimed && args[0].id === inputTask.id && args[0].status === 'running' && args[1].status === 'needs_input') { + inputReclaimed = true; + await advanceClaimGeneration(target, inputTask.id, 'needs_input'); + return false; + } + return target.replaceTaskForClaim(...args); + }; + } + if (property === 'replaceTaskForActiveClaim') { + return async (...args: Parameters): Promise => { + if (!handoffReclaimed && args[0].id === handoffTask.id && args[0].status === 'handoff') { + handoffReclaimed = true; + await advanceClaimGeneration(target, handoffTask.id, 'running'); + return false; + } + return target.replaceTaskForActiveClaim(...args); + }; + } + if (property === 'savePendingInput') { + return async (...args: Parameters): Promise => { + pendingInputWrites += 1; + return target.savePendingInput(...args); + }; + } + if (property === 'saveHandoff') { + return async (...args: Parameters): Promise => { + handoffWrites += 1; + return target.saveHandoff(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const service = new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('test-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } + ); + + const expectedConflict = { + code: 'conflict', + status: 409, + message: 'task claim generation changed concurrently' + }; + await expect(service.provideInput(inputTask.id, inputTask.userId, { kind: 'text', value: 'stale input' })) + .rejects.toMatchObject(expectedConflict); + await expect(service.requestHandoff(handoffTask.id, handoffTask.userId, 60)) + .rejects.toMatchObject(expectedConflict); + expect((await storage.getTask(inputTask.id))?.claimGeneration).toBe((inputClaim.task.claimGeneration ?? 0) + 1); + expect((await storage.getTask(handoffTask.id))?.claimGeneration).toBe((handoffClaim.task.claimGeneration ?? 0) + 1); + expect(pendingInputWrites).toBe(0); + expect(handoffWrites).toBe(0); + expect(await storage.takePendingInput(inputTask.id)).toBeUndefined(); + expect(await storage.getTask(handoffTask.id)).not.toHaveProperty('handoff'); + }); + it('lets either eligible machine claim queued work', async () => { const { repository, service } = setup(); await repository.savePool({ id: 'pool', visibility: 'platform', tags: {} }); diff --git a/control-plane/src/services/task-service.ts b/control-plane/src/services/task-service.ts index cb85be4..8a8c851 100644 --- a/control-plane/src/services/task-service.ts +++ b/control-plane/src/services/task-service.ts @@ -1,8 +1,8 @@ -import { conflict, deadlineExceeded, forbidden, notFound, taskCancelled, unauthorized, TalosError } from '../domain/errors.js'; +import { concurrentUpdate, conflict, deadlineExceeded, forbidden, notFound, taskCancelled, unauthorized, TalosError } from '../domain/errors.js'; import { timingSafeEqual } from 'node:crypto'; import { taskCreateSchema } from '../domain/schemas.js'; -import type { Lease, PublicTask, Task, TaskFinding, WebhookEvent } from '../domain/types.js'; -import type { Repository } from '../storage/repository.js'; +import type { Lease, MachineLeaseReservation, PublicTask, Task, TaskClaimRecoveryReason, TaskClaimGuard, TaskFinding, TaskRecoveryGuard, WebhookEvent } from '../domain/types.js'; +import type { Repository, TaskMaintenanceCursor } from '../storage/repository.js'; import { newId } from '../util/id.js'; import type { ProfileLockService } from './profile-lock.js'; import type { Scheduler } from './scheduler.js'; @@ -18,6 +18,13 @@ export interface TaskServiceOptions { logger?: Pick; } +const CLAIM_RECONCILIATION_BATCH_SIZE = 100; +const USER_TASK_CAS_ATTEMPTS = 3; +const TASK_MAINTENANCE_CURSOR_ID = 'task-claim-reconciliation' as const; +const ACTIVE_CLAIM_STATUSES: readonly Task['status'][] = ['claimed', 'running', 'needs_input', 'handoff', 'closing']; +const isNonEmptyString = (value: unknown): value is string => typeof value === 'string' && value.length > 0; +const isValidTimestamp = (value: unknown): value is string => isNonEmptyString(value) && Number.isFinite(Date.parse(value)); + export class TaskService { private readonly leaseSeconds: number; private readonly clock: () => number; @@ -69,11 +76,14 @@ export class TaskService { ...(data.profile_id === undefined ? {} : { profileId: data.profile_id }), ...(data.pool_id === undefined ? {} : { poolId: data.pool_id }), ...(requesterGroups.length === 0 ? {} : { requesterGroups: [...requesterGroups] }), - constraints: data.constraints, + constraints: data.constraints.deadline === undefined + ? data.constraints + : { ...data.constraints, deadline: new Date(data.constraints.deadline).toISOString() }, mode: data.mode, interaction, ...(data.callback === undefined ? {} : { callback: data.callback }), status: 'submitted', + taskVersion: 0, createdAt: now, updatedAt: now, findings: [], @@ -93,14 +103,10 @@ export class TaskService { try { const eligible = await this.scheduler.isEligible(candidate, machineId, candidate.userId, candidate.requesterGroups ?? []); if (eligible === undefined) continue; - const { machine } = eligible; - if (candidate.profileId !== undefined) { - const profile = await this.profiles.assertOwner(candidate.profileId, candidate.userId); - if (profile.machineId !== undefined && profile.machineId !== machine.id) continue; - await this.profiles.acquire(candidate.profileId, candidate.userId, candidate.id, now, machine.id, this.leaseSeconds); - } const expiresAt = new Date(now + this.leaseSeconds * 1000).toISOString(); const leaseToken = newId('lease'); + const claimId = newId('claim'); + const claimGeneration = (candidate.claimGeneration ?? 0) + 1; const task: Task = { ...candidate, status: 'claimed', @@ -108,14 +114,29 @@ export class TaskService { claimedAt: new Date(now).toISOString(), leaseExpiresAt: expiresAt, leaseToken, + claimId, + claimGeneration, + taskVersion: (candidate.taskVersion ?? 0) + 1, + claimCommitted: false, + claimReleased: false, + claimQueuePriority: candidate.queuePriority, queuePriority: undefined, workerId, machineId }; - await this.repository.saveTask(task); - await this.repository.saveMachine({ ...machine, activeLeases: machine.activeLeases + 1 }); - await this.emit(task, 'task.state_changed', { status: task.status }); - return { task, lease: { taskId: task.id, workerId, machineId, expiresAt }, leaseToken }; + const claimed = await this.repository.claimTask(task, candidate.claimGeneration ?? 0, candidate.taskVersion ?? 0); + if (claimed === undefined) continue; + if (!await this.ensureClaimProjections(claimed)) { + await this.abortClaim(claimed, now); + continue; + } + const committed = await this.replaceClaimedTask(claimed, { ...claimed, claimCommitted: true }); + if (!await this.verifyClaimProjections(committed)) { + await this.abortClaim(committed, now); + continue; + } + await this.emit(committed, 'task.state_changed', { status: committed.status }); + return { task: committed, lease: { taskId: committed.id, workerId, machineId, expiresAt }, leaseToken }; } catch (error) { if (error instanceof TalosError && error.code === 'conflict') continue; throw error; @@ -127,17 +148,19 @@ export class TaskService { public async heartbeat(taskId: string, workerId: string, leaseToken: string, extendSeconds: number): Promise { const task = await this.getWorkerTask(taskId, workerId, leaseToken); const now = this.clock(); - const nextStatus = task.status === 'claimed' ? 'running' : task.status; const updated: Task = { ...task, - status: nextStatus, + status: task.status === 'claimed' ? 'running' : task.status, updatedAt: new Date(now).toISOString(), leaseExpiresAt: new Date(now + extendSeconds * 1000).toISOString() }; - if (task.profileId !== undefined) await this.profiles.renew(task.profileId, task.id, now, extendSeconds); - await this.repository.saveTask(updated); + const persisted = await this.replaceActiveClaimedTask(task, updated); + if (!await this.ensureClaimProjections(persisted)) { + if (persisted.claimCommitted !== true) await this.abortClaim(persisted, now); + throw conflict('lease accounting could not be renewed'); + } if (task.status !== updated.status) await this.emit(updated, 'task.state_changed', { status: updated.status }); - return updated; + return persisted; } public async complete(taskId: string, workerId: string, leaseToken: string, status: 'completed' | 'failed', findings: readonly TaskFinding[], error?: { code: string; message: string }): Promise { @@ -149,7 +172,7 @@ export class TaskService { findings: [...findings], ...(error === undefined ? {} : { error }) }; - await this.repository.saveTask(updated); + await this.replaceClaimedTask(task, updated); await this.releaseLease(updated); await this.emit(updated, 'task.state_changed', { status }); if (status === 'completed') await this.emit(updated, 'task.completed', { status }); @@ -163,26 +186,33 @@ export class TaskService { updatedAt: new Date(this.clock()).toISOString(), artifacts: [...task.artifacts, artifact] }; - await this.repository.saveTask(updated); + await this.replaceClaimedTask(task, updated); return updated; } public async provideInput(id: string, userId: string, input: NonNullable): Promise { const task = await this.authorizedTask(id, userId); - if (task.interaction === 'interactive') throw conflict('interactive sessions do not accept task input'); - if (task.status !== 'needs_input') throw conflict('task is not waiting for input'); - const now = this.clock(); + this.assertTaskAcceptsInput(task); + const claimBinding = this.userClaimBinding(task); + const { persisted } = await this.replaceAuthorizedTask(id, userId, task, (current) => { + this.assertTaskAcceptsInput(current); + const now = this.clock(); + const currentLeaseExpiry = Date.parse(current.leaseExpiresAt ?? ''); + const extendedLeaseExpiry = now + this.leaseSeconds * 1000; + return { + ...current, + status: 'running', + updatedAt: new Date(now).toISOString(), + leaseExpiresAt: new Date(Math.max( + Number.isFinite(currentLeaseExpiry) ? currentLeaseExpiry : 0, + extendedLeaseExpiry + )).toISOString() + }; + }, { claimBinding }); await this.repository.savePendingInput(id, input); - const updated: Task = { - ...task, - status: 'running', - updatedAt: new Date(now).toISOString(), - leaseExpiresAt: task.workerId === undefined ? task.leaseExpiresAt : new Date(now + this.leaseSeconds * 1000).toISOString() - }; - if (task.workerId !== undefined && task.profileId !== undefined) await this.profiles.renew(task.profileId, task.id, now, this.leaseSeconds); - await this.repository.saveTask(updated); - await this.emit(updated, 'task.state_changed', { status: updated.status }); - return updated; + if (!await this.ensureClaimProjections(persisted)) throw conflict('lease accounting could not be renewed'); + await this.emit(persisted, 'task.state_changed', { status: persisted.status }); + return persisted; } public async needsInput(taskId: string, workerId: string, leaseToken: string): Promise { @@ -193,7 +223,7 @@ export class TaskService { status: 'needs_input', updatedAt: new Date(this.clock()).toISOString() }; - await this.repository.saveTask(updated); + await this.replaceClaimedTask(task, updated); await this.emit(updated, 'task.needs_input', { status: updated.status }); return updated; } @@ -206,61 +236,87 @@ export class TaskService { public async requestHandoff(id: string, userId: string, expiresInSeconds: number): Promise<{ handoff_url: string; expires: string }> { const task = await this.authorizedTask(id, userId); - if (task.interaction === 'interactive') throw conflict('interactive sessions do not support handoff'); - if (!['running', 'claimed'].includes(task.status)) throw conflict('task cannot request handoff in current state'); + this.assertTaskCanRequestHandoff(task); + const claimBinding = this.userClaimBinding(task); const expires = new Date(this.clock() + expiresInSeconds * 1000).toISOString(); const linkId = newId('handoff'); const url = `/v1/handoffs/${linkId}`; + const { persisted } = await this.replaceAuthorizedTask(id, userId, task, (current) => { + this.assertTaskCanRequestHandoff(current); + return { + ...current, + status: 'handoff', + updatedAt: new Date(this.clock()).toISOString(), + handoff: { url, expiresAt: expires } + }; + }, { claimBinding, requireActiveClaim: true }); await this.repository.saveHandoff({ id: linkId, taskId: id, userId, url, expiresAt: expires, used: false }); - const updated: Task = { - ...task, - status: 'handoff', - updatedAt: new Date(this.clock()).toISOString(), - handoff: { url, expiresAt: expires } - }; - await this.repository.saveTask(updated); - await this.emit(updated, 'task.handoff_requested', { handoff_url: url, expires }); + await this.emit(persisted, 'task.handoff_requested', { handoff_url: url, expires }); return { handoff_url: url, expires }; } public async cancel(id: string, userId: string): Promise { const task = await this.authorizedTask(id, userId); - if (task.interaction === 'interactive') throw conflict('interactive sessions must be closed through the session API'); - if (['completed', 'failed', 'cancelled'].includes(task.status)) throw conflict('task is already terminal'); - const updated: Task = { - ...task, - status: 'cancelled', - updatedAt: new Date(this.clock()).toISOString() - }; - await this.repository.saveTask(updated); - await this.releaseLease(updated); - await this.emit(updated, 'task.state_changed', { status: updated.status }); - return updated; + const { previous, persisted } = await this.replaceAuthorizedTask(id, userId, task, (current) => { + if (current.interaction === 'interactive') throw conflict('interactive sessions must be closed through the session API'); + if (['completed', 'failed', 'cancelled'].includes(current.status)) throw conflict('task is already terminal'); + return { + ...current, + status: 'cancelled', + updatedAt: new Date(this.clock()).toISOString() + }; + }); + if (previous.status !== 'submitted') await this.releaseLease(persisted); + const result = await this.authorizedTask(id, userId); + await this.emit(result, 'task.state_changed', { status: result.status }); + return result; } public async closeInteractive(id: string, userId: string): Promise { const task = await this.authorizedTask(id, userId); - if (task.interaction !== 'interactive') throw conflict('task is not an interactive session'); - if (['completed', 'failed', 'cancelled'].includes(task.status)) throw conflict('session is already terminal'); - const status = task.status === 'submitted' ? 'completed' : 'closing'; - const updated: Task = { - ...task, - status, - updatedAt: new Date(this.clock()).toISOString() - }; - await this.repository.saveTask(updated); - if (status === 'completed') await this.releaseLease(updated); - await this.emit(updated, 'task.state_changed', { status }); - if (status === 'completed') await this.emit(updated, 'task.completed', { status }); - return updated; + const { persisted } = await this.replaceAuthorizedTask(id, userId, task, (current) => { + if (current.interaction !== 'interactive') throw conflict('task is not an interactive session'); + if (['completed', 'failed', 'cancelled'].includes(current.status)) throw conflict('session is already terminal'); + return { + ...current, + status: current.status === 'submitted' ? 'completed' : 'closing', + updatedAt: new Date(this.clock()).toISOString() + }; + }); + await this.emit(persisted, 'task.state_changed', { status: persisted.status }); + if (persisted.status === 'completed') await this.emit(persisted, 'task.completed', { status: persisted.status }); + return persisted; } public async expireLeases(now = this.clock()): Promise { - const active = await this.repository.listTasks(); + await this.reconcileClaims(now); + const active = await this.repository.listExpirableTasks(now, CLAIM_RECONCILIATION_BATCH_SIZE); const expired: Task[] = []; for (const candidate of active) { + try { + const result = await this.expireTaskCandidate(candidate, now); + if (result !== undefined) expired.push(result); + } catch { + this.logger?.warn('task lease expiry failed', { taskId: candidate.id, error: 'claim_maintenance_failed' }); + } + } + await this.reconcileClaims(now); + return expired; + } + + private async expireTaskCandidate(candidate: Task, now: number): Promise { const current = await this.repository.getTask(candidate.id); - if (current?.kind === 'testing') continue; + if (current?.kind === 'testing' || current === undefined) return undefined; + if (await this.normalizeSubmittedDeadline(current, now)) return undefined; + const malformedReason = this.malformedClaimReason(current); + if ( + current.claimRecovery !== undefined || + malformedReason !== undefined || + (ACTIVE_CLAIM_STATUSES.includes(current.status) && current.claimId === undefined && current.claimGeneration === undefined) + ) { + await this.processTaskClaimMaintenance(current, now, malformedReason); + return undefined; + } if (current?.status === 'submitted' && current.constraints.deadline !== undefined && Date.parse(current.constraints.deadline) <= now) { const failed: Task = { ...current, @@ -268,9 +324,10 @@ export class TaskService { updatedAt: new Date(now).toISOString(), error: { code: 'deadline_exceeded', message: deadlineExceeded().message } }; - await this.repository.saveTask(failed); - await this.emit(failed, 'task.state_changed', { status: failed.status, error: failed.error }); - continue; + if (await this.repository.replaceSubmittedTask(failed, current.claimGeneration ?? 0, current.taskVersion ?? 0)) { + await this.emit(failed, 'task.state_changed', { status: failed.status, error: failed.error }); + } + return undefined; } if (current?.leaseExpiresAt !== undefined && Date.parse(current.leaseExpiresAt) <= now && ['claimed', 'running', 'closing'].includes(current.status)) { if (current.status === 'closing') { @@ -289,11 +346,11 @@ export class TaskService { pendingActionId: undefined, updatedAt: new Date(now).toISOString() }; - await this.repository.saveTask(completed); - await this.releaseLease(current); + if (!await this.tryReplaceClaimedTask(current, completed)) return undefined; + await this.releaseLease(completed); await this.emit(completed, 'task.state_changed', { status: completed.status }); await this.emit(completed, 'task.completed', { status: completed.status }); - continue; + return undefined; } const requeued: Task = { ...current, @@ -302,16 +359,330 @@ export class TaskService { leaseExpiresAt: undefined, leaseToken: undefined, workerId: undefined, - machineId: undefined, queuePriority: -1 }; + if (!await this.tryReplaceClaimedTask(current, requeued)) return undefined; if (current.interaction === 'interactive') await this.repository.requeueSessionAction(current.id); - await this.repository.saveTask(requeued); await this.releaseLease(current); - expired.push(requeued); + return requeued; } + return undefined; + } + + public async reconcileClaims(now = this.clock()): Promise { + let cursor = await this.getOrCreateTaskMaintenanceCursor(now); + let tasks = await this.repository.listTaskMaintenancePage(cursor, CLAIM_RECONCILIATION_BATCH_SIZE); + if (tasks.length === 0) { + const highWatermark = await this.repository.getTaskMaintenanceHighWatermark(); + const wrapped: TaskMaintenanceCursor = { + id: TASK_MAINTENANCE_CURSOR_ID, + version: cursor.version + 1, + cycleCutoffAt: highWatermark ?? new Date(now).toISOString(), + updatedAt: new Date(now).toISOString() + }; + if (!await this.repository.replaceTaskMaintenanceCursor(wrapped, cursor.version)) return; + cursor = wrapped; + tasks = await this.repository.listTaskMaintenancePage(cursor, CLAIM_RECONCILIATION_BATCH_SIZE); + if (tasks.length === 0) return; } - return expired; + for (const task of tasks) { + try { + await this.processTaskClaimMaintenance(task, now); + } catch { + this.logger?.warn('task claim reconciliation failed', { + taskId: task.id, + error: 'claim_maintenance_failed' + }); + } + } + const last = tasks[tasks.length - 1]!; + await this.repository.replaceTaskMaintenanceCursor({ + ...cursor, + version: cursor.version + 1, + afterCreatedAt: last.createdAt, + afterTaskId: last.id, + updatedAt: new Date(now).toISOString() + }, cursor.version); + } + + private async getOrCreateTaskMaintenanceCursor(now: number): Promise { + const existing = await this.repository.getTaskMaintenanceCursor(TASK_MAINTENANCE_CURSOR_ID); + if (existing !== undefined) return existing; + const highWatermark = await this.repository.getTaskMaintenanceHighWatermark(); + const cursor: TaskMaintenanceCursor = { + id: TASK_MAINTENANCE_CURSOR_ID, + version: 0, + cycleCutoffAt: highWatermark ?? new Date(now).toISOString(), + updatedAt: new Date(now).toISOString() + }; + if (await this.repository.createTaskMaintenanceCursor(cursor)) return cursor; + const concurrent = await this.repository.getTaskMaintenanceCursor(TASK_MAINTENANCE_CURSOR_ID); + if (concurrent === undefined) throw new Error('task maintenance cursor unavailable'); + return concurrent; + } + + private async processTaskClaimMaintenance(task: Task, now: number, knownMalformedReason?: TaskClaimRecoveryReason): Promise { + if (task.kind === 'testing') return; + if (await this.normalizeSubmittedDeadline(task, now)) return; + if (task.claimRecovery !== undefined) { + if (task.claimRecovery.kind === 'legacy' && task.claimRecovery.phase !== 'quarantined') { + await this.continueLegacyClaimRecovery(task, now); + } else if (task.claimRecovery.kind === 'malformed' && task.claimRecovery.phase !== 'quarantined') { + await this.continueMalformedClaimRecovery(task, now); + } + return; + } + const malformedReason = knownMalformedReason ?? this.malformedClaimReason(task); + if (malformedReason !== undefined) { + await this.beginMalformedClaimRecovery(task, malformedReason, now); + return; + } + if (ACTIVE_CLAIM_STATUSES.includes(task.status) && task.claimId === undefined && task.claimGeneration === undefined) { + await this.beginLegacyClaimRecovery(task, now); + return; + } + if (this.isActiveClaim(task)) { + const projectionsReady = await this.ensureClaimProjections(task); + if (!projectionsReady && task.claimCommitted !== true) await this.abortClaim(task, now); + else if (projectionsReady && task.claimCommitted !== true) { + await this.replaceClaimedTask(task, { ...task, claimCommitted: true }); + } + return; + } + if (task.claimId !== undefined && task.claimGeneration !== undefined) await this.releaseLease(task); + } + + private async normalizeSubmittedDeadline(task: Task, now: number): Promise { + const deadline = task.status === 'submitted' ? task.constraints.deadline : undefined; + if (deadline === undefined) return false; + const deadlineAt = Date.parse(deadline); + if (!Number.isFinite(deadlineAt)) { + const failed: Task = { + ...task, + status: 'failed', + updatedAt: new Date(now).toISOString(), + error: { code: 'invalid_deadline', message: 'task deadline is invalid' } + }; + if (await this.repository.replaceSubmittedTask(failed, task.claimGeneration ?? 0, task.taskVersion ?? 0)) { + await this.emit(failed, 'task.state_changed', { status: failed.status, error: failed.error }); + } + return true; + } + const canonicalDeadline = new Date(deadlineAt).toISOString(); + if (canonicalDeadline === deadline) return false; + await this.repository.replaceSubmittedTask({ + ...task, + constraints: { ...task.constraints, deadline: canonicalDeadline }, + updatedAt: new Date(now).toISOString() + }, task.claimGeneration ?? 0, task.taskVersion ?? 0); + return true; + } + + private malformedClaimReason(task: Task): TaskClaimRecoveryReason | undefined { + const hasClaimId = task.claimId !== undefined; + const hasClaimGeneration = task.claimGeneration !== undefined; + if (hasClaimId !== hasClaimGeneration) return 'partial_claim_identity'; + if (hasClaimId && !isNonEmptyString(task.claimId)) return 'invalid_claim_credentials'; + if (hasClaimGeneration && (!Number.isInteger(task.claimGeneration) || task.claimGeneration! <= 0)) { + return 'invalid_claim_generation'; + } + if (!ACTIVE_CLAIM_STATUSES.includes(task.status) || !hasClaimId) return undefined; + if ( + !isNonEmptyString(task.machineId) || + !isNonEmptyString(task.workerId) || + !isNonEmptyString(task.leaseToken) || + task.leaseExpiresAt === undefined + ) return 'active_claim_missing_credentials'; + if (!isValidTimestamp(task.leaseExpiresAt)) return 'invalid_lease_expiry'; + if (task.claimReleased === true) return 'active_claim_marked_released'; + return undefined; + } + + private recoveryGuard(task: Task): TaskRecoveryGuard { + return { + status: task.status, + taskVersion: task.taskVersion ?? 0, + updatedAt: task.updatedAt, + claimId: task.claimId, + claimGeneration: task.claimGeneration, + recoveryId: task.claimRecovery?.recoveryId, + recoveryPhase: task.claimRecovery?.phase + }; + } + + private async beginLegacyClaimRecovery(task: Task, now: number): Promise { + if (!ACTIVE_CLAIM_STATUSES.includes(task.status)) return; + const timestamp = new Date(now).toISOString(); + const recovering: Task = { + ...task, + updatedAt: timestamp, + claimRecovery: { + schemaVersion: 'talos.task-claim-recovery/v1', + recoveryId: newId('recovery'), + kind: 'legacy', + phase: 'draining', + sourceStatus: task.status, + ...(task.machineId === undefined ? {} : { sourceMachineId: task.machineId }), + ...(task.profileId === undefined ? {} : { sourceProfileId: task.profileId }), + restoredQueuePriority: task.claimQueuePriority ?? task.queuePriority ?? 0, + startedAt: timestamp, + updatedAt: timestamp + } + }; + if (!await this.repository.replaceTaskForRecovery(recovering, this.recoveryGuard(task))) return; + const persisted = await this.repository.getTask(task.id); + if (persisted !== undefined) await this.continueLegacyClaimRecovery(persisted, now); + } + + private async continueLegacyClaimRecovery(task: Task, now: number): Promise { + const recovery = task.claimRecovery; + if (recovery?.kind !== 'legacy' || recovery.phase === 'quarantined') return; + const timestamp = new Date(now).toISOString(); + if (recovery.phase === 'draining') { + if (recovery.sourceProfileId !== undefined) { + const profile = await this.repository.getProfile(recovery.sourceProfileId); + if ( + profile?.lockedByTaskId === task.id && + (profile.lockedByClaimId !== undefined || profile.lockedByClaimGeneration !== undefined) + ) { + await this.quarantineLegacyClaimRecovery(task, now); + return; + } + } + if (recovery.sourceStatus === 'closing') { + const pending = await this.repository.getPendingSessionAction(task.id); + if (pending !== undefined) { + await this.repository.finalizeSessionAction({ + actionId: pending.id, + taskId: task.id, + result: { error: { code: 'session_closed', message: 'session closed before the action completed' } }, + completedAt: timestamp + }, ['pending', 'dispatched']); + } + } + if (recovery.sourceMachineId !== undefined) { + await this.repository.releaseLegacyMachineLease(recovery.sourceMachineId, recovery.recoveryId, task.id); + } + if (recovery.sourceProfileId !== undefined) { + await this.repository.releaseLegacyProfileLease(recovery.sourceProfileId, task.id); + } + if (recovery.sourceStatus !== 'closing' && task.interaction === 'interactive') { + await this.repository.requeueSessionAction(task.id); + } + const finalizing: Task = { + ...task, + status: recovery.sourceStatus === 'closing' ? 'completed' : 'submitted', + updatedAt: timestamp, + workerId: undefined, + machineId: undefined, + leaseToken: undefined, + leaseExpiresAt: undefined, + claimCommitted: undefined, + claimReleased: undefined, + claimQueuePriority: undefined, + ...(recovery.sourceStatus === 'closing' + ? { pendingActionId: undefined } + : { queuePriority: recovery.restoredQueuePriority }), + claimRecovery: { ...recovery, phase: 'finalizing', updatedAt: timestamp } + }; + if (!await this.repository.replaceTaskForRecovery(finalizing, this.recoveryGuard(task))) return; + const persisted = await this.repository.getTask(task.id); + if (persisted !== undefined) await this.continueLegacyClaimRecovery(persisted, now); + return; + } + if (recovery.sourceMachineId !== undefined) { + await this.repository.clearLegacyMachineLeaseMarker(recovery.sourceMachineId, recovery.recoveryId); + } + await this.repository.replaceTaskForRecovery({ + ...task, + updatedAt: timestamp, + claimRecovery: undefined + }, this.recoveryGuard(task)); + } + + private async quarantineLegacyClaimRecovery(task: Task, now: number): Promise { + const recovery = task.claimRecovery; + if (recovery?.kind !== 'legacy') return; + const timestamp = new Date(now).toISOString(); + await this.repository.replaceTaskForRecovery({ + ...task, + status: 'failed', + updatedAt: timestamp, + workerId: undefined, + leaseToken: undefined, + leaseExpiresAt: undefined, + error: { code: 'claim_state_malformed', message: 'task claim state requires operator recovery' }, + claimRecovery: { + ...recovery, + kind: 'malformed', + phase: 'quarantined', + reasonCode: 'legacy_profile_identity_conflict', + updatedAt: timestamp + } + }, this.recoveryGuard(task)); + } + + private async beginMalformedClaimRecovery(task: Task, reason: TaskClaimRecoveryReason, now: number): Promise { + const timestamp = new Date(now).toISOString(); + const recovering: Task = { + ...task, + updatedAt: timestamp, + claimRecovery: { + schemaVersion: 'talos.task-claim-recovery/v1', + recoveryId: newId('recovery'), + kind: 'malformed', + phase: 'draining', + sourceStatus: task.status, + ...(task.machineId === undefined ? {} : { sourceMachineId: task.machineId }), + ...(task.profileId === undefined ? {} : { sourceProfileId: task.profileId }), + restoredQueuePriority: task.claimQueuePriority ?? task.queuePriority ?? 0, + reasonCode: reason, + startedAt: timestamp, + updatedAt: timestamp + } + }; + if (!await this.repository.replaceTaskForRecovery(recovering, this.recoveryGuard(task))) return; + const persisted = await this.repository.getTask(task.id); + if (persisted !== undefined) await this.continueMalformedClaimRecovery(persisted, now); + } + + private async continueMalformedClaimRecovery(task: Task, now: number): Promise { + const recovery = task.claimRecovery; + if (recovery?.kind !== 'malformed' || recovery.phase === 'quarantined') return; + if ( + isNonEmptyString(task.claimId) && + typeof task.claimGeneration === 'number' && + Number.isInteger(task.claimGeneration) && + task.claimGeneration > 0 + ) { + const reservation = { + taskId: task.id, + claimId: task.claimId, + claimGeneration: task.claimGeneration + }; + let machineReleased = false; + if (isNonEmptyString(recovery.sourceMachineId)) { + machineReleased = await this.repository.releaseMachineLease(recovery.sourceMachineId, reservation); + } + if (!machineReleased && isNonEmptyString(task.machineId) && task.machineId !== recovery.sourceMachineId) { + machineReleased = await this.repository.releaseMachineLease(task.machineId, reservation); + } + if (!machineReleased) await this.repository.releaseMachineLeaseReservation(reservation); + if (isNonEmptyString(recovery.sourceProfileId)) { + await this.repository.releaseProfileLease(recovery.sourceProfileId, reservation); + } + } + const timestamp = new Date(now).toISOString(); + await this.repository.replaceTaskForRecovery({ + ...task, + status: 'failed', + updatedAt: timestamp, + workerId: undefined, + leaseToken: undefined, + leaseExpiresAt: undefined, + error: { code: 'claim_state_malformed', message: 'task claim state requires operator recovery' }, + claimRecovery: { ...recovery, phase: 'quarantined', updatedAt: timestamp } + }, this.recoveryGuard(task)); } private async authorizedTask(id: string, userId: string): Promise { @@ -326,11 +697,24 @@ export class TaskService { const task = await this.repository.getTask(taskId); if (task === undefined) throw notFound('task not found'); if (task.kind === 'testing') throw conflict('testing tasks require the Testing Executor API'); - if (task.workerId !== workerId || task.leaseToken === undefined || !['claimed', 'running', 'needs_input', 'handoff', 'closing', 'cancelled'].includes(task.status)) throw unauthorized('worker does not own active lease'); + if (task.claimRecovery !== undefined || this.malformedClaimReason(task) !== undefined) { + throw unauthorized('worker does not own active lease'); + } + if ( + task.workerId !== workerId || + task.leaseToken === undefined || + task.claimId === undefined || + task.claimGeneration === undefined || + task.claimGeneration <= 0 + ) throw unauthorized('worker does not own active lease'); const expected = Buffer.from(task.leaseToken); const actual = Buffer.from(leaseToken); if (expected.length !== actual.length || !timingSafeEqual(expected, actual)) throw unauthorized('invalid lease token'); if (task.status === 'cancelled') throw taskCancelled(); + if ( + task.claimCommitted !== true || + !['claimed', 'running', 'needs_input', 'handoff', 'closing'].includes(task.status) + ) throw unauthorized('worker does not own active lease'); if (task.leaseExpiresAt !== undefined && Date.parse(task.leaseExpiresAt) <= this.clock() && !['needs_input', 'handoff'].includes(task.status)) throw unauthorized('lease expired'); return task; } @@ -347,6 +731,8 @@ export class TaskService { if (task === undefined) throw unauthorized('worker does not own action result'); if (task.kind === 'testing') throw conflict('testing tasks require the Testing Executor API'); if ( + task.claimRecovery !== undefined || + this.malformedClaimReason(task) !== undefined || task.workerId !== workerId || task.machineId === undefined || task.leaseToken === undefined @@ -369,12 +755,262 @@ export class TaskService { return this.getWorkerTask(taskId, workerId, leaseToken); } + private async replaceClaimedTask(current: Task, updated: Task): Promise { + if (!await this.tryReplaceClaimedTask(current, updated)) throw unauthorized('lease generation is no longer active'); + const persisted = await this.repository.getTask(current.id); + if (persisted === undefined) throw unauthorized('lease generation is no longer active'); + return persisted; + } + + private async replaceActiveClaimedTask(current: Task, updated: Task): Promise { + const guard = this.claimGuard(current); + if (current.leaseExpiresAt === undefined) throw unauthorized('lease generation is no longer active'); + if (!await this.repository.replaceTaskForActiveClaim(updated, { ...guard, leaseExpiresAt: current.leaseExpiresAt })) { + throw unauthorized('lease generation is no longer active'); + } + const persisted = await this.repository.getTask(current.id); + if (persisted === undefined) throw unauthorized('lease generation is no longer active'); + return persisted; + } + + private async tryReplaceClaimedTask(current: Task, updated: Task): Promise { + return this.repository.replaceTaskForClaim(updated, this.claimGuard(current)); + } + + private async replaceAuthorizedTask( + id: string, + userId: string, + initial: Task, + update: (current: Task) => Task, + options: { + claimBinding?: Pick; + requireActiveClaim?: boolean; + } = {} + ): Promise<{ previous: Task; persisted: Task }> { + let current = initial; + for (let attempt = 0; attempt < USER_TASK_CAS_ATTEMPTS; attempt += 1) { + if ( + options.claimBinding !== undefined && + ( + current.status === 'submitted' || + current.claimId !== options.claimBinding.claimId || + current.claimGeneration !== options.claimBinding.claimGeneration + ) + ) throw conflict('task claim generation changed concurrently'); + + const updated = update(current); + let replaced: boolean; + if (current.status === 'submitted') { + replaced = await this.repository.replaceSubmittedTask( + updated, + current.claimGeneration ?? 0, + current.taskVersion ?? 0 + ); + } else { + const guard = this.userClaimGuard(current); + if (options.requireActiveClaim === true) { + if (current.leaseExpiresAt === undefined) throw conflict('task claim is not active'); + replaced = await this.repository.replaceTaskForActiveClaim(updated, { ...guard, leaseExpiresAt: current.leaseExpiresAt }); + } else { + replaced = await this.repository.replaceTaskForClaim(updated, guard); + } + } + if (replaced) return { previous: current, persisted: await this.authorizedTask(id, userId) }; + const latest = await this.authorizedTask(id, userId); + if ( + options.requireActiveClaim === true && + latest.claimId === current.claimId && + latest.claimGeneration === current.claimGeneration && + (latest.taskVersion ?? 0) === (current.taskVersion ?? 0) && + latest.status === current.status && + latest.leaseExpiresAt === current.leaseExpiresAt + ) throw conflict('task claim is not active'); + current = latest; + } + throw concurrentUpdate(); + } + + private userClaimBinding(task: Task): Pick { + if (task.claimId === undefined || task.claimGeneration === undefined || task.claimGeneration <= 0) { + throw conflict('task claim is not active'); + } + return { claimId: task.claimId, claimGeneration: task.claimGeneration }; + } + + private userClaimGuard(task: Task): TaskClaimGuard { + if (task.claimRecovery !== undefined || this.malformedClaimReason(task) !== undefined) { + throw conflict('task claim state cannot be changed'); + } + return { ...this.userClaimBinding(task), taskVersion: task.taskVersion ?? 0, status: task.status }; + } + + private assertTaskAcceptsInput(task: Task): void { + if (task.interaction === 'interactive') throw conflict('interactive sessions do not accept task input'); + if (task.status !== 'needs_input') throw conflict('task is not waiting for input'); + } + + private assertTaskCanRequestHandoff(task: Task): void { + if (task.interaction === 'interactive') throw conflict('interactive sessions do not support handoff'); + if (!['running', 'claimed'].includes(task.status)) throw conflict('task cannot request handoff in current state'); + if (task.claimCommitted !== true) throw conflict('task claim is not active'); + } + + private claimGuard(task: Task): TaskClaimGuard { + if (task.claimId === undefined || task.claimGeneration === undefined || task.claimGeneration <= 0) { + throw unauthorized('lease generation is no longer active'); + } + return { claimId: task.claimId, claimGeneration: task.claimGeneration, taskVersion: task.taskVersion ?? 0, status: task.status }; + } + + private reservation(task: Task): MachineLeaseReservation { + const guard = this.claimGuard(task); + if (task.leaseExpiresAt === undefined) throw unauthorized('lease generation is no longer active'); + return { + taskId: task.id, + claimId: guard.claimId, + claimGeneration: guard.claimGeneration, + expiresAt: task.leaseExpiresAt + }; + } + + private async ensureClaimProjections(task: Task): Promise { + if (task.machineId === undefined) return false; + const requested = this.reservation(task); + const authoritative = await this.repository.getTask(task.id); + if (authoritative === undefined || !this.matchesActiveClaim(authoritative, requested) || authoritative.machineId !== task.machineId) return false; + const reservation = this.reservation(authoritative); + if (!await this.repository.reserveMachineLease(task.machineId, reservation)) return false; + await this.repository.renewMachineLease(task.machineId, reservation); + if (task.profileId !== undefined) { + if (!await this.acquireProfileProjection(task, reservation)) { + await this.repository.releaseMachineLease(task.machineId, reservation); + return false; + } + } + const committed = await this.repository.getTask(task.id); + if (this.matchesActiveClaim(committed, reservation) && committed?.machineId === task.machineId) return true; + await this.repository.releaseMachineLease(task.machineId, reservation); + if (task.profileId !== undefined) await this.profiles.release(task.profileId, reservation); + return false; + } + + private async acquireProfileProjection(task: Task, reservation: MachineLeaseReservation): Promise { + if (task.profileId === undefined || task.machineId === undefined) return false; + try { + await this.profiles.acquire(task.profileId, task.userId, task.machineId, reservation); + return true; + } catch (error) { + if (!(error instanceof TalosError) || error.code !== 'conflict') throw error; + } + + const profile = await this.repository.getProfile(task.profileId); + if ( + profile?.lockedByTaskId === undefined || + profile.lockedByClaimId === undefined || + profile.lockedByClaimGeneration === undefined + ) return false; + const previous = await this.repository.getTask(profile.lockedByTaskId); + if ( + previous === undefined || + previous.kind === 'testing' || + previous.profileId !== task.profileId || + previous.claimId !== profile.lockedByClaimId || + previous.claimGeneration !== profile.lockedByClaimGeneration || + previous.leaseExpiresAt === undefined || + !['claimed', 'running'].includes(previous.status) + ) return false; + const requeued: Task = { + ...previous, + status: 'submitted', + updatedAt: new Date(this.clock()).toISOString(), + leaseExpiresAt: undefined, + leaseToken: undefined, + workerId: undefined, + queuePriority: previous.claimQueuePriority, + claimQueuePriority: undefined + }; + if (!await this.repository.replaceTaskForExpiredClaim(requeued, { + ...this.claimGuard(previous), + leaseExpiresAt: previous.leaseExpiresAt + })) return false; + if (previous.interaction === 'interactive') await this.repository.requeueSessionAction(previous.id); + await this.releaseLease(previous); + try { + await this.profiles.acquire(task.profileId, task.userId, task.machineId, reservation); + return true; + } catch (error) { + if (error instanceof TalosError && error.code === 'conflict') return false; + throw error; + } + } + + private async verifyClaimProjections(task: Task): Promise { + if (task.machineId === undefined) return false; + const reservation = this.reservation(task); + const machine = await this.repository.getMachine(task.machineId); + if (machine === undefined) return false; + const machineCommitted = machine?.leaseReservations?.some((entry) => + entry.taskId === reservation.taskId && + entry.claimId === reservation.claimId && + entry.claimGeneration === reservation.claimGeneration + ) === true; + if (!machineCommitted || machine.activeLeases > machine.capacity || machine.activeLeases < (machine.leaseReservations?.length ?? 0)) return false; + if (task.profileId === undefined) return true; + const profile = await this.repository.getProfile(task.profileId); + return profile?.lockedByTaskId === reservation.taskId && + profile.lockedByClaimId === reservation.claimId && + profile.lockedByClaimGeneration === reservation.claimGeneration; + } + + private async abortClaim(task: Task, now = this.clock()): Promise { + const requeued: Task = { + ...task, + status: 'submitted', + updatedAt: new Date(now).toISOString(), + leaseExpiresAt: undefined, + leaseToken: undefined, + workerId: undefined, + claimCommitted: false, + queuePriority: task.claimQueuePriority, + claimQueuePriority: undefined + }; + if (!await this.tryReplaceClaimedTask(task, requeued)) return; + await this.releaseLease(task); + } + + private isActiveClaim(task: Task): boolean { + return this.matchesActiveClaim(task, task.claimId === undefined || task.claimGeneration === undefined + ? undefined + : { taskId: task.id, claimId: task.claimId, claimGeneration: task.claimGeneration }); + } + + private matchesActiveClaim(task: Task | undefined, reservation: Omit | undefined): boolean { + return task !== undefined && + reservation !== undefined && + ['claimed', 'running', 'needs_input', 'handoff', 'closing'].includes(task.status) && + task.claimId === reservation.claimId && + task.claimGeneration === reservation.claimGeneration; + } + private async releaseLease(task: Task): Promise { - if (task.machineId !== undefined) { - const machine = await this.repository.getMachine(task.machineId); - if (machine !== undefined) await this.repository.saveMachine({ ...machine, activeLeases: Math.max(0, machine.activeLeases - 1) }); + if (task.claimId === undefined || task.claimGeneration === undefined) return; + const reservation = { taskId: task.id, claimId: task.claimId, claimGeneration: task.claimGeneration }; + const current = await this.repository.getTask(task.id); + if (current?.claimId !== reservation.claimId || current.claimGeneration !== reservation.claimGeneration || current.claimReleased === true) return; + let machineReleased = false; + if (current.machineId !== undefined) { + machineReleased = await this.repository.releaseMachineLease(current.machineId, reservation); + } + if (!machineReleased && task.machineId !== undefined && task.machineId !== current.machineId) { + machineReleased = await this.repository.releaseMachineLease(task.machineId, reservation); } - if (task.profileId !== undefined) await this.profiles.release(task.profileId, task.id); + if (!machineReleased) await this.repository.releaseMachineLeaseReservation(reservation); + if (current.profileId !== undefined) await this.profiles.release(current.profileId, reservation); + await this.repository.replaceTaskForClaim({ + ...current, + claimCommitted: false, + claimReleased: true + }, this.claimGuard(current)); } private async emit(task: Task, type: WebhookEvent['type'], payload: Record): Promise { @@ -390,10 +1026,10 @@ export class TaskService { await this.repository.saveWebhook(event); const signed = this.signer.sign(event, this.clock()); if (this.onWebhook !== undefined) { - void this.onWebhook(event, signed, task.callback).catch((error: unknown) => { + void this.onWebhook(event, signed, task.callback).catch(() => { this.logger?.warn('webhook delivery failed', { eventId: event.id, - error: error instanceof Error ? error.message : 'unknown' + error: 'webhook_delivery_failed' }); }); } @@ -403,12 +1039,21 @@ export class TaskService { public toPublicTask(task: Task): PublicTask { const hidden = new Set([ 'leaseToken', + 'claimId', + 'claimGeneration', + 'taskVersion', + 'claimCommitted', + 'claimReleased', + 'claimQueuePriority', 'queuePriority', 'workerId', 'machineId', 'leaseExpiresAt', 'input', 'requesterGroups', + 'pendingActionId', + 'lastActionId', + 'claimRecovery', 'testing' ]); return { diff --git a/control-plane/src/services/webhook-dispatcher.test.ts b/control-plane/src/services/webhook-dispatcher.test.ts index 8bddf76..2dc8a7f 100644 --- a/control-plane/src/services/webhook-dispatcher.test.ts +++ b/control-plane/src/services/webhook-dispatcher.test.ts @@ -40,10 +40,19 @@ describe('WebhookDispatcher', () => { const event = { id: 'evt2', type: 'task.state_changed' as const, taskId: 'task', userId: 'u', timestamp: new Date().toISOString(), payload: {}, delivery: { status: 'pending' as const, attempts: 0 } }; await repository.saveWebhook(event); let attempts = 0; - const dispatcher = new (await import('./webhook-dispatcher.js')).WebhookDispatcher(repository, signer, { fetchImpl: async () => { attempts += 1; return new Response('', { status: 500 }); }, backoffMs: 0 }); + const dispatcher = new (await import('./webhook-dispatcher.js')).WebhookDispatcher(repository, signer, { + fetchImpl: async () => { + attempts += 1; + throw new Error('claim-secret-sentinel lease-token-sentinel'); + }, + backoffMs: 0 + }); await dispatcher.dispatch(event, 'http://localhost/hook'); expect(attempts).toBe(3); expect((await repository.getWebhook('evt2'))?.delivery.status).toBe('failed'); + expect((await repository.getWebhook('evt2'))?.delivery.lastError).toBe('webhook_delivery_failed'); + expect(JSON.stringify(await repository.getWebhook('evt2'))).not.toContain('claim-secret-sentinel'); + expect(JSON.stringify(await repository.getWebhook('evt2'))).not.toContain('lease-token-sentinel'); await expect(dispatcher.dispatch(event, 'file:///tmp/hook')).rejects.toThrow('scheme'); }); @@ -81,12 +90,16 @@ describe('WebhookDispatcher', () => { await repository.saveMachine({ id: 'machine', poolId: 'pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: hashWorkerToken('worker-token-123456') }); const signer = new WebhookSigner('webhook-secret-1234'); const received: string[] = []; + const receivedBodies: string[] = []; const server = createServer((request, response) => { const chunks: Buffer[] = []; request.on('data', (chunk) => chunks.push(Buffer.from(chunk))); request.on('end', () => { const body = Buffer.concat(chunks).toString(); - if (signer.verify({ id: 'unused', timestamp: String(request.headers['x-talos-webhook-timestamp']), signature: String(request.headers['x-talos-webhook-signature']), body })) received.push(JSON.parse(body).type as string); + if (signer.verify({ id: 'unused', timestamp: String(request.headers['x-talos-webhook-timestamp']), signature: String(request.headers['x-talos-webhook-signature']), body })) { + received.push(JSON.parse(body).type as string); + receivedBodies.push(body); + } response.end(); }); }); @@ -104,6 +117,25 @@ describe('WebhookDispatcher', () => { await new Promise((resolve) => setTimeout(resolve, 5)); } expect(received).toContain('task.completed'); + const stored = await repository.getTask(task.id); + for (const body of receivedBodies) { + for (const field of [ + 'claimId', + 'claimGeneration', + 'taskVersion', + 'claimCommitted', + 'claimReleased', + 'claimQueuePriority', + 'queuePriority', + 'workerId', + 'machineId', + 'leaseExpiresAt', + 'leaseToken', + 'claimRecovery' + ]) expect(body).not.toContain(`\"${field}\"`); + expect(body).not.toContain(claim.leaseToken); + expect(body).not.toContain(stored?.claimId); + } server.close(); }); }); diff --git a/control-plane/src/services/webhook-dispatcher.ts b/control-plane/src/services/webhook-dispatcher.ts index aa30bbd..b9318ca 100644 --- a/control-plane/src/services/webhook-dispatcher.ts +++ b/control-plane/src/services/webhook-dispatcher.ts @@ -70,8 +70,8 @@ export class WebhookDispatcher { if (!response.ok) throw new Error(`HTTP ${response.status}`); await this.update(event, { status: 'delivered', attempts: attempt, lastAttemptAt: new Date(this.clock()).toISOString() }); return; - } catch (error) { - lastError = error instanceof Error ? error.message : lastError; + } catch { + lastError = 'webhook_delivery_failed'; await this.update(event, { status: 'pending', attempts: attempt, lastAttemptAt: new Date(this.clock()).toISOString(), lastError }); if (attempt < this.retries) await new Promise((resolve) => setTimeout(resolve, this.backoffMs * attempt)); } diff --git a/control-plane/src/storage/memory-repository.test.ts b/control-plane/src/storage/memory-repository.test.ts index 7829fa0..4ef4a9d 100644 --- a/control-plane/src/storage/memory-repository.test.ts +++ b/control-plane/src/storage/memory-repository.test.ts @@ -16,8 +16,8 @@ describe('MemoryRepository owner listings', () => { ownerUserId: 'bob', tags: {} }); - await repository.saveProfile({ id: 'alice-profile', userId: 'alice' }); - await repository.saveProfile({ id: 'bob-profile', userId: 'bob' }); + await repository.createProfile({ id: 'alice-profile', userId: 'alice' }); + await repository.createProfile({ id: 'bob-profile', userId: 'bob' }); expect((await repository.listPoolsByOwner('alice')).map((pool) => pool.id)).toEqual([ 'alice-pool' diff --git a/control-plane/src/storage/memory-repository.ts b/control-plane/src/storage/memory-repository.ts index 8d9ac46..381526a 100644 --- a/control-plane/src/storage/memory-repository.ts +++ b/control-plane/src/storage/memory-repository.ts @@ -1,6 +1,6 @@ -import type { HandoffLink, Machine, PendingSessionAction, Pool, Profile, SessionActionResult, Task, TaskInput, WebhookEvent } from '../domain/types.js'; +import type { HandoffLink, Machine, MachineLeaseReservation, PendingSessionAction, Pool, Profile, SessionActionResult, Task, TaskActiveClaimGuard, TaskClaimGuard, TaskInput, TaskRecoveryGuard, WebhookEvent } from '../domain/types.js'; import type { TestingMachineReservationRecord, TestingRunRecord } from '../domain/testing-types.js'; -import type { Repository, TestingAttemptDispatchGuard, TestingAttemptMutationGuard } from './repository.js'; +import type { Repository, TaskMaintenanceCursor, TestingAttemptDispatchGuard, TestingAttemptMutationGuard } from './repository.js'; const isFutureTimestamp = (value: string | undefined, observedNow: number): boolean => { if (value === undefined) return false; @@ -8,6 +8,27 @@ const isFutureTimestamp = (value: string | undefined, observedNow: number): bool return Number.isFinite(parsed) && parsed > observedNow; }; +const isValidClaim = (task: Task, expectedClaimGeneration: number, expectedTaskVersion: number): boolean => + task.kind !== 'testing' && + task.status === 'claimed' && + task.claimId !== undefined && + task.claimGeneration === expectedClaimGeneration + 1 && + task.taskVersion === expectedTaskVersion + 1 && + task.claimGeneration > 0 && + task.workerId !== undefined && + task.machineId !== undefined && + task.leaseToken !== undefined && + task.leaseExpiresAt !== undefined; + +const expirableTaskTimestamp = (task: Task): string => { + const value = task.status === 'submitted' ? task.constraints.deadline : task.leaseExpiresAt; + return typeof value === 'string' ? value : ''; +}; + +const assertPositivePageLimit = (limit: number): void => { + if (!Number.isSafeInteger(limit) || limit <= 0) throw new RangeError('page limit must be a positive safe integer'); +}; + export class MemoryRepository implements Repository { private readonly tasks = new Map(); private readonly pools = new Map(); @@ -21,6 +42,10 @@ export class MemoryRepository implements Repository { private readonly testingRuns = new Map(); private readonly testingRunIdempotency = new Map(); private readonly testingMachineReservations = new Map(); + private readonly taskMaintenanceCursors = new Map(); + private readonly legacyMachineRecoveryMarkers = new Map>(); + + public constructor(private readonly clock: () => number = () => Date.now()) {} public async ping(): Promise {} @@ -34,9 +59,74 @@ export class MemoryRepository implements Repository { this.tasks.set(task.id, task); } + public async claimTask(task: Task, expectedClaimGeneration: number, expectedTaskVersion: number): Promise { + if (!isValidClaim(task, expectedClaimGeneration, expectedTaskVersion)) return undefined; + const current = this.tasks.get(task.id); + if ( + current?.status !== 'submitted' || + (current.claimGeneration ?? 0) !== expectedClaimGeneration || + (current.taskVersion ?? 0) !== expectedTaskVersion || + (current.claimId !== undefined && current.claimReleased !== true) + ) return undefined; + this.tasks.set(task.id, task); + return task; + } + + public async replaceTaskForClaim(task: Task, guard: TaskClaimGuard): Promise { + if (task.claimId !== guard.claimId || task.claimGeneration !== guard.claimGeneration) return false; + const current = this.tasks.get(task.id); + if ( + current?.claimId !== guard.claimId || + current.claimGeneration !== guard.claimGeneration || + (current.taskVersion ?? 0) !== guard.taskVersion || + current.status !== guard.status + ) return false; + this.tasks.set(task.id, { ...task, taskVersion: guard.taskVersion + 1 }); + return true; + } + + public async replaceTaskForActiveClaim(task: Task, guard: TaskActiveClaimGuard): Promise { + const current = this.tasks.get(task.id); + if (current?.leaseExpiresAt !== guard.leaseExpiresAt || !isFutureTimestamp(current.leaseExpiresAt, this.clock())) return false; + return this.replaceTaskForClaim(task, guard); + } + + public async replaceTaskForExpiredClaim(task: Task, guard: TaskActiveClaimGuard): Promise { + const current = this.tasks.get(task.id); + if (current?.leaseExpiresAt !== guard.leaseExpiresAt || isFutureTimestamp(current.leaseExpiresAt, this.clock())) return false; + return this.replaceTaskForClaim(task, guard); + } + + public async replaceSubmittedTask(task: Task, expectedClaimGeneration: number, expectedTaskVersion: number): Promise { + if ((task.claimGeneration ?? 0) !== expectedClaimGeneration) return false; + const current = this.tasks.get(task.id); + if (current?.status !== 'submitted' || (current.claimGeneration ?? 0) !== expectedClaimGeneration || (current.taskVersion ?? 0) !== expectedTaskVersion) return false; + this.tasks.set(task.id, { ...task, taskVersion: expectedTaskVersion + 1 }); + return true; + } + + public async replaceTaskForRecovery(task: Task, guard: TaskRecoveryGuard): Promise { + const current = this.tasks.get(task.id); + if ( + current?.status !== guard.status || + (current.taskVersion ?? 0) !== guard.taskVersion || + current.updatedAt !== guard.updatedAt || + current.claimId !== guard.claimId || + current.claimGeneration !== guard.claimGeneration || + current.claimRecovery?.recoveryId !== guard.recoveryId || + current.claimRecovery?.phase !== guard.recoveryPhase + ) return false; + this.tasks.set(task.id, { ...task, taskVersion: guard.taskVersion + 1 }); + return true; + } + public async listQueuedTasks(): Promise { return [...this.tasks.values()] - .filter((task) => task.status === 'submitted') + .filter((task) => + task.status === 'submitted' && + task.claimRecovery === undefined && + (task.claimId === undefined || task.claimReleased === true) + ) .sort( (a, b) => (a.queuePriority ?? 0) - (b.queuePriority ?? 0) || @@ -47,6 +137,69 @@ export class MemoryRepository implements Repository { return [...this.tasks.values()]; } + public async listExpirableTasks(now: number, limit: number): Promise { + assertPositivePageLimit(limit); + const timestamp = new Date(now).toISOString(); + const deadlineTasks = [...this.tasks.values()] + .filter((task) => + task.kind !== 'testing' && + task.status === 'submitted' && + task.constraints.deadline !== undefined && + task.constraints.deadline <= timestamp + ) + .sort((left, right) => expirableTaskTimestamp(left).localeCompare(expirableTaskTimestamp(right)) || left.id.localeCompare(right.id)) + .slice(0, limit); + const leaseTasks = [...this.tasks.values()] + .filter((task) => + task.kind !== 'testing' && + ['claimed', 'running', 'closing'].includes(task.status) && + task.leaseExpiresAt !== undefined && + task.leaseExpiresAt <= timestamp + ) + .sort((left, right) => expirableTaskTimestamp(left).localeCompare(expirableTaskTimestamp(right)) || left.id.localeCompare(right.id)) + .slice(0, limit); + return [...deadlineTasks, ...leaseTasks] + .sort((left, right) => expirableTaskTimestamp(left).localeCompare(expirableTaskTimestamp(right)) || left.id.localeCompare(right.id)) + .slice(0, limit); + } + + public async listTaskMaintenancePage(cursor: TaskMaintenanceCursor, limit: number): Promise { + return [...this.tasks.values()] + .filter((task) => + task.createdAt <= cursor.cycleCutoffAt && + ( + cursor.afterCreatedAt === undefined || + task.createdAt > cursor.afterCreatedAt || + (task.createdAt === cursor.afterCreatedAt && task.id > (cursor.afterTaskId ?? '')) + ) + ) + .sort((left, right) => left.createdAt.localeCompare(right.createdAt) || left.id.localeCompare(right.id)) + .slice(0, limit); + } + + public async getTaskMaintenanceHighWatermark(): Promise { + return [...this.tasks.values()] + .map((task) => task.createdAt) + .sort((left, right) => right.localeCompare(left))[0]; + } + + public async createTaskMaintenanceCursor(cursor: TaskMaintenanceCursor): Promise { + if (this.taskMaintenanceCursors.has(cursor.id)) return false; + this.taskMaintenanceCursors.set(cursor.id, cursor); + return true; + } + + public async getTaskMaintenanceCursor(id: TaskMaintenanceCursor['id']): Promise { + return this.taskMaintenanceCursors.get(id); + } + + public async replaceTaskMaintenanceCursor(cursor: TaskMaintenanceCursor, expectedVersion: number): Promise { + const current = this.taskMaintenanceCursors.get(cursor.id); + if (current?.version !== expectedVersion || cursor.version !== expectedVersion + 1) return false; + this.taskMaintenanceCursors.set(cursor.id, cursor); + return true; + } + public async getPool(id: string): Promise { return this.pools.get(id); } @@ -72,12 +225,156 @@ export class MemoryRepository implements Repository { this.machines.set(machine.id, machine); } + public async rotateMachineToken(machineId: string, expectedTokenHash: string, tokenHash: string): Promise { + const machine = this.machines.get(machineId); + if (machine === undefined || machine.workerTokenHash !== expectedTokenHash) return false; + this.machines.set(machineId, { ...machine, workerTokenHash: tokenHash }); + return true; + } + + public async reserveMachineLease(machineId: string, reservation: MachineLeaseReservation): Promise { + const machine = this.machines.get(machineId); + if (machine === undefined) return false; + const reservations = machine.leaseReservations ?? []; + if (reservations.some((entry) => + entry.taskId === reservation.taskId && + entry.claimId === reservation.claimId && + entry.claimGeneration === reservation.claimGeneration + )) return true; + if (reservations.some((entry) => entry.claimId === reservation.claimId)) return false; + if (!machine.online || machine.activeLeases >= machine.capacity) return false; + this.machines.set(machineId, { + ...machine, + activeLeases: machine.activeLeases + 1, + leaseReservations: [...reservations, reservation] + }); + return true; + } + + public async renewMachineLease(machineId: string, reservation: MachineLeaseReservation): Promise { + const machine = this.machines.get(machineId); + if (machine === undefined) return false; + const reservations = machine.leaseReservations ?? []; + const index = reservations.findIndex((entry) => entry.claimId === reservation.claimId && entry.claimGeneration === reservation.claimGeneration && entry.taskId === reservation.taskId); + if (index < 0) return false; + const next = [...reservations]; + const current = next[index]!; + next[index] = Date.parse(current.expiresAt) >= Date.parse(reservation.expiresAt) ? current : reservation; + this.machines.set(machineId, { ...machine, leaseReservations: next }); + return true; + } + + public async releaseMachineLease(machineId: string, reservation: Omit): Promise { + const machine = this.machines.get(machineId); + if (machine === undefined) return false; + const reservations = machine.leaseReservations ?? []; + const next = reservations.filter((entry) => !( + entry.claimId === reservation.claimId && + entry.claimGeneration === reservation.claimGeneration && + entry.taskId === reservation.taskId + )); + if (next.length === reservations.length) return false; + this.machines.set(machineId, { ...machine, activeLeases: machine.activeLeases - 1, leaseReservations: next }); + return true; + } + + public async releaseMachineLeaseReservation(reservation: Omit): Promise { + const machine = [...this.machines.values()].find((candidate) => + candidate.leaseReservations?.some((entry) => + entry.taskId === reservation.taskId && + entry.claimId === reservation.claimId && + entry.claimGeneration === reservation.claimGeneration + ) === true + ); + return machine === undefined ? false : this.releaseMachineLease(machine.id, reservation); + } + + public async releaseLegacyMachineLease(machineId: string, recoveryId: string, taskId: string): Promise { + const machine = this.machines.get(machineId); + if (machine === undefined) return false; + const markers = this.legacyMachineRecoveryMarkers.get(machineId) ?? new Map(); + const markedTaskId = markers.get(recoveryId); + if (markedTaskId !== undefined) return markedTaskId === taskId; + const reservationCount = machine.leaseReservations?.length ?? 0; + this.machines.set(machineId, { + ...machine, + activeLeases: Math.max(reservationCount, machine.activeLeases - 1) + }); + markers.set(recoveryId, taskId); + this.legacyMachineRecoveryMarkers.set(machineId, markers); + return true; + } + + public async clearLegacyMachineLeaseMarker(machineId: string, recoveryId: string): Promise { + const markers = this.legacyMachineRecoveryMarkers.get(machineId); + if (markers === undefined) return false; + const removed = markers.delete(recoveryId); + if (markers.size === 0) this.legacyMachineRecoveryMarkers.delete(machineId); + return removed; + } + public async getProfile(id: string): Promise { return this.profiles.get(id); } - public async saveProfile(profile: Profile): Promise { + public async createProfile(profile: Profile): Promise { + if (this.profiles.has(profile.id)) return false; this.profiles.set(profile.id, profile); + return true; + } + + public async acquireProfileLease(profileId: string, userId: string, machineId: string, reservation: MachineLeaseReservation): Promise { + const profile = this.profiles.get(profileId); + if (profile === undefined || profile.userId !== userId) return undefined; + const sameClaim = profile.lockedByClaimId === reservation.claimId && profile.lockedByClaimGeneration === reservation.claimGeneration; + if (!sameClaim && profile.lockedByTaskId !== undefined) return undefined; + const updated: Profile = { + ...profile, + machineId, + lockedByTaskId: reservation.taskId, + lockedByClaimId: reservation.claimId, + lockedByClaimGeneration: reservation.claimGeneration, + lockExpiresAt: sameClaim && Date.parse(profile.lockExpiresAt ?? '') >= Date.parse(reservation.expiresAt) + ? profile.lockExpiresAt + : reservation.expiresAt + }; + this.profiles.set(profileId, updated); + return updated; + } + + + public async releaseProfileLease(profileId: string, reservation: Omit): Promise { + const profile = this.profiles.get(profileId); + if ( + profile?.lockedByTaskId !== reservation.taskId || + profile.lockedByClaimId !== reservation.claimId || + profile.lockedByClaimGeneration !== reservation.claimGeneration + ) return false; + this.profiles.set(profileId, { + id: profile.id, + userId: profile.userId, + ...(profile.machineId === undefined ? {} : { machineId: profile.machineId }) + }); + return true; + } + + public async releaseLegacyProfileLease(profileId: string, taskId: string): Promise { + const profile = this.profiles.get(profileId); + if ( + profile?.lockedByTaskId !== taskId || + profile.lockedByClaimId !== undefined || + profile.lockedByClaimGeneration !== undefined + ) return false; + this.profiles.set(profileId, { + ...profile, + lockedByTaskId: undefined, + lockExpiresAt: undefined + }); + return true; + } + + public async listProfiles(): Promise { + return [...this.profiles.values()]; } public async listProfilesByUser(userId: string): Promise { diff --git a/control-plane/src/storage/mongo-repository.test.ts b/control-plane/src/storage/mongo-repository.test.ts index 64be295..c7bce16 100644 --- a/control-plane/src/storage/mongo-repository.test.ts +++ b/control-plane/src/storage/mongo-repository.test.ts @@ -133,6 +133,14 @@ class FakeMongoClient { public async close(): Promise {} } +const fakeMongoRepository = (): { client: FakeMongoClient; repository: MongoRepository } => { + const client = new FakeMongoClient(); + const repository = new MongoRepository('mongodb://unused', 'talos-test', { + client: client as unknown as MongoClient + }); + return { client, repository }; +}; + const makeTestingRun = async (): Promise => { const digest = `sha256:${'a'.repeat(64)}`; const reference = (schema: string, ref: string) => ({ schema, ref, digest }); @@ -209,11 +217,33 @@ const makeTestingRun = async (): Promise => { }; describe('MongoRepository testing run persistence', () => { - it('creates the unique index, maps duplicate inserts, and rejects stale CAS writes', async () => { - const client = new FakeMongoClient(); - const repository = new MongoRepository('mongodb://unused', 'talos-test', { - client: client as unknown as MongoClient + it('normalizes legacy null machine reservations as absent', async () => { + const { client, repository } = fakeMongoRepository(); + await client.database.collection('machines').insertOne({ + _id: 'legacy-machine', + id: 'legacy-machine', + poolId: 'pool-1', + tags: {}, + capacity: 1, + activeLeases: 0, + online: true, + workerTokenHash: 'hash', + leaseReservations: null }); + + expect(await repository.getMachine('legacy-machine')).toEqual({ + id: 'legacy-machine', + poolId: 'pool-1', + tags: {}, + capacity: 1, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + }); + + it('creates the unique index, maps duplicate inserts, and rejects stale CAS writes', async () => { + const { client, repository } = fakeMongoRepository(); await repository.initialize(); expect(client.database.collection('testing_runs').indexes).toContainEqual({ keys: { userId: 1, idempotencyKey: 1 }, diff --git a/control-plane/src/storage/mongo-repository.ts b/control-plane/src/storage/mongo-repository.ts index c2016c5..5284e4c 100644 --- a/control-plane/src/storage/mongo-repository.ts +++ b/control-plane/src/storage/mongo-repository.ts @@ -1,10 +1,16 @@ -import { MongoClient, type Collection, type Db, type MongoClientOptions } from 'mongodb'; -import type { HandoffLink, Machine, PendingSessionAction, Pool, Profile, SessionActionResult, Task, TaskInput, WebhookEvent } from '../domain/types.js'; +import { MongoClient, type Collection, type Db, type Document as MongoDriverDocument, type Filter, type MongoClientOptions, type UpdateFilter } from 'mongodb'; +import type { HandoffLink, Machine, MachineLeaseReservation, PendingSessionAction, Pool, Profile, SessionActionResult, Task, TaskActiveClaimGuard, TaskClaimGuard, TaskInput, TaskRecoveryGuard, WebhookEvent } from '../domain/types.js'; import type { TestingMachineReservationRecord, TestingRunRecord } from '../domain/testing-types.js'; -import type { Repository, TestingAttemptDispatchGuard, TestingAttemptMutationGuard } from './repository.js'; +import type { Repository, TaskMaintenanceCursor, TestingAttemptDispatchGuard, TestingAttemptMutationGuard } from './repository.js'; type Document = { _id: string; [key: string]: unknown }; +type MachineDocument = Omit & { + _id: string; + leaseReservations?: MachineLeaseReservation[]; + legacyLeaseRecoveryMarkers?: Array<{ recoveryId: string; taskId: string }>; +}; + const mongoDate = (input: unknown): Readonly> => ({ $convert: { input, to: 'date', onError: null, onNull: null } }); @@ -13,6 +19,27 @@ const afterDatabaseNow = (input: unknown): Readonly> => $gt: [mongoDate(input), '$$NOW'] }); +const atOrBeforeDatabaseNow = (input: unknown): Readonly> => ({ + $lte: [mongoDate(input), '$$NOW'] +}); + +const NON_TESTING_TASK_KINDS = ['browse', 'computer_use'] as const; +const ACTIVE_TASK_STATUSES = ['claimed', 'running', 'closing'] as const; +const TASK_LEASE_EXPIRY_INDEX = 'task-lease-expiry-v1'; +const TASK_DEADLINE_EXPIRY_INDEX = 'task-deadline-expiry-v1'; + +const assertPositivePageLimit = (limit: number): void => { + if (!Number.isSafeInteger(limit) || limit <= 0) throw new RangeError('page limit must be a positive safe integer'); +}; + +const expirableTaskTimestamp = (task: Task): string => { + const value = task.status === 'submitted' ? task.constraints.deadline : task.leaseExpiresAt; + return typeof value === 'string' ? value : ''; +}; + +const compareExpirableTasks = (left: Task, right: Task): number => + expirableTaskTimestamp(left).localeCompare(expirableTaskTimestamp(right)) || left.id.localeCompare(right.id); + export interface MongoRepositoryOptions { client?: MongoClient; clientOptions?: MongoClientOptions; @@ -23,7 +50,7 @@ export class MongoRepository implements Repository { private readonly database: Db; private readonly tasks: Collection; private readonly pools: Collection; - private readonly machines: Collection; + private readonly machines: Collection; private readonly profiles: Collection; private readonly handoffs: Collection; private readonly webhooks: Collection; @@ -32,13 +59,14 @@ export class MongoRepository implements Repository { private readonly actionResults: Collection; private readonly testingRuns: Collection; private readonly testingMachineReservations: Collection; + private readonly maintenance: Collection; public constructor(url: string, databaseName = 'talos', options: MongoRepositoryOptions = {}) { this.client = options.client ?? new MongoClient(url, { ...options.clientOptions, ignoreUndefined: true }); this.database = this.client.db(databaseName); this.tasks = this.database.collection('tasks'); this.pools = this.database.collection('pools'); - this.machines = this.database.collection('machines'); + this.machines = this.database.collection('machines'); this.profiles = this.database.collection('profiles'); this.handoffs = this.database.collection('handoffs'); this.webhooks = this.database.collection('webhooks'); @@ -47,15 +75,31 @@ export class MongoRepository implements Repository { this.actionResults = this.database.collection('action_results'); this.testingRuns = this.database.collection('testing_runs'); this.testingMachineReservations = this.database.collection('testing_machine_reservations'); + this.maintenance = this.database.collection('maintenance'); } public async initialize(): Promise { await this.client.connect(); await Promise.all([ this.tasks.createIndex({ status: 1, queuePriority: 1, createdAt: 1 }), + this.tasks.createIndex({ kind: 1, claimId: 1, status: 1, updatedAt: 1 }), + this.tasks.createIndex( + { kind: 1, status: 1, leaseExpiresAt: 1, _id: 1 }, + { name: TASK_LEASE_EXPIRY_INDEX } + ), + this.tasks.createIndex( + { kind: 1, status: 1, 'constraints.deadline': 1, _id: 1 }, + { name: TASK_DEADLINE_EXPIRY_INDEX } + ), + this.tasks.createIndex({ createdAt: 1, _id: 1 }), this.pools.createIndex({ ownerUserId: 1 }), this.profiles.createIndex({ userId: 1 }), this.machines.createIndex({ poolId: 1 }), + this.machines.createIndex({ + 'leaseReservations.taskId': 1, + 'leaseReservations.claimId': 1, + 'leaseReservations.claimGeneration': 1 + }), this.pendingActions.createIndex( { taskId: 1 }, { unique: true, partialFilterExpression: { state: { $in: ['pending', 'dispatched'] } } } @@ -84,8 +128,103 @@ export class MongoRepository implements Repository { await this.tasks.replaceOne({ _id: task.id }, { ...task, _id: task.id, queuePriority: task.queuePriority ?? 0 }, { upsert: true }); } + public async claimTask(task: Task, expectedClaimGeneration: number, expectedTaskVersion: number): Promise { + if (!isValidClaim(task, expectedClaimGeneration, expectedTaskVersion)) return undefined; + const filter = { + _id: task.id, + status: 'submitted', + $and: [ + claimGenerationFilter(expectedClaimGeneration), + taskVersionFilter(expectedTaskVersion), + { $or: [{ claimId: { $exists: false } }, { claimReleased: true }] } + ] + } satisfies Filter; + const document = await this.tasks.findOneAndReplace( + filter, + { ...task, _id: task.id, queuePriority: task.queuePriority ?? 0 }, + { returnDocument: 'after' } + ); + return document === null ? undefined : taskFromDocument(document); + } + + public async replaceTaskForClaim(task: Task, guard: TaskClaimGuard): Promise { + if (task.claimId !== guard.claimId || task.claimGeneration !== guard.claimGeneration) return false; + const result = await this.tasks.replaceOne( + { _id: task.id, status: guard.status, claimId: guard.claimId, claimGeneration: guard.claimGeneration, ...taskVersionFilter(guard.taskVersion) }, + { ...task, taskVersion: guard.taskVersion + 1, _id: task.id, queuePriority: task.queuePriority ?? 0 } + ); + return result.matchedCount === 1; + } + + public async replaceTaskForActiveClaim(task: Task, guard: TaskActiveClaimGuard): Promise { + if (task.claimId !== guard.claimId || task.claimGeneration !== guard.claimGeneration) return false; + const result = await this.tasks.replaceOne( + { + _id: task.id, + status: guard.status, + claimId: guard.claimId, + claimGeneration: guard.claimGeneration, + leaseExpiresAt: guard.leaseExpiresAt, + ...taskVersionFilter(guard.taskVersion), + $expr: afterDatabaseNow('$leaseExpiresAt') + }, + { ...task, taskVersion: guard.taskVersion + 1, _id: task.id, queuePriority: task.queuePriority ?? 0 } + ); + return result.matchedCount === 1; + } + + public async replaceTaskForExpiredClaim(task: Task, guard: TaskActiveClaimGuard): Promise { + if (task.claimId !== guard.claimId || task.claimGeneration !== guard.claimGeneration) return false; + const result = await this.tasks.replaceOne( + { + _id: task.id, + status: guard.status, + claimId: guard.claimId, + claimGeneration: guard.claimGeneration, + leaseExpiresAt: guard.leaseExpiresAt, + ...taskVersionFilter(guard.taskVersion), + $expr: atOrBeforeDatabaseNow('$leaseExpiresAt') + }, + { ...task, taskVersion: guard.taskVersion + 1, _id: task.id, queuePriority: task.queuePriority ?? 0 } + ); + return result.matchedCount === 1; + } + + public async replaceSubmittedTask(task: Task, expectedClaimGeneration: number, expectedTaskVersion: number): Promise { + if ((task.claimGeneration ?? 0) !== expectedClaimGeneration) return false; + const result = await this.tasks.replaceOne( + { _id: task.id, status: 'submitted', ...claimGenerationFilter(expectedClaimGeneration), ...taskVersionFilter(expectedTaskVersion) }, + { ...task, taskVersion: expectedTaskVersion + 1, _id: task.id, queuePriority: task.queuePriority ?? 0 } + ); + return result.modifiedCount === 1; + } + + public async replaceTaskForRecovery(task: Task, guard: TaskRecoveryGuard): Promise { + const filter = { + _id: task.id, + status: guard.status, + updatedAt: guard.updatedAt, + $and: [ + taskVersionFilter(guard.taskVersion), + optionalFieldFilter('claimId', guard.claimId), + optionalFieldFilter('claimGeneration', guard.claimGeneration), + optionalFieldFilter('claimRecovery.recoveryId', guard.recoveryId), + optionalFieldFilter('claimRecovery.phase', guard.recoveryPhase) + ] + } satisfies Filter; + const result = await this.tasks.replaceOne( + filter, + { ...task, taskVersion: guard.taskVersion + 1, _id: task.id, queuePriority: task.queuePriority ?? 0 } + ); + return result.matchedCount === 1; + } + public async listQueuedTasks(): Promise { - const documents = await this.tasks.find({ status: 'submitted' }).sort({ queuePriority: 1, createdAt: 1 }).toArray(); + const documents = await this.tasks.find({ + status: 'submitted', + claimRecovery: { $exists: false }, + $or: [{ claimId: { $exists: false } }, { claimReleased: true }] + }).sort({ queuePriority: 1, createdAt: 1 }).toArray(); return documents.map(taskFromDocument); } @@ -93,6 +232,73 @@ export class MongoRepository implements Repository { return (await this.tasks.find({}).toArray()).map(taskFromDocument); } + public async listExpirableTasks(now: number, limit: number): Promise { + assertPositivePageLimit(limit); + const timestamp = new Date(now).toISOString(); + const [deadlineDocuments, leaseDocuments] = await Promise.all([ + this.tasks.find({ + kind: { $in: NON_TESTING_TASK_KINDS }, + status: 'submitted', + 'constraints.deadline': { $lte: timestamp } + }).sort({ 'constraints.deadline': 1, _id: 1 }).limit(limit).toArray(), + this.tasks.find({ + kind: { $in: NON_TESTING_TASK_KINDS }, + status: { $in: ACTIVE_TASK_STATUSES }, + leaseExpiresAt: { $lte: timestamp } + }).sort({ leaseExpiresAt: 1, _id: 1 }).limit(limit).toArray() + ]); + return [...deadlineDocuments, ...leaseDocuments] + .map(taskFromDocument) + .sort(compareExpirableTasks) + .slice(0, limit); + } + + public async listTaskMaintenancePage(cursor: TaskMaintenanceCursor, limit: number): Promise { + const after = cursor.afterCreatedAt === undefined + ? {} + : { + $or: [ + { createdAt: { $gt: cursor.afterCreatedAt } }, + { createdAt: cursor.afterCreatedAt, _id: { $gt: cursor.afterTaskId ?? '' } } + ] + }; + const documents = await this.tasks.find({ + createdAt: { $lte: cursor.cycleCutoffAt }, + ...after + }).sort({ createdAt: 1, _id: 1 }).limit(limit).toArray(); + return documents.map(taskFromDocument); + } + + public async getTaskMaintenanceHighWatermark(): Promise { + const document = await this.tasks.find({}).sort({ createdAt: -1, _id: -1 }).limit(1).toArray(); + const createdAt = document[0]?.createdAt; + return typeof createdAt === 'string' ? createdAt : undefined; + } + + public async createTaskMaintenanceCursor(cursor: TaskMaintenanceCursor): Promise { + try { + await this.maintenance.insertOne({ ...cursor, _id: cursor.id }); + return true; + } catch (error) { + if (isDuplicateKeyError(error)) return false; + throw error; + } + } + + public async getTaskMaintenanceCursor(id: TaskMaintenanceCursor['id']): Promise { + const document = await this.maintenance.findOne({ _id: id }); + return document === null ? undefined : taskMaintenanceCursorFromDocument(document); + } + + public async replaceTaskMaintenanceCursor(cursor: TaskMaintenanceCursor, expectedVersion: number): Promise { + if (cursor.version !== expectedVersion + 1) return false; + const result = await this.maintenance.replaceOne( + { _id: cursor.id, version: expectedVersion }, + { ...cursor, _id: cursor.id } + ); + return result.matchedCount === 1; + } + public async getPool(id: string): Promise { const document = await this.pools.findOne({ _id: id }); return document === null ? undefined : poolFromDocument(document); @@ -116,7 +322,119 @@ export class MongoRepository implements Repository { } public async saveMachine(machine: Machine): Promise { - await this.machines.replaceOne({ _id: machine.id }, { ...machine, _id: machine.id }, { upsert: true }); + const { leaseReservations, ...fields } = machine; + const update: UpdateFilter = leaseReservations === undefined + ? { $set: fields, $unset: { leaseReservations: '' } } + : { $set: { ...fields, leaseReservations: [...leaseReservations] } }; + await this.machines.updateOne({ _id: machine.id }, update, { upsert: true }); + } + + public async rotateMachineToken(machineId: string, expectedTokenHash: string, tokenHash: string): Promise { + const result = await this.machines.updateOne( + { _id: machineId, workerTokenHash: expectedTokenHash }, + { $set: { workerTokenHash: tokenHash } } + ); + return result.matchedCount === 1; + } + + public async reserveMachineLease(machineId: string, reservation: MachineLeaseReservation): Promise { + const existingFilter = { + _id: machineId, + leaseReservations: { $elemMatch: claimReservationFilter(reservation) } + } satisfies Filter; + const existing = await this.machines.findOne(existingFilter); + if (existing !== null) return true; + const filter = { + _id: machineId, + online: true, + $expr: { $lt: ['$activeLeases', '$capacity'] }, + leaseReservations: { $not: { $elemMatch: { claimId: reservation.claimId } } } + } satisfies Filter; + const update = { + $inc: { activeLeases: 1 }, + $push: { leaseReservations: reservation } + } satisfies UpdateFilter; + const document = await this.machines.findOneAndUpdate(filter, update, { returnDocument: 'after' }); + if (document !== null) return true; + return await this.machines.findOne(existingFilter) !== null; + } + + public async renewMachineLease(machineId: string, reservation: MachineLeaseReservation): Promise { + const filter = { + _id: machineId, + leaseReservations: { $elemMatch: claimReservationFilter(reservation) } + } satisfies Filter; + const update = { + $max: { 'leaseReservations.$.expiresAt': reservation.expiresAt } + } satisfies UpdateFilter; + const result = await this.machines.updateOne(filter, update); + return result.matchedCount === 1; + } + + public async releaseMachineLease(machineId: string, reservation: Omit): Promise { + const filter = { + _id: machineId, + leaseReservations: { $elemMatch: claimReservationFilter(reservation) } + } satisfies Filter; + const update = { + $inc: { activeLeases: -1 }, + $pull: { leaseReservations: claimReservationFilter(reservation) } + } satisfies UpdateFilter; + const result = await this.machines.updateOne(filter, update); + return result.modifiedCount === 1; + } + + public async releaseMachineLeaseReservation(reservation: Omit): Promise { + const filter = { + leaseReservations: { $elemMatch: claimReservationFilter(reservation) } + } satisfies Filter; + const update = { + $inc: { activeLeases: -1 }, + $pull: { leaseReservations: claimReservationFilter(reservation) } + } satisfies UpdateFilter; + const result = await this.machines.updateOne(filter, update); + return result.modifiedCount === 1; + } + + public async releaseLegacyMachineLease(machineId: string, recoveryId: string, taskId: string): Promise { + const filter = { + _id: machineId, + legacyLeaseRecoveryMarkers: { $not: { $elemMatch: { recoveryId } } } + } satisfies Filter; + const pipeline: MongoDriverDocument[] = [{ + $set: { + activeLeases: { + $max: [ + { $size: { $ifNull: ['$leaseReservations', []] } }, + { $subtract: ['$activeLeases', 1] } + ] + }, + legacyLeaseRecoveryMarkers: { + $concatArrays: [ + { $ifNull: ['$legacyLeaseRecoveryMarkers', []] }, + [{ recoveryId, taskId }] + ] + } + } + }]; + const result = await this.machines.updateOne(filter, pipeline); + if (result.modifiedCount === 1) return true; + return await this.machines.findOne({ + _id: machineId, + legacyLeaseRecoveryMarkers: { $elemMatch: { recoveryId, taskId } } + }) !== null; + } + + public async clearLegacyMachineLeaseMarker(machineId: string, recoveryId: string): Promise { + const filter = { + _id: machineId, + legacyLeaseRecoveryMarkers: { $elemMatch: { recoveryId } } + } satisfies Filter; + const update = { + $pull: { legacyLeaseRecoveryMarkers: { recoveryId } } + } satisfies UpdateFilter; + const result = await this.machines.updateOne(filter, update); + return result.modifiedCount === 1; } public async getProfile(id: string): Promise { @@ -124,8 +442,64 @@ export class MongoRepository implements Repository { return document === null ? undefined : profileFromDocument(document); } - public async saveProfile(profile: Profile): Promise { - await this.profiles.replaceOne({ _id: profile.id }, { ...profile, _id: profile.id }, { upsert: true }); + public async createProfile(profile: Profile): Promise { + try { + await this.profiles.insertOne({ ...profile, _id: profile.id }); + return true; + } catch (error) { + if (isDuplicateKeyError(error)) return false; + throw error; + } + } + + public async acquireProfileLease(profileId: string, userId: string, machineId: string, reservation: MachineLeaseReservation): Promise { + const document = await this.profiles.findOneAndUpdate( + { + _id: profileId, + userId, + $or: [ + { lockedByClaimId: reservation.claimId, lockedByClaimGeneration: reservation.claimGeneration }, + { lockedByTaskId: { $exists: false } } + ] + }, + { + $set: { + machineId, + lockedByTaskId: reservation.taskId, + lockedByClaimId: reservation.claimId, + lockedByClaimGeneration: reservation.claimGeneration + }, + $max: { lockExpiresAt: reservation.expiresAt } + }, + { returnDocument: 'after' } + ); + return document === null ? undefined : profileFromDocument(document); + } + + + public async releaseProfileLease(profileId: string, reservation: Omit): Promise { + const result = await this.profiles.updateOne( + { _id: profileId, lockedByTaskId: reservation.taskId, lockedByClaimId: reservation.claimId, lockedByClaimGeneration: reservation.claimGeneration }, + { $unset: { lockedByTaskId: '', lockedByClaimId: '', lockedByClaimGeneration: '', lockExpiresAt: '' } } + ); + return result.modifiedCount === 1; + } + + public async releaseLegacyProfileLease(profileId: string, taskId: string): Promise { + const result = await this.profiles.updateOne( + { + _id: profileId, + lockedByTaskId: taskId, + lockedByClaimId: { $exists: false }, + lockedByClaimGeneration: { $exists: false } + }, + { $unset: { lockedByTaskId: '', lockExpiresAt: '' } } + ); + return result.modifiedCount === 1; + } + + public async listProfiles(): Promise { + return (await this.profiles.find({}).toArray()).map(profileFromDocument); } public async listProfilesByUser(userId: string): Promise { @@ -388,16 +762,58 @@ export class MongoRepository implements Repository { } } +const isValidClaim = (task: Task, expectedClaimGeneration: number, expectedTaskVersion: number): boolean => + task.kind !== 'testing' && + task.status === 'claimed' && + task.claimId !== undefined && + task.claimGeneration === expectedClaimGeneration + 1 && + task.taskVersion === expectedTaskVersion + 1 && + task.claimGeneration > 0 && + task.workerId !== undefined && + task.machineId !== undefined && + task.leaseToken !== undefined && + task.leaseExpiresAt !== undefined; + +const claimGenerationFilter = (expectedClaimGeneration: number) => + expectedClaimGeneration === 0 + ? { $or: [{ claimGeneration: 0 }, { claimGeneration: { $exists: false } }] } + : { claimGeneration: expectedClaimGeneration }; + +const taskVersionFilter = (expectedTaskVersion: number) => + expectedTaskVersion === 0 + ? { $or: [{ taskVersion: 0 }, { taskVersion: { $exists: false } }] } + : { taskVersion: expectedTaskVersion }; + +const optionalFieldFilter = (field: string, value: unknown): Filter => + value === undefined ? { [field]: { $exists: false } } : { [field]: value }; + +const claimReservationFilter = (reservation: Omit) => ({ + claimId: reservation.claimId, + claimGeneration: reservation.claimGeneration, + taskId: reservation.taskId +}) satisfies Filter; + const withoutId = (document: Document): Record => { return Object.fromEntries(Object.entries(document).filter(([key, value]) => key !== '_id' && value !== null)); }; const taskFromDocument = (document: Document): Task => ({ interaction: 'autonomous', + workerId: undefined, + leaseExpiresAt: undefined, + claimQueuePriority: undefined, ...withoutId(document) }) as unknown as Task; const poolFromDocument = (document: Document): Pool => withoutId(document) as unknown as Pool; -const machineFromDocument = (document: Document): Machine => withoutId(document) as unknown as Machine; +const machineFromDocument = ({ + _id: _documentId, + leaseReservations, + legacyLeaseRecoveryMarkers: _recoveryMarkers, + ...machine +}: MachineDocument): Machine => + leaseReservations == null + ? machine + : { ...machine, leaseReservations }; const profileFromDocument = (document: Document): Profile => withoutId(document) as unknown as Profile; const handoffFromDocument = (document: Document): HandoffLink => withoutId(document) as unknown as HandoffLink; const webhookFromDocument = (document: Document): WebhookEvent => withoutId(document) as unknown as WebhookEvent; @@ -412,6 +828,8 @@ const completedSessionActionResultFromDocument = (document: Document): SessionAc const testingRunFromDocument = (document: Document): TestingRunRecord => withoutId(document) as unknown as TestingRunRecord; const testingMachineReservationFromDocument = (document: Document): TestingMachineReservationRecord => withoutId(document) as unknown as TestingMachineReservationRecord; +const taskMaintenanceCursorFromDocument = (document: Document): TaskMaintenanceCursor => + withoutId(document) as unknown as TaskMaintenanceCursor; const isDuplicateKeyError = (error: unknown): boolean => typeof error === 'object' && error !== null && 'code' in error && error.code === 11000; diff --git a/control-plane/src/storage/repository-contract.test.ts b/control-plane/src/storage/repository-contract.test.ts index 946d768..a846c3b 100644 --- a/control-plane/src/storage/repository-contract.test.ts +++ b/control-plane/src/storage/repository-contract.test.ts @@ -1,10 +1,15 @@ import { describe, expect, it, beforeAll, afterAll } from 'vitest'; -import { MongoClient } from 'mongodb'; +import { MongoClient, type Document as MongoDocument } from 'mongodb'; import { MongoMemoryServer } from 'mongodb-memory-server'; import type { Repository } from './repository.js'; import { MemoryRepository } from './memory-repository.js'; import { MongoRepository } from './mongo-repository.js'; import type { BrowserTask, WebhookEvent } from '../domain/types.js'; +import { TaskService } from '../services/task-service.js'; +import { SessionService } from '../services/session-service.js'; +import { Scheduler } from '../services/scheduler.js'; +import { ProfileLockService } from '../services/profile-lock.js'; +import { WebhookSigner } from '../services/webhook-signer.js'; import { TestingRunService } from '../services/testing-run-service.js'; import { submitTestingRun } from '../test-support/testing-transport.js'; import { digestJson } from '@talos/testing-protocol'; @@ -18,6 +23,7 @@ import { testTestingExecutionDependencyReadiness } from '../test-support/testing interface Harness { repository: Repository; + restart: () => Promise; close: () => Promise; } @@ -27,9 +33,211 @@ const deferred = () => { return { promise, resolve }; }; +const barrierRepository = (repository: Repository, participants = 2): Repository => { + const barrier = deferred(); + let arrivals = 0; + return new Proxy(repository, { + get(target, property) { + if (property === 'listQueuedTasks') { + return async () => { + const queued = await target.listQueuedTasks(); + arrivals += 1; + if (arrivals === participants) barrier.resolve(); + await barrier.promise; + return queued; + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); +}; + +const profileAcquireRaceRepositories = ( + repository: Repository, + taskIds: readonly string[] +): { repositories: readonly Repository[]; arrivals: ReadonlySet } => { + const barrier = deferred(); + const arrivals = new Set(); + const repositories = taskIds.map((taskId) => new Proxy(repository, { + get(target, property) { + if (property === 'listQueuedTasks') { + return async () => (await target.listQueuedTasks()).filter((task) => task.id === taskId); + } + if (property === 'acquireProfileLease') { + return async (...args: Parameters): Promise>> => { + arrivals.add(args[3].taskId); + if (arrivals.size === taskIds.length) barrier.resolve(); + await barrier.promise; + return target.acquireProfileLease(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + })); + return { repositories, arrivals }; +}; + +const barrierBeforeRepositoryMethods = ( + repository: Repository, + methods: readonly (keyof Repository)[], + participants = methods.length +): Repository => { + const barrier = deferred(); + let arrivals = 0; + return new Proxy(repository, { + get(target, property) { + const value = Reflect.get(target, property); + if (typeof value === 'function' && methods.includes(property as keyof Repository)) { + return async (...args: unknown[]) => { + arrivals += 1; + if (arrivals === participants) barrier.resolve(); + await barrier.promise; + return Reflect.apply(value, target, args); + }; + } + return typeof value === 'function' ? value.bind(target) : value; + } + }); +}; + +type FaultBoundary = + | 'claim_task' + | 'machine_reservation' + | 'profile_lock' + | 'claim_commit' + | 'terminal_task' + | 'machine_release' + | 'profile_release' + | 'task_release' + | 'legacy_draining' + | 'legacy_machine_release' + | 'legacy_profile_release' + | 'legacy_finalizing' + | 'legacy_marker_clear' + | 'legacy_done' + | 'legacy_action_requeue' + | 'legacy_action_finalize'; + +const faultAfterBoundary = ( + repository: Repository, + boundary: FaultBoundary +): { repository: Repository; hitCount: () => number } => { + let injected = false; + let hits = 0; + const inject = (): never => { + injected = true; + hits += 1; + throw new Error(`injected fault after ${boundary}`); + }; + const faulted = new Proxy(repository, { + get(target, property) { + if (property === 'claimTask' && boundary === 'claim_task') { + return async (...args: Parameters): Promise>> => { + const result = await target.claimTask(...args); + if (!injected && result !== undefined) inject(); + return result; + }; + } + if (property === 'reserveMachineLease' && boundary === 'machine_reservation') { + return async (...args: Parameters): Promise => { + const result = await target.reserveMachineLease(...args); + if (!injected && result) inject(); + return result; + }; + } + if (property === 'acquireProfileLease' && boundary === 'profile_lock') { + return async (...args: Parameters): Promise>> => { + const result = await target.acquireProfileLease(...args); + if (!injected && result !== undefined) inject(); + return result; + }; + } + if (property === 'replaceTaskForClaim' && ['claim_commit', 'terminal_task', 'task_release'].includes(boundary)) { + return async (...args: Parameters): Promise => { + const result = await target.replaceTaskForClaim(...args); + const task = args[0]; + const matches = + (boundary === 'claim_commit' && task.status === 'claimed' && task.claimCommitted === true && task.claimReleased !== true) || + (boundary === 'terminal_task' && ['completed', 'failed', 'cancelled', 'submitted'].includes(task.status) && task.claimReleased !== true) || + (boundary === 'task_release' && task.claimReleased === true); + if (!injected && result && matches) inject(); + return result; + }; + } + if (property === 'releaseMachineLease' && boundary === 'machine_release') { + return async (...args: Parameters): Promise => { + const result = await target.releaseMachineLease(...args); + if (!injected && result) inject(); + return result; + }; + } + if (property === 'releaseProfileLease' && boundary === 'profile_release') { + return async (...args: Parameters): Promise => { + const result = await target.releaseProfileLease(...args); + if (!injected && result) inject(); + return result; + }; + } + if (property === 'replaceTaskForRecovery' && ['legacy_draining', 'legacy_finalizing', 'legacy_done'].includes(boundary)) { + return async (...args: Parameters): Promise => { + const result = await target.replaceTaskForRecovery(...args); + const [task, guard] = args; + const matches = + (boundary === 'legacy_draining' && task.claimRecovery?.phase === 'draining' && guard.recoveryId === undefined) || + (boundary === 'legacy_finalizing' && task.claimRecovery?.phase === 'finalizing') || + (boundary === 'legacy_done' && task.claimRecovery === undefined && guard.recoveryPhase === 'finalizing'); + if (!injected && result && matches) inject(); + return result; + }; + } + if (property === 'releaseLegacyMachineLease' && boundary === 'legacy_machine_release') { + return async (...args: Parameters): Promise => { + const result = await target.releaseLegacyMachineLease(...args); + if (!injected && result) inject(); + return result; + }; + } + if (property === 'releaseLegacyProfileLease' && boundary === 'legacy_profile_release') { + return async (...args: Parameters): Promise => { + const result = await target.releaseLegacyProfileLease(...args); + if (!injected && result) inject(); + return result; + }; + } + if (property === 'clearLegacyMachineLeaseMarker' && boundary === 'legacy_marker_clear') { + return async (...args: Parameters): Promise => { + const result = await target.clearLegacyMachineLeaseMarker(...args); + if (!injected && result) inject(); + return result; + }; + } + if (property === 'requeueSessionAction' && boundary === 'legacy_action_requeue') { + return async (...args: Parameters): Promise => { + await target.requeueSessionAction(...args); + if (!injected) inject(); + }; + } + if (property === 'finalizeSessionAction' && boundary === 'legacy_action_finalize') { + return async (...args: Parameters): Promise => { + const result = await target.finalizeSessionAction(...args); + if (!injected && result) inject(); + return result; + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + return { repository: faulted, hitCount: () => hits }; +}; + const MONGODB_MEMORY_SERVER_VERSION = '7.0.14'; const MONGODB_CONNECT_TIMEOUT_MS = 5_000; const MONGODB_CONTRACT_TEST_TIMEOUT_MS = 30_000; +const EXPIRY_QUERY_BATCH_SIZE = 100; +const EXPIRY_QUERY_BACKLOG_SIZE = 2_000; const mongodbClientOptions = { connectTimeoutMS: MONGODB_CONNECT_TIMEOUT_MS, serverSelectionTimeoutMS: MONGODB_CONNECT_TIMEOUT_MS @@ -42,7 +250,11 @@ const baseTask = (overrides: Partial = {}): BrowserTask => ({ const memoryHarness = async (): Promise => { const repository = new MemoryRepository(); - return { repository, close: () => repository.close() }; + return { + repository, + restart: async () => repository, + close: () => repository.close() + }; }; let mongoServer: MongoMemoryServer | undefined; @@ -71,9 +283,10 @@ afterAll(async () => { const mongoHarness = async (): Promise => { if (mongoUrl === undefined) throw new Error('Mongo contract setup did not provide a database URL'); - const client = new MongoClient(mongoUrl, mongodbClientOptions); + const url = mongoUrl; const databaseName = `talos_test_${Date.now()}_${Math.random().toString(16).slice(2)}`; - const repository = new MongoRepository(mongoUrl, databaseName, { client }); + let client = new MongoClient(url, mongodbClientOptions); + let repository = new MongoRepository(url, databaseName, { client }); try { await repository.initialize(); } catch (error) { @@ -88,6 +301,13 @@ const mongoHarness = async (): Promise => { } return { repository, + restart: async () => { + await repository.close(); + client = new MongoClient(url, mongodbClientOptions); + repository = new MongoRepository(url, databaseName, { client }); + await repository.initialize(); + return repository; + }, close: async () => { try { await client.db(databaseName).dropDatabase(); @@ -98,13 +318,1872 @@ const mongoHarness = async (): Promise => { }; }; +const taskService = (repository: Repository, clock = { value: 1_000 }): TaskService => new TaskService( + repository, + new Scheduler(repository), + new ProfileLockService(repository), + new WebhookSigner('repository-contract-webhook-secret'), + { clock: () => clock.value, leaseSeconds: 10 } +); + +const executionPlanContainsStage = (value: unknown, stage: string): boolean => { + if (Array.isArray(value)) return value.some((entry) => executionPlanContainsStage(entry, stage)); + if (typeof value !== 'object' || value === null) return false; + const document = value as Readonly>; + return typeof document.stage === 'string' && document.stage.toUpperCase() === stage.toUpperCase() || + Object.values(document).some((entry) => executionPlanContainsStage(entry, stage)); +}; + +const executionPlanContainsValue = (value: unknown, expected: string): boolean => { + if (value === expected) return true; + if (Array.isArray(value)) return value.some((entry) => executionPlanContainsValue(entry, expected)); + if (typeof value !== 'object' || value === null) return false; + return Object.values(value).some((entry) => executionPlanContainsValue(entry, expected)); +}; + +const normalizedMongoSort = (value: unknown): Readonly> => { + if (value instanceof Map) return Object.fromEntries(value); + return typeof value === 'object' && value !== null ? value as Readonly> : {}; +}; + const contractTests = (makeHarness: () => Promise): void => { + it('returns one stable bounded page across deadline and lease expiry sources', async () => { + const { repository, close } = await makeHarness(); + try { + const now = Date.parse('2025-01-01T00:00:10.000Z'); + await repository.saveTask(baseTask({ + id: 'expiry-deadline-b', + constraints: { deadline: '2025-01-01T00:00:01.000Z' } + })); + await repository.saveTask(baseTask({ + id: 'expiry-deadline-a', + constraints: { deadline: '2025-01-01T00:00:01.000Z' } + })); + await repository.saveTask(baseTask({ + id: 'expiry-lease', + status: 'running', + leaseExpiresAt: '2025-01-01T00:00:02.000Z' + })); + await repository.saveTask(baseTask({ + id: 'expiry-future', + status: 'claimed', + leaseExpiresAt: '2025-01-01T00:01:00.000Z' + })); + + expect((await repository.listExpirableTasks(now, 3)).map((task) => task.id)).toEqual([ + 'expiry-deadline-a', + 'expiry-deadline-b', + 'expiry-lease' + ]); + for (const invalidLimit of [0, -1, 1.5, Number.MAX_SAFE_INTEGER + 1]) { + await expect(repository.listExpirableTasks(now, invalidLimit)) + .rejects.toThrow('page limit must be a positive safe integer'); + } + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('drains malformed deadline pages without starving a later valid expiry', async () => { + const { repository, close } = await makeHarness(); + try { + const now = Date.parse('2025-01-01T00:01:00.000Z'); + const service = taskService(repository, { value: now }); + for (let index = 0; index <= EXPIRY_QUERY_BATCH_SIZE; index += 1) { + await repository.saveTask(baseTask({ + id: `malformed-deadline-${String(index).padStart(3, '0')}`, + constraints: { deadline: '' } + })); + } + await repository.saveTask(baseTask({ + id: 'valid-deadline-after-malformed-page', + constraints: { deadline: '2025-01-01T00:00:01.000Z' }, + createdAt: '2025-01-01T00:00:01.000Z' + })); + + await service.expireLeases(now); + + expect(await repository.getTask('valid-deadline-after-malformed-page')).toMatchObject({ + status: 'failed', + error: { code: 'deadline_exceeded' } + }); + expect((await repository.listTasks()).filter((task) => task.error?.code === 'invalid_deadline')) + .toHaveLength(EXPIRY_QUERY_BATCH_SIZE + 1); + const canonical = await service.createTask('user-1', { + kind: 'browse', + goal: 'canonical deadline', + constraints: { deadline: '2025-01-01T08:02:00.000+08:00' } + }); + expect(canonical.constraints.deadline).toBe('2025-01-01T00:02:00.000Z'); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('linearizes task claim, machine admission, and profile ownership', async () => { + const { repository, close } = await makeHarness(); + try { + await repository.savePool({ id: 'claim-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'claim-machine', poolId: 'claim-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + await repository.createProfile({ id: 'claim-profile', userId: 'user-1' }); + await repository.saveTask(baseTask({ id: 'claim-task', profileId: 'claim-profile' })); + const service = taskService(barrierRepository(repository)); + + const results = await Promise.allSettled([ + service.claim('worker-a', 'claim-machine'), + service.claim('worker-b', 'claim-machine') + ]); + const winners = results.filter((result): result is PromiseFulfilledResult>> => result.status === 'fulfilled'); + expect(winners).toHaveLength(1); + const winner = winners[0]!.value; + const stored = await repository.getTask('claim-task'); + const machine = await repository.getMachine('claim-machine'); + const profile = await repository.getProfile('claim-profile'); + expect(stored).toMatchObject({ + status: 'claimed', + workerId: winner.task.workerId, + leaseToken: winner.leaseToken, + claimId: winner.task.claimId, + claimGeneration: 1 + }); + expect(machine).toMatchObject({ activeLeases: 1 }); + expect(machine?.leaseReservations).toEqual([expect.objectContaining({ + taskId: stored?.id, + claimId: stored?.claimId, + claimGeneration: stored?.claimGeneration + })]); + expect(profile).toMatchObject({ + lockedByTaskId: stored?.id, + lockedByClaimId: stored?.claimId, + lockedByClaimGeneration: stored?.claimGeneration + }); + expect(await repository.releaseMachineLease('claim-machine', { taskId: 'claim-task', claimId: 'loser-claim', claimGeneration: 1 })).toBe(false); + expect(await repository.releaseProfileLease('claim-profile', { taskId: 'claim-task', claimId: 'loser-claim', claimGeneration: 1 })).toBe(false); + expect(await repository.getMachine('claim-machine')).toMatchObject({ activeLeases: 1 }); + expect(await repository.getProfile('claim-profile')).toMatchObject({ lockedByClaimId: stored?.claimId }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('atomically admits only one claim into the final machine slot', async () => { + const { repository, close } = await makeHarness(); + try { + await repository.savePool({ id: 'capacity-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'capacity-machine', poolId: 'capacity-pool', tags: {}, capacity: 2, activeLeases: 1, online: true, workerTokenHash: 'hash' }); + await repository.saveTask(baseTask({ id: 'capacity-task-a', createdAt: '2025-01-01T00:00:00.000Z' })); + await repository.saveTask(baseTask({ id: 'capacity-task-b', createdAt: '2025-01-01T00:00:01.000Z' })); + const service = taskService(barrierRepository(repository)); + + const results = await Promise.allSettled([ + service.claim('worker-a', 'capacity-machine'), + service.claim('worker-b', 'capacity-machine') + ]); + expect(results.filter((result) => result.status === 'fulfilled')).toHaveLength(1); + expect(await repository.getMachine('capacity-machine')).toMatchObject({ + activeLeases: 2, + leaseReservations: [expect.any(Object)] + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('recovers a committed claim after faults at every claim persistence boundary', async () => { + const harness = await makeHarness(); + let repository = harness.repository; + try { + const clock = { value: Date.now() }; + await repository.savePool({ id: 'fault-claim-pool', visibility: 'platform', tags: {} }); + const boundaries = ['claim_task', 'machine_reservation', 'profile_lock', 'claim_commit'] as const; + for (const [index, boundary] of boundaries.entries()) { + const suffix = `${index}-${boundary}`; + const machineId = `fault-claim-machine-${suffix}`; + const profileId = `fault-claim-profile-${suffix}`; + const taskId = `fault-claim-task-${suffix}`; + await repository.saveMachine({ + id: machineId, + poolId: 'fault-claim-pool', + tags: {}, + capacity: 1, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ id: profileId, userId: 'user-1' }); + await repository.saveTask(baseTask({ id: taskId, profileId })); + const fault = faultAfterBoundary(repository, boundary); + const faulted = taskService(fault.repository, clock); + + await expect(faulted.claim(`worker-${suffix}`, machineId, clock.value)) + .rejects.toThrow(`injected fault after ${boundary}`); + expect(fault.hitCount()).toBe(1); + repository = await harness.restart(); + const restarted = taskService(repository, clock); + await restarted.reconcileClaims(clock.value); + + const stored = await repository.getTask(taskId); + expect(stored).toMatchObject({ + status: 'claimed', + claimCommitted: true, + claimReleased: false, + machineId, + workerId: `worker-${suffix}` + }); + expect(await repository.getMachine(machineId)).toMatchObject({ + activeLeases: 1, + leaseReservations: [{ + taskId, + claimId: stored?.claimId, + claimGeneration: stored?.claimGeneration + }] + }); + expect(await repository.getProfile(profileId)).toMatchObject({ + lockedByTaskId: taskId, + lockedByClaimId: stored?.claimId, + lockedByClaimGeneration: stored?.claimGeneration + }); + await expect(restarted.claim(`duplicate-${suffix}`, machineId, clock.value)) + .rejects.toMatchObject({ code: 'not_found' }); + } + } finally { + await harness.close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('recovers a lost claim response through expiry and a fenced generation N+1 reclaim', async () => { + const harness = await makeHarness(); + let repository = harness.repository; + try { + const clock = { value: 1_000 }; + await repository.savePool({ id: 'lost-response-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ + id: 'lost-response-machine', + poolId: 'lost-response-pool', + tags: {}, + capacity: 1, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ id: 'lost-response-profile', userId: 'user-1' }); + await repository.saveTask(baseTask({ id: 'lost-response-task', profileId: 'lost-response-profile' })); + const first = await taskService(repository, clock).claim('lost-response-worker', 'lost-response-machine', clock.value); + + repository = await harness.restart(); + const restarted = taskService(repository, clock); + await expect(restarted.claim('duplicate-worker', 'lost-response-machine', clock.value)) + .rejects.toMatchObject({ code: 'not_found' }); + + clock.value = 12_000; + expect(await restarted.expireLeases(clock.value)).toHaveLength(1); + expect(await repository.getTask('lost-response-task')).toMatchObject({ + status: 'submitted', + claimGeneration: first.task.claimGeneration, + claimReleased: true + }); + expect(await repository.getMachine('lost-response-machine')).toMatchObject({ + activeLeases: 0, + leaseReservations: [] + }); + expect((await repository.getProfile('lost-response-profile'))?.lockedByTaskId).toBeUndefined(); + + const second = await restarted.claim('replacement-worker', 'lost-response-machine', clock.value); + expect(second.task.claimGeneration).toBe((first.task.claimGeneration ?? 0) + 1); + expect(second.task.claimId).not.toBe(first.task.claimId); + await expect(restarted.heartbeat('lost-response-task', 'lost-response-worker', first.leaseToken, 10)) + .rejects.toMatchObject({ code: 'unauthorized' }); + await expect(restarted.complete('lost-response-task', 'lost-response-worker', first.leaseToken, 'completed', [])) + .rejects.toMatchObject({ code: 'unauthorized' }); + expect(await repository.getMachine('lost-response-machine')).toMatchObject({ + activeLeases: 1, + leaseReservations: [expect.objectContaining({ + taskId: second.task.id, + claimId: second.task.claimId, + claimGeneration: second.task.claimGeneration + })] + }); + expect(await repository.getProfile('lost-response-profile')).toMatchObject({ + lockedByTaskId: second.task.id, + lockedByClaimId: second.task.claimId, + lockedByClaimGeneration: second.task.claimGeneration + }); + } finally { + await harness.close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('converges exact release after faults without changing unrelated leases or profiles', async () => { + const harness = await makeHarness(); + let repository = harness.repository; + try { + const clock = { value: Date.now() }; + await repository.savePool({ id: 'fault-release-pool', visibility: 'platform', tags: {} }); + const boundaries = ['terminal_task', 'machine_release', 'profile_release', 'task_release'] as const; + for (const [index, boundary] of boundaries.entries()) { + const suffix = `${index}-${boundary}`; + const machineId = `fault-release-machine-${suffix}`; + const profileId = `fault-release-profile-${suffix}`; + const unrelatedProfileId = `unrelated-profile-${suffix}`; + const taskId = `fault-release-task-${suffix}`; + const unrelated = { + taskId: `unrelated-task-${suffix}`, + claimId: `unrelated-claim-${suffix}`, + claimGeneration: 9, + expiresAt: new Date(clock.value + 60_000).toISOString() + }; + await repository.saveMachine({ + id: machineId, + poolId: 'fault-release-pool', + tags: {}, + capacity: 2, + activeLeases: 1, + leaseReservations: [unrelated], + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ id: profileId, userId: 'user-1' }); + await repository.createProfile({ + id: unrelatedProfileId, + userId: 'user-1', + machineId, + lockedByTaskId: unrelated.taskId, + lockedByClaimId: unrelated.claimId, + lockedByClaimGeneration: unrelated.claimGeneration, + lockExpiresAt: unrelated.expiresAt + }); + await repository.saveTask(baseTask({ id: taskId, profileId })); + const claimed = await taskService(repository, clock).claim(`worker-${suffix}`, machineId, clock.value); + const fault = faultAfterBoundary(repository, boundary); + const faulted = taskService(fault.repository, clock); + + await expect(faulted.complete(taskId, `worker-${suffix}`, claimed.leaseToken, 'completed', [])) + .rejects.toThrow(`injected fault after ${boundary}`); + expect(fault.hitCount()).toBe(1); + repository = await harness.restart(); + await taskService(repository, clock).reconcileClaims(clock.value); + + expect(await repository.getTask(taskId)).toMatchObject({ status: 'completed', claimReleased: true }); + expect(await repository.getMachine(machineId)).toMatchObject({ + activeLeases: 1, + leaseReservations: [unrelated] + }); + expect((await repository.getProfile(profileId))?.lockedByTaskId).toBeUndefined(); + expect(await repository.getProfile(unrelatedProfileId)).toMatchObject({ + lockedByTaskId: unrelated.taskId, + lockedByClaimId: unrelated.claimId, + lockedByClaimGeneration: unrelated.claimGeneration + }); + } + } finally { + await harness.close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('converges legacy recovery after faults at every durable side-effect boundary', async () => { + const harness = await makeHarness(); + let repository = harness.repository; + try { + const clock = { value: Date.now() }; + const boundaries = [ + 'legacy_draining', + 'legacy_machine_release', + 'legacy_profile_release', + 'legacy_finalizing', + 'legacy_marker_clear', + 'legacy_done' + ] as const; + for (const [index, boundary] of boundaries.entries()) { + const suffix = `${index}-${boundary}`; + const machineId = `legacy-fault-machine-${suffix}`; + const profileId = `legacy-fault-profile-${suffix}`; + const taskId = `legacy-fault-task-${suffix}`; + await repository.saveMachine({ + id: machineId, + poolId: 'pool', + tags: {}, + capacity: 1, + activeLeases: 1, + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ + id: profileId, + userId: 'user-1', + machineId, + lockedByTaskId: taskId, + lockExpiresAt: new Date(clock.value + 60_000).toISOString() + }); + await repository.saveTask(baseTask({ + id: taskId, + status: 'running', + profileId, + machineId, + workerId: `legacy-worker-${suffix}`, + leaseToken: `legacy-token-${suffix}`, + leaseExpiresAt: new Date(clock.value + 60_000).toISOString(), + queuePriority: index + })); + + const fault = faultAfterBoundary(repository, boundary); + await taskService(fault.repository, clock).reconcileClaims(clock.value); + expect(fault.hitCount()).toBe(1); + const interrupted = await repository.getTask(taskId); + if (boundary === 'legacy_done') { + expect(interrupted?.claimRecovery).toBeUndefined(); + } else { + expect(interrupted?.claimRecovery?.phase).toBe( + ['legacy_finalizing', 'legacy_marker_clear'].includes(boundary) ? 'finalizing' : 'draining' + ); + } + expect(await repository.getMachine(machineId)).toMatchObject({ + activeLeases: ['legacy_draining'].includes(boundary) ? 1 : 0 + }); + expect((await repository.getProfile(profileId))?.lockedByTaskId).toBe( + ['legacy_draining', 'legacy_machine_release'].includes(boundary) ? taskId : undefined + ); + + repository = await harness.restart(); + const restarted = taskService(repository, clock); + await restarted.reconcileClaims(clock.value); + expect(await repository.getTask(taskId)).toMatchObject({ + status: 'submitted', + queuePriority: index + }); + expect((await repository.getTask(taskId))?.claimRecovery).toBeUndefined(); + expect(await repository.getMachine(machineId)).toMatchObject({ activeLeases: 0 }); + expect((await repository.getProfile(profileId))?.lockedByTaskId).toBeUndefined(); + + await restarted.reconcileClaims(clock.value); + expect(await repository.getMachine(machineId)).toMatchObject({ activeLeases: 0 }); + } + } finally { + await harness.close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('retries legacy interactive action side effects after restart', async () => { + const harness = await makeHarness(); + let repository = harness.repository; + try { + const clock = { value: Date.now() }; + await repository.saveTask(baseTask({ + id: 'legacy-action-requeue-task', + interaction: 'interactive', + status: 'running', + workerId: 'legacy-worker', + leaseToken: 'legacy-token', + pendingActionId: 'legacy-action-requeue' + })); + await repository.enqueueSessionAction({ + id: 'legacy-action-requeue', + taskId: 'legacy-action-requeue-task', + action: { type: 'navigate', url: 'https://example.com/retry' }, + state: 'pending', + createdAt: '2025-01-01T00:00:00.000Z' + }); + await repository.takePendingSessionAction('legacy-action-requeue-task'); + const requeueFault = faultAfterBoundary(repository, 'legacy_action_requeue'); + await taskService(requeueFault.repository, clock).reconcileClaims(clock.value); + expect(requeueFault.hitCount()).toBe(1); + expect((await repository.getTask('legacy-action-requeue-task'))?.claimRecovery?.phase).toBe('draining'); + expect(await repository.getPendingSessionAction('legacy-action-requeue-task')).toMatchObject({ + id: 'legacy-action-requeue', + state: 'pending' + }); + + repository = await harness.restart(); + await taskService(repository, clock).reconcileClaims(clock.value); + expect(await repository.getTask('legacy-action-requeue-task')).toMatchObject({ status: 'submitted' }); + expect((await repository.getTask('legacy-action-requeue-task'))?.claimRecovery).toBeUndefined(); + expect(await repository.getPendingSessionAction('legacy-action-requeue-task')).toMatchObject({ state: 'pending' }); + + await repository.saveTask(baseTask({ + id: 'legacy-action-closing-task', + interaction: 'interactive', + status: 'closing', + workerId: 'legacy-worker', + leaseToken: 'legacy-token', + pendingActionId: 'legacy-action-close', + createdAt: '2025-01-01T00:00:01.000Z' + })); + await repository.enqueueSessionAction({ + id: 'legacy-action-close', + taskId: 'legacy-action-closing-task', + action: { type: 'navigate', url: 'https://example.com/close' }, + state: 'pending', + createdAt: '2025-01-01T00:00:01.000Z' + }); + await repository.takePendingSessionAction('legacy-action-closing-task'); + const finalizeFault = faultAfterBoundary(repository, 'legacy_action_finalize'); + await taskService(finalizeFault.repository, clock).reconcileClaims(clock.value); + expect(finalizeFault.hitCount()).toBe(1); + expect((await repository.getTask('legacy-action-closing-task'))?.claimRecovery?.phase).toBe('draining'); + expect(await repository.getSessionActionResult('legacy-action-close')).toMatchObject({ + result: { error: { code: 'session_closed' } } + }); + + repository = await harness.restart(); + await taskService(repository, clock).reconcileClaims(clock.value); + expect(await repository.getTask('legacy-action-closing-task')).toMatchObject({ status: 'completed' }); + expect((await repository.getTask('legacy-action-closing-task'))?.claimRecovery).toBeUndefined(); + expect(await repository.getPendingSessionAction('legacy-action-closing-task')).toBeUndefined(); + expect(await repository.getSessionActionResult('legacy-action-close')).toMatchObject({ + result: { error: { code: 'session_closed' } } + }); + } finally { + await harness.close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('allows exactly one of two different tasks to own the same profile', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: Date.now() }; + await repository.savePool({ id: 'profile-race-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ + id: 'profile-race-machine', + poolId: 'profile-race-pool', + tags: {}, + capacity: 2, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ id: 'shared-race-profile', userId: 'user-1' }); + await repository.saveTask(baseTask({ id: 'profile-race-task-a', profileId: 'shared-race-profile' })); + await repository.saveTask(baseTask({ + id: 'profile-race-task-b', + profileId: 'shared-race-profile', + createdAt: '2025-01-01T00:00:01.000Z' + })); + const race = profileAcquireRaceRepositories(repository, ['profile-race-task-a', 'profile-race-task-b']); + const firstService = taskService(race.repositories[0]!, clock); + const secondService = taskService(race.repositories[1]!, clock); + + const results = await Promise.allSettled([ + firstService.claim('profile-race-worker-a', 'profile-race-machine', clock.value), + secondService.claim('profile-race-worker-b', 'profile-race-machine', clock.value) + ]); + const winners = results.filter((result) => result.status === 'fulfilled'); + expect(winners).toHaveLength(1); + expect([...race.arrivals].sort()).toEqual(['profile-race-task-a', 'profile-race-task-b']); + const tasks = await repository.listTasks(); + const claimed = tasks.filter((task) => task.id.startsWith('profile-race-task-') && task.status === 'claimed'); + const requeued = tasks.filter((task) => task.id.startsWith('profile-race-task-') && task.status === 'submitted'); + expect(claimed).toHaveLength(1); + expect(requeued).toHaveLength(1); + const machine = await repository.getMachine('profile-race-machine'); + expect(machine).toMatchObject({ activeLeases: 1 }); + expect(machine?.leaseReservations).toEqual([expect.objectContaining({ + taskId: claimed[0]?.id, + claimId: claimed[0]?.claimId, + claimGeneration: claimed[0]?.claimGeneration + })]); + expect(requeued[0]).toMatchObject({ claimReleased: true, claimCommitted: false }); + expect(await repository.getProfile('shared-race-profile')).toMatchObject({ + lockedByTaskId: claimed[0]?.id, + lockedByClaimId: claimed[0]?.claimId, + lockedByClaimGeneration: claimed[0]?.claimGeneration + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('rotates machine tokens without replacing lease accounting', async () => { + const { repository, close } = await makeHarness(); + try { + const unrelatedReservation = { taskId: 'task-unrelated', claimId: 'claim-unrelated', claimGeneration: 7, expiresAt: '2026-09-04T12:01:00.000Z' }; + const targetReservation = { taskId: 'task-target', claimId: 'claim-target', claimGeneration: 1, expiresAt: '2026-09-04T12:02:00.000Z' }; + await repository.saveMachine({ + id: 'machine-accounting', + poolId: 'pool', + tags: { os: 'macos' }, + capacity: 2, + activeLeases: 1, + leaseReservations: [unrelatedReservation], + online: true, + workerTokenHash: 'old-hash' + }); + + const reserveRace = barrierBeforeRepositoryMethods(repository, ['rotateMachineToken', 'reserveMachineLease']); + expect(await Promise.all([ + reserveRace.rotateMachineToken('machine-accounting', 'old-hash', 'rotated-hash'), + reserveRace.reserveMachineLease('machine-accounting', targetReservation) + ])).toEqual([true, true]); + expect(await repository.rotateMachineToken('machine-accounting', 'old-hash', 'stale-hash')).toBe(false); + expect(await repository.getMachine('machine-accounting')).toEqual({ + id: 'machine-accounting', + poolId: 'pool', + tags: { os: 'macos' }, + capacity: 2, + online: true, + activeLeases: 2, + workerTokenHash: 'rotated-hash', + leaseReservations: [unrelatedReservation, targetReservation] + }); + + const releaseRace = barrierBeforeRepositoryMethods(repository, ['rotateMachineToken', 'releaseMachineLease']); + expect(await Promise.all([ + releaseRace.rotateMachineToken('machine-accounting', 'rotated-hash', 'final-hash'), + releaseRace.releaseMachineLease('machine-accounting', { + taskId: targetReservation.taskId, + claimId: targetReservation.claimId, + claimGeneration: targetReservation.claimGeneration + }) + ])).toEqual([true, true]); + expect(await repository.getMachine('machine-accounting')).toEqual({ + id: 'machine-accounting', + poolId: 'pool', + tags: { os: 'macos' }, + capacity: 2, + online: true, + activeLeases: 1, + workerTokenHash: 'final-hash', + leaseReservations: [unrelatedReservation] + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('refuses a lease response when a committed projection disappears before verification', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: Date.now() }; + await repository.savePool({ id: 'projection-loss-pool', visibility: 'platform', tags: {} }); + for (const projection of ['machine', 'profile'] as const) { + const taskId = `projection-loss-${projection}-task`; + const machineId = `projection-loss-${projection}-machine`; + const profileId = `projection-loss-${projection}-profile`; + await repository.saveMachine({ id: machineId, poolId: 'projection-loss-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + await repository.createProfile({ id: profileId, userId: 'user-1' }); + await repository.saveTask(baseTask({ id: taskId, profileId })); + let injections = 0; + const projectionLossRepository = new Proxy(repository, { + get(target, property) { + if (property === 'replaceTaskForClaim') { + return async (...args: Parameters): Promise => { + const committed = await target.replaceTaskForClaim(...args); + const task = args[0]; + if (committed && injections === 0 && task.claimCommitted === true && task.claimReleased !== true) { + injections += 1; + const reservation = { taskId: task.id, claimId: task.claimId!, claimGeneration: task.claimGeneration! }; + if (projection === 'machine') await target.releaseMachineLease(machineId, reservation); + else await target.releaseProfileLease(profileId, reservation); + } + return committed; + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + + await expect(taskService(projectionLossRepository, clock).claim(`projection-loss-${projection}-worker`, machineId, clock.value)) + .rejects.toMatchObject({ code: 'not_found' }); + expect(injections).toBe(1); + expect(await repository.getTask(taskId)).toMatchObject({ + status: 'submitted', + claimGeneration: 1, + claimCommitted: false, + claimReleased: true + }); + expect(await repository.getMachine(machineId)).toMatchObject({ activeLeases: 0, leaseReservations: [] }); + expect(await repository.getProfile(profileId)).not.toHaveProperty('lockedByTaskId'); + + const reclaimed = await taskService(repository, clock).claim(`projection-loss-${projection}-replacement`, machineId, clock.value); + expect(reclaimed.task).toMatchObject({ status: 'claimed', claimGeneration: 2, claimCommitted: true }); + expect(await repository.getMachine(machineId)).toMatchObject({ + activeLeases: 1, + leaseReservations: [{ + taskId, + claimId: reclaimed.task.claimId, + claimGeneration: reclaimed.task.claimGeneration, + expiresAt: reclaimed.task.leaseExpiresAt + }] + }); + expect(await repository.getProfile(profileId)).toMatchObject({ + lockedByTaskId: taskId, + lockedByClaimId: reclaimed.task.claimId, + lockedByClaimGeneration: reclaimed.task.claimGeneration + }); + } + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('keeps machine reservation mutations idempotent and monotonic', async () => { + const { repository, close } = await makeHarness(); + try { + await repository.saveMachine({ + id: 'machine-reservation-mutations', + poolId: 'pool', + tags: {}, + capacity: 2, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + const initial = { + taskId: 'task-a', + claimId: 'claim-a', + claimGeneration: 1, + expiresAt: '2026-09-04T12:01:00.000Z' + }; + + expect(await repository.reserveMachineLease('machine-reservation-mutations', initial)).toBe(true); + expect(await repository.reserveMachineLease('machine-reservation-mutations', initial)).toBe(true); + expect(await repository.reserveMachineLease('machine-reservation-mutations', { + ...initial, + taskId: 'task-with-conflicting-claim-id' + })).toBe(false); + expect(await repository.getMachine('machine-reservation-mutations')).toMatchObject({ + activeLeases: 1, + leaseReservations: [initial] + }); + + expect(await repository.renewMachineLease('machine-reservation-mutations', { + ...initial, + expiresAt: '2026-09-04T12:00:00.000Z' + })).toBe(true); + expect(await repository.getMachine('machine-reservation-mutations')).toMatchObject({ + leaseReservations: [initial] + }); + + const extended = { ...initial, expiresAt: '2026-09-04T12:02:00.000Z' }; + expect(await repository.renewMachineLease('machine-reservation-mutations', extended)).toBe(true); + expect(await repository.getMachine('machine-reservation-mutations')).toMatchObject({ + activeLeases: 1, + leaseReservations: [extended] + }); + + const identity = { taskId: initial.taskId, claimId: initial.claimId, claimGeneration: initial.claimGeneration }; + expect(await repository.releaseMachineLease('machine-reservation-mutations', identity)).toBe(true); + expect(await repository.releaseMachineLease('machine-reservation-mutations', identity)).toBe(false); + expect(await repository.getMachine('machine-reservation-mutations')).toMatchObject({ + activeLeases: 0, + leaseReservations: [] + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('does not let profile creation clear an active claim lock', async () => { + const { repository, close } = await makeHarness(); + try { + expect(await repository.createProfile({ id: 'profile-insert', userId: 'user-1' })).toBe(true); + const reservation = { taskId: 'task-a', claimId: 'claim-a', claimGeneration: 1, expiresAt: '2026-09-07T12:01:00.000Z' }; + expect(await repository.acquireProfileLease('profile-insert', 'user-1', 'machine-a', reservation)).toBeDefined(); + expect(await repository.createProfile({ id: 'profile-insert', userId: 'user-1', machineId: 'machine-b' })).toBe(false); + expect(await repository.getProfile('profile-insert')).toMatchObject({ + machineId: 'machine-a', + lockedByTaskId: 'task-a', + lockedByClaimId: 'claim-a' + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('atomically selects one winner for concurrent profile creation', async () => { + const { repository, close } = await makeHarness(); + try { + const candidates = [ + { id: 'profile-create-race', userId: 'user-a', machineId: 'machine-a' }, + { id: 'profile-create-race', userId: 'user-b', machineId: 'machine-b' } + ] as const; + const results = await Promise.all(candidates.map((profile) => repository.createProfile(profile))); + expect(results.filter(Boolean)).toHaveLength(1); + expect(await repository.getProfile('profile-create-race')).toEqual(candidates[results.indexOf(true)]); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('serializes profile takeover after authoritative task expiry', async () => { + const { repository, close } = await makeHarness(); + try { + const now = Date.now(); + await repository.createProfile({ id: 'profile-renewing', userId: 'user-1' }); + const submitted = baseTask({ id: 'task-renewing', profileId: 'profile-renewing' }); + await repository.saveTask(submitted); + const active = (await repository.claimTask({ + ...submitted, + status: 'claimed', + workerId: 'worker-active', + machineId: 'machine-a', + leaseToken: 'lease-active', + leaseExpiresAt: new Date(now + 60_000).toISOString(), + claimId: 'claim-active', + claimGeneration: 1, + taskVersion: 1, + claimCommitted: true, + claimedAt: new Date(now - 60_000).toISOString(), + updatedAt: new Date(now - 60_000).toISOString() + }, 0, 0))!; + const staleProjection = { + taskId: active.id, + claimId: active.claimId!, + claimGeneration: active.claimGeneration!, + expiresAt: new Date(now - 60_000).toISOString() + }; + const takeover = { taskId: 'task-new', claimId: 'claim-new', claimGeneration: 1, expiresAt: new Date(now + 120_000).toISOString() }; + expect(await repository.acquireProfileLease('profile-renewing', 'user-1', 'machine-a', staleProjection)).toBeDefined(); + + const renewedDeadline = new Date(now + 180_000).toISOString(); + const renewalCommitted = deferred(); + const releaseRenewal = deferred(); + const renewal = (async () => { + const result = await repository.replaceTaskForActiveClaim( + { ...active, status: 'running', leaseExpiresAt: renewedDeadline }, + { + claimId: active.claimId!, + claimGeneration: active.claimGeneration!, + taskVersion: active.taskVersion!, + status: active.status, + leaseExpiresAt: active.leaseExpiresAt! + } + ); + renewalCommitted.resolve(); + await releaseRenewal.promise; + return result; + })(); + + await renewalCommitted.promise; + try { + expect(await repository.acquireProfileLease('profile-renewing', 'user-1', 'machine-b', takeover)).toBeUndefined(); + const renewed = (await repository.getTask(active.id))!; + expect(await repository.replaceTaskForExpiredClaim({ + ...renewed, + status: 'submitted', + leaseExpiresAt: undefined, + leaseToken: undefined, + workerId: undefined, + queuePriority: -1 + }, { + claimId: renewed.claimId!, + claimGeneration: renewed.claimGeneration!, + taskVersion: renewed.taskVersion!, + status: renewed.status, + leaseExpiresAt: renewed.leaseExpiresAt! + })).toBe(false); + expect(await repository.getProfile('profile-renewing')).toMatchObject({ + lockedByTaskId: active.id, + lockedByClaimId: active.claimId + }); + } finally { + releaseRenewal.resolve(); + } + expect(await renewal).toBe(true); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('takes over a profile only after requeueing its expired task generation', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: Date.now() }; + await repository.savePool({ id: 'takeover-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'takeover-machine', poolId: 'takeover-pool', tags: {}, capacity: 2, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + await repository.createProfile({ id: 'takeover-profile', userId: 'user-1' }); + const previous = baseTask({ id: 'takeover-previous', profileId: 'takeover-profile', queuePriority: 7 }); + await repository.saveTask(previous); + const expired = (await repository.claimTask({ + ...previous, + status: 'claimed', + workerId: 'worker-previous', + machineId: 'takeover-machine', + leaseToken: 'lease-previous', + leaseExpiresAt: '2000-01-01T00:00:00.000Z', + claimId: 'claim-previous', + claimGeneration: 1, + taskVersion: 1, + claimCommitted: true, + claimQueuePriority: previous.queuePriority, + queuePriority: undefined, + claimedAt: '1999-12-31T23:59:00.000Z', + updatedAt: '1999-12-31T23:59:00.000Z' + }, 0, 0))!; + const expiredReservation = { + taskId: expired.id, + claimId: expired.claimId!, + claimGeneration: expired.claimGeneration!, + expiresAt: expired.leaseExpiresAt! + }; + expect(await repository.reserveMachineLease('takeover-machine', expiredReservation)).toBe(true); + expect(await repository.acquireProfileLease('takeover-profile', 'user-1', 'takeover-machine', expiredReservation)).toBeDefined(); + await repository.saveTask(baseTask({ id: 'takeover-next', profileId: 'takeover-profile', createdAt: new Date(clock.value).toISOString(), updatedAt: new Date(clock.value).toISOString() })); + + const claimed = await taskService(repository, clock).claim('worker-next', 'takeover-machine', clock.value); + expect(claimed.task.id).toBe('takeover-next'); + expect(await repository.getTask(expired.id)).toMatchObject({ + status: 'submitted', + claimReleased: true, + workerId: undefined, + leaseExpiresAt: undefined, + queuePriority: previous.queuePriority, + claimQueuePriority: undefined + }); + expect(await repository.getMachine('takeover-machine')).toMatchObject({ activeLeases: 1 }); + expect(await repository.getProfile('takeover-profile')).toMatchObject({ + lockedByTaskId: claimed.task.id, + lockedByClaimId: claimed.task.claimId + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('fences stale lifecycle mutations with the exact task version', async () => { + const { repository, close } = await makeHarness(); + try { + await repository.savePool({ id: 'version-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'version-machine', poolId: 'version-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + await repository.saveTask(baseTask({ id: 'version-task' })); + const claimed = (await taskService(repository).claim('worker-a', 'version-machine')).task; + const guard = { + claimId: claimed.claimId!, + claimGeneration: claimed.claimGeneration!, + taskVersion: claimed.taskVersion!, + status: claimed.status + }; + + expect(await repository.replaceTaskForClaim({ ...claimed, status: 'running' }, guard)).toBe(true); + expect(await repository.replaceTaskForClaim({ ...claimed, status: 'failed' }, guard)).toBe(false); + expect(await repository.getTask(claimed.id)).toMatchObject({ status: 'running', taskVersion: guard.taskVersion + 1 }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('rejects a delayed initial claimant after a later generation was requeued', async () => { + const { repository, close } = await makeHarness(); + try { + const submitted = baseTask({ id: 'delayed-claim-task' }); + await repository.saveTask(submitted); + const delayedClaim = { + ...submitted, + status: 'claimed' as const, + workerId: 'worker-delayed', + machineId: 'machine-delayed', + leaseToken: 'lease-delayed', + leaseExpiresAt: '2026-09-08T12:01:00.000Z', + claimId: 'claim-delayed', + claimGeneration: 1, + taskVersion: 1, + claimCommitted: false, + claimReleased: false, + claimedAt: '2026-09-08T12:00:00.000Z', + updatedAt: '2026-09-08T12:00:00.000Z' + }; + const winner = await repository.claimTask({ + ...delayedClaim, + workerId: 'worker-winner', + machineId: 'machine-winner', + leaseToken: 'lease-winner', + claimId: 'claim-winner' + }, 0, 0); + expect(winner).toBeDefined(); + expect(await repository.replaceTaskForClaim({ + ...winner!, + status: 'submitted', + workerId: undefined, + machineId: undefined, + leaseToken: undefined, + leaseExpiresAt: undefined, + claimCommitted: false, + claimReleased: true, + queuePriority: -1 + }, { + claimId: winner!.claimId!, + claimGeneration: winner!.claimGeneration!, + taskVersion: winner!.taskVersion!, + status: winner!.status + })).toBe(true); + + expect(await repository.claimTask(delayedClaim, 0, 0)).toBeUndefined(); + expect(await repository.getTask(submitted.id)).toMatchObject({ + status: 'submitted', + claimId: 'claim-winner', + claimGeneration: 1, + taskVersion: 2, + claimReleased: true + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('fences a delayed initial claimant after a complete generation N+1 reclaim', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: 1_000 }; + await repository.savePool({ id: 'delayed-reclaim-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ + id: 'delayed-reclaim-machine', + poolId: 'delayed-reclaim-pool', + tags: {}, + capacity: 2, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ id: 'delayed-reclaim-profile', userId: 'user-1' }); + await repository.saveTask(baseTask({ + id: 'delayed-reclaim-task', + profileId: 'delayed-reclaim-profile' + })); + const delayedAtCas = deferred(); + const resumeDelayed = deferred(); + const delayedRepository = new Proxy(repository, { + get(target, property) { + if (property === 'claimTask') { + return async (...args: Parameters): Promise>> => { + delayedAtCas.resolve(); + await resumeDelayed.promise; + return target.claimTask(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const delayed = taskService(delayedRepository, clock) + .claim('delayed-worker', 'delayed-reclaim-machine', clock.value); + await delayedAtCas.promise; + + const authority = taskService(repository, clock); + const generationN = await authority.claim('generation-n-worker', 'delayed-reclaim-machine', clock.value); + clock.value = 12_000; + expect(await authority.expireLeases(clock.value)).toHaveLength(1); + const generationNPlusOne = await authority.claim('generation-n-plus-one-worker', 'delayed-reclaim-machine', clock.value); + expect(generationNPlusOne.task.claimGeneration).toBe((generationN.task.claimGeneration ?? 0) + 1); + + resumeDelayed.resolve(); + await expect(delayed).rejects.toMatchObject({ code: 'not_found' }); + expect(await repository.getTask('delayed-reclaim-task')).toMatchObject({ + status: 'claimed', + workerId: 'generation-n-plus-one-worker', + claimId: generationNPlusOne.task.claimId, + claimGeneration: generationNPlusOne.task.claimGeneration + }); + expect(await repository.getMachine('delayed-reclaim-machine')).toMatchObject({ + activeLeases: 1, + leaseReservations: [expect.objectContaining({ + taskId: generationNPlusOne.task.id, + claimId: generationNPlusOne.task.claimId, + claimGeneration: generationNPlusOne.task.claimGeneration + })] + }); + expect(await repository.getProfile('delayed-reclaim-profile')).toMatchObject({ + lockedByTaskId: generationNPlusOne.task.id, + lockedByClaimId: generationNPlusOne.task.claimId, + lockedByClaimGeneration: generationNPlusOne.task.claimGeneration + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('rejects a heartbeat mutation that reaches the claim CAS after lease expiry', async () => { + const { repository, close } = await makeHarness(); + try { + const admittedAt = Date.parse('1999-12-31T23:59:59.999Z'); + const submitted = baseTask({ id: 'expired-heartbeat-task' }); + await repository.saveTask(submitted); + const claimed = await repository.claimTask({ + ...submitted, + status: 'claimed', + workerId: 'worker-a', + machineId: 'machine-a', + leaseToken: 'lease-a', + leaseExpiresAt: '2000-01-01T00:00:00.000Z', + claimId: 'claim-a', + claimGeneration: 1, + taskVersion: 1, + claimCommitted: true, + claimedAt: '1999-12-31T23:59:50.000Z', + updatedAt: '1999-12-31T23:59:50.000Z' + }, 0, 0); + expect(claimed).toBeDefined(); + + const heartbeatRead = deferred(); + const expiryBarrier = deferred(); + const renewal = (async () => { + const current = (await repository.getTask('expired-heartbeat-task'))!; + expect(Date.parse(current.leaseExpiresAt!)).toBeGreaterThan(admittedAt); + heartbeatRead.resolve(); + await expiryBarrier.promise; + return repository.replaceTaskForActiveClaim( + { ...current, status: 'running', leaseExpiresAt: '2000-01-01T00:01:00.000Z' }, + { + claimId: current.claimId!, + claimGeneration: current.claimGeneration!, + taskVersion: current.taskVersion!, + status: current.status, + leaseExpiresAt: current.leaseExpiresAt! + } + ); + })(); + + await heartbeatRead.promise; + expiryBarrier.resolve(); + expect(await renewal).toBe(false); + expect(await repository.getTask('expired-heartbeat-task')).toMatchObject({ + status: 'claimed', + leaseExpiresAt: '2000-01-01T00:00:00.000Z', + taskVersion: 1 + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('retries an authorized cancellation after heartbeat wins the first task CAS', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: Date.now() }; + await repository.savePool({ id: 'cancel-race-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'cancel-race-machine', poolId: 'cancel-race-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + const cancelCasReached = deferred(); + const resumeCancelCas = deferred(); + let cancelAttempts = 0; + let leaseReleases = 0; + const raced = new Proxy(repository, { + get(target, property) { + if (property === 'replaceTaskForClaim') { + return async (...args: Parameters): Promise => { + if (args[0].status === 'cancelled' && args[0].claimReleased !== true) { + cancelAttempts += 1; + if (cancelAttempts === 1) { + cancelCasReached.resolve(); + await resumeCancelCas.promise; + } + } + return target.replaceTaskForClaim(...args); + }; + } + if (property === 'releaseMachineLease') { + return async (...args: Parameters): Promise => { + leaseReleases += 1; + return target.releaseMachineLease(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const service = taskService(raced, clock); + const task = await service.createTask('cancel-race-user', { kind: 'browse', goal: 'cancel race' }); + const claim = await service.claim('cancel-race-worker', 'cancel-race-machine', clock.value); + + const cancellation = service.cancel(task.id, task.userId); + await cancelCasReached.promise; + await service.heartbeat(task.id, 'cancel-race-worker', claim.leaseToken, 30); + resumeCancelCas.resolve(); + + await expect(cancellation).resolves.toMatchObject({ status: 'cancelled', claimReleased: true }); + expect(cancelAttempts).toBe(2); + expect(leaseReleases).toBe(1); + expect(await repository.getMachine('cancel-race-machine')).toMatchObject({ + activeLeases: 0, + leaseReservations: [] + }); + expect((await repository.listWebhooks()).filter((event) => event.payload.status === 'cancelled')).toHaveLength(1); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('retries interactive close after heartbeat and preserves pending-action teardown', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: Date.now() }; + await repository.savePool({ id: 'close-race-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'close-race-machine', poolId: 'close-race-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + const closeCasReached = deferred(); + const resumeCloseCas = deferred(); + let closeAttempts = 0; + const raced = new Proxy(repository, { + get(target, property) { + if (property === 'replaceTaskForClaim') { + return async (...args: Parameters): Promise => { + if (args[0].status === 'closing') { + closeAttempts += 1; + if (closeAttempts === 1) { + closeCasReached.resolve(); + await resumeCloseCas.promise; + } + } + return target.replaceTaskForClaim(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const tasks = taskService(raced, clock); + const sessions = new SessionService(tasks, raced, { clock: () => clock.value }); + const session = await sessions.create('close-race-user', { mode: 'act', constraints: {} }); + const claim = await tasks.claim('close-race-worker', 'close-race-machine', clock.value); + const action = await sessions.sendAction(session.id, 'close-race-user', { + type: 'screenshot', + format: 'png', + quality: 80 + }, 0); + + const closing = sessions.close(session.id, 'close-race-user'); + await closeCasReached.promise; + await tasks.heartbeat(session.id, 'close-race-worker', claim.leaseToken, 30); + resumeCloseCas.resolve(); + + await expect(closing).resolves.toMatchObject({ status: 'closing', lastActionId: action.action_id }); + expect(closeAttempts).toBe(2); + await expect(sessions.getAction(session.id, action.action_id, 'close-race-user', 0)).resolves.toMatchObject({ + status: 'completed', + result: { error: { code: 'session_closed' } } + }); + await expect(sessions.pollWorkerAction(session.id, 'close-race-worker', claim.leaseToken)) + .resolves.toEqual({ closing: true }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('reconciles interrupted projections and fences a requeued generation', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: 1_000 }; + await repository.savePool({ id: 'reconcile-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ id: 'reconcile-machine', poolId: 'reconcile-pool', tags: {}, capacity: 1, activeLeases: 0, online: true, workerTokenHash: 'hash' }); + await repository.createProfile({ id: 'reconcile-profile', userId: 'user-1' }); + const submitted = baseTask({ id: 'reconcile-task', profileId: 'reconcile-profile' }); + await repository.saveTask(submitted); + const interrupted = await repository.claimTask({ + ...submitted, + status: 'claimed', + workerId: 'worker-a', + machineId: 'reconcile-machine', + leaseToken: 'lease-a', + leaseExpiresAt: '1970-01-01T00:00:11.000Z', + claimId: 'claim-a', + claimGeneration: 1, + taskVersion: 1, + claimCommitted: false, + claimedAt: '1970-01-01T00:00:01.000Z', + updatedAt: '1970-01-01T00:00:01.000Z' + }, 0, 0); + expect(interrupted).toBeDefined(); + const service = taskService(repository, clock); + await service.reconcileClaims(); + expect(await repository.getMachine('reconcile-machine')).toMatchObject({ activeLeases: 1 }); + expect(await repository.getProfile('reconcile-profile')).toMatchObject({ lockedByClaimId: 'claim-a' }); + + const claimed = (await repository.getTask('reconcile-task'))!; + expect(await repository.replaceTaskForClaim({ + ...claimed, + status: 'submitted', + workerId: undefined, + machineId: undefined, + leaseToken: undefined, + leaseExpiresAt: undefined, + queuePriority: -1 + }, { claimId: 'claim-a', claimGeneration: 1, taskVersion: claimed.taskVersion ?? 0, status: 'claimed' })).toBe(true); + + const restarted = taskService(repository, clock); + await restarted.reconcileClaims(); + expect(await repository.getMachine('reconcile-machine')).toMatchObject({ activeLeases: 0, leaseReservations: [] }); + expect(await repository.getProfile('reconcile-profile')).not.toHaveProperty('lockedByTaskId'); + + clock.value = 12_000; + const reclaimed = await restarted.claim('worker-b', 'reconcile-machine', 12_000); + expect(reclaimed.task.claimGeneration).toBe(2); + await expect(restarted.heartbeat('reconcile-task', 'worker-a', 'lease-a', 10)).rejects.toMatchObject({ code: 'unauthorized' }); + expect(await repository.getMachine('reconcile-machine')).toMatchObject({ activeLeases: 1 }); + expect(await repository.getProfile('reconcile-profile')).toMatchObject({ lockedByClaimId: reclaimed.task.claimId }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('does not let a stale generation reconciler remove generation N+1 projections', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: Date.now() - 30_000 }; + await repository.savePool({ id: 'stale-reconcile-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ + id: 'stale-reconcile-machine', + poolId: 'stale-reconcile-pool', + tags: {}, + capacity: 2, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ id: 'stale-reconcile-profile', userId: 'user-1' }); + await repository.saveTask(baseTask({ id: 'stale-reconcile-task', profileId: 'stale-reconcile-profile' })); + const authority = taskService(repository, clock); + const generationN = await authority.claim('stale-reconcile-worker-n', 'stale-reconcile-machine', clock.value); + + const staleAtReservation = deferred(); + const resumeStale = deferred(); + const staleReservationPersisted = deferred(); + const returnStaleReservation = deferred(); + let paused = false; + let staleReservationResult: boolean | undefined; + let staleReservationError: unknown; + const staleRepository = new Proxy(repository, { + get(target, property) { + if (property === 'reserveMachineLease') { + return async (...args: Parameters): Promise => { + const reservation = args[1]; + if (!paused && reservation.claimId === generationN.task.claimId) { + paused = true; + staleAtReservation.resolve(); + await resumeStale.promise; + try { + staleReservationResult = await target.reserveMachineLease(...args); + } catch (error) { + staleReservationError = error; + } finally { + staleReservationPersisted.resolve(); + } + await returnStaleReservation.promise; + if (staleReservationError !== undefined) throw staleReservationError; + return staleReservationResult ?? false; + } + return target.reserveMachineLease(...args); + }; + } + const value = Reflect.get(target, property); + return typeof value === 'function' ? value.bind(target) : value; + } + }); + const staleReconciliation = taskService(staleRepository, { value: Date.now() }).reconcileClaims(Date.now()); + await staleAtReservation.promise; + + clock.value = Date.now(); + expect(await authority.expireLeases(clock.value)).toHaveLength(1); + clock.value = Date.now() + 600_000; + const generationNPlusOne = await authority.claim('stale-reconcile-worker-n-plus-one', 'stale-reconcile-machine', clock.value); + expect(generationNPlusOne.task.claimGeneration).toBe((generationN.task.claimGeneration ?? 0) + 1); + const cursorBeforeStaleResume = await repository.getTaskMaintenanceCursor('task-claim-reconciliation'); + expect(cursorBeforeStaleResume).toMatchObject({ id: 'task-claim-reconciliation' }); + if (cursorBeforeStaleResume === undefined) throw new Error('authoritative maintenance cursor missing'); + expect(cursorBeforeStaleResume.version).toBeGreaterThan(0); + expect(await repository.getMachine('stale-reconcile-machine')).toMatchObject({ + activeLeases: 1, + leaseReservations: [{ + taskId: 'stale-reconcile-task', + claimId: generationNPlusOne.task.claimId, + claimGeneration: generationNPlusOne.task.claimGeneration, + expiresAt: generationNPlusOne.task.leaseExpiresAt + }] + }); + + resumeStale.resolve(); + await staleReservationPersisted.promise; + let temporaryMachine: Awaited>; + let observationError: unknown; + try { + temporaryMachine = await repository.getMachine('stale-reconcile-machine'); + } catch (error) { + observationError = error; + } finally { + returnStaleReservation.resolve(); + } + await staleReconciliation; + if (observationError !== undefined) throw observationError; + expect(staleReservationError).toBeUndefined(); + expect(staleReservationResult).toBe(true); + expect(temporaryMachine).toMatchObject({ activeLeases: 2 }); + expect(temporaryMachine?.leaseReservations).toHaveLength(2); + expect(temporaryMachine?.leaseReservations).toEqual(expect.arrayContaining([ + { + taskId: 'stale-reconcile-task', + claimId: generationN.task.claimId, + claimGeneration: generationN.task.claimGeneration, + expiresAt: generationN.task.leaseExpiresAt + }, + { + taskId: 'stale-reconcile-task', + claimId: generationNPlusOne.task.claimId, + claimGeneration: generationNPlusOne.task.claimGeneration, + expiresAt: generationNPlusOne.task.leaseExpiresAt + } + ])); + + expect(await repository.getTask('stale-reconcile-task')).toMatchObject({ + status: 'claimed', + claimId: generationNPlusOne.task.claimId, + claimGeneration: generationNPlusOne.task.claimGeneration, + workerId: 'stale-reconcile-worker-n-plus-one' + }); + expect(await repository.getMachine('stale-reconcile-machine')).toMatchObject({ + activeLeases: 1, + leaseReservations: [{ + taskId: 'stale-reconcile-task', + claimId: generationNPlusOne.task.claimId, + claimGeneration: generationNPlusOne.task.claimGeneration, + expiresAt: generationNPlusOne.task.leaseExpiresAt + }] + }); + expect(await repository.getProfile('stale-reconcile-profile')).toMatchObject({ + lockedByTaskId: 'stale-reconcile-task', + lockedByClaimId: generationNPlusOne.task.claimId, + lockedByClaimGeneration: generationNPlusOne.task.claimGeneration + }); + expect(await repository.getTaskMaintenanceCursor('task-claim-reconciliation')) + .toEqual(cursorBeforeStaleResume); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('fences generation N action operations after N+1 reclaims the original action', async () => { + const { repository, close } = await makeHarness(); + try { + const clock = { value: Date.now() - 30_000 }; + await repository.savePool({ id: 'action-reclaim-pool', visibility: 'platform', tags: {} }); + await repository.saveMachine({ + id: 'action-reclaim-machine', + poolId: 'action-reclaim-pool', + tags: {}, + capacity: 1, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ id: 'action-reclaim-profile', userId: 'user-1' }); + const tasks = taskService(repository, clock); + const sessions = new SessionService(tasks, repository, { clock: () => clock.value }); + const session = await sessions.create('user-1', { + profile_id: 'action-reclaim-profile', + mode: 'act', + constraints: {} + }); + const generationN = await tasks.claim('action-reclaim-worker-n', 'action-reclaim-machine', clock.value); + const pending = await sessions.sendAction(session.id, 'user-1', { type: 'wait', milliseconds: 1 }, 0); + expect((await sessions.pollWorkerAction(session.id, 'action-reclaim-worker-n', generationN.leaseToken)).action?.id) + .toBe(pending.action_id); + + clock.value = Date.now(); + expect(await tasks.expireLeases(clock.value)).toHaveLength(1); + const generationNPlusOne = await tasks.claim('action-reclaim-worker-n-plus-one', 'action-reclaim-machine', clock.value); + expect(generationNPlusOne.task.claimGeneration).toBe((generationN.task.claimGeneration ?? 0) + 1); + + await expect(sessions.pollWorkerAction(session.id, 'action-reclaim-worker-n', generationN.leaseToken)) + .rejects.toMatchObject({ code: 'unauthorized' }); + await expect(sessions.saveWorkerResult( + session.id, + pending.action_id, + 'action-reclaim-worker-n', + generationN.leaseToken, + { value: 'stale' }, + 'action-reclaim-machine' + )).rejects.toMatchObject({ code: 'unauthorized' }); + expect(await repository.getSessionActionResult(pending.action_id)).toBeUndefined(); + expect(await repository.getPendingSessionAction(session.id)).toMatchObject({ + id: pending.action_id, + state: 'pending' + }); + + expect((await sessions.pollWorkerAction(session.id, 'action-reclaim-worker-n-plus-one', generationNPlusOne.leaseToken)).action?.id) + .toBe(pending.action_id); + await sessions.saveWorkerResult( + session.id, + pending.action_id, + 'action-reclaim-worker-n-plus-one', + generationNPlusOne.leaseToken, + { value: 'generation-n-plus-one' }, + 'action-reclaim-machine' + ); + await expect(sessions.getAction(session.id, pending.action_id, 'user-1', 0)).resolves.toEqual({ + action_id: pending.action_id, + status: 'completed', + result: { value: 'generation-n-plus-one' } + }); + expect(await repository.getMachine('action-reclaim-machine')).toMatchObject({ + activeLeases: 1, + leaseReservations: [expect.objectContaining({ + taskId: session.id, + claimId: generationNPlusOne.task.claimId, + claimGeneration: generationNPlusOne.task.claimGeneration + })] + }); + expect(await repository.getProfile('action-reclaim-profile')).toMatchObject({ + lockedByTaskId: session.id, + lockedByClaimId: generationNPlusOne.task.claimId, + lockedByClaimGeneration: generationNPlusOne.task.claimGeneration + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('recovers legacy machine and profile accounting without changing modern reservations', async () => { + const { repository, close } = await makeHarness(); + try { + const modernReservation = { + taskId: 'modern-task', + claimId: 'modern-claim', + claimGeneration: 3, + expiresAt: '2026-09-08T12:10:00.000Z' + }; + await repository.saveMachine({ + id: 'legacy-machine', + poolId: 'pool', + tags: {}, + capacity: 2, + activeLeases: 2, + leaseReservations: [modernReservation], + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ + id: 'legacy-profile', + userId: 'user-1', + machineId: 'legacy-machine', + lockedByTaskId: 'legacy-task', + lockExpiresAt: '2026-09-08T12:10:00.000Z' + }); + await repository.saveTask(baseTask({ + id: 'legacy-task', + status: 'running', + machineId: 'legacy-machine', + profileId: 'legacy-profile', + workerId: 'legacy-worker', + leaseToken: 'legacy-token', + leaseExpiresAt: '2026-09-08T12:10:00.000Z', + queuePriority: 7 + })); + + const service = taskService(repository); + await service.reconcileClaims(); + expect(await repository.getTask('legacy-task')).toMatchObject({ + status: 'submitted', + queuePriority: 7 + }); + expect((await repository.getTask('legacy-task'))?.claimRecovery).toBeUndefined(); + expect((await repository.getTask('legacy-task'))?.machineId).toBeUndefined(); + expect(await repository.getMachine('legacy-machine')).toMatchObject({ + activeLeases: 1, + leaseReservations: [modernReservation] + }); + expect(await repository.getProfile('legacy-profile')).toMatchObject({ + machineId: 'legacy-machine' + }); + expect((await repository.getProfile('legacy-profile'))?.lockedByTaskId).toBeUndefined(); + + await service.reconcileClaims(); + expect(await repository.getMachine('legacy-machine')).toMatchObject({ + activeLeases: 1, + leaseReservations: [modernReservation] + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('requeues legacy interactive actions but terminalizes legacy closing actions', async () => { + const { repository, close } = await makeHarness(); + try { + await repository.saveTask(baseTask({ + id: 'legacy-interactive', + interaction: 'interactive', + status: 'running', + workerId: 'legacy-worker', + leaseToken: 'legacy-token', + pendingActionId: 'action-requeue', + createdAt: '2025-01-01T00:00:00.000Z' + })); + await repository.saveTask(baseTask({ + id: 'legacy-closing', + interaction: 'interactive', + status: 'closing', + workerId: 'legacy-worker', + leaseToken: 'legacy-token', + pendingActionId: 'action-close', + createdAt: '2025-01-01T00:00:01.000Z' + })); + await repository.enqueueSessionAction({ + id: 'action-requeue', + taskId: 'legacy-interactive', + action: { type: 'navigate', url: 'https://example.com/requeue' }, + state: 'pending', + createdAt: '2025-01-01T00:00:00.000Z' + }); + await repository.enqueueSessionAction({ + id: 'action-close', + taskId: 'legacy-closing', + action: { type: 'navigate', url: 'https://example.com/close' }, + state: 'pending', + createdAt: '2025-01-01T00:00:01.000Z' + }); + await repository.takePendingSessionAction('legacy-interactive'); + await repository.takePendingSessionAction('legacy-closing'); + + await taskService(repository).reconcileClaims(); + expect(await repository.getTask('legacy-interactive')).toMatchObject({ status: 'submitted' }); + expect(await repository.getPendingSessionAction('legacy-interactive')).toMatchObject({ + id: 'action-requeue', + state: 'pending' + }); + expect(await repository.getTask('legacy-closing')).toMatchObject({ status: 'completed' }); + expect((await repository.getTask('legacy-closing'))?.pendingActionId).toBeUndefined(); + expect(await repository.getPendingSessionAction('legacy-closing')).toBeUndefined(); + expect(await repository.getSessionActionResult('action-close')).toMatchObject({ + taskId: 'legacy-closing', + result: { error: { code: 'session_closed' } } + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('quarantines malformed claims without preventing later legacy recovery', async () => { + const { repository, close } = await makeHarness(); + try { + await repository.saveTask(baseTask({ + id: 'malformed-claim', + status: 'claimed', + claimId: 'malformed-id', + claimGeneration: 0, + machineId: 'unknown-machine', + workerId: 'malformed-worker', + leaseToken: 'malformed-token', + leaseExpiresAt: '1970-01-01T00:00:00.000Z', + createdAt: '2025-01-01T00:00:00.000Z' + })); + await repository.saveTask(baseTask({ + id: 'invalid-expiry', + status: 'claimed', + claimId: 'invalid-expiry-claim', + claimGeneration: 1, + machineId: 'unknown-machine', + workerId: 'invalid-expiry-worker', + leaseToken: 'invalid-expiry-token', + leaseExpiresAt: 'not-a-timestamp', + createdAt: '2025-01-01T00:00:00.500Z' + })); + await repository.saveTask(baseTask({ + id: 'invalid-credentials', + status: 'claimed', + claimId: '', + claimGeneration: 1, + machineId: 'unknown-machine', + workerId: 'invalid-credentials-worker', + leaseToken: 'invalid-credentials-token', + leaseExpiresAt: '2026-09-08T12:10:00.000Z', + createdAt: '2025-01-01T00:00:00.750Z' + })); + await repository.saveTask(baseTask({ + id: 'later-legacy', + status: 'running', + workerId: 'legacy-worker', + leaseToken: 'legacy-token', + createdAt: '2025-01-01T00:00:01.000Z' + })); + + await taskService(repository).reconcileClaims(); + expect(await repository.getTask('malformed-claim')).toMatchObject({ + status: 'failed', + error: { code: 'claim_state_malformed' }, + claimRecovery: { + kind: 'malformed', + phase: 'quarantined', + reasonCode: 'invalid_claim_generation' + } + }); + expect(await repository.getTask('invalid-expiry')).toMatchObject({ + status: 'failed', + claimRecovery: { reasonCode: 'invalid_lease_expiry', phase: 'quarantined' } + }); + expect(await repository.getTask('invalid-credentials')).toMatchObject({ + status: 'failed', + claimRecovery: { reasonCode: 'invalid_claim_credentials', phase: 'quarantined' } + }); + expect(await repository.getTask('later-legacy')).toMatchObject({ status: 'submitted' }); + expect((await repository.getTask('later-legacy'))?.claimRecovery).toBeUndefined(); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('drains exact projections even when a malformed active claim says it was released', async () => { + const { repository, close } = await makeHarness(); + try { + const reservation = { + taskId: 'false-released-task', + claimId: 'false-released-claim', + claimGeneration: 4, + expiresAt: '2026-09-08T12:10:00.000Z' + }; + await repository.saveMachine({ + id: 'false-released-machine', + poolId: 'pool', + tags: {}, + capacity: 1, + activeLeases: 1, + leaseReservations: [reservation], + online: true, + workerTokenHash: 'hash' + }); + await repository.createProfile({ + id: 'false-released-profile', + userId: 'user-1', + machineId: 'false-released-machine', + lockedByTaskId: reservation.taskId, + lockedByClaimId: reservation.claimId, + lockedByClaimGeneration: reservation.claimGeneration, + lockExpiresAt: reservation.expiresAt + }); + await repository.saveTask(baseTask({ + id: reservation.taskId, + status: 'claimed', + profileId: 'false-released-profile', + machineId: 'false-released-machine', + workerId: 'false-released-worker', + leaseToken: 'false-released-token', + leaseExpiresAt: reservation.expiresAt, + claimId: reservation.claimId, + claimGeneration: reservation.claimGeneration, + claimCommitted: true, + claimReleased: true + })); + + const service = taskService(repository); + await service.reconcileClaims(); + expect(await repository.getTask(reservation.taskId)).toMatchObject({ + status: 'failed', + claimRecovery: { + kind: 'malformed', + phase: 'quarantined', + reasonCode: 'active_claim_marked_released' + } + }); + expect(await repository.getMachine('false-released-machine')).toMatchObject({ + activeLeases: 0, + leaseReservations: [] + }); + expect((await repository.getProfile('false-released-profile'))?.lockedByTaskId).toBeUndefined(); + + await service.reconcileClaims(); + expect(await repository.getMachine('false-released-machine')).toMatchObject({ + activeLeases: 0, + leaseReservations: [] + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('advances durable maintenance past one full batch', async () => { + const { repository, close } = await makeHarness(); + try { + for (let index = 0; index < 100; index += 1) { + await repository.saveTask(baseTask({ + id: `healthy-${String(index).padStart(3, '0')}`, + createdAt: '2025-01-01T00:00:00.000Z' + })); + } + await repository.saveTask(baseTask({ + id: 'legacy-after-full-batch', + status: 'running', + workerId: 'legacy-worker', + leaseToken: 'legacy-token', + createdAt: '2025-01-01T00:00:01.000Z' + })); + const service = taskService(repository); + + await service.reconcileClaims(); + expect(await repository.getTask('legacy-after-full-batch')).toMatchObject({ status: 'running' }); + await service.reconcileClaims(); + expect(await repository.getTask('legacy-after-full-batch')).toMatchObject({ status: 'submitted' }); + expect((await repository.getTask('legacy-after-full-batch'))?.claimRecovery).toBeUndefined(); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('reconciles a reservation when the task machine pointer is stale', async () => { + const { repository, close } = await makeHarness(); + try { + await repository.saveMachine({ + id: 'reservation-machine', + poolId: 'pool', + tags: {}, + capacity: 1, + activeLeases: 0, + online: true, + workerTokenHash: 'hash' + }); + const submitted = baseTask({ id: 'stale-machine-pointer-task' }); + await repository.saveTask(submitted); + const claimed = (await repository.claimTask({ + ...submitted, + status: 'claimed', + workerId: 'worker-a', + machineId: 'stale-machine-pointer', + leaseToken: 'lease-a', + leaseExpiresAt: '1970-01-01T00:00:11.000Z', + claimId: 'claim-a', + claimGeneration: 1, + taskVersion: 1, + claimCommitted: true, + claimedAt: '1970-01-01T00:00:01.000Z', + updatedAt: '1970-01-01T00:00:01.000Z' + }, 0, 0))!; + const reservation = { + taskId: claimed.id, + claimId: claimed.claimId!, + claimGeneration: claimed.claimGeneration!, + expiresAt: claimed.leaseExpiresAt! + }; + expect(await repository.reserveMachineLease('reservation-machine', reservation)).toBe(true); + expect(await repository.replaceTaskForClaim({ + ...claimed, + status: 'submitted', + workerId: undefined, + leaseToken: undefined, + leaseExpiresAt: undefined, + queuePriority: -1 + }, { + claimId: claimed.claimId!, + claimGeneration: claimed.claimGeneration!, + taskVersion: claimed.taskVersion!, + status: claimed.status + })).toBe(true); + + const service = taskService(repository); + await service.reconcileClaims(); + expect(await repository.getMachine('reservation-machine')).toMatchObject({ + activeLeases: 0, + leaseReservations: [] + }); + expect(await repository.getTask(claimed.id)).toMatchObject({ claimReleased: true }); + + await service.reconcileClaims(); + expect(await repository.getMachine('reservation-machine')).toMatchObject({ + activeLeases: 0, + leaseReservations: [] + }); + } finally { + await close(); + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + it('round-trips registry entities and task state', async () => { const { repository, close } = await makeHarness(); try { await repository.savePool({ id: 'pool-1', visibility: 'org', ownerUserId: 'user-1', sharedWithGroups: ['eng'], tags: { os: 'linux' } }); await repository.saveMachine({ id: 'machine-1', poolId: 'pool-1', tags: { browser: true }, capacity: 2, activeLeases: 0, online: true, workerTokenHash: 'hash' }); - await repository.saveProfile({ id: 'profile-1', userId: 'user-1', machineId: 'machine-1' }); + await repository.createProfile({ id: 'profile-1', userId: 'user-1', machineId: 'machine-1' }); const task = baseTask({ profileId: 'profile-1', poolId: 'pool-1' }); await repository.saveTask(task); await repository.saveHandoff({ id: 'handoff-1', taskId: task.id, userId: task.userId, url: '/v1/handoffs/handoff-1', expiresAt: '2025-01-01T00:10:00.000Z', used: false }); @@ -495,4 +2574,126 @@ describe('Repository contract: mongo', () => { contractTests(mongoHarness); undefinedLeaseTest(mongoHarness); testingRunContractTest(mongoHarness); + + it('continues the durable maintenance cursor after a real repository reconnect', async () => { + if (mongoUrl === undefined) throw new Error('Mongo contract setup did not provide a database URL'); + const databaseName = `talos_restart_${Date.now()}_${Math.random().toString(16).slice(2)}`; + const firstClient = new MongoClient(mongoUrl, mongodbClientOptions); + const first = new MongoRepository(mongoUrl, databaseName, { client: firstClient }); + let firstClosed = false; + let secondClient: MongoClient | undefined; + let second: MongoRepository | undefined; + try { + await first.initialize(); + for (let index = 0; index < 100; index += 1) { + await first.saveTask(baseTask({ + id: `restart-healthy-${String(index).padStart(3, '0')}`, + createdAt: '2025-01-01T00:00:00.000Z' + })); + } + await first.saveTask(baseTask({ + id: 'restart-legacy-after-batch', + status: 'running', + workerId: 'legacy-worker', + leaseToken: 'legacy-token', + createdAt: '2025-01-01T00:00:01.000Z' + })); + await taskService(first).reconcileClaims(); + expect(await first.getTask('restart-legacy-after-batch')).toMatchObject({ status: 'running' }); + + await first.close(); + firstClosed = true; + secondClient = new MongoClient(mongoUrl, mongodbClientOptions); + second = new MongoRepository(mongoUrl, databaseName, { client: secondClient }); + await second.initialize(); + await taskService(second).reconcileClaims(); + + expect(await second.getTask('restart-legacy-after-batch')).toMatchObject({ status: 'submitted' }); + expect((await second.getTask('restart-legacy-after-batch'))?.claimRecovery).toBeUndefined(); + } finally { + if (!firstClosed) await first.close(); + if (secondClient !== undefined) { + try { + await secondClient.db(databaseName).dropDatabase(); + } finally { + await second?.close(); + } + } + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); + + it('bounds expiry query work to the requested pages instead of the expired backlog', async () => { + if (mongoUrl === undefined) throw new Error('Mongo contract setup did not provide a database URL'); + const databaseName = `talos_expiry_plan_${Date.now()}_${Math.random().toString(16).slice(2)}`; + const client = new MongoClient(mongoUrl, { ...mongodbClientOptions, monitorCommands: true }); + const repository = new MongoRepository(mongoUrl, databaseName, { client }); + const findCommands: MongoDocument[] = []; + client.on('commandStarted', (event) => { + if (event.commandName === 'find' && event.command.find === 'tasks') { + findCommands.push({ + find: event.command.find, + filter: event.command.filter, + sort: event.command.sort, + limit: event.command.limit + }); + } + }); + try { + await repository.initialize(); + const expiryBase = Date.parse('2025-01-01T00:00:00.000Z'); + const documents = Array.from({ length: EXPIRY_QUERY_BACKLOG_SIZE }, (_, index) => { + const submitted = index % 2 === 0; + const id = `expiry-plan-${String(index).padStart(4, '0')}`; + const expiredAt = new Date(expiryBase + (EXPIRY_QUERY_BACKLOG_SIZE - index) * 1_000).toISOString(); + return { + ...baseTask({ + id, + kind: index % 4 < 2 ? 'browse' : 'computer_use', + status: submitted ? 'submitted' : (['claimed', 'running', 'closing'] as const)[index % 3], + constraints: submitted ? { deadline: expiredAt } : {}, + ...(submitted ? {} : { leaseExpiresAt: expiredAt }) + }), + _id: id + }; + }); + await client.db(databaseName).collection<{ _id: string; [key: string]: unknown }>('tasks').insertMany(documents); + + expect(await repository.listExpirableTasks(Date.parse('2026-01-01T00:00:00.000Z'), EXPIRY_QUERY_BATCH_SIZE)) + .toHaveLength(EXPIRY_QUERY_BATCH_SIZE); + expect(findCommands).toHaveLength(2); + + const expectedQueries = [ + { sortField: 'constraints.deadline', indexName: 'task-deadline-expiry-v1' }, + { sortField: 'leaseExpiresAt', indexName: 'task-lease-expiry-v1' } + ] as const; + for (const expected of expectedQueries) { + const command = findCommands.find((candidate) => { + const filter = candidate.filter; + return typeof filter === 'object' && filter !== null && Object.hasOwn(filter, expected.sortField); + }); + if (command === undefined) throw new Error(`missing captured ${expected.sortField} expiry query`); + expect(command).toMatchObject({ find: 'tasks', limit: EXPIRY_QUERY_BATCH_SIZE }); + expect(normalizedMongoSort(command.sort)).toEqual({ [expected.sortField]: 1, _id: 1 }); + const explanation = await client.db(databaseName).command({ explain: command, verbosity: 'executionStats' }); + const executionStats = explanation.executionStats as { + nReturned: number; + totalDocsExamined: number; + totalKeysExamined: number; + executionStages: unknown; + }; + expect(executionStats.nReturned).toBe(EXPIRY_QUERY_BATCH_SIZE); + expect(executionStats.totalDocsExamined).toBeLessThanOrEqual(EXPIRY_QUERY_BATCH_SIZE * 8); + expect(executionStats.totalKeysExamined).toBeLessThanOrEqual(EXPIRY_QUERY_BATCH_SIZE * 8); + expect(executionPlanContainsStage(executionStats.executionStages, 'SORT')).toBe(false); + expect(executionPlanContainsStage(executionStats.executionStages, 'COLLSCAN')).toBe(false); + expect(executionPlanContainsValue(executionStats.executionStages, expected.indexName)).toBe(true); + } + } finally { + try { + await client.db(databaseName).dropDatabase(); + } finally { + await repository.close(); + } + } + }, MONGODB_CONTRACT_TEST_TIMEOUT_MS); }); diff --git a/control-plane/src/storage/repository.ts b/control-plane/src/storage/repository.ts index 00cabb9..b39470c 100644 --- a/control-plane/src/storage/repository.ts +++ b/control-plane/src/storage/repository.ts @@ -1,4 +1,4 @@ -import type { HandoffLink, Machine, PendingSessionAction, Pool, Profile, SessionActionResult, Task, TaskInput, WebhookEvent } from '../domain/types.js'; +import type { HandoffLink, Machine, MachineLeaseReservation, PendingSessionAction, Pool, Profile, SessionActionResult, Task, TaskActiveClaimGuard, TaskClaimGuard, TaskInput, TaskRecoveryGuard, WebhookEvent } from '../domain/types.js'; import type { TestingAttemptStatus, TestingMachineReservationRecord, TestingRunRecord } from '../domain/testing-types.js'; export interface TestingAttemptMutationGuard { @@ -16,21 +16,53 @@ export interface TestingAttemptDispatchGuard extends TestingAttemptMutationGuard readonly dispatchAuthorizationExpiresAt: string; } +export interface TaskMaintenanceCursor { + readonly id: 'task-claim-reconciliation'; + readonly version: number; + readonly cycleCutoffAt: string; + readonly afterCreatedAt?: string; + readonly afterTaskId?: string; + readonly updatedAt: string; +} + export interface Repository { ping(): Promise; close(): Promise; getTask(id: string): Promise; saveTask(task: Task): Promise; + claimTask(task: Task, expectedClaimGeneration: number, expectedTaskVersion: number): Promise; + replaceTaskForClaim(task: Task, guard: TaskClaimGuard): Promise; + replaceTaskForActiveClaim(task: Task, guard: TaskActiveClaimGuard): Promise; + replaceTaskForExpiredClaim(task: Task, guard: TaskActiveClaimGuard): Promise; + replaceSubmittedTask(task: Task, expectedClaimGeneration: number, expectedTaskVersion: number): Promise; + replaceTaskForRecovery(task: Task, guard: TaskRecoveryGuard): Promise; + getTaskMaintenanceHighWatermark(): Promise; + listTaskMaintenancePage(cursor: TaskMaintenanceCursor, limit: number): Promise; + createTaskMaintenanceCursor(cursor: TaskMaintenanceCursor): Promise; + getTaskMaintenanceCursor(id: TaskMaintenanceCursor['id']): Promise; + replaceTaskMaintenanceCursor(cursor: TaskMaintenanceCursor, expectedVersion: number): Promise; listQueuedTasks(): Promise; listTasks(): Promise; + listExpirableTasks(now: number, limit: number): Promise; getPool(id: string): Promise; savePool(pool: Pool): Promise; listPoolsByOwner(ownerUserId: string): Promise; listMachines(poolId?: string): Promise; getMachine(id: string): Promise; saveMachine(machine: Machine): Promise; + rotateMachineToken(machineId: string, expectedTokenHash: string, tokenHash: string): Promise; + reserveMachineLease(machineId: string, reservation: MachineLeaseReservation): Promise; + renewMachineLease(machineId: string, reservation: MachineLeaseReservation): Promise; + releaseMachineLease(machineId: string, reservation: Omit): Promise; + releaseMachineLeaseReservation(reservation: Omit): Promise; + releaseLegacyMachineLease(machineId: string, recoveryId: string, taskId: string): Promise; + clearLegacyMachineLeaseMarker(machineId: string, recoveryId: string): Promise; getProfile(id: string): Promise; - saveProfile(profile: Profile): Promise; + createProfile(profile: Profile): Promise; + acquireProfileLease(profileId: string, userId: string, machineId: string, reservation: MachineLeaseReservation): Promise; + releaseProfileLease(profileId: string, reservation: Omit): Promise; + releaseLegacyProfileLease(profileId: string, taskId: string): Promise; + listProfiles(): Promise; listProfilesByUser(userId: string): Promise; saveHandoff(link: HandoffLink): Promise; getHandoff(id: string): Promise; diff --git a/specs/talos-openapi.yaml b/specs/talos-openapi.yaml index 9c1a505..b3c7b9e 100644 --- a/specs/talos-openapi.yaml +++ b/specs/talos-openapi.yaml @@ -1655,6 +1655,7 @@ components: properties: { kind: { type: string, enum: [choice, text, otp] }, value: { type: string } } Task: type: object + additionalProperties: false required: [id, userId, kind, goal, interaction, status, createdAt, updatedAt, findings, artifacts] properties: id: { type: string } @@ -1686,6 +1687,7 @@ components: constraints: { type: object, properties: { budget: { type: number, minimum: 0 }, deadline: { type: string, format: date-time }, requirements: { type: object, additionalProperties: { oneOf: [{ type: string }, { type: boolean }] } } } } Session: type: object + additionalProperties: false required: [id, status, mode, createdAt, updatedAt] properties: id: { type: string } @@ -1742,6 +1744,7 @@ components: properties: { id: { type: string }, name: { type: string }, contentType: { type: string }, size: { type: integer }, uri: { type: string, format: uri }, createdAt: { type: string, format: date-time } } ClaimResponse: type: object + additionalProperties: false required: [task, lease, leaseToken] properties: task: { $ref: '#/components/schemas/Task' } @@ -1749,6 +1752,7 @@ components: leaseToken: { type: string, minLength: 1 } Lease: type: object + additionalProperties: false required: [taskId, workerId, machineId, expiresAt] properties: { taskId: { type: string }, workerId: { type: string }, machineId: { type: string }, expiresAt: { type: string, format: date-time } } WorkerCredentials: @@ -1897,6 +1901,7 @@ components: properties: { error: { $ref: '#/components/schemas/PublicErrorDetail' } } PublicErrorDetail: type: object + additionalProperties: false required: [code, message, retryable] properties: code: { type: string } @@ -1921,6 +1926,7 @@ components: snapshot_digest: { $ref: '#/components/schemas/Sha256Digest' } WebhookEvent: type: object + additionalProperties: false required: [id, type, taskId, userId, timestamp, payload] properties: id: { type: string } @@ -1937,24 +1943,29 @@ components: delivery: { $ref: '#/components/schemas/WebhookDelivery' } StateChangedPayload: type: object + additionalProperties: false required: [status] properties: status: { type: string, enum: [submitted, claimed, running, needs_input, handoff, closing, completed, failed, cancelled] } error: { $ref: '#/components/schemas/TaskError' } NeedsInputPayload: type: object + additionalProperties: false required: [status] properties: { status: { type: string, const: needs_input } } HandoffRequestedPayload: type: object + additionalProperties: false required: [handoff_url, expires] properties: { handoff_url: { type: string }, expires: { type: string, format: date-time } } CompletedPayload: type: object + additionalProperties: false required: [status] properties: { status: { type: string, const: completed } } WebhookDelivery: type: object + additionalProperties: false required: [status, attempts] properties: status: { type: string, enum: [pending, delivered, failed] } @@ -1963,6 +1974,7 @@ components: lastError: { type: string } TaskError: type: object + additionalProperties: false required: [code, message] properties: { code: { type: string }, message: { type: string } } x-webhook-signing: