Thales Maciel
fc72fcd3a1
Add explicit workspace secrets across the CLI, SDK, and MCP, with create-time secret definitions and per-call secret-to-env mapping for exec, shell open, and service start. Persist only safe secret metadata in workspace records, materialize secret files under /run/pyro-secrets, and redact secret values from exec output, shell reads, service logs, and surfaced errors. Fix the remaining real-guest shell gap by shipping bundled guest init alongside the guest agent and patching both into guest-backed workspace rootfs images before boot. The new init mounts devpts so PTY shells work on Firecracker guests, while reset continues to recreate the sandbox and re-materialize secrets from stored task-local secret material. 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 with secrets, secret-backed exec, shell, service, reset, and delete.
250 lines
10 KiB
Python
250 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import io
|
|
import tomllib
|
|
from contextlib import redirect_stdout
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
from pyro_mcp import Pyro, __version__
|
|
from pyro_mcp.cli import _build_parser
|
|
from pyro_mcp.contract import (
|
|
PUBLIC_CLI_COMMANDS,
|
|
PUBLIC_CLI_DEMO_SUBCOMMANDS,
|
|
PUBLIC_CLI_ENV_SUBCOMMANDS,
|
|
PUBLIC_CLI_RUN_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_CREATE_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_DIFF_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_EXEC_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_EXPORT_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_RESET_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SERVICE_LIST_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SERVICE_LOGS_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SERVICE_START_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SERVICE_STATUS_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SERVICE_STOP_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SERVICE_SUBCOMMANDS,
|
|
PUBLIC_CLI_WORKSPACE_SHELL_CLOSE_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SHELL_OPEN_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SHELL_READ_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SHELL_SIGNAL_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SHELL_SUBCOMMANDS,
|
|
PUBLIC_CLI_WORKSPACE_SHELL_WRITE_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SNAPSHOT_CREATE_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SNAPSHOT_DELETE_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SNAPSHOT_LIST_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SNAPSHOT_SUBCOMMANDS,
|
|
PUBLIC_CLI_WORKSPACE_SUBCOMMANDS,
|
|
PUBLIC_CLI_WORKSPACE_SYNC_PUSH_FLAGS,
|
|
PUBLIC_CLI_WORKSPACE_SYNC_SUBCOMMANDS,
|
|
PUBLIC_MCP_TOOLS,
|
|
PUBLIC_SDK_METHODS,
|
|
)
|
|
from pyro_mcp.vm_manager import VmManager
|
|
from pyro_mcp.vm_network import TapNetworkManager
|
|
|
|
|
|
def _subparser_choice(parser: argparse.ArgumentParser, name: str) -> argparse.ArgumentParser:
|
|
subparsers = getattr(parser, "_subparsers", None)
|
|
if subparsers is None:
|
|
raise AssertionError("parser does not define subparsers")
|
|
group_actions = cast(list[Any], subparsers._group_actions) # noqa: SLF001
|
|
if not group_actions:
|
|
raise AssertionError("parser subparsers are empty")
|
|
choices = cast(dict[str, argparse.ArgumentParser], group_actions[0].choices)
|
|
return choices[name]
|
|
|
|
|
|
def test_public_sdk_methods_exist() -> None:
|
|
assert tuple(sorted(PUBLIC_SDK_METHODS)) == PUBLIC_SDK_METHODS
|
|
for method_name in PUBLIC_SDK_METHODS:
|
|
assert hasattr(Pyro, method_name), method_name
|
|
|
|
|
|
def test_public_cli_help_lists_commands_and_run_flags() -> None:
|
|
parser = _build_parser()
|
|
help_text = parser.format_help()
|
|
assert "--version" in help_text
|
|
for command_name in PUBLIC_CLI_COMMANDS:
|
|
assert command_name in help_text
|
|
|
|
run_parser = _build_parser()
|
|
run_help = run_parser.parse_args(["run", "debian:12-base", "--", "true"])
|
|
assert run_help.command == "run"
|
|
assert run_help.environment == "debian:12-base"
|
|
assert run_help.vcpu_count == 1
|
|
assert run_help.mem_mib == 1024
|
|
|
|
run_help_text = _subparser_choice(parser, "run").format_help()
|
|
for flag in PUBLIC_CLI_RUN_FLAGS:
|
|
assert flag in run_help_text
|
|
|
|
env_help_text = _subparser_choice(parser, "env").format_help()
|
|
for subcommand_name in PUBLIC_CLI_ENV_SUBCOMMANDS:
|
|
assert subcommand_name in env_help_text
|
|
|
|
workspace_help_text = _subparser_choice(parser, "workspace").format_help()
|
|
for subcommand_name in PUBLIC_CLI_WORKSPACE_SUBCOMMANDS:
|
|
assert subcommand_name in workspace_help_text
|
|
workspace_create_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"), "create"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_CREATE_FLAGS:
|
|
assert flag in workspace_create_help_text
|
|
workspace_exec_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"), "exec"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_EXEC_FLAGS:
|
|
assert flag in workspace_exec_help_text
|
|
workspace_sync_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"),
|
|
"sync",
|
|
).format_help()
|
|
for subcommand_name in PUBLIC_CLI_WORKSPACE_SYNC_SUBCOMMANDS:
|
|
assert subcommand_name in workspace_sync_help_text
|
|
workspace_sync_push_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "sync"), "push"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SYNC_PUSH_FLAGS:
|
|
assert flag in workspace_sync_push_help_text
|
|
workspace_export_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"), "export"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_EXPORT_FLAGS:
|
|
assert flag in workspace_export_help_text
|
|
workspace_diff_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"), "diff"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_DIFF_FLAGS:
|
|
assert flag in workspace_diff_help_text
|
|
workspace_snapshot_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"),
|
|
"snapshot",
|
|
).format_help()
|
|
for subcommand_name in PUBLIC_CLI_WORKSPACE_SNAPSHOT_SUBCOMMANDS:
|
|
assert subcommand_name in workspace_snapshot_help_text
|
|
workspace_snapshot_create_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "snapshot"),
|
|
"create",
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SNAPSHOT_CREATE_FLAGS:
|
|
assert flag in workspace_snapshot_create_help_text
|
|
workspace_snapshot_list_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "snapshot"),
|
|
"list",
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SNAPSHOT_LIST_FLAGS:
|
|
assert flag in workspace_snapshot_list_help_text
|
|
workspace_snapshot_delete_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "snapshot"),
|
|
"delete",
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SNAPSHOT_DELETE_FLAGS:
|
|
assert flag in workspace_snapshot_delete_help_text
|
|
workspace_reset_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"), "reset"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_RESET_FLAGS:
|
|
assert flag in workspace_reset_help_text
|
|
workspace_shell_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"),
|
|
"shell",
|
|
).format_help()
|
|
for subcommand_name in PUBLIC_CLI_WORKSPACE_SHELL_SUBCOMMANDS:
|
|
assert subcommand_name in workspace_shell_help_text
|
|
workspace_shell_open_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "open"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SHELL_OPEN_FLAGS:
|
|
assert flag in workspace_shell_open_help_text
|
|
workspace_shell_read_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "read"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SHELL_READ_FLAGS:
|
|
assert flag in workspace_shell_read_help_text
|
|
workspace_shell_write_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "write"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SHELL_WRITE_FLAGS:
|
|
assert flag in workspace_shell_write_help_text
|
|
workspace_shell_signal_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "signal"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SHELL_SIGNAL_FLAGS:
|
|
assert flag in workspace_shell_signal_help_text
|
|
workspace_shell_close_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "close"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SHELL_CLOSE_FLAGS:
|
|
assert flag in workspace_shell_close_help_text
|
|
workspace_service_help_text = _subparser_choice(
|
|
_subparser_choice(parser, "workspace"),
|
|
"service",
|
|
).format_help()
|
|
for subcommand_name in PUBLIC_CLI_WORKSPACE_SERVICE_SUBCOMMANDS:
|
|
assert subcommand_name in workspace_service_help_text
|
|
workspace_service_start_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "service"), "start"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SERVICE_START_FLAGS:
|
|
assert flag in workspace_service_start_help_text
|
|
workspace_service_list_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "service"), "list"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SERVICE_LIST_FLAGS:
|
|
assert flag in workspace_service_list_help_text
|
|
workspace_service_status_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "service"), "status"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SERVICE_STATUS_FLAGS:
|
|
assert flag in workspace_service_status_help_text
|
|
workspace_service_logs_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "service"), "logs"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SERVICE_LOGS_FLAGS:
|
|
assert flag in workspace_service_logs_help_text
|
|
workspace_service_stop_help_text = _subparser_choice(
|
|
_subparser_choice(_subparser_choice(parser, "workspace"), "service"), "stop"
|
|
).format_help()
|
|
for flag in PUBLIC_CLI_WORKSPACE_SERVICE_STOP_FLAGS:
|
|
assert flag in workspace_service_stop_help_text
|
|
|
|
demo_help_text = _subparser_choice(parser, "demo").format_help()
|
|
for subcommand_name in PUBLIC_CLI_DEMO_SUBCOMMANDS:
|
|
assert subcommand_name in demo_help_text
|
|
|
|
|
|
def test_public_cli_version_matches_package_version() -> None:
|
|
parser = _build_parser()
|
|
stdout = io.StringIO()
|
|
with pytest.raises(SystemExit, match="0"), redirect_stdout(stdout):
|
|
parser.parse_args(["--version"])
|
|
assert stdout.getvalue().strip().endswith(f" {__version__}")
|
|
|
|
|
|
def test_public_mcp_tools_match_contract(tmp_path: Path) -> None:
|
|
pyro = Pyro(
|
|
manager=VmManager(
|
|
backend_name="mock",
|
|
base_dir=tmp_path / "vms",
|
|
network_manager=TapNetworkManager(enabled=False),
|
|
)
|
|
)
|
|
|
|
async def _run() -> tuple[str, ...]:
|
|
server = pyro.create_server()
|
|
tools = await server.list_tools()
|
|
return tuple(sorted(tool.name for tool in tools))
|
|
|
|
tool_names = asyncio.run(_run())
|
|
assert tool_names == tuple(sorted(PUBLIC_MCP_TOOLS))
|
|
|
|
|
|
def test_pyproject_exposes_single_public_cli_script() -> None:
|
|
pyproject = tomllib.loads(Path("pyproject.toml").read_text(encoding="utf-8"))
|
|
scripts = pyproject["project"]["scripts"]
|
|
assert scripts == {"pyro": "pyro_mcp.cli:main"}
|