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()

tests/test_task_scheduler.py

@@ -9,7 +9,6 @@ import pytest
 from cloud.config import CloudConfig
 from cloud.pool import DevicePool
 from cloud.scheduler import (
-    AssignmentStrategy,
     FIFO_MATCH_STRATEGY_NAME,
     QueueFullError,
     ScheduledTask,
@@ -35,11 +34,15 @@ def _config(**overrides) -> CloudConfig:
     return CloudConfig(**base)
 
 
-def _device(device_id: str, *, status: str = "idle", driver_type: str = "wda") -> Device:
+def _device(
+    device_id: str, *, status: str = "idle", driver_type: str = "wda"
+) -> Device:
     return Device(id=device_id, status=status, driver_type=driver_type)  # type: ignore[arg-type]
 
 
-def _pool_with_devices(tmp_path, *devices: Device, host_id: str = "host-local") -> DevicePool:
+def _pool_with_devices(
+    tmp_path, *devices: Device, host_id: str = "host-local"
+) -> DevicePool:
     pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())
     pool.sync_host_devices(host_id, list(devices))
     return pool
@@ -134,6 +137,11 @@ def test_two_tasks_assigned_in_submission_order_with_one_device(tmp_path) -> Non
     assert first_task.status == "assigned"
     assert second_task.status == "queued"
 
+    assert scheduler.assign() == []
+    second_task = pool.store.get_task(second_id)
+    assert second_task is not None
+    assert second_task.status == "queued"
+
 
 def test_unknown_strategy_raises_at_init(tmp_path) -> None:
     pool = _pool_with_devices(tmp_path)