feat(cloud-scheduler): reserve devices atomically

tests/test_cloud_repository_contract.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import os
-from datetime import UTC, datetime
+from datetime import UTC, datetime, timedelta
 from pathlib import Path
 from typing import get_protocol_members
 from uuid import uuid4
@@ -278,3 +278,162 @@ def test_task_attempt_history_is_ordered_and_complete(database_url: str) -> None
         assert attempts[1].failure_reason == "execution failed"
     finally:
         database.close()
+
+
+def test_atomic_assignment_creates_lease_attempt_and_reservation(
+    database_url: str,
+) -> None:
+    database = CloudDatabase(database_url)
+    host_id = _unique_id("assignment-host")
+    device_id = _unique_id("assignment-device")
+    task_id = _unique_id("assignment-task")
+    now = datetime(2026, 7, 12, 2, 0, tzinfo=UTC)
+    lease_expires_at = now + timedelta(minutes=1)
+
+    try:
+        database.repository.upsert_host(host_id, address=None, last_seen_at=now)
+        database.repository.replace_host_devices(
+            host_id,
+            [_device(device_id, host_id)],
+        )
+        database.repository.enqueue_task(
+            ScheduledTask(
+                id=task_id,
+                goal="open settings",
+                workflow_definition_id=None,
+                constraints=TaskConstraints(),
+                created_at=now,
+            )
+        )
+
+        assignment = database.repository.assign_task(
+            task_id=task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id="lease-1",
+            lease_expires_at=lease_expires_at,
+            now=now,
+        )
+
+        assert assignment is not None
+        assert assignment.attempt == 1
+        assert assignment.lease_id == "lease-1"
+        task = database.repository.get_task(task_id)
+        assert task is not None
+        assert task.status == "assigned"
+        assert task.attempt_count == 1
+        assert task.lease_expires_at == lease_expires_at
+        assert device_id in database.repository.list_reserved_device_ids(now=now)
+        attempts = database.repository.list_task_attempts(task_id)
+        assert len(attempts) == 1
+        assert attempts[0].status == "assigned"
+        assert attempts[0].lease_id == "lease-1"
+    finally:
+        database.close()
+
+
+def test_active_assignment_blocks_reuse_of_stale_idle_snapshot(
+    database_url: str,
+) -> None:
+    database = CloudDatabase(database_url)
+    host_id = _unique_id("reservation-host")
+    device_id = _unique_id("reservation-device")
+    first_task_id = _unique_id("reservation-task")
+    second_task_id = _unique_id("reservation-task")
+    now = datetime(2026, 7, 12, 3, 0, tzinfo=UTC)
+
+    try:
+        database.repository.upsert_host(host_id, address=None, last_seen_at=now)
+        database.repository.replace_host_devices(
+            host_id,
+            [_device(device_id, host_id)],
+        )
+        for task_id in (first_task_id, second_task_id):
+            database.repository.enqueue_task(
+                ScheduledTask(
+                    id=task_id,
+                    goal="run task",
+                    workflow_definition_id=None,
+                    constraints=TaskConstraints(),
+                    created_at=now,
+                )
+            )
+
+        first = database.repository.assign_task(
+            task_id=first_task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id="lease-active",
+            lease_expires_at=now + timedelta(minutes=1),
+            now=now,
+        )
+        second = database.repository.assign_task(
+            task_id=second_task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id="lease-blocked",
+            lease_expires_at=now + timedelta(minutes=1),
+            now=now,
+        )
+
+        assert first is not None
+        assert second is None
+        blocked_task = database.repository.get_task(second_task_id)
+        assert blocked_task is not None
+        assert blocked_task.status == "queued"
+        assert database.repository.list_task_attempts(second_task_id) == []
+    finally:
+        database.close()
+
+
+def test_expired_assignment_no_longer_reserves_device(database_url: str) -> None:
+    database = CloudDatabase(database_url)
+    host_id = _unique_id("expired-host")
+    device_id = _unique_id("expired-device")
+    first_task_id = _unique_id("expired-task")
+    second_task_id = _unique_id("expired-task")
+    assigned_at = datetime(2026, 7, 12, 4, 0, tzinfo=UTC)
+    after_expiry = assigned_at + timedelta(minutes=2)
+
+    try:
+        database.repository.upsert_host(
+            host_id,
+            address=None,
+            last_seen_at=assigned_at,
+        )
+        database.repository.replace_host_devices(
+            host_id,
+            [_device(device_id, host_id)],
+        )
+        for task_id in (first_task_id, second_task_id):
+            database.repository.enqueue_task(
+                ScheduledTask(
+                    id=task_id,
+                    goal="run task",
+                    workflow_definition_id=None,
+                    constraints=TaskConstraints(),
+                    created_at=assigned_at,
+                )
+            )
+        assert database.repository.assign_task(
+            task_id=first_task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id="lease-expired",
+            lease_expires_at=assigned_at + timedelta(minutes=1),
+            now=assigned_at,
+        )
+
+        assert device_id not in database.repository.list_reserved_device_ids(
+            now=after_expiry
+        )
+        assert database.repository.assign_task(
+            task_id=second_task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id="lease-new",
+            lease_expires_at=after_expiry + timedelta(minutes=1),
+            now=after_expiry,
+        )
+    finally:
+        database.close()