feat(cloud-scheduler): reserve devices atomically

openspec/changes/cloud-control-plane-integration/tasks.md

@@ -18,7 +18,7 @@
 ## 3. Lease-Backed Scheduling
 
 - [x] 3.1 Extend scheduled-task persistence with attempt count, lease id/expiry, terminal result, failure reason, and auditable attempt records.
-- [ ] 3.2 Implement atomic queued-task assignment and device reservation while excluding devices with active assignments even when snapshots report idle.
+- [x] 3.2 Implement atomic queued-task assignment and device reservation while excluding devices with active assignments even when snapshots report idle.
 - [ ] 3.3 Implement owning-host claim that atomically transitions one assigned attempt to dispatched under its active lease.
 - [ ] 3.4 Implement lease renewal with host/task/attempt ownership validation and conflict responses for stale leases.
 - [ ] 3.5 Implement idempotent terminal result recording and reservation release for active leases.

packages/cloud-platform/cloud/config.py

@@ -25,6 +25,9 @@ class CloudConfig:
     # are rejected rather than letting the backlog grow without limit.
     max_queue_depth: int = 100
 
+    # Duration of each scheduler-created assignment lease.
+    lease_duration_seconds: float = 60.0
+
     # Strategy name looked up in the AssignmentStrategy registry (default fifo_match).
     default_assignment_strategy: str = "fifo_match"
 

packages/cloud-platform/cloud/scheduler.py

@@ -9,7 +9,7 @@ by ``driver/registry.py``, ``perception/provider.py``, and ``workflow/conditions
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable
 from uuid import uuid4
 
@@ -157,32 +157,38 @@ class TaskScheduler:
             return assignments
 
         devices = self.pool.list_devices()
-        assigned_device_ids: set[str] = set()
+        now = utc_now()
+        reserved_device_ids = self.store.list_reserved_device_ids(now=now)
         for task in queued:
             candidates = [
                 device
                 for device in devices
-                if device.device_id not in assigned_device_ids
+                if device.device_id not in reserved_device_ids
                 and device.status == "idle"
                 and _matches(device, task.constraints)
             ]
             selected = strategy.select(task, candidates)
             if selected is None:
                 continue
-            assigned_device_ids.add(selected.device_id)
-            self.store.update_task(
-                task.id,
-                status="assigned",
-                assigned_device_id=selected.device_id,
-                assigned_host_id=selected.host_id,
+            leased = self.store.assign_task(
+                task_id=task.id,
+                host_id=selected.host_id,
+                device_id=selected.device_id,
+                lease_id=uuid4().hex,
+                lease_expires_at=now
+                + timedelta(seconds=self.config.lease_duration_seconds),
+                now=now,
             )
+            if leased is None:
+                continue
+            reserved_device_ids.add(selected.device_id)
             assignments.append(
                 Assignment(
-                    task_id=task.id,
-                    device_id=selected.device_id,
-                    host_id=selected.host_id,
-                    goal=task.goal,
-                    workflow_definition_id=task.workflow_definition_id,
+                    task_id=leased.task_id,
+                    device_id=leased.device_id,
+                    host_id=leased.host_id,
+                    goal=leased.goal,
+                    workflow_definition_id=leased.workflow_definition_id,
                 )
             )
         return assignments

packages/cloud-platform/cloud/sql_repository.py

@@ -207,6 +207,89 @@ class SQLAlchemyCloudRepository:
             row = session.get(PluginRow, name)
         return _plugin_from_row(row) if row else None
 
+    def list_reserved_device_ids(self, *, now: datetime) -> set[str]:
+        with self._sessions() as session:
+            device_ids = session.scalars(
+                select(ScheduledTaskRow.assigned_device_id).where(
+                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
+                    ScheduledTaskRow.assigned_device_id.is_not(None),
+                    ScheduledTaskRow.lease_expires_at.is_not(None),
+                    ScheduledTaskRow.lease_expires_at > _iso(now),
+                )
+            ).all()
+        return {device_id for device_id in device_ids if device_id is not None}
+
+    def assign_task(
+        self,
+        *,
+        task_id: str,
+        host_id: str,
+        device_id: str,
+        lease_id: str,
+        lease_expires_at: datetime,
+        now: datetime,
+    ) -> Any | None:
+        with self._sessions.begin() as session:
+            task_statement = select(ScheduledTaskRow).where(
+                ScheduledTaskRow.id == task_id
+            )
+            device_statement = select(PooledDeviceRow).where(
+                PooledDeviceRow.host_id == host_id,
+                PooledDeviceRow.device_id == device_id,
+            )
+            if self.engine.dialect.name == "postgresql":
+                task_statement = task_statement.with_for_update(skip_locked=True)
+                device_statement = device_statement.with_for_update()
+
+            task = session.scalars(task_statement).first()
+            if task is None or task.status != "queued":
+                return None
+            device = session.scalars(device_statement).first()
+            if device is None or device.status != "idle":
+                return None
+
+            active_reservation = session.scalar(
+                select(ScheduledTaskRow.id)
+                .where(
+                    ScheduledTaskRow.id != task_id,
+                    ScheduledTaskRow.assigned_device_id == device_id,
+                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
+                    ScheduledTaskRow.lease_expires_at.is_not(None),
+                    ScheduledTaskRow.lease_expires_at > _iso(now),
+                )
+                .limit(1)
+            )
+            if active_reservation is not None:
+                return None
+
+            attempt = task.attempt_count + 1
+            task.status = "assigned"
+            task.assigned_host_id = host_id
+            task.assigned_device_id = device_id
+            task.attempt_count = attempt
+            task.lease_id = lease_id
+            task.lease_expires_at = _iso(lease_expires_at)
+            task.failure_reason = None
+            task.result_json = None
+            task.updated_at = _iso(now)
+            session.add(
+                TaskAttemptRow(
+                    task_id=task.id,
+                    attempt=attempt,
+                    lease_id=lease_id,
+                    host_id=host_id,
+                    device_id=device_id,
+                    status="assigned",
+                    lease_expires_at=_iso(lease_expires_at),
+                    created_at=_iso(now),
+                    completed_at=None,
+                    failure_reason=None,
+                    result_json=None,
+                )
+            )
+            session.flush()
+            return _leased_assignment_from_row(task)
+
     def list_task_attempts(self, task_id: str) -> list[Any]:
         with self._sessions() as session:
             rows = session.scalars(
@@ -311,6 +394,29 @@ def _task_attempt_from_row(row: TaskAttemptRow) -> Any:
     )
 
 
+def _leased_assignment_from_row(row: ScheduledTaskRow) -> Any:
+    from cloud.repository import LeasedAssignment
+
+    lease_expires_at = _parse_dt(row.lease_expires_at)
+    if (
+        row.lease_id is None
+        or lease_expires_at is None
+        or row.assigned_host_id is None
+        or row.assigned_device_id is None
+    ):
+        raise ValueError(f"task {row.id!r} does not contain a complete lease")
+    return LeasedAssignment(
+        task_id=row.id,
+        attempt=row.attempt_count,
+        lease_id=row.lease_id,
+        lease_expires_at=lease_expires_at,
+        host_id=row.assigned_host_id,
+        device_id=row.assigned_device_id,
+        goal=row.goal,
+        workflow_definition_id=row.workflow_definition_id,
+    )
+
+
 def _parse_json_object(value: str | None) -> dict[str, Any] | None:
     if value is None:
         return None

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)