From 0a9392b7ead41def824314932984fec393e50ac8 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Sun, 12 Jul 2026 17:15:09 +0800 Subject: [PATCH] feat(cloud-store): persist task lease history --- .../cloud-control-plane-integration/tasks.md | 2 +- packages/cloud-platform/cloud/repository.py | 3 +- packages/cloud-platform/cloud/scheduler.py | 12 ++- .../cloud-platform/cloud/sql_repository.py | 63 +++++++++++++- tests/test_cloud_repository_contract.py | 82 +++++++++++++++++++ 5 files changed, 155 insertions(+), 7 deletions(-) diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index b7bd39d..a526e3a 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -17,7 +17,7 @@ ## 3. Lease-Backed Scheduling -- [ ] 3.1 Extend scheduled-task persistence with attempt count, lease id/expiry, terminal result, failure reason, and auditable attempt records. +- [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. - [ ] 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. diff --git a/packages/cloud-platform/cloud/repository.py b/packages/cloud-platform/cloud/repository.py index f02b1f3..cce943b 100644 --- a/packages/cloud-platform/cloud/repository.py +++ b/packages/cloud-platform/cloud/repository.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass from datetime import datetime -from typing import TYPE_CHECKING, Literal, Protocol +from typing import TYPE_CHECKING, Any, Literal, Protocol if TYPE_CHECKING: from cloud.plugins import PluginManifest @@ -27,6 +27,7 @@ class TaskAttemptRecord: created_at: datetime completed_at: datetime | None = None failure_reason: str | None = None + terminal_result: dict[str, Any] | None = None @dataclass(frozen=True) diff --git a/packages/cloud-platform/cloud/scheduler.py b/packages/cloud-platform/cloud/scheduler.py index f72e162..be4d63d 100644 --- a/packages/cloud-platform/cloud/scheduler.py +++ b/packages/cloud-platform/cloud/scheduler.py @@ -10,7 +10,7 @@ from __future__ import annotations from dataclasses import dataclass, field from datetime import datetime -from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable from uuid import uuid4 from core.models import utc_now @@ -44,6 +44,12 @@ class ScheduledTask: status: ScheduledTaskStatus = "queued" assigned_device_id: str | None = None assigned_host_id: str | None = None + attempt_count: int = 0 + lease_id: str | None = None + lease_expires_at: datetime | None = None + terminal_result: dict[str, Any] | None = None + failure_reason: str | None = None + updated_at: datetime | None = None created_at: datetime = field(default_factory=utc_now) @@ -107,7 +113,9 @@ class TaskScheduler: self.pool = pool self.store = store self.config = config - self._strategies = dict(strategies) if strategies is not None else dict(DEFAULT_STRATEGIES) + self._strategies = ( + dict(strategies) if strategies is not None else dict(DEFAULT_STRATEGIES) + ) if config.default_assignment_strategy not in self._strategies: raise UnknownAssignmentStrategyError( f"unknown assignment strategy {config.default_assignment_strategy!r}; " diff --git a/packages/cloud-platform/cloud/sql_repository.py b/packages/cloud-platform/cloud/sql_repository.py index e8736f9..66b96cd 100644 --- a/packages/cloud-platform/cloud/sql_repository.py +++ b/packages/cloud-platform/cloud/sql_repository.py @@ -6,7 +6,7 @@ from datetime import datetime from typing import Any from sqlalchemy import Engine, delete, func, select -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import sessionmaker from cloud.db_models import ( Base, @@ -14,6 +14,7 @@ from cloud.db_models import ( PluginRow, PooledDeviceRow, ScheduledTaskRow, + TaskAttemptRow, ) from core.models import utc_now @@ -119,6 +120,18 @@ class SQLAlchemyCloudRepository: status=task.status, assigned_device_id=task.assigned_device_id, assigned_host_id=task.assigned_host_id, + attempt_count=task.attempt_count, + lease_id=task.lease_id, + lease_expires_at=( + _iso(task.lease_expires_at) if task.lease_expires_at else None + ), + failure_reason=task.failure_reason, + result_json=( + json.dumps(task.terminal_result, ensure_ascii=False) + if task.terminal_result is not None + else None + ), + updated_at=_iso(task.updated_at) if task.updated_at else None, created_at=_iso(task.created_at), ) ) @@ -194,6 +207,15 @@ class SQLAlchemyCloudRepository: row = session.get(PluginRow, name) return _plugin_from_row(row) if row else None + def list_task_attempts(self, task_id: str) -> list[Any]: + with self._sessions() as session: + rows = session.scalars( + select(TaskAttemptRow) + .where(TaskAttemptRow.task_id == task_id) + .order_by(TaskAttemptRow.attempt) + ).all() + return [_task_attempt_from_row(row) for row in rows] + def health_check(self) -> None: with self._sessions() as session: session.execute(select(1)) @@ -230,7 +252,7 @@ def _device_from_row(row: PooledDeviceRow) -> Any: try: tags = list(json.loads(row.capability_tags_json)) - except (TypeError, ValueError): + except TypeError, ValueError: tags = [] return PooledDevice( device_id=row.device_id, @@ -247,8 +269,9 @@ def _task_from_row(row: ScheduledTaskRow) -> Any: try: constraints_data = json.loads(row.constraints_json) - except (TypeError, ValueError): + except TypeError, ValueError: constraints_data = {} + terminal_result = _parse_json_object(row.result_json) return ScheduledTask( id=row.id, goal=row.goal, @@ -260,10 +283,44 @@ def _task_from_row(row: ScheduledTaskRow) -> Any: status=row.status, assigned_device_id=row.assigned_device_id, assigned_host_id=row.assigned_host_id, + attempt_count=row.attempt_count, + lease_id=row.lease_id, + lease_expires_at=_parse_dt(row.lease_expires_at), + terminal_result=terminal_result, + failure_reason=row.failure_reason, + updated_at=_parse_dt(row.updated_at), created_at=_parse_dt(row.created_at) or utc_now(), ) +def _task_attempt_from_row(row: TaskAttemptRow) -> Any: + from cloud.repository import TaskAttemptRecord + + return TaskAttemptRecord( + task_id=row.task_id, + attempt=row.attempt, + lease_id=row.lease_id, + host_id=row.host_id, + device_id=row.device_id, + status=row.status, + lease_expires_at=_parse_dt(row.lease_expires_at) or utc_now(), + created_at=_parse_dt(row.created_at) or utc_now(), + completed_at=_parse_dt(row.completed_at), + failure_reason=row.failure_reason, + terminal_result=_parse_json_object(row.result_json), + ) + + +def _parse_json_object(value: str | None) -> dict[str, Any] | None: + if value is None: + return None + try: + parsed = json.loads(value) + except TypeError, ValueError: + return None + return parsed if isinstance(parsed, dict) else None + + def _plugin_from_row(row: PluginRow) -> tuple[Any, bool]: from cloud.plugins import PluginManifest diff --git a/tests/test_cloud_repository_contract.py b/tests/test_cloud_repository_contract.py index cacd6ea..09de4bc 100644 --- a/tests/test_cloud_repository_contract.py +++ b/tests/test_cloud_repository_contract.py @@ -8,8 +8,10 @@ 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, LeasedAssignment, TaskAttemptRecord @@ -196,3 +198,83 @@ def test_repository_state_survives_application_restart(database_url: str) -> Non 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()