from __future__ import annotations import os from datetime import UTC, datetime, timedelta from pathlib import Path from typing import get_protocol_members from uuid import uuid4 import pytest from sqlalchemy import event from sqlalchemy.orm import Session from cloud.database import CloudDatabase from cloud.db_models import TaskAttemptRow from cloud.plugins import PluginManifest from cloud.pool import PooledDevice from cloud.repository import ( CloudRepository, DeviceEnrollmentConflictError, EnrollmentTokenConflictError, HostEnrollmentConflictError, LeasedAssignment, TaskAttemptRecord, ) from cloud.scheduler import ScheduledTask, TaskConstraints @pytest.fixture( params=[ pytest.param("sqlite", id="sqlite"), pytest.param("postgresql", id="postgresql", marks=pytest.mark.integration), ] ) def database_url(request: pytest.FixtureRequest, tmp_path: Path) -> str: if request.param == "sqlite": return f"sqlite:///{(tmp_path / 'contract.sqlite3').as_posix()}" url = os.getenv("TEST_POSTGRES_URL") if not url: pytest.skip("TEST_POSTGRES_URL is required for PostgreSQL contract tests") return url def _unique_id(prefix: str) -> str: return f"{prefix}-{uuid4().hex}" def _device(device_id: str, host_id: str) -> PooledDevice: return PooledDevice( device_id=device_id, host_id=host_id, driver_type="wda", status="idle", capability_tags=["ios", "physical"], synced_at=datetime(2026, 7, 12, tzinfo=UTC), ) def test_cloud_repository_exposes_crud_and_atomic_lease_operations() -> None: members = get_protocol_members(CloudRepository) assert { "enroll_host", "authenticate_enrolled_host", "revoke_enrolled_host", "is_enrollment_managed_host", "enroll_device", "get_device_enrollment", "list_device_enrollments", "upsert_host", "replace_host_devices", "list_hosts", "list_devices", "enqueue_task", "get_task", "list_tasks", "count_tasks", "save_plugin", "assign_task", "claim_assignment", "renew_lease", "record_task_result", "reap_expired_leases", "list_task_attempts", "record_planner_decision", "prune_planner_decision_log", "list_planner_decisions", "health_check", "close", } <= members def test_repository_transfer_records_are_immutable() -> None: assert TaskAttemptRecord.__dataclass_params__.frozen is True assert LeasedAssignment.__dataclass_params__.frozen is True def test_host_enrollment_is_idempotent_and_token_is_one_time( database_url: str, ) -> None: database = CloudDatabase(database_url) repository = database.repository enrolled_at = datetime(2026, 7, 13, 4, 0, tzinfo=UTC) host_id = _unique_id("managed-host") agent_instance_id = _unique_id("agent-instance") credential_digest = _unique_id("credential-digest") enrollment_digest = _unique_id("enrollment-digest") try: created = repository.enroll_host( host_id=host_id, agent_instance_id=agent_instance_id, credential_digest=credential_digest, enrollment_token_digest=enrollment_digest, display_name="Edge Mac", enrolled_at=enrolled_at, ) assert created.host_id == host_id assert created.agent_instance_id == agent_instance_id assert created.display_name == "Edge Mac" assert created.enrolled_at == enrolled_at assert repository.is_enrollment_managed_host(host_id) is True assert repository.authenticate_enrolled_host(credential_digest) == host_id retried = repository.enroll_host( host_id=_unique_id("ignored-host"), agent_instance_id=agent_instance_id, credential_digest=credential_digest, enrollment_token_digest=enrollment_digest, display_name="Renamed Edge Mac", enrolled_at=enrolled_at + timedelta(minutes=1), ) assert retried == created with pytest.raises(HostEnrollmentConflictError): repository.enroll_host( host_id=_unique_id("host"), agent_instance_id=agent_instance_id, credential_digest=_unique_id("different-credential"), enrollment_token_digest=enrollment_digest, display_name=None, enrolled_at=enrolled_at, ) with pytest.raises(EnrollmentTokenConflictError): repository.enroll_host( host_id=_unique_id("host"), agent_instance_id=_unique_id("different-instance"), credential_digest=_unique_id("credential"), enrollment_token_digest=enrollment_digest, display_name=None, enrolled_at=enrolled_at, ) assert repository.revoke_enrolled_host( host_id, revoked_at=enrolled_at + timedelta(hours=1), ) assert repository.authenticate_enrolled_host(credential_digest) is None assert repository.get_host(host_id) is not None finally: database.close() def test_self_service_enrollment_is_idempotent_with_null_token_digest( database_url: str, ) -> None: database = CloudDatabase(database_url) repository = database.repository enrolled_at = datetime(2026, 7, 13, 6, 0, tzinfo=UTC) host_id = _unique_id("self-service-host") agent_instance_id = _unique_id("self-service-instance") credential_digest = _unique_id("self-service-credential") try: created = repository.enroll_host( host_id=host_id, agent_instance_id=agent_instance_id, credential_digest=credential_digest, enrollment_token_digest=None, display_name="Self-Service Host", enrolled_at=enrolled_at, ) assert created.host_id == host_id retried = repository.enroll_host( host_id=_unique_id("ignored-host"), agent_instance_id=agent_instance_id, credential_digest=credential_digest, enrollment_token_digest=None, display_name="Renamed Self-Service Host", enrolled_at=enrolled_at + timedelta(minutes=1), ) assert retried == created finally: database.close() def test_multiple_self_service_hosts_coexist_without_token_conflict( database_url: str, ) -> None: database = CloudDatabase(database_url) repository = database.repository enrolled_at = datetime(2026, 7, 13, 7, 0, tzinfo=UTC) credential_a = _unique_id("self-service-credential-a") credential_b = _unique_id("self-service-credential-b") try: first = repository.enroll_host( host_id=_unique_id("self-service-host-a"), agent_instance_id=_unique_id("self-service-instance-a"), credential_digest=credential_a, enrollment_token_digest=None, display_name=None, enrolled_at=enrolled_at, ) second = repository.enroll_host( host_id=_unique_id("self-service-host-b"), agent_instance_id=_unique_id("self-service-instance-b"), credential_digest=credential_b, enrollment_token_digest=None, display_name=None, enrolled_at=enrolled_at, ) assert first.host_id != second.host_id assert repository.authenticate_enrolled_host(credential_a) == first.host_id assert repository.authenticate_enrolled_host(credential_b) == second.host_id finally: database.close() def test_device_enrollment_is_host_scoped_and_idempotent( database_url: str, ) -> None: database = CloudDatabase(database_url) repository = database.repository enrolled_at = datetime(2026, 7, 13, 5, 0, tzinfo=UTC) host_a = _unique_id("host-a") host_b = _unique_id("host-b") local_device_id = _unique_id("local-device") try: device_a = repository.enroll_device( device_id=_unique_id("cloud-device"), host_id=host_a, local_device_id=local_device_id, driver_type="wda", name="iPhone", capability_tags=["ios"], enrolled_at=enrolled_at, ) retried = repository.enroll_device( device_id=_unique_id("ignored-device"), host_id=host_a, local_device_id=local_device_id, driver_type="wda", name="Renamed iPhone", capability_tags=["ios", "physical"], enrolled_at=enrolled_at + timedelta(minutes=1), ) assert retried.device_id == device_a.device_id assert retried.name == "Renamed iPhone" assert retried.capability_tags == ["ios", "physical"] assert repository.get_device_enrollment(device_a.device_id) == retried assert repository.list_device_enrollments(host_a) == [retried] device_b = repository.enroll_device( device_id=_unique_id("cloud-device"), host_id=host_b, local_device_id=local_device_id, driver_type="wda", name="Moved iPhone", capability_tags=[], enrolled_at=enrolled_at, ) assert device_b.device_id != device_a.device_id assert device_b.host_id == host_b with pytest.raises(DeviceEnrollmentConflictError): repository.enroll_device( device_id=_unique_id("device"), host_id=host_a, local_device_id=local_device_id, driver_type="android", name=None, capability_tags=[], enrolled_at=enrolled_at, ) finally: database.close() def test_repository_crud_contract(database_url: str) -> None: database = CloudDatabase(database_url) repository = database.repository host_id = _unique_id("host") device_id = _unique_id("device") task_id = _unique_id("task") plugin_name = _unique_id("plugin") seen_at = datetime(2026, 7, 12, 1, 2, 3, tzinfo=UTC) try: repository.upsert_host( host_id, address="127.0.0.1:9000", last_seen_at=seen_at, ) repository.replace_host_devices(host_id, [_device(device_id, host_id)]) repository.enqueue_task( ScheduledTask( id=task_id, goal="open settings", workflow_definition_id=None, constraints=TaskConstraints( driver_type="wda", capability_tags=["ios"], ), created_at=seen_at, ) ) repository.save_plugin( PluginManifest( name=plugin_name, version="1.0.0", entry_point_kind="tool", target="cloud.store:CloudStore", ), wired=False, ) assert repository.get_host(host_id) is not None assert repository.get_device(device_id) == _device(device_id, host_id) assert repository.get_task(task_id) is not None assert repository.get_plugin(plugin_name) is not None repository.health_check() finally: database.close() def test_repository_can_atomically_transfer_explicit_device_takeover( database_url: str, ) -> None: database = CloudDatabase(database_url) repository = database.repository old_host_id = _unique_id("old-host") new_host_id = _unique_id("new-host") device_id = _unique_id("shared-device") try: repository.replace_host_devices( old_host_id, [_device(device_id, old_host_id)], ) repository.replace_host_devices( new_host_id, [_device(device_id, new_host_id)], allow_device_takeover=True, ) matching = [ device for device in repository.list_devices() if device.device_id == device_id ] assert len(matching) == 1 assert matching[0].host_id == new_host_id finally: database.close() def test_failed_snapshot_transaction_rolls_back(database_url: str) -> None: database = CloudDatabase(database_url) repository = database.repository host_id = _unique_id("rollback-host") original_device_id = _unique_id("original-device") duplicate_device_id = _unique_id("duplicate-device") try: repository.upsert_host( host_id, address=None, last_seen_at=datetime(2026, 7, 12, tzinfo=UTC), ) repository.replace_host_devices( host_id, [_device(original_device_id, host_id)], ) def fail_before_insert( _connection: object, _cursor: object, statement: str, _parameters: object, _context: object, _executemany: bool, ) -> None: if statement.startswith("INSERT INTO pooled_devices"): raise RuntimeError("injected snapshot write failure") event.listen(database.engine, "before_cursor_execute", fail_before_insert) try: with pytest.raises(RuntimeError, match="injected snapshot write failure"): repository.replace_host_devices( host_id, [_device(duplicate_device_id, host_id)], ) finally: event.remove(database.engine, "before_cursor_execute", fail_before_insert) devices = [ device for device in repository.list_devices() if device.host_id == host_id ] assert [device.device_id for device in devices] == [original_device_id] repository.health_check() finally: database.close() def test_repository_state_survives_application_restart(database_url: str) -> None: host_id = _unique_id("restart-host") device_id = _unique_id("restart-device") first_process = CloudDatabase(database_url) try: first_process.repository.upsert_host( host_id, address="host.internal", last_seen_at=datetime(2026, 7, 12, tzinfo=UTC), ) first_process.repository.replace_host_devices( host_id, [_device(device_id, host_id)], ) finally: first_process.close() restarted_process = CloudDatabase(database_url) try: host = restarted_process.repository.get_host(host_id) device = restarted_process.repository.get_device(device_id) assert host is not None assert host.address == "host.internal" assert device == _device(device_id, host_id) finally: restarted_process.close() def test_task_lease_and_terminal_fields_round_trip(database_url: str) -> None: database = CloudDatabase(database_url) task_id = _unique_id("lease-task") created_at = datetime(2026, 7, 12, 1, 0, tzinfo=UTC) lease_expires_at = datetime(2026, 7, 12, 1, 5, tzinfo=UTC) updated_at = datetime(2026, 7, 12, 1, 1, tzinfo=UTC) task = ScheduledTask( id=task_id, goal="capture diagnostics", workflow_definition_id=None, constraints=TaskConstraints(), status="failed", assigned_device_id="device-a", assigned_host_id="host-a", attempt_count=2, lease_id="lease-a", lease_expires_at=lease_expires_at, terminal_result={"steps": 3, "status": "failed"}, failure_reason="device disconnected", updated_at=updated_at, created_at=created_at, ) try: database.repository.enqueue_task(task) assert database.repository.get_task(task_id) == task finally: database.close() def test_task_attempt_history_is_ordered_and_complete(database_url: str) -> None: database = CloudDatabase(database_url) task_id = _unique_id("attempt-task") created_at = datetime(2026, 7, 12, 1, 0, tzinfo=UTC) lease_expires_at = datetime(2026, 7, 12, 1, 5, tzinfo=UTC) try: with Session(database.engine) as session, session.begin(): session.add_all( [ TaskAttemptRow( task_id=task_id, attempt=2, lease_id="lease-2", host_id="host-b", device_id="device-b", status="failed", lease_expires_at=lease_expires_at.isoformat(), created_at=created_at.isoformat(), completed_at=lease_expires_at.isoformat(), failure_reason="execution failed", result_json='{"exit_code": 1}', ), TaskAttemptRow( task_id=task_id, attempt=1, lease_id="lease-1", host_id="host-a", device_id="device-a", status="expired", lease_expires_at=lease_expires_at.isoformat(), created_at=created_at.isoformat(), completed_at=lease_expires_at.isoformat(), failure_reason="lease expired", result_json=None, ), ] ) attempts = database.repository.list_task_attempts(task_id) assert [attempt.attempt for attempt in attempts] == [1, 2] assert attempts[0].status == "expired" assert attempts[1].terminal_result == {"exit_code": 1} assert attempts[1].failure_reason == "execution failed" finally: database.close() def test_list_tasks_returns_empty_when_repository_has_no_tasks( database_url: str, ) -> None: database = CloudDatabase(database_url) try: assert database.repository.list_tasks() == [] assert database.repository.count_tasks() == 0 assert database.repository.list_tasks(status="queued") == [] assert database.repository.count_tasks(status="queued") == 0 finally: database.close() def test_list_tasks_returns_most_recent_first_with_optional_status_filter( database_url: str, ) -> None: database = CloudDatabase(database_url) base = datetime(2026, 7, 12, 6, 0, tzinfo=UTC) queued_ids = [_unique_id("list-task") for _ in range(2)] failed_ids = [_unique_id("list-task") for _ in range(2)] try: for index, task_id in enumerate(queued_ids): database.repository.enqueue_task( ScheduledTask( id=task_id, goal="queued goal", workflow_definition_id=None, constraints=TaskConstraints(), status="queued", created_at=base + timedelta(seconds=index), ) ) for index, task_id in enumerate(failed_ids): database.repository.enqueue_task( ScheduledTask( id=task_id, goal="failed goal", workflow_definition_id=None, constraints=TaskConstraints(), status="failed", failure_reason="boom", created_at=base + timedelta(seconds=10 + index), ) ) unfiltered = database.repository.list_tasks() assert [task.id for task in unfiltered] == ( list(reversed(failed_ids)) + list(reversed(queued_ids)) ) assert database.repository.count_tasks() == 4 queued = database.repository.list_tasks(status="queued") assert [task.id for task in queued] == list(reversed(queued_ids)) assert all(task.status == "queued" for task in queued) assert database.repository.count_tasks(status="queued") == 2 failed = database.repository.list_tasks(status="failed") assert [task.id for task in failed] == list(reversed(failed_ids)) assert database.repository.count_tasks(status="failed") == 2 # A status with no matches returns an empty page and zero count. assert database.repository.list_tasks(status="done") == [] assert database.repository.count_tasks(status="done") == 0 finally: database.close() def test_list_tasks_pagination_bounds(database_url: str) -> None: database = CloudDatabase(database_url) base = datetime(2026, 7, 12, 7, 0, tzinfo=UTC) task_ids = [_unique_id("page-task") for _ in range(4)] try: for index, task_id in enumerate(task_ids): database.repository.enqueue_task( ScheduledTask( id=task_id, goal=f"goal-{index}", workflow_definition_id=None, constraints=TaskConstraints(), created_at=base + timedelta(seconds=index), ) ) # Most-recent-first ordering means page 1 returns the newest two ids. page_one = database.repository.list_tasks(limit=2, offset=0) assert [task.id for task in page_one] == [task_ids[3], task_ids[2]] page_two = database.repository.list_tasks(limit=2, offset=2) assert [task.id for task in page_two] == [task_ids[1], task_ids[0]] # An offset past the end of the result set returns an empty page, # not an error — the caller is expected to consult count_tasks(). assert database.repository.list_tasks(limit=10, offset=100) == [] 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() def test_owning_host_claims_one_active_assignment(database_url: str) -> None: database = CloudDatabase(database_url) host_id = _unique_id("claim-host") device_id = _unique_id("claim-device") task_id = _unique_id("claim-task") now = datetime(2026, 7, 12, 5, 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)], ) database.repository.enqueue_task( ScheduledTask( id=task_id, goal="claim me", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="claim-lease", lease_expires_at=now + timedelta(minutes=1), now=now, ) assignment = database.repository.claim_assignment( host_id=host_id, now=now + timedelta(seconds=1), ) assert assignment is not None assert assignment.task_id == task_id assert assignment.lease_id == "claim-lease" task = database.repository.get_task(task_id) assert task is not None assert task.status == "dispatched" attempts = database.repository.list_task_attempts(task_id) assert len(attempts) == 1 assert attempts[0].status == "dispatched" assert database.repository.claim_assignment(host_id=host_id, now=now) is None finally: database.close() def test_foreign_host_cannot_claim_assignment(database_url: str) -> None: database = CloudDatabase(database_url) host_id = _unique_id("owner-host") device_id = _unique_id("owner-device") task_id = _unique_id("owner-task") now = datetime(2026, 7, 12, 6, 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)], ) database.repository.enqueue_task( ScheduledTask( id=task_id, goal="owner only", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="owner-lease", lease_expires_at=now + timedelta(minutes=1), now=now, ) assert ( database.repository.claim_assignment( host_id=_unique_id("foreign-host"), now=now, ) is None ) task = database.repository.get_task(task_id) assert task is not None assert task.status == "assigned" finally: database.close() def test_expired_assignment_cannot_be_claimed(database_url: str) -> None: database = CloudDatabase(database_url) host_id = _unique_id("expired-claim-host") device_id = _unique_id("expired-claim-device") task_id = _unique_id("expired-claim-task") now = datetime(2026, 7, 12, 7, 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)], ) database.repository.enqueue_task( ScheduledTask( id=task_id, goal="too late", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="expired-claim-lease", lease_expires_at=now + timedelta(seconds=1), now=now, ) assert ( database.repository.claim_assignment( host_id=host_id, now=now + timedelta(seconds=2), ) is None ) task = database.repository.get_task(task_id) assert task is not None assert task.status == "assigned" finally: database.close() def test_active_lease_renews_for_owning_host(database_url: str) -> None: database = CloudDatabase(database_url) host_id = _unique_id("renew-host") device_id = _unique_id("renew-device") task_id = _unique_id("renew-task") now = datetime(2026, 7, 12, 8, 0, tzinfo=UTC) initial_expiry = now + timedelta(minutes=1) renewed_expiry = now + timedelta(minutes=2) 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="renew me", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="renew-lease", lease_expires_at=initial_expiry, now=now, ) database.repository.claim_assignment(host_id=host_id, now=now) status = database.repository.renew_lease( task_id=task_id, attempt=1, lease_id="renew-lease", host_id=host_id, lease_expires_at=renewed_expiry, now=now + timedelta(seconds=30), ) assert status == "renewed" task = database.repository.get_task(task_id) assert task is not None assert task.lease_expires_at == renewed_expiry attempts = database.repository.list_task_attempts(task_id) assert attempts[0].lease_expires_at == renewed_expiry finally: database.close() @pytest.mark.parametrize( ("attempt", "lease_id", "host_id"), [ (2, "lease-current", "owner"), (1, "lease-stale", "owner"), (1, "lease-current", "foreign"), ], ) def test_stale_or_foreign_lease_renewal_conflicts( database_url: str, attempt: int, lease_id: str, host_id: str, ) -> None: database = CloudDatabase(database_url) owner_id = _unique_id("renew-owner") device_id = _unique_id("renew-conflict-device") task_id = _unique_id("renew-conflict-task") now = datetime(2026, 7, 12, 9, 0, tzinfo=UTC) initial_expiry = now + timedelta(minutes=1) try: database.repository.upsert_host(owner_id, address=None, last_seen_at=now) database.repository.replace_host_devices( owner_id, [_device(device_id, owner_id)], ) database.repository.enqueue_task( ScheduledTask( id=task_id, goal="do not renew", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=owner_id, device_id=device_id, lease_id="lease-current", lease_expires_at=initial_expiry, now=now, ) status = database.repository.renew_lease( task_id=task_id, attempt=attempt, lease_id=lease_id, host_id=owner_id if host_id == "owner" else _unique_id("foreign"), lease_expires_at=now + timedelta(minutes=2), now=now + timedelta(seconds=30), ) assert status == "conflict" task = database.repository.get_task(task_id) assert task is not None assert task.lease_expires_at == initial_expiry finally: database.close() def test_expired_or_missing_lease_cannot_be_renewed(database_url: str) -> None: database = CloudDatabase(database_url) host_id = _unique_id("expired-renew-host") device_id = _unique_id("expired-renew-device") task_id = _unique_id("expired-renew-task") now = datetime(2026, 7, 12, 10, 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)], ) database.repository.enqueue_task( ScheduledTask( id=task_id, goal="expired renewal", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="expired-renew-lease", lease_expires_at=now + timedelta(seconds=1), now=now, ) assert ( database.repository.renew_lease( task_id=task_id, attempt=1, lease_id="expired-renew-lease", host_id=host_id, lease_expires_at=now + timedelta(minutes=2), now=now + timedelta(seconds=2), ) == "expired" ) assert ( database.repository.renew_lease( task_id=_unique_id("missing-task"), attempt=1, lease_id="missing-lease", host_id=host_id, lease_expires_at=now + timedelta(minutes=2), now=now, ) == "not_found" ) finally: database.close() def test_terminal_result_is_recorded_idempotently_and_releases_reservation( database_url: str, ) -> None: database = CloudDatabase(database_url) host_id = _unique_id("result-host") device_id = _unique_id("result-device") task_id = _unique_id("result-task") now = datetime(2026, 7, 12, 11, 0, tzinfo=UTC) completed_at = now + timedelta(seconds=30) result = {"steps": 4, "summary": "completed"} 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="complete me", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="result-lease", lease_expires_at=now + timedelta(minutes=1), now=now, ) database.repository.claim_assignment(host_id=host_id, now=now) first_status = database.repository.record_task_result( task_id=task_id, attempt=1, lease_id="result-lease", host_id=host_id, status="done", failure_reason=None, terminal_result=result, completed_at=completed_at, ) repeated_status = database.repository.record_task_result( task_id=task_id, attempt=1, lease_id="result-lease", host_id=host_id, status="done", failure_reason=None, terminal_result=result, completed_at=completed_at + timedelta(seconds=1), ) assert first_status == "recorded" assert repeated_status == "already_recorded" task = database.repository.get_task(task_id) assert task is not None assert task.status == "done" assert task.terminal_result == result attempts = database.repository.list_task_attempts(task_id) assert attempts[0].status == "done" assert attempts[0].terminal_result == result assert attempts[0].completed_at == completed_at assert device_id not in database.repository.list_reserved_device_ids( now=completed_at ) finally: database.close() def test_conflicting_terminal_result_cannot_overwrite_recorded_outcome( database_url: str, ) -> None: database = CloudDatabase(database_url) host_id = _unique_id("result-conflict-host") device_id = _unique_id("result-conflict-device") task_id = _unique_id("result-conflict-task") now = datetime(2026, 7, 12, 12, 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)], ) database.repository.enqueue_task( ScheduledTask( id=task_id, goal="stable outcome", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="stable-lease", lease_expires_at=now + timedelta(minutes=1), now=now, ) assert ( database.repository.record_task_result( task_id=task_id, attempt=1, lease_id="stable-lease", host_id=host_id, status="failed", failure_reason="device offline", terminal_result={"retryable": True}, completed_at=now + timedelta(seconds=10), ) == "recorded" ) assert ( database.repository.record_task_result( task_id=task_id, attempt=1, lease_id="stable-lease", host_id=host_id, status="done", failure_reason=None, terminal_result={"retryable": False}, completed_at=now + timedelta(seconds=20), ) == "conflict" ) task = database.repository.get_task(task_id) assert task is not None assert task.status == "failed" assert task.failure_reason == "device offline" assert task.terminal_result == {"retryable": True} finally: database.close() @pytest.mark.parametrize( ("attempt", "lease_id", "host_kind", "report_delay"), [ (2, "active-lease", "owner", 10), (1, "stale-lease", "owner", 10), (1, "active-lease", "foreign", 10), (1, "active-lease", "owner", 61), ], ) def test_stale_foreign_or_expired_result_is_rejected( database_url: str, attempt: int, lease_id: str, host_kind: str, report_delay: int, ) -> None: database = CloudDatabase(database_url) owner_id = _unique_id("result-owner") device_id = _unique_id("result-stale-device") task_id = _unique_id("result-stale-task") now = datetime(2026, 7, 12, 13, 0, tzinfo=UTC) try: database.repository.upsert_host(owner_id, address=None, last_seen_at=now) database.repository.replace_host_devices( owner_id, [_device(device_id, owner_id)], ) database.repository.enqueue_task( ScheduledTask( id=task_id, goal="reject stale result", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=owner_id, device_id=device_id, lease_id="active-lease", lease_expires_at=now + timedelta(minutes=1), now=now, ) status = database.repository.record_task_result( task_id=task_id, attempt=attempt, lease_id=lease_id, host_id=owner_id if host_kind == "owner" else _unique_id("foreign"), status="done", failure_reason=None, terminal_result={"ignored": True}, completed_at=now + timedelta(seconds=report_delay), ) assert status == "conflict" task = database.repository.get_task(task_id) assert task is not None assert task.status == "assigned" assert task.terminal_result is None finally: database.close() @pytest.mark.parametrize("claimed", [False, True], ids=["assigned", "dispatched"]) def test_expired_lease_requeues_with_auditable_history( database_url: str, claimed: bool, ) -> None: database = CloudDatabase(database_url) host_id = _unique_id("requeue-host") device_id = _unique_id("requeue-device") task_id = _unique_id("requeue-task") now = datetime(2026, 7, 12, 14, 0, tzinfo=UTC) expired_at = now + timedelta(seconds=10) reaped_at = expired_at + timedelta(seconds=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="retry after expiry", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="expired-attempt-1", lease_expires_at=expired_at, now=now, ) if claimed: database.repository.claim_assignment(host_id=host_id, now=now) reaped_task_ids = database.repository.reap_expired_leases( now=reaped_at, max_attempts=2, ) assert task_id in reaped_task_ids task = database.repository.get_task(task_id) assert task is not None assert task.status == "queued" assert task.attempt_count == 1 assert task.assigned_host_id is None assert task.assigned_device_id is None assert task.lease_id is None assert task.lease_expires_at is None assert task.failure_reason is None attempts = database.repository.list_task_attempts(task_id) assert len(attempts) == 1 assert attempts[0].status == "expired" assert attempts[0].failure_reason == "lease expired" assert attempts[0].completed_at == reaped_at assert device_id not in database.repository.list_reserved_device_ids( now=reaped_at ) second_assignment = database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="active-attempt-2", lease_expires_at=reaped_at + timedelta(minutes=1), now=reaped_at, ) assert second_assignment is not None assert second_assignment.attempt == 2 assert [ attempt.status for attempt in database.repository.list_task_attempts(task_id) ] == ["expired", "assigned"] finally: database.close() def test_expired_lease_fails_at_attempt_limit(database_url: str) -> None: database = CloudDatabase(database_url) host_id = _unique_id("limit-host") device_id = _unique_id("limit-device") task_id = _unique_id("limit-task") now = datetime(2026, 7, 12, 15, 0, tzinfo=UTC) expired_at = now + timedelta(seconds=10) reaped_at = expired_at + timedelta(seconds=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="fail after expiry", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="final-attempt", lease_expires_at=expired_at, now=now, ) reaped_task_ids = database.repository.reap_expired_leases( now=reaped_at, max_attempts=1, ) assert task_id in reaped_task_ids task = database.repository.get_task(task_id) assert task is not None assert task.status == "failed" assert task.failure_reason == "lease expired after 1 attempts" assert task.lease_id is None assert task.lease_expires_at is None attempts = database.repository.list_task_attempts(task_id) assert attempts[0].status == "expired" assert device_id not in database.repository.list_reserved_device_ids( now=reaped_at ) assert ( database.repository.reap_expired_leases( now=reaped_at, max_attempts=1, ) == [] ) finally: database.close() def test_unexpired_lease_is_not_reaped(database_url: str) -> None: database = CloudDatabase(database_url) host_id = _unique_id("active-reaper-host") device_id = _unique_id("active-reaper-device") task_id = _unique_id("active-reaper-task") now = datetime(2026, 7, 12, 16, 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)], ) database.repository.enqueue_task( ScheduledTask( id=task_id, goal="stay active", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="active-reaper-lease", lease_expires_at=now + timedelta(minutes=1), now=now, ) assert ( database.repository.reap_expired_leases( now=now + timedelta(seconds=30), max_attempts=2, ) == [] ) task = database.repository.get_task(task_id) assert task is not None assert task.status == "assigned" assert device_id in database.repository.list_reserved_device_ids(now=now) finally: database.close() def test_task_lifecycle_logs_structured_identifiers( database_url: str, monkeypatch ) -> None: import cloud.sql_repository as repository_module database = CloudDatabase(database_url) now = datetime(2026, 7, 12, 20, 0, tzinfo=UTC) host_id = _unique_id("log-host") device_id = _unique_id("log-device") task_id = _unique_id("log-task") events: list[dict[str, object]] = [] def record_info(_message: str, *, extra: dict[str, object]) -> None: events.append(extra) monkeypatch.setattr(repository_module.logger, "info", record_info) 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="sensitive typed text", workflow_definition_id=None, constraints=TaskConstraints(), created_at=now, ) ) database.repository.assign_task( task_id=task_id, host_id=host_id, device_id=device_id, lease_id="log-lease", lease_expires_at=now + timedelta(minutes=1), now=now, ) database.repository.claim_assignment(host_id=host_id, now=now) database.repository.record_task_result( task_id=task_id, attempt=1, lease_id="log-lease", host_id=host_id, status="done", failure_reason=None, terminal_result={"screenshot": "secret-image"}, completed_at=now + timedelta(seconds=1), ) assert [event["event"] for event in events] == [ "assigned", "claimed", "completed", ] assert all(event["task_id"] == task_id for event in events) assert all(event["host_id"] == host_id for event in events) assert all(event["device_id"] == device_id for event in events) assert all(event["correlation_id"] for event in events) assert "sensitive typed text" not in repr(events) assert "secret-image" not in repr(events) finally: database.close() def test_record_planner_decision_assigns_incrementing_step_index( database_url: str, ) -> None: database = CloudDatabase(database_url) task_id = _unique_id("planner-task") other_task_id = _unique_id("planner-task-other") host_id = _unique_id("planner-host") now = datetime(2026, 7, 14, 0, 0, tzinfo=UTC) try: step1 = database.repository.record_planner_decision( host_id=host_id, task_id=task_id, attempt=1, system_prompt="system", user_prompt="prompt-1", tool_name="tap", arguments_json='{"x": 1}', now=now, ) step2 = database.repository.record_planner_decision( host_id=host_id, task_id=task_id, attempt=1, system_prompt="system", user_prompt="prompt-2", tool_name="swipe", arguments_json='{"y": 2}', now=now + timedelta(seconds=1), ) step3 = database.repository.record_planner_decision( host_id=host_id, task_id=task_id, attempt=1, system_prompt="system", user_prompt="prompt-3", tool_name="wait", arguments_json="{}", now=now + timedelta(seconds=2), ) assert [step1, step2, step3] == [1, 2, 3] # Different (task_id, attempt) gets its own counter starting at 1. other_step = database.repository.record_planner_decision( host_id=host_id, task_id=other_task_id, attempt=1, system_prompt="system", user_prompt="other", tool_name="tap", arguments_json="{}", now=now, ) assert other_step == 1 # Different attempt on the same task also gets its own counter. attempt2_step = database.repository.record_planner_decision( host_id=host_id, task_id=task_id, attempt=2, system_prompt="system", user_prompt="retry", tool_name="tap", arguments_json="{}", now=now, ) assert attempt2_step == 1 # Verify stored rows. decisions = database.repository.list_planner_decisions( task_id=task_id, attempt=1 ) assert len(decisions) == 3 assert [d.step_index for d in decisions] == [1, 2, 3] assert [d.user_prompt for d in decisions] == [ "prompt-1", "prompt-2", "prompt-3", ] assert [d.tool_name for d in decisions] == ["tap", "swipe", "wait"] assert decisions[0].arguments_json == '{"x": 1}' finally: database.close() def test_prune_planner_decision_log_deletes_only_old_terminal_tasks( database_url: str, ) -> None: database = CloudDatabase(database_url) old_terminal_task = _unique_id("old-terminal") in_flight_task = _unique_id("in-flight") recent_terminal_task = _unique_id("recent-terminal") host_id = _unique_id("prune-host") now = datetime(2026, 7, 14, 12, 0, tzinfo=UTC) try: with Session(database.engine) as session, session.begin(): # Old terminal task (completed > window ago). session.add( TaskAttemptRow( task_id=old_terminal_task, attempt=1, lease_id="lease-old", host_id=host_id, device_id="device-old", status="done", lease_expires_at=(now - timedelta(days=10)).isoformat(), created_at=(now - timedelta(days=11)).isoformat(), completed_at=(now - timedelta(days=10)).isoformat(), failure_reason=None, result_json=None, ) ) # In-flight task (also old, but NOT terminal). session.add( TaskAttemptRow( task_id=in_flight_task, attempt=1, lease_id="lease-flight", host_id=host_id, device_id="device-flight", status="dispatched", lease_expires_at=(now - timedelta(days=10)).isoformat(), created_at=(now - timedelta(days=11)).isoformat(), completed_at=None, failure_reason=None, result_json=None, ) ) # Recently terminal task (within window). session.add( TaskAttemptRow( task_id=recent_terminal_task, attempt=1, lease_id="lease-recent", host_id=host_id, device_id="device-recent", status="done", lease_expires_at=(now - timedelta(hours=1)).isoformat(), created_at=(now - timedelta(hours=2)).isoformat(), completed_at=(now - timedelta(hours=1)).isoformat(), failure_reason=None, result_json=None, ) ) # Seed decision log rows for all three tasks. for task_id in [old_terminal_task, in_flight_task, recent_terminal_task]: database.repository.record_planner_decision( host_id=host_id, task_id=task_id, attempt=1, system_prompt="s", user_prompt="u", tool_name="tap", arguments_json="{}", now=now - timedelta(days=11), ) # 7-day window. deleted = database.repository.prune_planner_decision_log( now=now, prune_after_terminal_seconds=7 * 86_400, ) assert deleted == 1 # Old terminal task's rows are gone. assert ( database.repository.list_planner_decisions( task_id=old_terminal_task, attempt=1 ) == [] ) # In-flight and recent terminal rows survive. assert ( len( database.repository.list_planner_decisions( task_id=in_flight_task, attempt=1 ) ) == 1 ) assert ( len( database.repository.list_planner_decisions( task_id=recent_terminal_task, attempt=1 ) ) == 1 ) finally: database.close()