from datetime import UTC, datetime from .compute_provider import ComputeProvider, ComputeResult from .errors import AppError from .models import ( Computer, ComputerMetric, ComputerNetworkPolicy, ComputerOperation, ComputerSnapshot, ComputerTemplate, ) COMPUTER_STATES = { "unconfigured", "pending", "provisioning", "starting", "running", "idle", "suspending", "suspended", "stopping", "stopped", "restarting", "recovering", "failed", "destroying", "destroyed", } TRANSITIONAL_STATES = { "pending", "provisioning", "starting", "suspending", "stopping", "restarting", "recovering", "destroying", } class ComputerProvisioningService: """The only control-plane component allowed to invoke a hypervisor gateway.""" def __init__(self, provider: ComputeProvider): self.provider = provider async def create_computer( self, computer: Computer, template: ComputerTemplate, network_policy: ComputerNetworkPolicy, operation: ComputerOperation, ) -> None: if not self.provider.configured or template.status != "available": self._blocked( computer, operation, "COMPUTE_PROVIDER_UNCONFIGURED", "A real provider and provider-backed OS template must be configured.", ) return computer.status = "provisioning" computer.compute_provider = self.provider.name operation.status = "running" operation.started_at = datetime.now(UTC) spec = { "computer_id": str(computer.id), "template_ref": template.provider_template_ref, "os_family": computer.os_family, "cpu_cores": computer.cpu_cores, "memory_mb": computer.memory_mb, "disk_gb": computer.disk_gb, "gpu": computer.gpu, "persistent_disk": computer.persistent_disk, "system_privilege": computer.system_privilege, "browser": computer.browser, "network_policy": self._network_spec(network_policy), } try: result = await self.provider.create(spec, idempotency_key=operation.idempotency_key) await self._verify_and_apply(computer, operation, result) except AppError as error: self._uncertain(computer, operation, error) async def lifecycle(self, computer: Computer, operation: ComputerOperation) -> None: if ( operation.operation == "destroy" and computer.status == "unconfigured" and not computer.provider_resource_id ): now = datetime.now(UTC) computer.status = "destroyed" computer.desired_state = "destroyed" computer.stopped_at = now computer.version += 1 operation.status = "succeeded" operation.started_at = now operation.completed_at = now operation.verification = { "observed_status": "destroyed", "provider_resource_absent": True, } return if not self.provider.configured or not computer.provider_resource_id: self._blocked( computer, operation, "COMPUTE_PROVIDER_UNCONFIGURED", "This computer has no configured virtual-machine provider resource.", ) return desired = { "start": ("starting", "running"), "stop": ("stopping", "stopped"), "restart": ("restarting", "running"), "pause": ("suspending", "suspended"), "resume": ("starting", "running"), "destroy": ("destroying", "destroyed"), }[operation.operation] computer.status, computer.desired_state = desired operation.status = "running" operation.started_at = datetime.now(UTC) method_name = "delete" if operation.operation == "destroy" else operation.operation try: method = getattr(self.provider, method_name) result = await method( computer.provider_resource_id, idempotency_key=operation.idempotency_key ) await self._verify_and_apply(computer, operation, result) except AppError as error: self._uncertain(computer, operation, error) async def create_snapshot( self, computer: Computer, snapshot: ComputerSnapshot, operation: ComputerOperation, ) -> None: if not self.provider.configured or not computer.provider_resource_id: snapshot.status = "failed" self._blocked( computer, operation, "COMPUTE_PROVIDER_UNCONFIGURED", "This computer has no configured virtual-machine provider resource.", ) return operation.status = "running" operation.started_at = datetime.now(UTC) try: result = await self.provider.snapshot( computer.provider_resource_id, {"name": snapshot.name, "includes_memory": snapshot.includes_memory}, idempotency_key=operation.idempotency_key, ) snapshot.provider_snapshot_id = result.provider_snapshot_id snapshot.status = "ready" if result.status == "ready" else "creating" operation.provider_request_id = result.provider_request_id operation.status = "succeeded" if snapshot.status == "ready" else "running" operation.verification = { "provider_snapshot_id": result.provider_snapshot_id, "observed_status": result.status, } if operation.status == "succeeded": operation.completed_at = datetime.now(UTC) except AppError as error: snapshot.status = "failed" self._uncertain(computer, operation, error) async def restore_snapshot( self, computer: Computer, snapshot: ComputerSnapshot, operation: ComputerOperation, ) -> None: if ( not self.provider.configured or not computer.provider_resource_id or not snapshot.provider_snapshot_id ): self._blocked( computer, operation, "COMPUTE_PROVIDER_UNCONFIGURED", "The computer or snapshot has no configured provider resource.", ) return computer.status = "recovering" snapshot.status = "restoring" operation.status = "running" operation.started_at = datetime.now(UTC) try: result = await self.provider.restore( computer.provider_resource_id, snapshot.provider_snapshot_id, idempotency_key=operation.idempotency_key, ) await self._verify_and_apply(computer, operation, result) snapshot.status = "ready" if operation.status == "succeeded" else "restoring" except AppError as error: snapshot.status = "failed" self._uncertain(computer, operation, error) async def resize_computer( self, computer: Computer, operation: ComputerOperation, spec: dict ) -> None: if not self.provider.configured or not computer.provider_resource_id: self._blocked( computer, operation, "COMPUTE_PROVIDER_UNCONFIGURED", "This computer has no configured virtual-machine provider resource.", ) return operation.status = "running" operation.started_at = datetime.now(UTC) try: result = await self.provider.resize( computer.provider_resource_id, spec, idempotency_key=operation.idempotency_key ) await self._verify_and_apply(computer, operation, result) if operation.status == "succeeded": computer.cpu_cores = spec["cpu_cores"] computer.memory_mb = spec["memory_mb"] computer.disk_gb = spec["disk_gb"] except AppError as error: self._uncertain(computer, operation, error) async def refresh(self, computer: Computer, operation: ComputerOperation) -> None: if not self.provider.configured or not computer.provider_resource_id: self._blocked( computer, operation, "COMPUTE_PROVIDER_UNCONFIGURED", "This computer has no configured virtual-machine provider resource.", ) return operation.status = "running" operation.started_at = datetime.now(UTC) try: result = await self.provider.get_status(computer.provider_resource_id) self._apply_result(computer, result) operation.provider_request_id = result.provider_request_id operation.status = "succeeded" operation.verification = {"observed_status": result.status} operation.completed_at = datetime.now(UTC) except AppError as error: self._uncertain(computer, operation, error) async def collect_metrics(self, computer: Computer) -> ComputerMetric: if not self.provider.configured or not computer.provider_resource_id: raise AppError( "COMPUTE_PROVIDER_UNCONFIGURED", "This computer has no configured virtual-machine provider resource.", 503, ) result = await self.provider.get_metrics(computer.provider_resource_id) metrics = result.metrics return ComputerMetric( workspace_id=computer.workspace_id, computer_id=computer.id, cpu_percent=int(metrics.get("cpu_percent", 0)), memory_used_mb=int(metrics.get("memory_used_mb", 0)), disk_used_gb=int(metrics.get("disk_used_gb", 0)), network_rx_bytes=int(metrics.get("network_rx_bytes", 0)), network_tx_bytes=int(metrics.get("network_tx_bytes", 0)), top_processes=list(metrics.get("top_processes", []))[:20], ) async def _verify_and_apply( self, computer: Computer, operation: ComputerOperation, result: ComputeResult ) -> None: computer.provider_resource_id = result.provider_resource_id operation.provider_request_id = result.provider_request_id observed = await self.provider.get_status(result.provider_resource_id) if observed.provider_resource_id != result.provider_resource_id: raise AppError( "COMPUTE_PROVIDER_VERIFICATION_FAILED", "The provider returned inconsistent computer identity evidence.", 502, ) self._apply_result(computer, observed) operation.verification = { "provider_resource_id": observed.provider_resource_id, "observed_status": observed.status, } operation.status = "running" if observed.status in TRANSITIONAL_STATES else "succeeded" if operation.status == "succeeded": operation.completed_at = datetime.now(UTC) @staticmethod def _apply_result(computer: Computer, result: ComputeResult) -> None: if result.status not in COMPUTER_STATES - {"unconfigured"}: raise AppError( "COMPUTE_PROVIDER_RESPONSE_INVALID", "The provider returned an unknown computer state.", 502, ) if computer.provider_resource_id and computer.provider_resource_id != result.provider_resource_id: raise AppError( "COMPUTE_PROVIDER_IDENTITY_MISMATCH", "The provider returned a different computer identity.", 502, ) now = datetime.now(UTC) computer.provider_resource_id = result.provider_resource_id computer.status = result.status computer.ip_address = result.ip_address computer.control_endpoint = result.control_endpoint computer.browser_status = result.browser_status computer.last_seen_at = now computer.provider_error_code = None computer.provider_error_message = None computer.version += 1 if result.status in {"running", "idle"}: computer.started_at = computer.started_at or now computer.stopped_at = None elif result.status in {"stopped", "destroyed"}: computer.stopped_at = now @staticmethod def _network_spec(policy: ComputerNetworkPolicy) -> dict: return { "internet_access": policy.internet_access, "lan_access": policy.lan_access, "platform_api_access": policy.platform_api_access, "other_agent_networks": policy.other_agent_networks, "host_management_access": False, "allowed_domains": policy.allowed_domains, "blocked_domains": policy.blocked_domains, } @staticmethod def _blocked( computer: Computer, operation: ComputerOperation, code: str, message: str ) -> None: computer.status = "unconfigured" computer.provider_error_code = code computer.provider_error_message = message computer.version += 1 operation.status = "blocked" operation.error_code = code operation.error_message = message operation.completed_at = datetime.now(UTC) @staticmethod def _uncertain( computer: Computer, operation: ComputerOperation, error: AppError ) -> None: computer.status = "recovering" computer.provider_error_code = error.code computer.provider_error_message = error.message computer.version += 1 operation.status = "unknown" operation.error_code = error.code operation.error_message = error.message operation.completed_at = datetime.now(UTC)