from __future__ import annotations import asyncio from datetime import UTC, datetime from cloud.internal_api.models import AssignmentModel, TerminalResultResponse from host_agent.assignment import AssignmentExecutionResult from host_agent.processor import AssignmentProcessor from host_agent.status import AgentStatusTracker def _assignment() -> AssignmentModel: return AssignmentModel( task_id="task-a", attempt=1, lease_id="lease-a", lease_expires_at=datetime(2026, 7, 12, tzinfo=UTC), host_id="host-a", device_id="device-a", goal="open settings", ) def test_processor_reports_success_with_execution_metadata() -> None: async def scenario() -> None: reports: list[dict[str, object]] = [] class SuccessfulExecutor: async def run(self, assignment): return AssignmentExecutionResult( status="done", failure_reason="ignored success detail", metadata={"runtime_task_id": "runtime-a"}, ) class RecordingClient: async def report_result(self, assignment, **kwargs): reports.append(kwargs) return TerminalResultResponse(status="recorded") result = await AssignmentProcessor( RecordingClient(), # type: ignore[arg-type] SuccessfulExecutor(), ).process(_assignment()) assert result.report_status == "recorded" assert reports == [ { "status": "done", "failure_reason": None, "result": {"runtime_task_id": "runtime-a"}, } ] asyncio.run(scenario()) def test_processor_preserves_runtime_failure_reason() -> None: async def scenario() -> None: reports: list[dict[str, object]] = [] class FailedExecutor: async def run(self, assignment): return AssignmentExecutionResult( status="failed", failure_reason="planner unavailable", metadata={"runtime_status": "failed"}, ) class RecordingClient: async def report_result(self, assignment, **kwargs): reports.append(kwargs) return TerminalResultResponse(status="already_recorded") result = await AssignmentProcessor( RecordingClient(), # type: ignore[arg-type] FailedExecutor(), ).process(_assignment()) assert result.report_status == "already_recorded" assert result.execution.failure_reason == "planner unavailable" assert reports[0] == { "status": "failed", "failure_reason": "planner unavailable", "result": {"runtime_status": "failed"}, } asyncio.run(scenario()) def test_processor_reports_cancelled_status_with_reason() -> None: async def scenario() -> None: reports: list[dict[str, object]] = [] class CancelledExecutor: async def run(self, assignment): return AssignmentExecutionResult( status="cancelled", failure_reason="cancellation requested by control plane", metadata={"runtime_status": "cancelled"}, ) class RecordingClient: async def report_result(self, assignment, **kwargs): reports.append(kwargs) return TerminalResultResponse(status="recorded") result = await AssignmentProcessor( RecordingClient(), # type: ignore[arg-type] CancelledExecutor(), ).process(_assignment()) assert result.report_status == "recorded" assert reports[0] == { "status": "cancelled", "failure_reason": "cancellation requested by control plane", "result": {"runtime_status": "cancelled"}, } asyncio.run(scenario()) def test_status_tracker_sees_started_then_finished_even_on_raise() -> None: async def scenario() -> None: tracker = AgentStatusTracker() snapshots: list[dict[str, object]] = [] class RaisingExecutor: async def run(self, assignment): snapshots.append(tracker.snapshot()) raise RuntimeError("executor exploded") class RecordingClient: async def report_result(self, assignment, **kwargs): return TerminalResultResponse(status="recorded") processor = AssignmentProcessor( RecordingClient(), # type: ignore[arg-type] RaisingExecutor(), status_tracker=tracker, ) try: await processor.process(_assignment()) except RuntimeError: pass assert snapshots[0]["current_assignment"] is not None assert snapshots[0]["current_assignment"]["task_id"] == "task-a" assert tracker.snapshot()["current_assignment"] is None asyncio.run(scenario()) def test_on_result_receives_assignment_and_result_and_swallows_exceptions() -> None: async def scenario() -> None: received: list[tuple[object, object]] = [] class SuccessfulExecutor: async def run(self, assignment): return AssignmentExecutionResult( status="done", failure_reason=None, metadata={}, ) class RecordingClient: async def report_result(self, assignment, **kwargs): return TerminalResultResponse(status="recorded") def on_result(assignment, result) -> None: received.append((assignment, result)) raise RuntimeError("history recording exploded") assignment = _assignment() processor = AssignmentProcessor( RecordingClient(), # type: ignore[arg-type] SuccessfulExecutor(), on_result=on_result, ) result = await processor.process(assignment) assert received == [(assignment, result)] asyncio.run(scenario())