mcp-virtual-computer 0.2.3__py3-none-any.whl

kilntainers/cli.py

@@ -0,0 +1,375 @@
+"""CLI argument parsing and main entry point."""
+
+import argparse
+import os
+import signal
+import sys
+import threading
+from typing import NoReturn
+
+from kilntainers.auth import BearerTokenMiddleware
+from kilntainers.backends import (
+    get_available_backend_names,
+    get_backend_class,
+)
+from kilntainers.computers import validate_computer_id
+from kilntainers.config import BackendConfig, ServerConfig, env_flag
+from kilntainers.errors import BackendError
+from kilntainers.server import create_server
+
+# Sentinel for detecting unset HTTP-only arguments
+_UNSET = object()
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the CLI argument parser.
+
+    Returns:
+        An ArgumentParser with all kilntainers arguments organized into groups.
+    """
+    parser = argparse.ArgumentParser(
+        prog="mcp-virtual-computer",
+        description=(
+            "MCP server providing one persistent, visually rendered local "
+            "Docker computer for LLM agents."
+        ),
+        usage="%(prog)s [-h] [--transport {stdio,http}] [...]",
+    )
+
+    # Get available backend names for choices
+    available_backends = get_available_backend_names()
+
+    # --- Core parameters ---
+    core = parser.add_argument_group("core options")
+    core.add_argument(
+        "--transport",
+        default="stdio",
+        choices=["stdio", "http"],
+        help="MCP transport (default: stdio)",
+    )
+    core.add_argument(
+        "--host",
+        default=_UNSET,
+        help="HTTP bind address (default: 127.0.0.1, HTTP mode only)",
+    )
+    core.add_argument(
+        "--port",
+        type=int,
+        default=_UNSET,
+        help="HTTP listen port (default: 8435, HTTP mode only)",
+    )
+    core.add_argument(
+        "--timeout",
+        type=int,
+        default=120,
+        help="Default exec timeout in seconds (default: 120)",
+    )
+    core.add_argument(
+        "--output-limit",
+        type=int,
+        default=2_097_152,
+        help="Max combined stdout+stderr bytes per exec (default: 2097152 = 2 MiB)",
+    )
+    core.add_argument(
+        "--session-timeout",
+        type=int,
+        default=_UNSET,
+        help="Idle session timeout in seconds (default: 300, HTTP mode only)",
+    )
+    core.add_argument(
+        "--auth-token",
+        default=os.getenv("KILNTAINERS_AUTH_TOKEN"),
+        help=(
+            "Static bearer token for the /mcp HTTP route "
+            "(default: KILNTAINERS_AUTH_TOKEN)"
+        ),
+    )
+    core.add_argument(
+        "--allow-unauthenticated-http",
+        action="store_true",
+        default=False,
+        help=(
+            "Allow a non-loopback HTTP listener without authentication. "
+            "Only use behind a trusted private network or auth proxy."
+        ),
+    )
+    core.add_argument(
+        "--shell",
+        default="/bin/bash",
+        help="Shell binary for command mode (e.g., /bin/bash, ash). Default: /bin/bash.",
+    )
+    core.add_argument(
+        "--network",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Enable network access in sandboxes (default: enabled)",
+    )
+
+    # --- Tool description ---
+    desc = parser.add_argument_group("tool description")
+    desc.add_argument(
+        "--tool-instruction-override",
+        default=None,
+        help="Replace the entire terminal_execute tool description",
+    )
+    desc.add_argument(
+        "--extended-tool-instruction",
+        default=None,
+        help="Append to the backend's default tool description",
+    )
+
+    # --- Backend-specific parameters (delegated to each backend) ---
+    for name in available_backends:
+        group = parser.add_argument_group(f"{name} backend options")
+        try:
+            backend_cls = get_backend_class(name)
+            backend_cls.add_cli_arguments(group)
+        except BackendError:
+            # Keep --help available if the local Docker adapter cannot initialize.
+            pass
+
+    return parser
+
+
+def build_configs(
+    args: argparse.Namespace,
+) -> tuple[ServerConfig, BackendConfig]:
+    """Build config dataclasses from parsed arguments.
+
+    This function maps flat CLI arguments to the typed config objects
+    consumed by the server and backend layers. Server config is built
+    here; backend config is delegated to the backend class.
+
+    Args:
+        args: Parsed command-line arguments from argparse.
+
+    Returns:
+        A tuple of (ServerConfig, BackendConfig).
+    """
+    # Handle HTTP-only args that may be _UNSET
+    host = "127.0.0.1" if args.host is _UNSET else args.host
+    port = 8435 if args.port is _UNSET else args.port
+    session_timeout = 300 if args.session_timeout is _UNSET else args.session_timeout
+
+    server_config = ServerConfig(
+        transport=args.transport,
+        host=host,
+        port=port,
+        default_timeout=args.timeout,
+        output_limit=args.output_limit,
+        computer_id=os.getenv("COMPUTER_ID", ""),
+        desktop_environment=env_flag("DESKTOP_ENVIRONMENT", default=False),
+        network_access=env_flag("NETWORK_ACCESS", default=args.network),
+        tool_instruction_override=args.tool_instruction_override,
+        extended_tool_instruction=args.extended_tool_instruction,
+        session_timeout=session_timeout,
+        auth_token=args.auth_token,
+        allow_unauthenticated_http=args.allow_unauthenticated_http,
+    )
+
+    # Delegate backend config construction to the backend class
+    backend_cls = get_backend_class("docker")
+    backend_config = backend_cls.config_from_args(args)
+
+    return server_config, backend_config
+
+
+def _startup_error(message: str) -> NoReturn:
+    """Write an error message to stderr and exit with code 1.
+
+    Used for all startup/configuration errors.
+
+    Args:
+        message: The error message to display.
+
+    Raises:
+        SystemExit: Always exits with code 1.
+    """
+    sys.stderr.write(f"mcp-virtual-computer: error: {message}\n")
+    sys.exit(1)
+
+
+def validate_config(server_config: ServerConfig) -> None:
+    """Validate configuration constraints that span multiple parameters.
+
+    Raises SystemExit with a descriptive message on failure.
+    Individual argument type validation is handled by argparse.
+    Cross-cutting constraints are checked here.
+
+    Args:
+        server_config: The server configuration to validate.
+
+    Raises:
+        SystemExit: If validation fails, with code 1 and an error message.
+    """
+    # HTTP-only parameters in stdio mode
+    # When called from main(), _parsed_args contains the actual parsed args
+    # When called from tests, we rely on comparing config values to defaults
+    if server_config.transport == "stdio":
+        # Check if user explicitly passed HTTP-only args
+        # We check if values differ from defaults, indicating explicit setting
+        if server_config.host != "127.0.0.1":
+            _startup_error(
+                "--host is only valid with --transport http. "
+                "In stdio mode, there is no HTTP server to bind."
+            )
+        if server_config.port != 8435:
+            _startup_error(
+                "--port is only valid with --transport http. "
+                "In stdio mode, there is no HTTP server to bind."
+            )
+        if server_config.session_timeout != 300:
+            _startup_error(
+                "--session-timeout is only valid with --transport http. "
+                "In stdio mode, the session lives as long as the process."
+            )
+
+    # Mutual exclusivity: tool description params
+    if (
+        server_config.tool_instruction_override is not None
+        and server_config.extended_tool_instruction is not None
+    ):
+        _startup_error(
+            "Cannot use both --tool-instruction-override and "
+            "--extended-tool-instruction. Use override to replace "
+            "the description entirely, or extended to append to "
+            "the backend default."
+        )
+
+    # Timeout must be positive
+    if server_config.default_timeout < 1:
+        _startup_error("--timeout must be at least 1 second.")
+
+    # Output limit must be positive
+    if server_config.output_limit < 1:
+        _startup_error("--output-limit must be at least 1 byte.")
+
+    if not server_config.computer_id:
+        _startup_error("COMPUTER_ID is required (example: COMPUTER_ID=agent-workstation).")
+    try:
+        validate_computer_id(server_config.computer_id)
+    except BackendError as error:
+        _startup_error(str(error).replace("computer_id", "COMPUTER_ID"))
+
+    if (
+        server_config.transport == "http"
+        and server_config.host not in {"127.0.0.1", "localhost", "::1"}
+        and not server_config.auth_token
+        and not server_config.allow_unauthenticated_http
+    ):
+        _startup_error(
+            "A non-loopback HTTP listener can execute arbitrary sandbox commands. "
+            "Set KILNTAINERS_AUTH_TOKEN/--auth-token, or explicitly pass "
+            "--allow-unauthenticated-http behind a trusted private network."
+        )
+
+
+async def _async_main(
+    server_config: ServerConfig,
+    backend_config: BackendConfig,
+    backend_name: str,
+) -> None:
+    """Async startup: build server, run.
+
+    This function performs all async startup operations:
+    - Creates the backend (validation happens lazily)
+    - Creates the MCP server
+    - Runs the transport (blocking until shutdown)
+
+    Args:
+        server_config: Server configuration.
+        backend_config: Backend configuration.
+        backend_name: Name of the backend to use.
+
+    Raises:
+        SystemExit: If server creation fails.
+    """
+    # Create backend (validation happens lazily on first terminal_execute)
+    backend_class = get_backend_class(backend_name)
+    backend = backend_class(backend_config)
+
+    # Create the MCP server (assembles tool description, registers tool)
+    try:
+        mcp = create_server(backend, server_config)
+    except BackendError as e:
+        _startup_error(str(e))
+
+    # Run the transport (blocks until shutdown)
+    transport = "stdio" if server_config.transport == "stdio" else "streamable-http"
+    mcp.run(transport=transport)
+
+
+def main() -> None:
+    """CLI entry point. Parses args, configures, and runs the server.
+
+    This is the main entry point for the kilntainers command. It:
+    1. Parses CLI arguments
+    2. Builds configuration objects
+    3. Validates configuration constraints
+    4. Creates the backend (validation happens lazily on first exec)
+    5. Creates and runs the MCP server
+
+    Never returns normally (exits on KeyboardInterrupt or server shutdown).
+    """
+    parser = build_parser()
+    args = parser.parse_args()
+
+    server_config, backend_config = build_configs(args)
+    validate_config(server_config)
+
+    # Create backend (validation happens lazily on first terminal_execute)
+    backend_name = "docker"
+    backend_class = get_backend_class(backend_name)
+    backend_class.prepare_runtime()
+    backend = backend_class(backend_config)
+
+    # Create the MCP server (assembles tool description, registers tool)
+    try:
+        mcp = create_server(backend, server_config)
+    except BackendError as e:
+        _startup_error(str(e))
+
+    # Run the transport (blocks until shutdown)
+    # mcp.run() manages its own event loop, so we call it directly
+    transport = "stdio" if server_config.transport == "stdio" else "streamable-http"
+
+    # Register SIGTERM handler to convert to SIGINT for clean shutdown
+    # FastMCP handles SIGINT (Ctrl+C) gracefully, so we redirect SIGTERM to the same path
+    def _handle_sigterm(signum: int, frame: object) -> None:
+        """
+        Convert SIGTERM to SIGINT for clean shutdown (triggers mcp library graceful shutdown).
+
+        Watchdog timmer needed as `mcp` library doesn't exit on SIGTERM.
+
+        `mcp` library is adding sigterm support, but not in a release yet.
+        """
+        # Schedule forced exit as fallback in case graceful shutdown hangs.
+        # Uses a daemon thread so it won't block normal exit if shutdown succeeds.
+        timer = threading.Timer(5.0, lambda: os._exit(0))
+        timer.daemon = True
+        timer.start()
+        os.kill(os.getpid(), signal.SIGINT)
+
+    signal.signal(signal.SIGTERM, _handle_sigterm)
+    if hasattr(signal, "SIGBREAK"):
+        signal.signal(signal.SIGBREAK, _handle_sigterm)
+
+    try:
+        if server_config.transport == "http" and server_config.auth_token:
+            import uvicorn
+
+            app = mcp.streamable_http_app()
+            app.add_middleware(
+                BearerTokenMiddleware,  # ty: ignore[invalid-argument-type]
+                token=server_config.auth_token,
+            )
+            uvicorn.run(
+                app,
+                host=server_config.host,
+                port=server_config.port,
+                log_level="info",
+            )
+        else:
+            mcp.run(transport=transport)
+    except KeyboardInterrupt:
+        pass  # Clean exit on Ctrl+C

kilntainers/computers.py

@@ -0,0 +1,284 @@
+"""Named computer registry shared by MCP sessions and dashboard tools."""
+
+import asyncio
+import re
+import secrets
+from dataclasses import dataclass
+
+from kilntainers.backends.base import Backend, ComputerInfo, Sandbox
+from kilntainers.errors import BackendError
+
+_COMPUTER_ID_RE = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$")
+_ADJECTIVES = (
+    "amber",
+    "brisk",
+    "calm",
+    "clever",
+    "coral",
+    "crisp",
+    "gentle",
+    "lucky",
+    "quiet",
+    "rapid",
+    "silver",
+    "steady",
+)
+_NOUNS = (
+    "badger",
+    "comet",
+    "falcon",
+    "gecko",
+    "heron",
+    "lynx",
+    "otter",
+    "panda",
+    "raven",
+    "tiger",
+    "whale",
+    "wolf",
+)
+
+
+def random_computer_id() -> str:
+    """Generate a readable, collision-resistant provider-safe slug."""
+    return (
+        f"{secrets.choice(_ADJECTIVES)}-{secrets.choice(_NOUNS)}-{secrets.token_hex(2)}"
+    )
+
+
+def validate_computer_id(computer_id: str) -> str:
+    """Validate and return a stable computer ID.
+
+    IDs intentionally use a conservative lowercase Docker-safe slug format.
+    """
+    value = computer_id.strip()
+    if not _COMPUTER_ID_RE.fullmatch(value):
+        raise BackendError(
+            "computer_id must be 1-63 lowercase letters, numbers, or hyphens; "
+            "it must start and end with a letter or number."
+        )
+    return value
+
+
+@dataclass(slots=True)
+class _ComputerRecord:
+    sandbox: Sandbox
+    temporary: bool
+    owners: int = 0
+
+
+class ComputerRegistry:
+    """Coordinate named sandboxes across MCP sessions in one server process.
+
+    The Docker backend also discovers computers provider-side, so permanent
+    records survive a server restart and can be reattached by ID.
+    """
+
+    def __init__(self, backend: Backend) -> None:
+        self.backend = backend
+        self._records: dict[str, _ComputerRecord] = {}
+        self._lock = asyncio.Lock()
+
+    @staticmethod
+    def _tag(sandbox: Sandbox, computer_id: str, temporary: bool) -> Sandbox:
+        """Apply registry metadata used by legacy Sandbox base properties."""
+        setattr(sandbox, "_managed_computer_id", computer_id)
+        setattr(sandbox, "_managed_temporary", temporary)
+        return sandbox
+
+    async def acquire(
+        self,
+        computer_id: str | None,
+        *,
+        temporary: bool,
+        add_owner: bool,
+    ) -> tuple[str, Sandbox]:
+        """Attach to or create a computer and optionally add an owner ref."""
+        requested_id = (
+            random_computer_id()
+            if computer_id is None
+            else validate_computer_id(computer_id)
+        )
+
+        async with self._lock:
+            record = self._records.get(requested_id)
+            if record is not None:
+                if record.temporary != temporary:
+                    raise BackendError(
+                        f"Computer '{requested_id}' already exists with "
+                        f"temporary={str(record.temporary).lower()}; lifecycle "
+                        "mode cannot be changed without a factory reset or delete."
+                    )
+                if add_owner:
+                    record.owners += 1
+                return requested_id, record.sandbox
+
+            sandbox = await self.backend.attach_sandbox(requested_id)
+            if sandbox is not None:
+                actual_temporary = sandbox.temporary
+                if actual_temporary != temporary:
+                    raise BackendError(
+                        f"Computer '{requested_id}' already exists with "
+                        f"temporary={str(actual_temporary).lower()}; requested "
+                        f"temporary={str(temporary).lower()}."
+                    )
+            else:
+                sandbox = await self.backend.create_sandbox(
+                    computer_id=requested_id,
+                    temporary=temporary,
+                )
+                # Legacy backends keep their original Sandbox interface. These
+                # attributes let the base properties expose registry semantics
+                # without wrapping the instance or breaking backend-specific APIs.
+                self._tag(sandbox, requested_id, temporary)
+                actual_temporary = temporary
+
+            self._records[requested_id] = _ComputerRecord(
+                sandbox=sandbox,
+                temporary=actual_temporary,
+                owners=1 if add_owner else 0,
+            )
+            return requested_id, sandbox
+
+    async def get_owned(self, computer_id: str) -> Sandbox | None:
+        """Return a locally attached sandbox without changing owner refs."""
+        async with self._lock:
+            record = self._records.get(computer_id)
+            return record.sandbox if record is not None else None
+
+    def peek(self, computer_id: str) -> Sandbox | None:
+        """Return a local sandbox for synchronous status properties."""
+        record = self._records.get(computer_id)
+        return record.sandbox if record is not None else None
+
+    async def release(self, computer_id: str) -> None:
+        """Release one owner and remove an unowned temporary computer."""
+        sandbox: Sandbox | None = None
+        async with self._lock:
+            record = self._records.get(computer_id)
+            if record is None:
+                return
+            record.owners = max(0, record.owners - 1)
+            if record.temporary and record.owners == 0:
+                sandbox = record.sandbox
+                del self._records[computer_id]
+
+        if sandbox is not None:
+            await sandbox.stop()
+
+    async def list(self) -> list[ComputerInfo]:
+        """Return a de-duplicated provider and in-process inventory."""
+        provider_items = await self.backend.list_computers()
+        by_id = {item.computer_id: item for item in provider_items}
+
+        async with self._lock:
+            for computer_id, record in self._records.items():
+                by_id.setdefault(
+                    computer_id,
+                    ComputerInfo(
+                        computer_id=computer_id,
+                        sandbox_id=record.sandbox.sandbox_id,
+                        backend=self.backend.__class__.__name__,
+                        state="running",
+                        temporary=record.temporary,
+                    ),
+                )
+        return sorted(by_id.values(), key=lambda item: item.computer_id)
+
+    async def restart(self, computer_id: str) -> Sandbox:
+        """Restart a computer while preserving its filesystem."""
+        computer_id = validate_computer_id(computer_id)
+        async with self._lock:
+            record = self._records.get(computer_id)
+            replacement = await self.backend.restart_computer(computer_id)
+            if replacement is None:
+                if record is None:
+                    raise BackendError(f"Computer '{computer_id}' was not found.")
+                await record.sandbox.stop()
+                replacement = await self.backend.create_sandbox(
+                    computer_id=computer_id,
+                    temporary=record.temporary,
+                )
+                self._tag(replacement, computer_id, record.temporary)
+            if record is None:
+                record = _ComputerRecord(
+                    replacement,
+                    replacement.temporary,
+                    owners=0,
+                )
+                self._records[computer_id] = record
+            else:
+                record.sandbox = replacement
+            return replacement
+
+    async def factory_reset(self, computer_id: str) -> Sandbox:
+        """Delete a computer's writable state and recreate its base image."""
+        computer_id = validate_computer_id(computer_id)
+        async with self._lock:
+            record = self._records.get(computer_id)
+            replacement = await self.backend.factory_reset_computer(computer_id)
+            if replacement is None:
+                if record is None:
+                    raise BackendError(f"Computer '{computer_id}' was not found.")
+                await record.sandbox.stop()
+                replacement = await self.backend.create_sandbox(
+                    computer_id=computer_id,
+                    temporary=record.temporary,
+                )
+                self._tag(replacement, computer_id, record.temporary)
+            if record is None:
+                self._records[computer_id] = _ComputerRecord(
+                    replacement,
+                    replacement.temporary,
+                    owners=0,
+                )
+            else:
+                record.sandbox = replacement
+            return replacement
+
+    async def set_network_access(self, computer_id: str, enabled: bool) -> Sandbox:
+        """Change real network access and refresh the attached sandbox handle."""
+        computer_id = validate_computer_id(computer_id)
+        async with self._lock:
+            record = self._records.get(computer_id)
+            if record is None:
+                raise BackendError(f"Computer '{computer_id}' was not found.")
+            replacement = await self.backend.set_network_access(computer_id, enabled)
+            if replacement is None:
+                raise BackendError("This backend cannot change network access at runtime.")
+            self._tag(replacement, computer_id, record.temporary)
+            record.sandbox = replacement
+            return replacement
+
+    async def switch_desktop_environment(
+        self,
+        computer_id: str,
+        enabled: bool,
+    ) -> Sandbox:
+        """Switch virtual/Xfce mode while preserving the whole computer."""
+        computer_id = validate_computer_id(computer_id)
+        async with self._lock:
+            record = self._records.get(computer_id)
+            if record is None:
+                raise BackendError(f"Computer '{computer_id}' was not found.")
+            updated = await self.backend.switch_desktop_environment(
+                computer_id,
+                enabled,
+            )
+            if updated is None:
+                raise BackendError("This backend cannot switch desktop mode at runtime.")
+            self._tag(updated, computer_id, record.temporary)
+            record.sandbox = updated
+            return updated
+
+    async def delete(self, computer_id: str) -> None:
+        """Permanently delete a managed computer."""
+        computer_id = validate_computer_id(computer_id)
+        async with self._lock:
+            record = self._records.pop(computer_id, None)
+            deleted = await self.backend.delete_computer(computer_id)
+            if deleted:
+                return
+            if record is None:
+                raise BackendError(f"Computer '{computer_id}' was not found.")
+            await record.sandbox.stop()