From 23a2dfb3307f2e6c46b44f74b3465d8de7a9f9cb Mon Sep 17 00:00:00 2001 From: Thales Maciel Date: Sat, 7 Mar 2026 16:28:28 -0300 Subject: [PATCH] Unify public UX around pyro CLI and Pyro facade --- AGENTS.md | 12 +- Makefile | 38 ++-- README.md | 264 ++++++++++++++------------ pyproject.toml | 8 +- src/pyro_mcp/__init__.py | 3 +- src/pyro_mcp/api.py | 179 +++++++++++++++++ src/pyro_mcp/cli.py | 103 ++++++++++ src/pyro_mcp/demo.py | 25 ++- src/pyro_mcp/ollama_demo.py | 102 +++++++--- src/pyro_mcp/runtime_network_check.py | 35 ++-- src/pyro_mcp/server.py | 63 +----- src/pyro_mcp/vm_manager.py | 48 ++++- tests/test_api.py | 85 +++++++++ tests/test_cli.py | 89 +++++++++ tests/test_demo.py | 70 ++++--- tests/test_ollama_demo.py | 137 ++++++------- tests/test_runtime_network_check.py | 40 ++-- tests/test_server.py | 21 +- tests/test_vm_manager.py | 21 +- 19 files changed, 936 insertions(+), 407 deletions(-) create mode 100644 src/pyro_mcp/api.py create mode 100644 src/pyro_mcp/cli.py create mode 100644 tests/test_api.py create mode 100644 tests/test_cli.py diff --git a/AGENTS.md b/AGENTS.md index fc3e694..a6ed039 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -11,6 +11,8 @@ This repository ships `pyro-mcp`, an MCP-compatible package for ephemeral VM lif - Use `uv` for all Python environment and command execution. - Run `make setup` after cloning. - Run `make check` before opening a PR. +- Public user-facing CLI is `pyro`. +- Public Python SDK entrypoint is `from pyro_mcp import Pyro`. - Use `make runtime-bundle` to regenerate the packaged runtime bundle from `runtime_sources/`. - Use `make runtime-materialize` to build real runtime inputs into `build/runtime_sources/`. - Use `make runtime-fetch-binaries`, `make runtime-build-kernel-real`, and `make runtime-build-rootfs-real` if you need to debug the real-source pipeline step by step. @@ -33,11 +35,13 @@ These checks run in pre-commit hooks and should all pass locally. ## Key API Contract -- Public factory: `pyro_mcp.create_server()` -- Runtime diagnostics CLI: `pyro-mcp-doctor` -- Runtime bundle build CLI: `pyro-mcp-runtime-build` +- Public SDK facade: `pyro_mcp.Pyro` +- Public MCP factory: `pyro_mcp.create_server()` +- Public CLI: `pyro` - Current bundled runtime is guest-capable for VM boot, guest exec, and guest networking; check `make doctor` for runtime capabilities. -- Lifecycle tools: +- Primary tool: + - `vm_run` +- Advanced lifecycle tools: - `vm_list_profiles` - `vm_create` - `vm_start` diff --git a/Makefile b/Makefile index cc21003..75e9582 100644 --- a/Makefile +++ b/Makefile @@ -60,66 +60,66 @@ test: check: lint typecheck test demo: - uv run python examples/static_tool_demo.py + uv run pyro demo network-demo: - PYRO_VM_ENABLE_NETWORK=1 uv run python examples/static_tool_demo.py + uv run pyro demo --network doctor: - uv run pyro-mcp-doctor + uv run pyro doctor ollama: ollama-demo ollama-demo: - PYRO_VM_ENABLE_NETWORK=1 uv run pyro-mcp-ollama-demo --base-url "$(OLLAMA_BASE_URL)" --model "$(OLLAMA_MODEL)" $(OLLAMA_DEMO_FLAGS) + uv run pyro demo ollama --base-url "$(OLLAMA_BASE_URL)" --model "$(OLLAMA_MODEL)" $(OLLAMA_DEMO_FLAGS) run-server: - uv run pyro-mcp-server + uv run pyro mcp serve install-hooks: uv run pre-commit install runtime-binaries: - uv run pyro-mcp-runtime-build stage-binaries --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build stage-binaries --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-kernel: - uv run pyro-mcp-runtime-build stage-kernel --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build stage-kernel --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-rootfs: - uv run pyro-mcp-runtime-build stage-rootfs --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build stage-rootfs --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-agent: - uv run pyro-mcp-runtime-build stage-agent --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build stage-agent --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-validate: - uv run pyro-mcp-runtime-build validate --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build validate --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-manifest: - uv run pyro-mcp-runtime-build manifest --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build manifest --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-sync: - uv run pyro-mcp-runtime-build sync --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build sync --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-bundle: - uv run pyro-mcp-runtime-build bundle --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build bundle --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-fetch-binaries: - uv run pyro-mcp-runtime-build fetch-binaries --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build fetch-binaries --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-build-kernel-real: - uv run pyro-mcp-runtime-build build-kernel --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build build-kernel --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-build-rootfs-real: - uv run pyro-mcp-runtime-build build-rootfs --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build build-rootfs --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-materialize: - uv run pyro-mcp-runtime-build materialize --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" + uv run python -m pyro_mcp.runtime_build materialize --platform "$(RUNTIME_PLATFORM)" --source-dir "$(RUNTIME_SOURCE_DIR)" --build-dir "$(RUNTIME_BUILD_DIR)" --bundle-dir "$(RUNTIME_BUNDLE_DIR)" --materialized-dir "$(RUNTIME_MATERIALIZED_DIR)" runtime-boot-check: - uv run pyro-mcp-runtime-boot-check + uv run python -m pyro_mcp.runtime_boot_check runtime-network-check: - uv run pyro-mcp-runtime-network-check + uv run python -m pyro_mcp.runtime_network_check runtime-clean: rm -rf "$(RUNTIME_BUILD_DIR)" "$(RUNTIME_MATERIALIZED_DIR)" diff --git a/README.md b/README.md index 80c1cbf..d81e559 100644 --- a/README.md +++ b/README.md @@ -1,170 +1,200 @@ # pyro-mcp -`pyro-mcp` is an MCP-compatible tool package for running ephemeral development environments with a VM lifecycle API. +`pyro-mcp` is a Firecracker-backed sandbox for coding agents. -## v0.1.0 Capabilities +It exposes the same runtime in two public forms: -- Split lifecycle tools for coding agents: `vm_list_profiles`, `vm_create`, `vm_start`, `vm_exec`, `vm_stop`, `vm_delete`, `vm_status`, `vm_network_info`, `vm_reap_expired`. -- Standard environment profiles: -- `debian-base`: minimal Debian shell/core Unix tools. -- `debian-git`: Debian base with Git preinstalled. -- `debian-build`: Debian Git profile with common build tooling. -- Explicit sizing contract for agents (`vcpu_count`, `mem_mib`) with guardrails. -- Strict ephemerality for command execution (`vm_exec` auto-deletes VM on completion). -- Ollama demo that asks an LLM to clone a small public Git repository through lifecycle tools. +- a `pyro` CLI +- a Python SDK via `from pyro_mcp import Pyro` -## Runtime +It also ships an MCP server so LLM clients can use the same VM runtime through tools. -The package includes a bundled Linux x86_64 runtime payload: -- Firecracker binary -- Jailer binary -- Profile artifacts for `debian-base`, `debian-git`, and `debian-build` +## Public UX -No system Firecracker installation is required for basic usage. +Primary install/run path: -Current status: -- The bundled runtime is real, not shim-based. -- `doctor` reports real guest capability flags for VM boot, guest exec, and guest networking. -- `vm_exec` now runs in `guest_vsock` mode when the VM is started from the bundled runtime. -- Networking still requires host privileges for TAP/NAT setup; see the networking section below. +```bash +uvx --from pyro-mcp pyro mcp serve +``` -Host requirements still apply: -- Linux host -- `/dev/kvm` available for full virtualization mode +Installed package path: + +```bash +pyro mcp serve +``` + +The public user-facing interface is `pyro` and `Pyro`. +`Makefile` targets are contributor conveniences for this repository and are not the primary product UX. + +## Capabilities + +- Firecracker microVM execution with bundled runtime artifacts +- standard profiles: + - `debian-base` + - `debian-git` + - `debian-build` +- high-level one-shot execution via `vm_run` / `Pyro.run_in_vm(...)` +- low-level lifecycle control when needed: + - `vm_create` + - `vm_start` + - `vm_exec` + - `vm_stop` + - `vm_delete` + - `vm_status` + - `vm_network_info` + - `vm_reap_expired` +- outbound guest networking with explicit opt-in ## Requirements +- Linux host +- `/dev/kvm` - Python 3.12+ -- `uv` -- Optional for Ollama demo: local Ollama server and `llama:3.2-3b` model. +- host privilege for TAP/NAT setup when using guest networking -## Setup +The current implementation uses `sudo -n` for `ip`, `nft`, and `iptables` when networked runs are requested. + +## CLI + +Start the MCP server: ```bash -make setup +pyro mcp serve ``` -## Build runtime bundle +Run one command in an ephemeral VM: ```bash -make runtime-bundle +pyro run --profile debian-git --vcpu-count 1 --mem-mib 1024 -- git --version ``` -This builds the packaged runtime bundle from `runtime_sources/` and syncs the result into `src/pyro_mcp/runtime_bundle/`. -For real artifacts, first materialize upstream sources into `build/runtime_sources/`. - -Available staged targets: -- `make runtime-binaries` -- `make runtime-kernel` -- `make runtime-rootfs` -- `make runtime-agent` -- `make runtime-validate` -- `make runtime-manifest` -- `make runtime-sync` -- `make runtime-clean` - -Available real-runtime targets: -- `make runtime-fetch-binaries` -- `make runtime-build-kernel-real` -- `make runtime-build-rootfs-real` -- `make runtime-materialize` -- `make runtime-boot-check` -- `make runtime-network-check` - -Notes: -- the real-source path depends on `docker`, outbound access to GitHub and Debian snapshot mirrors, and enough disk for kernel/rootfs builds -- `make runtime-boot-check` validates that the bundled runtime can boot a real microVM -- `make runtime-network-check` validates outbound internet access from inside the guest by cloning `https://github.com/octocat/Hello-World.git` - -## Run deterministic lifecycle demo +Run with outbound internet enabled: ```bash -make demo +pyro run --profile debian-git --vcpu-count 1 --mem-mib 1024 --network -- \ + "git clone --depth 1 https://github.com/octocat/Hello-World.git hello-world && git -C hello-world rev-parse --is-inside-work-tree" ``` -The demo creates a VM, starts it, runs a command, and returns structured output. -If the VM was started with networking enabled, it uses an internet probe. -Otherwise it runs `git --version`. - -To run the deterministic demo with guest networking enabled: +Show runtime and host diagnostics: ```bash -make network-demo +pyro doctor ``` -## Runtime doctor +Run the deterministic demo: ```bash -make doctor +pyro demo +pyro demo --network ``` -This prints bundled runtime paths, profile availability, checksum validation status, runtime capability flags, KVM host checks, and host networking diagnostics. - -## Networking - -- Host-side network allocation and diagnostics are implemented. -- The MCP server exposes `vm_network_info` for per-VM network metadata. -- Primary network-enabled entrypoints: - -```bash -make network-demo -make ollama-demo -``` - -- Network setup requires host privilege to manage TAP/NAT state. -- The current implementation auto-uses `sudo -n` for `ip`, `nft`, and `iptables` commands when available. -- Manual opt-in for other commands is still available with: - -```bash -PYRO_VM_ENABLE_NETWORK=1 make demo -``` - -- To validate real guest egress directly: - -```bash -make runtime-network-check -``` - -## Run Ollama lifecycle demo +Run the Ollama demo: ```bash ollama serve ollama pull llama:3.2-3b -make ollama-demo +pyro demo ollama ``` -Defaults are configured in `Makefile`. -The demo streams lifecycle progress logs and ends with a short text summary. -`make ollama-demo` now enables guest networking by default. -The command it asks the model to run is a small public repository clone: +Verbose Ollama logs: ```bash -rm -rf hello-world && git clone --depth 1 https://github.com/octocat/Hello-World.git hello-world >/dev/null && git -C hello-world rev-parse --is-inside-work-tree -``` -By default it omits log values; to include prompt content, tool args, and tool results use: - -```bash -make ollama-demo OLLAMA_DEMO_FLAGS=-v +pyro demo ollama -v ``` -## Run MCP server +## Python SDK -```bash -make run-server +```python +from pyro_mcp import Pyro + +pyro = Pyro() +result = pyro.run_in_vm( + profile="debian-git", + command="git --version", + vcpu_count=1, + mem_mib=1024, + timeout_seconds=30, + network=False, +) +print(result["stdout"]) ``` -## Quality checks +Lower-level lifecycle control remains available: + +```python +from pyro_mcp import Pyro + +pyro = Pyro() +created = pyro.create_vm( + profile="debian-git", + vcpu_count=1, + mem_mib=1024, + ttl_seconds=600, + network=True, +) +vm_id = created["vm_id"] +pyro.start_vm(vm_id) +result = pyro.exec_vm(vm_id, command="git --version", timeout_seconds=30) +print(result["stdout"]) +``` + +## MCP Tools + +Primary agent-facing tool: + +- `vm_run(profile, command, vcpu_count, mem_mib, timeout_seconds=30, ttl_seconds=600, network=false)` + +Advanced lifecycle tools: + +- `vm_list_profiles()` +- `vm_create(profile, vcpu_count, mem_mib, ttl_seconds=600, network=false)` +- `vm_start(vm_id)` +- `vm_exec(vm_id, command, timeout_seconds=30)` +- `vm_stop(vm_id)` +- `vm_delete(vm_id)` +- `vm_status(vm_id)` +- `vm_network_info(vm_id)` +- `vm_reap_expired()` + +## Runtime + +The package ships a bundled Linux x86_64 runtime payload with: + +- Firecracker +- Jailer +- guest kernel +- guest agent +- profile rootfs images + +No system Firecracker installation is required. + +Runtime diagnostics: ```bash +pyro doctor +``` + +The doctor report includes: + +- runtime integrity +- component versions +- capability flags +- KVM availability +- host networking prerequisites + +## Contributor Workflow + +For work inside this repository: + +```bash +make help +make setup make check ``` -Includes `ruff`, `mypy`, and `pytest` with coverage threshold. +Runtime build and validation helpers remain available through `make`, including: -## Pre-commit - -```bash -make install-hooks -``` - -Hooks execute the same lint/type/test gates. +- `make runtime-bundle` +- `make runtime-materialize` +- `make runtime-boot-check` +- `make runtime-network-check` diff --git a/pyproject.toml b/pyproject.toml index c49634e..8af43ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,13 +12,7 @@ dependencies = [ ] [project.scripts] -pyro-mcp-server = "pyro_mcp.server:main" -pyro-mcp-demo = "pyro_mcp.demo:main" -pyro-mcp-ollama-demo = "pyro_mcp.ollama_demo:main" -pyro-mcp-doctor = "pyro_mcp.doctor:main" -pyro-mcp-runtime-build = "pyro_mcp.runtime_build:main" -pyro-mcp-runtime-boot-check = "pyro_mcp.runtime_boot_check:main" -pyro-mcp-runtime-network-check = "pyro_mcp.runtime_network_check:main" +pyro = "pyro_mcp.cli:main" [build-system] requires = ["hatchling"] diff --git a/src/pyro_mcp/__init__.py b/src/pyro_mcp/__init__.py index 47b9b49..ec32f55 100644 --- a/src/pyro_mcp/__init__.py +++ b/src/pyro_mcp/__init__.py @@ -1,6 +1,7 @@ """Public package surface for pyro_mcp.""" +from pyro_mcp.api import Pyro from pyro_mcp.server import create_server from pyro_mcp.vm_manager import VmManager -__all__ = ["VmManager", "create_server"] +__all__ = ["Pyro", "VmManager", "create_server"] diff --git a/src/pyro_mcp/api.py b/src/pyro_mcp/api.py new file mode 100644 index 0000000..6814f9d --- /dev/null +++ b/src/pyro_mcp/api.py @@ -0,0 +1,179 @@ +"""Public facade shared by the Python SDK and MCP server.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from mcp.server.fastmcp import FastMCP + +from pyro_mcp.vm_manager import VmManager + + +class Pyro: + """High-level facade over the ephemeral VM runtime.""" + + def __init__( + self, + manager: VmManager | None = None, + *, + backend_name: str | None = None, + base_dir: Path | None = None, + artifacts_dir: Path | None = None, + max_active_vms: int = 4, + ) -> None: + self._manager = manager or VmManager( + backend_name=backend_name, + base_dir=base_dir, + artifacts_dir=artifacts_dir, + max_active_vms=max_active_vms, + ) + + @property + def manager(self) -> VmManager: + return self._manager + + def list_profiles(self) -> list[dict[str, object]]: + return self._manager.list_profiles() + + def create_vm( + self, + *, + profile: str, + vcpu_count: int, + mem_mib: int, + ttl_seconds: int = 600, + network: bool = False, + ) -> dict[str, Any]: + return self._manager.create_vm( + profile=profile, + vcpu_count=vcpu_count, + mem_mib=mem_mib, + ttl_seconds=ttl_seconds, + network=network, + ) + + def start_vm(self, vm_id: str) -> dict[str, Any]: + return self._manager.start_vm(vm_id) + + def exec_vm(self, vm_id: str, *, command: str, timeout_seconds: int = 30) -> dict[str, Any]: + return self._manager.exec_vm(vm_id, command=command, timeout_seconds=timeout_seconds) + + def stop_vm(self, vm_id: str) -> dict[str, Any]: + return self._manager.stop_vm(vm_id) + + def delete_vm(self, vm_id: str) -> dict[str, Any]: + return self._manager.delete_vm(vm_id) + + def status_vm(self, vm_id: str) -> dict[str, Any]: + return self._manager.status_vm(vm_id) + + def network_info_vm(self, vm_id: str) -> dict[str, Any]: + return self._manager.network_info_vm(vm_id) + + def reap_expired(self) -> dict[str, Any]: + return self._manager.reap_expired() + + def run_in_vm( + self, + *, + profile: str, + command: str, + vcpu_count: int, + mem_mib: int, + timeout_seconds: int = 30, + ttl_seconds: int = 600, + network: bool = False, + ) -> dict[str, Any]: + return self._manager.run_vm( + profile=profile, + command=command, + vcpu_count=vcpu_count, + mem_mib=mem_mib, + timeout_seconds=timeout_seconds, + ttl_seconds=ttl_seconds, + network=network, + ) + + def create_server(self) -> FastMCP: + server = FastMCP(name="pyro_mcp") + + @server.tool() + async def vm_run( + profile: str, + command: str, + vcpu_count: int, + mem_mib: int, + timeout_seconds: int = 30, + ttl_seconds: int = 600, + network: bool = False, + ) -> dict[str, Any]: + """Create, start, execute, and clean up an ephemeral VM.""" + return self.run_in_vm( + profile=profile, + command=command, + vcpu_count=vcpu_count, + mem_mib=mem_mib, + timeout_seconds=timeout_seconds, + ttl_seconds=ttl_seconds, + network=network, + ) + + @server.tool() + async def vm_list_profiles() -> list[dict[str, object]]: + """List standard environment profiles and package highlights.""" + return self.list_profiles() + + @server.tool() + async def vm_create( + profile: str, + vcpu_count: int, + mem_mib: int, + ttl_seconds: int = 600, + network: bool = False, + ) -> dict[str, Any]: + """Create an ephemeral VM record with profile and resource sizing.""" + return self.create_vm( + profile=profile, + vcpu_count=vcpu_count, + mem_mib=mem_mib, + ttl_seconds=ttl_seconds, + network=network, + ) + + @server.tool() + async def vm_start(vm_id: str) -> dict[str, Any]: + """Start a created VM and transition it into a command-ready state.""" + return self.start_vm(vm_id) + + @server.tool() + async def vm_exec(vm_id: str, command: str, timeout_seconds: int = 30) -> dict[str, Any]: + """Run one non-interactive command and auto-clean the VM.""" + return self.exec_vm(vm_id, command=command, timeout_seconds=timeout_seconds) + + @server.tool() + async def vm_stop(vm_id: str) -> dict[str, Any]: + """Stop a running VM.""" + return self.stop_vm(vm_id) + + @server.tool() + async def vm_delete(vm_id: str) -> dict[str, Any]: + """Delete a VM and its runtime artifacts.""" + return self.delete_vm(vm_id) + + @server.tool() + async def vm_status(vm_id: str) -> dict[str, Any]: + """Get the current state and metadata for a VM.""" + return self.status_vm(vm_id) + + @server.tool() + async def vm_network_info(vm_id: str) -> dict[str, Any]: + """Get the current network configuration assigned to a VM.""" + return self.network_info_vm(vm_id) + + @server.tool() + async def vm_reap_expired() -> dict[str, Any]: + """Delete VMs whose TTL has expired.""" + return self.reap_expired() + + return server diff --git a/src/pyro_mcp/cli.py b/src/pyro_mcp/cli.py new file mode 100644 index 0000000..6bac51a --- /dev/null +++ b/src/pyro_mcp/cli.py @@ -0,0 +1,103 @@ +"""Public CLI for pyro-mcp.""" + +from __future__ import annotations + +import argparse +import json +from typing import Any + +from pyro_mcp.api import Pyro +from pyro_mcp.demo import run_demo +from pyro_mcp.ollama_demo import DEFAULT_OLLAMA_BASE_URL, DEFAULT_OLLAMA_MODEL, run_ollama_tool_demo +from pyro_mcp.runtime import DEFAULT_PLATFORM, doctor_report + + +def _print_json(payload: dict[str, Any]) -> None: + print(json.dumps(payload, indent=2, sort_keys=True)) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="pyro CLI for ephemeral Firecracker VMs.") + subparsers = parser.add_subparsers(dest="command", required=True) + + mcp_parser = subparsers.add_parser("mcp", help="Run the MCP server.") + mcp_subparsers = mcp_parser.add_subparsers(dest="mcp_command", required=True) + mcp_subparsers.add_parser("serve", help="Run the MCP server over stdio.") + + run_parser = subparsers.add_parser("run", help="Run one command inside an ephemeral VM.") + run_parser.add_argument("--profile", required=True) + run_parser.add_argument("--vcpu-count", type=int, required=True) + run_parser.add_argument("--mem-mib", type=int, required=True) + run_parser.add_argument("--timeout-seconds", type=int, default=30) + run_parser.add_argument("--ttl-seconds", type=int, default=600) + run_parser.add_argument("--network", action="store_true") + run_parser.add_argument("command_args", nargs=argparse.REMAINDER) + + doctor_parser = subparsers.add_parser("doctor", help="Inspect runtime and host diagnostics.") + doctor_parser.add_argument("--platform", default=DEFAULT_PLATFORM) + + demo_parser = subparsers.add_parser("demo", help="Run built-in demos.") + demo_subparsers = demo_parser.add_subparsers(dest="demo_command") + demo_parser.add_argument("--network", action="store_true") + ollama_parser = demo_subparsers.add_parser("ollama", help="Run the Ollama MCP demo.") + ollama_parser.add_argument("--base-url", default=DEFAULT_OLLAMA_BASE_URL) + ollama_parser.add_argument("--model", default=DEFAULT_OLLAMA_MODEL) + ollama_parser.add_argument("-v", "--verbose", action="store_true") + + return parser + + +def _require_command(command_args: list[str]) -> str: + if command_args and command_args[0] == "--": + command_args = command_args[1:] + if not command_args: + raise ValueError("command is required after `pyro run --`") + return " ".join(command_args) + + +def main() -> None: + args = _build_parser().parse_args() + if args.command == "mcp": + Pyro().create_server().run(transport="stdio") + return + if args.command == "run": + command = _require_command(args.command_args) + result = Pyro().run_in_vm( + profile=args.profile, + command=command, + vcpu_count=args.vcpu_count, + mem_mib=args.mem_mib, + timeout_seconds=args.timeout_seconds, + ttl_seconds=args.ttl_seconds, + network=args.network, + ) + _print_json(result) + return + if args.command == "doctor": + _print_json(doctor_report(platform=args.platform)) + return + if args.command == "demo" and args.demo_command == "ollama": + try: + result = run_ollama_tool_demo( + base_url=args.base_url, + model=args.model, + verbose=args.verbose, + log=lambda message: print(message, flush=True), + ) + except Exception as exc: # noqa: BLE001 + print(f"[error] {exc}", flush=True) + raise SystemExit(1) from exc + exec_result = result["exec_result"] + if not isinstance(exec_result, dict): + raise RuntimeError("demo produced invalid execution result") + print( + f"[summary] exit_code={int(exec_result.get('exit_code', -1))} " + f"fallback_used={bool(result.get('fallback_used'))} " + f"execution_mode={str(exec_result.get('execution_mode', 'unknown'))}", + flush=True, + ) + if args.verbose: + print(f"[summary] stdout={str(exec_result.get('stdout', '')).strip()}", flush=True) + return + result = run_demo(network=bool(args.network)) + _print_json(result) diff --git a/src/pyro_mcp/demo.py b/src/pyro_mcp/demo.py index 833f8c5..196491d 100644 --- a/src/pyro_mcp/demo.py +++ b/src/pyro_mcp/demo.py @@ -5,7 +5,7 @@ from __future__ import annotations import json from typing import Any -from pyro_mcp.vm_manager import VmManager +from pyro_mcp.api import Pyro INTERNET_PROBE_COMMAND = ( 'python3 -c "import urllib.request; ' @@ -20,15 +20,22 @@ def _demo_command(status: dict[str, Any]) -> str: return "git --version" -def run_demo() -> dict[str, Any]: +def run_demo(*, network: bool = False) -> dict[str, Any]: """Create/start/exec/delete a VM and return command output.""" - manager = VmManager() - created = manager.create_vm(profile="debian-git", vcpu_count=1, mem_mib=512, ttl_seconds=600) - vm_id = str(created["vm_id"]) - manager.start_vm(vm_id) - status = manager.status_vm(vm_id) - executed = manager.exec_vm(vm_id, command=_demo_command(status), timeout_seconds=30) - return executed + pyro = Pyro() + status = { + "network_enabled": network, + "execution_mode": "guest_vsock" if network else "host_compat", + } + return pyro.run_in_vm( + profile="debian-git", + command=_demo_command(status), + vcpu_count=1, + mem_mib=512, + timeout_seconds=30, + ttl_seconds=600, + network=network, + ) def main() -> None: diff --git a/src/pyro_mcp/ollama_demo.py b/src/pyro_mcp/ollama_demo.py index 449d69f..4c0e729 100644 --- a/src/pyro_mcp/ollama_demo.py +++ b/src/pyro_mcp/ollama_demo.py @@ -9,9 +9,9 @@ import urllib.request from collections.abc import Callable from typing import Any, Final, cast -from pyro_mcp.vm_manager import VmManager +from pyro_mcp.api import Pyro -__all__ = ["VmManager", "run_ollama_tool_demo"] +__all__ = ["Pyro", "run_ollama_tool_demo"] DEFAULT_OLLAMA_BASE_URL: Final[str] = "http://localhost:11434/v1" DEFAULT_OLLAMA_MODEL: Final[str] = "llama:3.2-3b" @@ -24,6 +24,27 @@ NETWORK_PROOF_COMMAND: Final[str] = ( ) TOOL_SPECS: Final[list[dict[str, Any]]] = [ + { + "type": "function", + "function": { + "name": "vm_run", + "description": "Create, start, execute, and clean up an ephemeral VM in one call.", + "parameters": { + "type": "object", + "properties": { + "profile": {"type": "string"}, + "command": {"type": "string"}, + "vcpu_count": {"type": "integer"}, + "mem_mib": {"type": "integer"}, + "timeout_seconds": {"type": "integer"}, + "ttl_seconds": {"type": "integer"}, + "network": {"type": "boolean"}, + }, + "required": ["profile", "command", "vcpu_count", "mem_mib"], + "additionalProperties": False, + }, + }, + }, { "type": "function", "function": { @@ -48,6 +69,7 @@ TOOL_SPECS: Final[list[dict[str, Any]]] = [ "vcpu_count": {"type": "integer"}, "mem_mib": {"type": "integer"}, "ttl_seconds": {"type": "integer"}, + "network": {"type": "boolean"}, }, "required": ["profile", "vcpu_count", "mem_mib"], "additionalProperties": False, @@ -170,35 +192,55 @@ def _require_int(arguments: dict[str, Any], key: str) -> int: raise ValueError(f"{key} must be an integer") +def _require_bool(arguments: dict[str, Any], key: str, *, default: bool = False) -> bool: + value = arguments.get(key, default) + if isinstance(value, bool): + return value + raise ValueError(f"{key} must be a boolean") + + def _dispatch_tool_call( - manager: VmManager, tool_name: str, arguments: dict[str, Any] + pyro: Pyro, tool_name: str, arguments: dict[str, Any] ) -> dict[str, Any]: + if tool_name == "vm_run": + ttl_seconds = arguments.get("ttl_seconds", 600) + timeout_seconds = arguments.get("timeout_seconds", 30) + return pyro.run_in_vm( + profile=_require_str(arguments, "profile"), + command=_require_str(arguments, "command"), + vcpu_count=_require_int(arguments, "vcpu_count"), + mem_mib=_require_int(arguments, "mem_mib"), + timeout_seconds=_require_int({"timeout_seconds": timeout_seconds}, "timeout_seconds"), + ttl_seconds=_require_int({"ttl_seconds": ttl_seconds}, "ttl_seconds"), + network=_require_bool(arguments, "network", default=False), + ) if tool_name == "vm_list_profiles": - return {"profiles": manager.list_profiles()} + return {"profiles": pyro.list_profiles()} if tool_name == "vm_create": ttl_seconds = arguments.get("ttl_seconds", 600) - return manager.create_vm( + return pyro.create_vm( profile=_require_str(arguments, "profile"), vcpu_count=_require_int(arguments, "vcpu_count"), mem_mib=_require_int(arguments, "mem_mib"), ttl_seconds=_require_int({"ttl_seconds": ttl_seconds}, "ttl_seconds"), + network=_require_bool(arguments, "network", default=False), ) if tool_name == "vm_start": - return manager.start_vm(_require_str(arguments, "vm_id")) + return pyro.start_vm(_require_str(arguments, "vm_id")) if tool_name == "vm_exec": timeout_seconds = arguments.get("timeout_seconds", 30) vm_id = _require_str(arguments, "vm_id") - status = manager.status_vm(vm_id) + status = pyro.status_vm(vm_id) state = status.get("state") if state in {"created", "stopped"}: - manager.start_vm(vm_id) - return manager.exec_vm( + pyro.start_vm(vm_id) + return pyro.exec_vm( vm_id, command=_require_str(arguments, "command"), timeout_seconds=_require_int({"timeout_seconds": timeout_seconds}, "timeout_seconds"), ) if tool_name == "vm_status": - return manager.status_vm(_require_str(arguments, "vm_id")) + return pyro.status_vm(_require_str(arguments, "vm_id")) raise RuntimeError(f"unexpected tool requested by model: {tool_name!r}") @@ -212,11 +254,16 @@ def _format_tool_error(tool_name: str, arguments: dict[str, Any], exc: Exception } -def _run_direct_lifecycle_fallback(manager: VmManager) -> dict[str, Any]: - created = manager.create_vm(profile="debian-git", vcpu_count=1, mem_mib=512, ttl_seconds=600) - vm_id = str(created["vm_id"]) - manager.start_vm(vm_id) - return manager.exec_vm(vm_id, command=NETWORK_PROOF_COMMAND, timeout_seconds=60) +def _run_direct_lifecycle_fallback(pyro: Pyro) -> dict[str, Any]: + return pyro.run_in_vm( + profile="debian-git", + command=NETWORK_PROOF_COMMAND, + vcpu_count=1, + mem_mib=512, + timeout_seconds=60, + ttl_seconds=600, + network=True, + ) def _is_vm_id_placeholder(value: str) -> bool: @@ -264,6 +311,7 @@ def run_ollama_tool_demo( base_url: str = DEFAULT_OLLAMA_BASE_URL, model: str = DEFAULT_OLLAMA_MODEL, *, + pyro: Pyro | None = None, strict: bool = True, verbose: bool = False, log: Callable[[str], None] | None = None, @@ -271,15 +319,15 @@ def run_ollama_tool_demo( """Ask Ollama to run git version check in an ephemeral VM through lifecycle tools.""" emit = log or (lambda _: None) emit(f"[ollama] starting tool demo with model={model}") - manager = VmManager() + pyro_client = pyro or Pyro() messages: list[dict[str, Any]] = [ { "role": "user", "content": ( - "Use the lifecycle tools to prove outbound internet access in an ephemeral VM.\n" - "Required order: vm_list_profiles -> vm_create -> vm_start -> vm_exec.\n" - "Use profile `debian-git`, choose adequate vCPU/memory, and pass the `vm_id` " - "returned by vm_create into vm_start/vm_exec.\n" + "Use the VM tools to prove outbound internet access in an ephemeral VM.\n" + "Prefer `vm_run` unless a lower-level lifecycle step is strictly necessary.\n" + "Use profile `debian-git`, choose adequate vCPU/memory, " + "and set `network` to true.\n" f"Run this exact command: `{NETWORK_PROOF_COMMAND}`.\n" f"Success means the clone completes and the command prints `true`.\n" "If a tool returns an error, fix arguments and retry." @@ -350,7 +398,7 @@ def run_ollama_tool_demo( else: emit(f"[tool] calling {tool_name}") try: - result = _dispatch_tool_call(manager, tool_name, arguments) + result = _dispatch_tool_call(pyro_client, tool_name, arguments) success = True emit(f"[tool] {tool_name} succeeded") if tool_name == "vm_create": @@ -382,26 +430,26 @@ def run_ollama_tool_demo( } ) else: - raise RuntimeError("tool-calling loop exceeded maximum rounds") + raise RuntimeError("tool-calling loop exceeded maximum rounds") exec_event = next( ( event for event in reversed(tool_events) - if event.get("tool_name") == "vm_exec" and bool(event.get("success")) + if event.get("tool_name") in {"vm_exec", "vm_run"} and bool(event.get("success")) ), None, ) fallback_used = False if exec_event is None: if strict: - raise RuntimeError("demo did not execute a successful vm_exec") - emit("[fallback] model did not complete vm_exec; running direct lifecycle command") - exec_result = _run_direct_lifecycle_fallback(manager) + raise RuntimeError("demo did not execute a successful vm_run or vm_exec") + emit("[fallback] model did not complete vm_run; running direct lifecycle command") + exec_result = _run_direct_lifecycle_fallback(pyro_client) fallback_used = True tool_events.append( { - "tool_name": "vm_exec_fallback", + "tool_name": "vm_run_fallback", "arguments": {"command": NETWORK_PROOF_COMMAND}, "result": exec_result, "success": True, diff --git a/src/pyro_mcp/runtime_network_check.py b/src/pyro_mcp/runtime_network_check.py index 609853c..fcadef0 100644 --- a/src/pyro_mcp/runtime_network_check.py +++ b/src/pyro_mcp/runtime_network_check.py @@ -6,8 +6,7 @@ import argparse from dataclasses import dataclass from pathlib import Path -from pyro_mcp.vm_manager import VmManager -from pyro_mcp.vm_network import TapNetworkManager +from pyro_mcp.api import Pyro NETWORK_CHECK_COMMAND = ( "rm -rf hello-world " @@ -36,32 +35,24 @@ def run_network_check( timeout_seconds: int = 120, base_dir: Path | None = None, ) -> NetworkCheckResult: # pragma: no cover - integration helper - manager = VmManager( - base_dir=base_dir, - network_manager=TapNetworkManager(enabled=True), - ) - created = manager.create_vm( + pyro = Pyro(base_dir=base_dir) + result = pyro.run_in_vm( profile=profile, + command=NETWORK_CHECK_COMMAND, vcpu_count=vcpu_count, mem_mib=mem_mib, - ttl_seconds=ttl_seconds, - ) - vm_id = str(created["vm_id"]) - manager.start_vm(vm_id) - status = manager.status_vm(vm_id) - executed = manager.exec_vm( - vm_id, - command=NETWORK_CHECK_COMMAND, timeout_seconds=timeout_seconds, + ttl_seconds=ttl_seconds, + network=True, ) return NetworkCheckResult( - vm_id=vm_id, - execution_mode=str(executed["execution_mode"]), - network_enabled=bool(status["network_enabled"]), - exit_code=int(executed["exit_code"]), - stdout=str(executed["stdout"]), - stderr=str(executed["stderr"]), - cleanup=dict(executed["cleanup"]), + vm_id=str(result["vm_id"]), + execution_mode=str(result["execution_mode"]), + network_enabled=True, + exit_code=int(result["exit_code"]), + stdout=str(result["stdout"]), + stderr=str(result["stderr"]), + cleanup=dict(result["cleanup"]), ) diff --git a/src/pyro_mcp/server.py b/src/pyro_mcp/server.py index 8794f4a..bd094e0 100644 --- a/src/pyro_mcp/server.py +++ b/src/pyro_mcp/server.py @@ -2,74 +2,15 @@ from __future__ import annotations -from typing import Any - from mcp.server.fastmcp import FastMCP +from pyro_mcp.api import Pyro from pyro_mcp.vm_manager import VmManager def create_server(manager: VmManager | None = None) -> FastMCP: """Create and return a configured MCP server instance.""" - vm_manager = manager or VmManager() - server = FastMCP(name="pyro_mcp") - - @server.tool() - async def vm_list_profiles() -> list[dict[str, object]]: - """List standard environment profiles and package highlights.""" - return vm_manager.list_profiles() - - @server.tool() - async def vm_create( - profile: str, - vcpu_count: int, - mem_mib: int, - ttl_seconds: int = 600, - ) -> dict[str, Any]: - """Create an ephemeral VM record with profile and resource sizing.""" - return vm_manager.create_vm( - profile=profile, - vcpu_count=vcpu_count, - mem_mib=mem_mib, - ttl_seconds=ttl_seconds, - ) - - @server.tool() - async def vm_start(vm_id: str) -> dict[str, Any]: - """Start a created VM and transition it into a command-ready state.""" - return vm_manager.start_vm(vm_id) - - @server.tool() - async def vm_exec(vm_id: str, command: str, timeout_seconds: int = 30) -> dict[str, Any]: - """Run one non-interactive command and auto-clean the VM.""" - return vm_manager.exec_vm(vm_id, command=command, timeout_seconds=timeout_seconds) - - @server.tool() - async def vm_stop(vm_id: str) -> dict[str, Any]: - """Stop a running VM.""" - return vm_manager.stop_vm(vm_id) - - @server.tool() - async def vm_delete(vm_id: str) -> dict[str, Any]: - """Delete a VM and its runtime artifacts.""" - return vm_manager.delete_vm(vm_id) - - @server.tool() - async def vm_status(vm_id: str) -> dict[str, Any]: - """Get the current state and metadata for a VM.""" - return vm_manager.status_vm(vm_id) - - @server.tool() - async def vm_network_info(vm_id: str) -> dict[str, Any]: - """Get the current network configuration assigned to a VM.""" - return vm_manager.network_info_vm(vm_id) - - @server.tool() - async def vm_reap_expired() -> dict[str, Any]: - """Delete VMs whose TTL has expired.""" - return vm_manager.reap_expired() - - return server + return Pyro(manager=manager).create_server() def main() -> None: diff --git a/src/pyro_mcp/vm_manager.py b/src/pyro_mcp/vm_manager.py index a4b6ef6..eb74dca 100644 --- a/src/pyro_mcp/vm_manager.py +++ b/src/pyro_mcp/vm_manager.py @@ -40,6 +40,7 @@ class VmInstance: expires_at: float workdir: Path state: VmState = "created" + network_requested: bool = False firecracker_pid: int | None = None last_error: str | None = None metadata: dict[str, str] = field(default_factory=dict) @@ -165,7 +166,7 @@ class FirecrackerBackend(VmBackend): # pragma: no cover rootfs_copy = instance.workdir / "rootfs.ext4" shutil.copy2(artifacts.rootfs_image, rootfs_copy) instance.metadata["rootfs_image"] = str(rootfs_copy) - if self._network_manager.enabled: + if instance.network_requested: network = self._network_manager.allocate(instance.vm_id) instance.network = network instance.metadata.update(self._network_manager.to_metadata(network)) @@ -342,7 +343,12 @@ class VmManager: reason="mock backend does not boot a guest", ) self._max_active_vms = max_active_vms - self._network_manager = network_manager or TapNetworkManager() + if network_manager is not None: + self._network_manager = network_manager + elif self._backend_name == "firecracker": + self._network_manager = TapNetworkManager(enabled=True) + else: + self._network_manager = TapNetworkManager(enabled=False) self._lock = threading.Lock() self._instances: dict[str, VmInstance] = {} self._base_dir.mkdir(parents=True, exist_ok=True) @@ -367,7 +373,13 @@ class VmManager: return list_profiles() def create_vm( - self, *, profile: str, vcpu_count: int, mem_mib: int, ttl_seconds: int + self, + *, + profile: str, + vcpu_count: int, + mem_mib: int, + ttl_seconds: int, + network: bool = False, ) -> dict[str, Any]: self._validate_limits(vcpu_count=vcpu_count, mem_mib=mem_mib, ttl_seconds=ttl_seconds) get_profile(profile) @@ -389,11 +401,41 @@ class VmManager: created_at=now, expires_at=now + ttl_seconds, workdir=self._base_dir / vm_id, + network_requested=network, ) self._backend.create(instance) self._instances[vm_id] = instance return self._serialize(instance) + def run_vm( + self, + *, + profile: str, + command: str, + vcpu_count: int, + mem_mib: int, + timeout_seconds: int = 30, + ttl_seconds: int = 600, + network: bool = False, + ) -> dict[str, Any]: + created = self.create_vm( + profile=profile, + vcpu_count=vcpu_count, + mem_mib=mem_mib, + ttl_seconds=ttl_seconds, + network=network, + ) + vm_id = str(created["vm_id"]) + try: + self.start_vm(vm_id) + return self.exec_vm(vm_id, command=command, timeout_seconds=timeout_seconds) + except Exception: + try: + self.delete_vm(vm_id, reason="run_vm_error_cleanup") + except ValueError: + pass + raise + def start_vm(self, vm_id: str) -> dict[str, Any]: with self._lock: instance = self._get_instance_locked(vm_id) diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..9448bf3 --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any, cast + +from pyro_mcp.api import Pyro +from pyro_mcp.vm_manager import VmManager +from pyro_mcp.vm_network import TapNetworkManager + + +def test_pyro_run_in_vm_delegates_to_manager(tmp_path: Path) -> None: + pyro = Pyro( + manager=VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + ) + result = pyro.run_in_vm( + profile="debian-base", + command="printf 'ok\\n'", + vcpu_count=1, + mem_mib=512, + timeout_seconds=30, + ttl_seconds=600, + network=False, + ) + assert int(result["exit_code"]) == 0 + assert str(result["stdout"]) == "ok\n" + + +def test_pyro_create_server_registers_vm_run(tmp_path: Path) -> None: + pyro = Pyro( + manager=VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + ) + + async def _run() -> list[str]: + server = pyro.create_server() + tools = await server.list_tools() + return sorted(tool.name for tool in tools) + + tool_names = asyncio.run(_run()) + assert "vm_run" in tool_names + assert "vm_create" in tool_names + + +def test_pyro_vm_run_tool_executes(tmp_path: Path) -> None: + pyro = Pyro( + manager=VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + ) + + def _extract_structured(raw_result: object) -> dict[str, Any]: + if not isinstance(raw_result, tuple) or len(raw_result) != 2: + raise TypeError("unexpected call_tool result shape") + _, structured = raw_result + if not isinstance(structured, dict): + raise TypeError("expected structured dictionary result") + return cast(dict[str, Any], structured) + + async def _run() -> dict[str, Any]: + server = pyro.create_server() + return _extract_structured( + await server.call_tool( + "vm_run", + { + "profile": "debian-base", + "command": "printf 'ok\\n'", + "vcpu_count": 1, + "mem_mib": 512, + "network": False, + }, + ) + ) + + result = asyncio.run(_run()) + assert int(result["exit_code"]) == 0 diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..cbd8ac6 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +import argparse +import json +from typing import Any + +import pytest + +import pyro_mcp.cli as cli + + +def test_cli_run_prints_json( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + class StubPyro: + def run_in_vm(self, **kwargs: Any) -> dict[str, Any]: + assert kwargs["network"] is True + assert kwargs["command"] == "echo hi" + return {"exit_code": 0, "stdout": "hi\n"} + + class StubParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="run", + profile="debian-git", + vcpu_count=1, + mem_mib=512, + timeout_seconds=30, + ttl_seconds=600, + network=True, + command_args=["--", "echo", "hi"], + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: StubParser()) + monkeypatch.setattr(cli, "Pyro", StubPyro) + cli.main() + output = json.loads(capsys.readouterr().out) + assert output["exit_code"] == 0 + + +def test_cli_doctor_prints_json( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + class StubParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace(command="doctor", platform="linux-x86_64") + + monkeypatch.setattr(cli, "_build_parser", lambda: StubParser()) + monkeypatch.setattr( + cli, + "doctor_report", + lambda platform: {"platform": platform, "runtime_ok": True}, + ) + cli.main() + output = json.loads(capsys.readouterr().out) + assert output["runtime_ok"] is True + + +def test_cli_demo_ollama_prints_summary( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + class StubParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="demo", + demo_command="ollama", + base_url="http://localhost:11434/v1", + model="llama3.2:3b", + verbose=False, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: StubParser()) + monkeypatch.setattr( + cli, + "run_ollama_tool_demo", + lambda **kwargs: { + "exec_result": {"exit_code": 0, "execution_mode": "guest_vsock", "stdout": "true\n"}, + "fallback_used": False, + }, + ) + cli.main() + output = capsys.readouterr().out + assert "[summary] exit_code=0 fallback_used=False execution_mode=guest_vsock" in output + + +def test_cli_requires_run_command() -> None: + with pytest.raises(ValueError, match="command is required"): + cli._require_command([]) diff --git a/tests/test_demo.py b/tests/test_demo.py index 6947042..e640817 100644 --- a/tests/test_demo.py +++ b/tests/test_demo.py @@ -11,56 +11,55 @@ import pyro_mcp.demo as demo_module def test_run_demo_happy_path(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[tuple[str, dict[str, Any]]] = [] - class StubManager: + class StubPyro: def __init__(self) -> None: pass - def create_vm( + def run_in_vm( self, *, profile: str, + command: str, vcpu_count: int, mem_mib: int, + timeout_seconds: int, ttl_seconds: int, - ) -> dict[str, str]: + network: bool, + ) -> dict[str, Any]: calls.append( ( - "create_vm", + "run_in_vm", { "profile": profile, + "command": command, "vcpu_count": vcpu_count, "mem_mib": mem_mib, + "timeout_seconds": timeout_seconds, "ttl_seconds": ttl_seconds, + "network": network, }, ) ) - return {"vm_id": "vm-1"} + return {"vm_id": "vm-1", "stdout": "git version 2.x", "exit_code": 0} - def start_vm(self, vm_id: str) -> dict[str, str]: - calls.append(("start_vm", {"vm_id": vm_id})) - return {"vm_id": vm_id} - - def status_vm(self, vm_id: str) -> dict[str, Any]: - calls.append(("status_vm", {"vm_id": vm_id})) - return {"vm_id": vm_id, "network_enabled": False, "execution_mode": "host_compat"} - - def exec_vm(self, vm_id: str, *, command: str, timeout_seconds: int) -> dict[str, Any]: - calls.append( - ( - "exec_vm", - {"vm_id": vm_id, "command": command, "timeout_seconds": timeout_seconds}, - ) - ) - return {"vm_id": vm_id, "stdout": "git version 2.x", "exit_code": 0} - - monkeypatch.setattr(demo_module, "VmManager", StubManager) + monkeypatch.setattr(demo_module, "Pyro", StubPyro) result = demo_module.run_demo() assert result["exit_code"] == 0 - assert calls[0][0] == "create_vm" - assert calls[1] == ("start_vm", {"vm_id": "vm-1"}) - assert calls[2] == ("status_vm", {"vm_id": "vm-1"}) - assert calls[3][0] == "exec_vm" + assert calls == [ + ( + "run_in_vm", + { + "profile": "debian-git", + "command": "git --version", + "vcpu_count": 1, + "mem_mib": 512, + "timeout_seconds": 30, + "ttl_seconds": 600, + "network": False, + }, + ) + ] def test_demo_command_prefers_network_probe_for_guest_vsock() -> None: @@ -79,3 +78,20 @@ def test_main_prints_json( demo_module.main() rendered = json.loads(capsys.readouterr().out) assert rendered["exit_code"] == 0 + + +def test_run_demo_network_uses_probe(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + class StubPyro: + def __init__(self) -> None: + pass + + def run_in_vm(self, **kwargs: Any) -> dict[str, Any]: + captured.update(kwargs) + return {"exit_code": 0} + + monkeypatch.setattr(demo_module, "Pyro", StubPyro) + demo_module.run_demo(network=True) + assert "https://example.com" in str(captured["command"]) + assert captured["network"] is True diff --git a/tests/test_ollama_demo.py b/tests/test_ollama_demo.py index 82b1235..086bebe 100644 --- a/tests/test_ollama_demo.py +++ b/tests/test_ollama_demo.py @@ -10,16 +10,17 @@ from typing import Any import pytest import pyro_mcp.ollama_demo as ollama_demo +from pyro_mcp.api import Pyro as RealPyro from pyro_mcp.vm_manager import VmManager as RealVmManager @pytest.fixture(autouse=True) def _mock_vm_manager_for_tests(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - class TestVmManager(RealVmManager): + class TestPyro(RealPyro): def __init__(self) -> None: - super().__init__(backend_name="mock", base_dir=tmp_path / "vms") + super().__init__(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms")) - monkeypatch.setattr(ollama_demo, "VmManager", TestVmManager) + monkeypatch.setattr(ollama_demo, "Pyro", TestPyro) def _stepwise_model_response(payload: dict[str, Any], step: int) -> dict[str, Any]: @@ -46,55 +47,14 @@ def _stepwise_model_response(payload: dict[str, Any], step: int) -> dict[str, An { "id": "2", "function": { - "name": "vm_create", - "arguments": json.dumps( - {"profile": "debian-git", "vcpu_count": 1, "mem_mib": 512} - ), - }, - } - ], - } - } - ] - } - if step == 3: - vm_id = json.loads(payload["messages"][-1]["content"])["vm_id"] - return { - "choices": [ - { - "message": { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "3", - "function": { - "name": "vm_start", - "arguments": json.dumps({"vm_id": vm_id}), - }, - } - ], - } - } - ] - } - if step == 4: - vm_id = json.loads(payload["messages"][-1]["content"])["vm_id"] - return { - "choices": [ - { - "message": { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "4", - "function": { - "name": "vm_exec", + "name": "vm_run", "arguments": json.dumps( { - "vm_id": vm_id, + "profile": "debian-git", "command": "printf 'true\\n'", + "vcpu_count": 1, + "mem_mib": 512, + "network": True, } ), }, @@ -127,12 +87,12 @@ def test_run_ollama_tool_demo_happy_path(monkeypatch: pytest.MonkeyPatch) -> Non assert result["fallback_used"] is False assert str(result["exec_result"]["stdout"]).strip() == "true" assert result["final_response"] == "Executed git command in ephemeral VM." - assert len(result["tool_events"]) == 4 + assert len(result["tool_events"]) == 2 assert any(line == "[model] input user" for line in logs) assert any(line == "[model] output assistant" for line in logs) - assert any("[model] tool_call vm_exec" in line for line in logs) - assert any(line == "[tool] calling vm_exec" for line in logs) - assert any(line == "[tool] result vm_exec" for line in logs) + assert any("[model] tool_call vm_run" in line for line in logs) + assert any(line == "[tool] calling vm_run" for line in logs) + assert any(line == "[tool] result vm_run" for line in logs) def test_run_ollama_tool_demo_recovers_from_bad_vm_id( @@ -256,12 +216,12 @@ def test_run_ollama_tool_demo_resolves_vm_id_placeholder( def test_dispatch_tool_call_vm_exec_autostarts_created_vm(tmp_path: Path) -> None: - manager = RealVmManager(backend_name="mock", base_dir=tmp_path / "vms") - created = manager.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=60) + pyro = RealPyro(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms")) + created = pyro.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=60) vm_id = str(created["vm_id"]) executed = ollama_demo._dispatch_tool_call( - manager, + pyro, "vm_exec", {"vm_id": vm_id, "command": "printf 'git version\\n'", "timeout_seconds": "30"}, ) @@ -275,7 +235,7 @@ def test_run_ollama_tool_demo_raises_without_vm_exec(monkeypatch: pytest.MonkeyP return {"choices": [{"message": {"role": "assistant", "content": "No tools"}}]} monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion) - with pytest.raises(RuntimeError, match="did not execute a successful vm_exec"): + with pytest.raises(RuntimeError, match="did not execute a successful vm_run or vm_exec"): ollama_demo.run_ollama_tool_demo() @@ -286,16 +246,16 @@ def test_run_ollama_tool_demo_uses_fallback_when_not_strict( del base_url, payload return {"choices": [{"message": {"role": "assistant", "content": "No tools"}}]} - class TestVmManager(RealVmManager): + class TestPyro(RealPyro): def __init__(self) -> None: - super().__init__(backend_name="mock", base_dir=tmp_path / "vms") + super().__init__(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms")) monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion) - monkeypatch.setattr(ollama_demo, "VmManager", TestVmManager) + monkeypatch.setattr(ollama_demo, "Pyro", TestPyro) monkeypatch.setattr( ollama_demo, "_run_direct_lifecycle_fallback", - lambda manager: { + lambda pyro: { "vm_id": "vm-1", "command": ollama_demo.NETWORK_PROOF_COMMAND, "stdout": "true\n", @@ -332,7 +292,7 @@ def test_run_ollama_tool_demo_verbose_logs_values(monkeypatch: pytest.MonkeyPatc assert str(result["exec_result"]["stdout"]).strip() == "true" assert any("[model] input user:" in line for line in logs) assert any("[model] tool_call vm_list_profiles args={}" in line for line in logs) - assert any("[tool] result vm_exec " in line for line in logs) + assert any("[tool] result vm_run " in line for line in logs) @pytest.mark.parametrize( @@ -397,7 +357,7 @@ def test_run_ollama_tool_demo_exec_result_validation( "message": { "role": "assistant", "tool_calls": [ - {"id": "1", "function": {"name": "vm_exec", "arguments": "{}"}} + {"id": "1", "function": {"name": "vm_run", "arguments": "{}"}} ], } } @@ -410,9 +370,9 @@ def test_run_ollama_tool_demo_exec_result_validation( del base_url, payload return responses.pop(0) - def fake_dispatch(manager: Any, tool_name: str, arguments: dict[str, Any]) -> Any: - del manager, arguments - if tool_name == "vm_exec": + def fake_dispatch(pyro: Any, tool_name: str, arguments: dict[str, Any]) -> Any: + del pyro, arguments + if tool_name == "vm_run": return exec_result return {"ok": True} @@ -423,27 +383,46 @@ def test_run_ollama_tool_demo_exec_result_validation( def test_dispatch_tool_call_coverage(tmp_path: Path) -> None: - manager = RealVmManager(backend_name="mock", base_dir=tmp_path / "vms") - profiles = ollama_demo._dispatch_tool_call(manager, "vm_list_profiles", {}) + pyro = RealPyro(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms")) + profiles = ollama_demo._dispatch_tool_call(pyro, "vm_list_profiles", {}) assert "profiles" in profiles created = ollama_demo._dispatch_tool_call( - manager, + pyro, "vm_create", - {"profile": "debian-base", "vcpu_count": "1", "mem_mib": "512", "ttl_seconds": "60"}, + { + "profile": "debian-base", + "vcpu_count": "1", + "mem_mib": "512", + "ttl_seconds": "60", + "network": False, + }, ) vm_id = str(created["vm_id"]) - started = ollama_demo._dispatch_tool_call(manager, "vm_start", {"vm_id": vm_id}) + started = ollama_demo._dispatch_tool_call(pyro, "vm_start", {"vm_id": vm_id}) assert started["state"] == "started" - status = ollama_demo._dispatch_tool_call(manager, "vm_status", {"vm_id": vm_id}) + status = ollama_demo._dispatch_tool_call(pyro, "vm_status", {"vm_id": vm_id}) assert status["vm_id"] == vm_id executed = ollama_demo._dispatch_tool_call( - manager, + pyro, "vm_exec", {"vm_id": vm_id, "command": "printf 'true\\n'", "timeout_seconds": "30"}, ) assert int(executed["exit_code"]) == 0 + executed_run = ollama_demo._dispatch_tool_call( + pyro, + "vm_run", + { + "profile": "debian-base", + "command": "printf 'true\\n'", + "vcpu_count": "1", + "mem_mib": "512", + "timeout_seconds": "30", + "network": False, + }, + ) + assert int(executed_run["exit_code"]) == 0 with pytest.raises(RuntimeError, match="unexpected tool requested by model"): - ollama_demo._dispatch_tool_call(manager, "nope", {}) + ollama_demo._dispatch_tool_call(pyro, "nope", {}) def test_format_tool_error() -> None: @@ -483,6 +462,16 @@ def test_require_int_validation(value: Any) -> None: ollama_demo._require_int({"k": value}, "k") +@pytest.mark.parametrize(("arguments", "expected"), [({}, False), ({"k": True}, True)]) +def test_require_bool(arguments: dict[str, Any], expected: bool) -> None: + assert ollama_demo._require_bool(arguments, "k", default=False) is expected + + +def test_require_bool_validation() -> None: + with pytest.raises(ValueError, match="must be a boolean"): + ollama_demo._require_bool({"k": "true"}, "k") + + def test_post_chat_completion_success(monkeypatch: pytest.MonkeyPatch) -> None: class StubResponse: def __enter__(self) -> StubResponse: diff --git a/tests/test_runtime_network_check.py b/tests/test_runtime_network_check.py index 88a3684..44fe7dd 100644 --- a/tests/test_runtime_network_check.py +++ b/tests/test_runtime_network_check.py @@ -3,48 +3,38 @@ from __future__ import annotations import pytest import pyro_mcp.runtime_network_check as runtime_network_check -from pyro_mcp.vm_network import TapNetworkManager def test_network_check_uses_network_enabled_manager(monkeypatch: pytest.MonkeyPatch) -> None: observed: dict[str, object] = {} - class StubManager: + class StubPyro: def __init__(self, **kwargs: object) -> None: observed.update(kwargs) - def create_vm(self, **kwargs: object) -> dict[str, object]: - observed["create_kwargs"] = kwargs - return {"vm_id": "vm123"} - - def start_vm(self, vm_id: str) -> dict[str, object]: - observed["started_vm_id"] = vm_id - return {"state": "started"} - - def status_vm(self, vm_id: str) -> dict[str, object]: - observed["status_vm_id"] = vm_id - return {"network_enabled": True} - - def exec_vm(self, vm_id: str, *, command: str, timeout_seconds: int) -> dict[str, object]: - observed["exec_vm_id"] = vm_id - observed["command"] = command - observed["timeout_seconds"] = timeout_seconds + def run_in_vm(self, **kwargs: object) -> dict[str, object]: + observed["run_kwargs"] = kwargs return { + "vm_id": "vm123", "execution_mode": "guest_vsock", "exit_code": 0, "stdout": "true\n", "stderr": "", - "cleanup": {"deleted": True, "vm_id": vm_id, "reason": "post_exec_cleanup"}, + "cleanup": {"deleted": True, "vm_id": "vm123", "reason": "post_exec_cleanup"}, } - monkeypatch.setattr(runtime_network_check, "VmManager", StubManager) + monkeypatch.setattr(runtime_network_check, "Pyro", StubPyro) result = runtime_network_check.run_network_check() - network_manager = observed["network_manager"] - assert isinstance(network_manager, TapNetworkManager) - assert network_manager.enabled is True - assert observed["command"] == runtime_network_check.NETWORK_CHECK_COMMAND - assert observed["timeout_seconds"] == 120 + assert observed["run_kwargs"] == { + "profile": "debian-git", + "command": runtime_network_check.NETWORK_CHECK_COMMAND, + "vcpu_count": 1, + "mem_mib": 1024, + "timeout_seconds": 120, + "ttl_seconds": 600, + "network": True, + } assert result.execution_mode == "guest_vsock" assert result.network_enabled is True assert result.exit_code == 0 diff --git a/tests/test_server.py b/tests/test_server.py index ba38313..bddddb2 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -29,10 +29,11 @@ def test_create_server_registers_vm_tools(tmp_path: Path) -> None: assert "vm_exec" in tool_names assert "vm_list_profiles" in tool_names assert "vm_network_info" in tool_names + assert "vm_run" in tool_names assert "vm_status" in tool_names -def test_vm_tools_lifecycle_round_trip(tmp_path: Path) -> None: +def test_vm_run_round_trip(tmp_path: Path) -> None: manager = VmManager( backend_name="mock", base_dir=tmp_path / "vms", @@ -49,17 +50,17 @@ def test_vm_tools_lifecycle_round_trip(tmp_path: Path) -> None: async def _run() -> dict[str, Any]: server = create_server(manager=manager) - created = _extract_structured( - await server.call_tool( - "vm_create", - {"profile": "debian-git", "vcpu_count": 1, "mem_mib": 512, "ttl_seconds": 600}, - ) - ) - vm_id = str(created["vm_id"]) - await server.call_tool("vm_start", {"vm_id": vm_id}) executed = _extract_structured( await server.call_tool( - "vm_exec", {"vm_id": vm_id, "command": "printf 'git version 2.0\\n'"} + "vm_run", + { + "profile": "debian-git", + "command": "printf 'git version 2.0\\n'", + "vcpu_count": 1, + "mem_mib": 512, + "ttl_seconds": 600, + "network": False, + }, ) ) return executed diff --git a/tests/test_vm_manager.py b/tests/test_vm_manager.py index e32c057..c5a3c98 100644 --- a/tests/test_vm_manager.py +++ b/tests/test_vm_manager.py @@ -107,7 +107,7 @@ def test_vm_manager_reaps_started_vm(tmp_path: Path) -> None: ({"vcpu_count": 1, "mem_mib": 512, "ttl_seconds": 30}, "ttl_seconds must be between"), ], ) -def test_vm_manager_validates_limits(tmp_path: Path, kwargs: dict[str, int], msg: str) -> None: +def test_vm_manager_validates_limits(tmp_path: Path, kwargs: dict[str, Any], msg: str) -> None: manager = VmManager( backend_name="mock", base_dir=tmp_path / "vms", @@ -188,6 +188,25 @@ def test_vm_manager_network_info(tmp_path: Path) -> None: assert info["network_enabled"] is False +def test_vm_manager_run_vm(tmp_path: Path) -> None: + manager = VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + result = manager.run_vm( + profile="debian-base", + command="printf 'ok\\n'", + vcpu_count=1, + mem_mib=512, + timeout_seconds=30, + ttl_seconds=600, + network=False, + ) + assert int(result["exit_code"]) == 0 + assert str(result["stdout"]) == "ok\n" + + def test_vm_manager_firecracker_backend_path( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: