from __future__ import annotations from datetime import UTC, datetime from types import SimpleNamespace from cloud.internal_api.models import AssignmentModel from core.models import Task from host_agent.assignment import AssignmentExecutor from host_agent.execution import ExecutionFactories def _assignment(**overrides) -> AssignmentModel: values = { "task_id": "cloud-task", "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", } values.update(overrides) return AssignmentModel(**values) def test_goal_assignment_executes_through_task_runner() -> None: received: list[Task] = [] created: list[tuple[str, str | None, int | None]] = [] class FakeTaskRunner: def run(self, task: Task) -> Task: received.append(task) task.status = "completed" return task class FakeMetadataStore: def create_task( self, task: Task, *, source_task_id: str | None = None, source_attempt: int | None = None, ) -> None: created.append((task.id, source_task_id, source_attempt)) factories = ExecutionFactories( task_runner_factory=lambda: FakeTaskRunner(), # type: ignore[arg-type,return-value] workflow_runner_factory=lambda: object(), # type: ignore[arg-type,return-value] workflow_store=object(), # type: ignore[arg-type] metadata_store=FakeMetadataStore(), # type: ignore[arg-type] ) result = AssignmentExecutor(factories).execute(_assignment()) assert result.status == "done" assert received[0].goal == "open settings" assert received[0].device_id == "device-a" assert result.metadata["runtime_task_id"] == received[0].id assert created == [(received[0].id, "cloud-task", 1)] def test_goal_assignment_preserves_runtime_failure_reason() -> None: class FakeTaskRunner: def run(self, task: Task) -> Task: task.status = "failed" task.failure_reason = "planner unavailable" return task factories = ExecutionFactories( task_runner_factory=lambda: FakeTaskRunner(), # type: ignore[arg-type,return-value] workflow_runner_factory=lambda: object(), # type: ignore[arg-type,return-value] workflow_store=object(), # type: ignore[arg-type] ) result = AssignmentExecutor(factories).execute(_assignment()) assert result.status == "failed" assert result.failure_reason == "planner unavailable" def test_goal_assignment_maps_cancellation_stop_to_cancelled_status() -> None: # First should_stop() call is Executor.execute()'s pre-flight check (must pass # through so the runner is actually invoked); the runner's own loop then stops. calls = {"count": 0} def should_stop() -> bool: calls["count"] += 1 return calls["count"] > 1 class FakeTaskRunner: def run(self, task: Task, *, should_stop=None, stop_reason=None) -> Task: assert should_stop is not None and should_stop() assert stop_reason is not None task.status = "cancelled" task.failure_reason = stop_reason() return task factories = ExecutionFactories( task_runner_factory=lambda: FakeTaskRunner(), # type: ignore[arg-type,return-value] workflow_runner_factory=lambda: object(), # type: ignore[arg-type,return-value] workflow_store=object(), # type: ignore[arg-type] ) result = AssignmentExecutor(factories).execute( _assignment(), should_stop=should_stop, stop_reason=lambda: "cancellation requested by control plane", ) assert result.status == "cancelled" assert result.failure_reason == "cancellation requested by control plane" def test_workflow_assignment_loads_and_executes_definition() -> None: definition = object() calls: list[tuple[object, str]] = [] class FakeWorkflowStore: def get_definition(self, definition_id: str): return definition if definition_id == "workflow-a" else None class FakeWorkflowRunner: def run(self, loaded_definition, device_id: str): calls.append((loaded_definition, device_id)) return SimpleNamespace(id="run-a", status="completed") factories = ExecutionFactories( task_runner_factory=lambda: object(), # type: ignore[arg-type,return-value] workflow_runner_factory=lambda: FakeWorkflowRunner(), # type: ignore[arg-type,return-value] workflow_store=FakeWorkflowStore(), # type: ignore[arg-type] ) result = AssignmentExecutor(factories).execute( _assignment(goal=None, workflow_definition_id="workflow-a") ) assert result.status == "done" assert calls == [(definition, "device-a")] assert result.metadata == { "workflow_run_id": "run-a", "workflow_status": "completed", } def test_workflow_assignment_maps_cancellation_stop_to_cancelled_status() -> None: calls = {"count": 0} def should_stop() -> bool: calls["count"] += 1 return calls["count"] > 1 class FakeWorkflowStore: def get_definition(self, definition_id: str): return object() if definition_id == "workflow-a" else None class FakeWorkflowRunner: def run( self, loaded_definition, device_id: str, *, should_stop=None, stop_reason=None, ): assert should_stop is not None and should_stop() assert stop_reason is not None return SimpleNamespace( id="run-a", status="cancelled", failure_reason=stop_reason() ) factories = ExecutionFactories( task_runner_factory=lambda: object(), # type: ignore[arg-type,return-value] workflow_runner_factory=lambda: FakeWorkflowRunner(), # type: ignore[arg-type,return-value] workflow_store=FakeWorkflowStore(), # type: ignore[arg-type] ) result = AssignmentExecutor(factories).execute( _assignment(goal=None, workflow_definition_id="workflow-a"), should_stop=should_stop, stop_reason=lambda: "cancellation requested by control plane", ) assert result.status == "cancelled" assert result.metadata == { "workflow_run_id": "run-a", "workflow_status": "cancelled", } def test_execute_fails_fast_when_mcp_session_holds_device() -> None: """Cloud assignment arriving for a device currently held by an MCP session must fail immediately rather than fight for the device.""" from host_agent.mcp_lock import McpBusyTracker tracker = McpBusyTracker() tracker.acquire("phone-1", "sess-mcp") executor = AssignmentExecutor( _build_factories(), mcp_busy_tracker=tracker, ) assignment = _assignment(device_id="phone-1") result = executor.execute(assignment) assert result.status == "failed" assert "MCP" in (result.failure_reason or "") def test_execute_skips_check_when_tracker_is_none() -> None: """Default backward-compat: no tracker → no fail-fast.""" executor = AssignmentExecutor(_build_factories()) # Without a real workflow store / task runner this test verifies the # entry-point path doesn't raise on the mcp_busy check. # We use a goal + a mock runner factory so execute() runs through. assignment = _assignment() result = executor.execute(assignment) # Should run through normally (not fail on MCP check) assert result.status == "done" def _build_factories() -> ExecutionFactories: """Shared factory fixture used by MCP-hold tests.""" received: list[Task] = [] class FakeTaskRunner: def run(self, task: Task) -> Task: received.append(task) task.status = "completed" return task class FakeMetadataStore: def create_task( self, task: Task, *, source_task_id: str | None = None, source_attempt: int | None = None, ) -> None: pass return ExecutionFactories( task_runner_factory=lambda: FakeTaskRunner(), # type: ignore[arg-type,return-value] workflow_runner_factory=lambda: object(), # type: ignore[arg-type,return-value] workflow_store=object(), # type: ignore[arg-type] metadata_store=FakeMetadataStore(), # type: ignore[arg-type] ) def test_unknown_workflow_fails_without_running() -> None: class FakeWorkflowStore: def get_definition(self, definition_id: str): return None factories = ExecutionFactories( task_runner_factory=lambda: object(), # type: ignore[arg-type,return-value] workflow_runner_factory=lambda: object(), # type: ignore[arg-type,return-value] workflow_store=FakeWorkflowStore(), # type: ignore[arg-type] ) result = AssignmentExecutor(factories).execute( _assignment(goal=None, workflow_definition_id="missing") ) assert result.status == "failed" assert result.failure_reason == "unknown workflow definition 'missing'"