Thales Maciel
84a7e18d4d
Complete the 2.6.0 workspace milestone by adding explicit host-out export and immutable-baseline diff across the CLI, Python SDK, and MCP server. Capture a baseline archive at workspace creation, export live /workspace paths through the guest agent, and compute structured whole-workspace diffs on the host without affecting command logs or shell state. The docs, roadmap, bundled guest agent, and workspace example now reflect the new create -> sync -> diff -> export workflow. Validation: uv lock, UV_CACHE_DIR=.uv-cache make check, UV_CACHE_DIR=.uv-cache make dist-check, and a real guest-backed Firecracker smoke covering workspace create, sync push, diff, export, and delete.
359 lines
12 KiB
Python
359 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
import pyro_mcp.server as server_module
|
|
from pyro_mcp.server import create_server
|
|
from pyro_mcp.vm_manager import VmManager
|
|
from pyro_mcp.vm_network import TapNetworkManager
|
|
|
|
|
|
def test_create_server_registers_vm_tools(tmp_path: Path) -> None:
|
|
manager = VmManager(
|
|
backend_name="mock",
|
|
base_dir=tmp_path / "vms",
|
|
network_manager=TapNetworkManager(enabled=False),
|
|
)
|
|
|
|
async def _run() -> list[str]:
|
|
server = create_server(manager=manager)
|
|
tools = await server.list_tools()
|
|
return sorted(tool.name for tool in tools)
|
|
|
|
tool_names = asyncio.run(_run())
|
|
assert "vm_create" in tool_names
|
|
assert "vm_exec" in tool_names
|
|
assert "vm_list_environments" in tool_names
|
|
assert "vm_network_info" in tool_names
|
|
assert "vm_run" in tool_names
|
|
assert "vm_status" in tool_names
|
|
assert "workspace_create" in tool_names
|
|
assert "workspace_diff" in tool_names
|
|
assert "workspace_export" in tool_names
|
|
assert "workspace_logs" in tool_names
|
|
assert "workspace_sync_push" in tool_names
|
|
assert "shell_open" in tool_names
|
|
assert "shell_read" in tool_names
|
|
assert "shell_write" in tool_names
|
|
assert "shell_signal" in tool_names
|
|
assert "shell_close" in tool_names
|
|
|
|
|
|
def test_vm_run_round_trip(tmp_path: Path) -> None:
|
|
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 = create_server(manager=manager)
|
|
executed = _extract_structured(
|
|
await server.call_tool(
|
|
"vm_run",
|
|
{
|
|
"environment": "debian:12",
|
|
"command": "printf 'git version 2.0\\n'",
|
|
"ttl_seconds": 600,
|
|
"network": False,
|
|
"allow_host_compat": True,
|
|
},
|
|
)
|
|
)
|
|
return executed
|
|
|
|
executed = asyncio.run(_run())
|
|
assert int(executed["exit_code"]) == 0
|
|
assert "git version" in str(executed["stdout"])
|
|
|
|
|
|
def test_vm_tools_status_stop_delete_and_reap(tmp_path: Path) -> None:
|
|
manager = VmManager(
|
|
backend_name="mock",
|
|
base_dir=tmp_path / "vms",
|
|
network_manager=TapNetworkManager(enabled=False),
|
|
)
|
|
manager.MIN_TTL_SECONDS = 1
|
|
|
|
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() -> tuple[
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
list[dict[str, object]],
|
|
dict[str, Any],
|
|
]:
|
|
server = create_server(manager=manager)
|
|
environments_raw = await server.call_tool("vm_list_environments", {})
|
|
if not isinstance(environments_raw, tuple) or len(environments_raw) != 2:
|
|
raise TypeError("unexpected environments result")
|
|
_, environments_structured = environments_raw
|
|
if not isinstance(environments_structured, dict):
|
|
raise TypeError("environments tool should return a dictionary")
|
|
raw_environments = environments_structured.get("result")
|
|
if not isinstance(raw_environments, list):
|
|
raise TypeError("environments tool did not contain a result list")
|
|
created = _extract_structured(
|
|
await server.call_tool(
|
|
"vm_create",
|
|
{
|
|
"environment": "debian:12-base",
|
|
"ttl_seconds": 600,
|
|
"allow_host_compat": True,
|
|
},
|
|
)
|
|
)
|
|
vm_id = str(created["vm_id"])
|
|
await server.call_tool("vm_start", {"vm_id": vm_id})
|
|
status = _extract_structured(await server.call_tool("vm_status", {"vm_id": vm_id}))
|
|
network = _extract_structured(await server.call_tool("vm_network_info", {"vm_id": vm_id}))
|
|
stopped = _extract_structured(await server.call_tool("vm_stop", {"vm_id": vm_id}))
|
|
deleted = _extract_structured(await server.call_tool("vm_delete", {"vm_id": vm_id}))
|
|
|
|
expiring = _extract_structured(
|
|
await server.call_tool(
|
|
"vm_create",
|
|
{
|
|
"environment": "debian:12-base",
|
|
"ttl_seconds": 1,
|
|
"allow_host_compat": True,
|
|
},
|
|
)
|
|
)
|
|
expiring_id = str(expiring["vm_id"])
|
|
manager._instances[expiring_id].expires_at = 0.0 # noqa: SLF001
|
|
reaped = _extract_structured(await server.call_tool("vm_reap_expired", {}))
|
|
return (
|
|
status,
|
|
network,
|
|
stopped,
|
|
deleted,
|
|
cast(list[dict[str, object]], raw_environments),
|
|
reaped,
|
|
)
|
|
|
|
status, network, stopped, deleted, environments, reaped = asyncio.run(_run())
|
|
assert status["state"] == "started"
|
|
assert network["network_enabled"] is False
|
|
assert stopped["state"] == "stopped"
|
|
assert bool(deleted["deleted"]) is True
|
|
assert environments[0]["name"] == "debian:12"
|
|
assert int(reaped["count"]) == 1
|
|
|
|
|
|
def test_server_main_runs_stdio_transport(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
called: dict[str, str] = {}
|
|
|
|
class StubServer:
|
|
def run(self, transport: str) -> None:
|
|
called["transport"] = transport
|
|
|
|
monkeypatch.setattr(server_module, "create_server", lambda: StubServer())
|
|
server_module.main()
|
|
assert called == {"transport": "stdio"}
|
|
|
|
|
|
def test_workspace_tools_round_trip(tmp_path: Path) -> None:
|
|
manager = VmManager(
|
|
backend_name="mock",
|
|
base_dir=tmp_path / "vms",
|
|
network_manager=TapNetworkManager(enabled=False),
|
|
)
|
|
source_dir = tmp_path / "seed"
|
|
source_dir.mkdir()
|
|
(source_dir / "note.txt").write_text("ok\n", encoding="utf-8")
|
|
|
|
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() -> tuple[
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
dict[str, Any],
|
|
]:
|
|
server = create_server(manager=manager)
|
|
created = _extract_structured(
|
|
await server.call_tool(
|
|
"workspace_create",
|
|
{
|
|
"environment": "debian:12-base",
|
|
"allow_host_compat": True,
|
|
"seed_path": str(source_dir),
|
|
},
|
|
)
|
|
)
|
|
workspace_id = str(created["workspace_id"])
|
|
update_dir = tmp_path / "update"
|
|
update_dir.mkdir()
|
|
(update_dir / "more.txt").write_text("more\n", encoding="utf-8")
|
|
synced = _extract_structured(
|
|
await server.call_tool(
|
|
"workspace_sync_push",
|
|
{
|
|
"workspace_id": workspace_id,
|
|
"source_path": str(update_dir),
|
|
"dest": "subdir",
|
|
},
|
|
)
|
|
)
|
|
executed = _extract_structured(
|
|
await server.call_tool(
|
|
"workspace_exec",
|
|
{
|
|
"workspace_id": workspace_id,
|
|
"command": "cat subdir/more.txt",
|
|
},
|
|
)
|
|
)
|
|
diffed = _extract_structured(
|
|
await server.call_tool("workspace_diff", {"workspace_id": workspace_id})
|
|
)
|
|
export_path = tmp_path / "exported-more.txt"
|
|
exported = _extract_structured(
|
|
await server.call_tool(
|
|
"workspace_export",
|
|
{
|
|
"workspace_id": workspace_id,
|
|
"path": "subdir/more.txt",
|
|
"output_path": str(export_path),
|
|
},
|
|
)
|
|
)
|
|
opened = _extract_structured(
|
|
await server.call_tool("shell_open", {"workspace_id": workspace_id})
|
|
)
|
|
shell_id = str(opened["shell_id"])
|
|
written = _extract_structured(
|
|
await server.call_tool(
|
|
"shell_write",
|
|
{
|
|
"workspace_id": workspace_id,
|
|
"shell_id": shell_id,
|
|
"input": "pwd",
|
|
},
|
|
)
|
|
)
|
|
read = _extract_structured(
|
|
await server.call_tool(
|
|
"shell_read",
|
|
{
|
|
"workspace_id": workspace_id,
|
|
"shell_id": shell_id,
|
|
},
|
|
)
|
|
)
|
|
deadline = time.time() + 5
|
|
while "/workspace" not in str(read["output"]) and time.time() < deadline:
|
|
read = _extract_structured(
|
|
await server.call_tool(
|
|
"shell_read",
|
|
{
|
|
"workspace_id": workspace_id,
|
|
"shell_id": shell_id,
|
|
"cursor": 0,
|
|
},
|
|
)
|
|
)
|
|
await asyncio.sleep(0.05)
|
|
signaled = _extract_structured(
|
|
await server.call_tool(
|
|
"shell_signal",
|
|
{
|
|
"workspace_id": workspace_id,
|
|
"shell_id": shell_id,
|
|
},
|
|
)
|
|
)
|
|
closed = _extract_structured(
|
|
await server.call_tool(
|
|
"shell_close",
|
|
{
|
|
"workspace_id": workspace_id,
|
|
"shell_id": shell_id,
|
|
},
|
|
)
|
|
)
|
|
logs = _extract_structured(
|
|
await server.call_tool("workspace_logs", {"workspace_id": workspace_id})
|
|
)
|
|
deleted = _extract_structured(
|
|
await server.call_tool("workspace_delete", {"workspace_id": workspace_id})
|
|
)
|
|
return (
|
|
created,
|
|
synced,
|
|
executed,
|
|
diffed,
|
|
exported,
|
|
opened,
|
|
written,
|
|
read,
|
|
signaled,
|
|
closed,
|
|
logs,
|
|
deleted,
|
|
)
|
|
|
|
(
|
|
created,
|
|
synced,
|
|
executed,
|
|
diffed,
|
|
exported,
|
|
opened,
|
|
written,
|
|
read,
|
|
signaled,
|
|
closed,
|
|
logs,
|
|
deleted,
|
|
) = asyncio.run(_run())
|
|
assert created["state"] == "started"
|
|
assert created["workspace_seed"]["mode"] == "directory"
|
|
assert synced["workspace_sync"]["destination"] == "/workspace/subdir"
|
|
assert executed["stdout"] == "more\n"
|
|
assert diffed["changed"] is True
|
|
assert exported["artifact_type"] == "file"
|
|
assert Path(str(exported["output_path"])).read_text(encoding="utf-8") == "more\n"
|
|
assert opened["state"] == "running"
|
|
assert written["input_length"] == 3
|
|
assert "/workspace" in read["output"]
|
|
assert signaled["signal"] == "INT"
|
|
assert closed["closed"] is True
|
|
assert logs["count"] == 1
|
|
assert deleted["deleted"] is True
|