trycua--cua/libs/python/cua-sandbox/cua_sandbox/runtime/base.py

"""Abstract runtime — starts a VM or container from an Image spec and returns connection info."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Optional

from cua_sandbox.image import Image

if TYPE_CHECKING:
    pass


@dataclass
class CheckpointInfo:
    """Metadata for a saved checkpoint / snapshot of a sandbox."""

    name: str
    runtime_type: str
    created_at: float = 0.0
    metadata: dict = field(default_factory=dict)


@dataclass
class RuntimeInfo:
    """Connection info returned after a runtime spins up."""

    host: str
    api_port: int
    vnc_port: Optional[int] = None
    api_key: Optional[str] = None
    container_id: Optional[str] = None
    name: Optional[str] = None
    qmp_port: Optional[int] = None  # Set when QMP transport should be used
    environment: Optional[str] = None  # OS type hint for QMP transport
    agent_type: Optional[str] = None  # e.g. "osworld" for OSWorld Flask server
    guest_server_port: int = 8000  # Port the guest server listens on
    ssh_port: Optional[int] = None  # SSH port on the guest
    ssh_username: Optional[str] = None
    ssh_password: Optional[str] = None
    ssh_key_filename: Optional[str] = None
    vnc_host: Optional[str] = None  # VNC host (if different from host, e.g. localhost for Tart)
    vnc_password: Optional[str] = None
    grpc_port: Optional[int] = (
        None  # Set when gRPC transport should be used (emulator gRPC service)
    )


class Runtime(ABC):
    """Base class for local runtimes that create VMs/containers from an Image."""

    @abstractmethod
    async def start(self, image: Image, name: str, **opts) -> RuntimeInfo:
        """Start a VM/container and return connection info."""

    @abstractmethod
    async def stop(self, name: str) -> None:
        """Stop and clean up a VM/container."""

    @abstractmethod
    async def is_ready(self, info: RuntimeInfo, timeout: float = 120) -> bool:
        """Wait until the computer-server inside is reachable."""

    async def suspend(self, name: str) -> None:
        """Suspend (pause/save state of) a running VM."""
        raise NotImplementedError(f"{type(self).__name__} does not support suspend()")

    async def resume(self, image: "Image", name: str, **opts) -> RuntimeInfo:
        """Resume a suspended VM and return its RuntimeInfo."""
        raise NotImplementedError(f"{type(self).__name__} does not support resume()")

    async def list(self) -> list[dict]:
        """List known VMs managed by this runtime.

        Returns a list of dicts with at minimum: name, status.
        """
        raise NotImplementedError(f"{type(self).__name__} does not support list()")

    # ── Checkpoint / fork primitives ─────────────────────────────────────────

    async def ensure_base(self, image: "Image", base_name: str) -> CheckpointInfo:
        """Pull/build image into a stopped base sandbox if not already present.

        Idempotent. First call does the full pull; subsequent calls return
        immediately. The base sandbox is never started — it only serves as a
        fork source for ephemeral VMs.
        """
        raise NotImplementedError(f"{type(self).__name__} does not support ensure_base()")

    async def fork(self, base_name: str, new_name: str, **opts) -> None:
        """Create a new sandbox from a base/checkpoint using CoW where possible."""
        raise NotImplementedError(f"{type(self).__name__} does not support fork()")

    async def checkpoint(self, name: str, checkpoint_name: str, **opts) -> CheckpointInfo:
        """Capture the current state of a running sandbox as a named checkpoint."""
        raise NotImplementedError(f"{type(self).__name__} does not support checkpoint()")

    async def list_checkpoints(self) -> list[CheckpointInfo]:
        raise NotImplementedError(f"{type(self).__name__} does not support list_checkpoints()")

    async def delete_checkpoint(self, checkpoint_name: str) -> None:
        raise NotImplementedError(f"{type(self).__name__} does not support delete_checkpoint()")