diff --git a/CHANGELOG.md b/CHANGELOG.md index b813849..e1ff83e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,16 @@ All notable user-visible changes to `pyro-mcp` are documented here. +## 2.8.0 + +- Added explicit named workspace snapshots across the CLI, Python SDK, and MCP server with + `pyro workspace snapshot *`, `Pyro.create_snapshot()` / `list_snapshots()` / + `delete_snapshot()`, and the matching `snapshot_*` MCP tools. +- Added `pyro workspace reset` and `Pyro.reset_workspace()` so a workspace can recreate its full + sandbox from the immutable baseline or one named snapshot while keeping the same identity. +- Made reset a full-sandbox recovery path that clears command history, shells, and services while + preserving the workspace spec, named snapshots, and immutable baseline. + ## 2.7.0 - Added first-class workspace services across the CLI, Python SDK, and MCP server with diff --git a/README.md b/README.md index ba89cc6..518ed0f 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ It exposes the same runtime in three public forms: - First run transcript: [docs/first-run.md](docs/first-run.md) - Terminal walkthrough GIF: [docs/assets/first-run.gif](docs/assets/first-run.gif) - PyPI package: [pypi.org/project/pyro-mcp](https://pypi.org/project/pyro-mcp/) -- What's new in 2.7.0: [CHANGELOG.md#270](CHANGELOG.md#270) +- What's new in 2.8.0: [CHANGELOG.md#280](CHANGELOG.md#280) - Host requirements: [docs/host-requirements.md](docs/host-requirements.md) - Integration targets: [docs/integrations.md](docs/integrations.md) - Public contract: [docs/public-contract.md](docs/public-contract.md) @@ -57,7 +57,7 @@ What success looks like: ```bash Platform: linux-x86_64 Runtime: PASS -Catalog version: 2.7.0 +Catalog version: 2.8.0 ... [pull] phase=install environment=debian:12 [pull] phase=ready environment=debian:12 @@ -79,6 +79,8 @@ After the quickstart works: - create a persistent workspace with `uvx --from pyro-mcp pyro workspace create debian:12 --seed-path ./repo` - update a live workspace from the host with `uvx --from pyro-mcp pyro workspace sync push WORKSPACE_ID ./changes` - diff the live workspace against its create-time baseline with `uvx --from pyro-mcp pyro workspace diff WORKSPACE_ID` +- capture a checkpoint with `uvx --from pyro-mcp pyro workspace snapshot create WORKSPACE_ID checkpoint` +- reset a broken workspace with `uvx --from pyro-mcp pyro workspace reset WORKSPACE_ID --snapshot checkpoint` - export a changed file or directory with `uvx --from pyro-mcp pyro workspace export WORKSPACE_ID note.txt --output ./note.txt` - open a persistent interactive shell with `uvx --from pyro-mcp pyro workspace shell open WORKSPACE_ID` - start long-running workspace services with `uvx --from pyro-mcp pyro workspace service start WORKSPACE_ID app --ready-file .ready -- sh -lc 'touch .ready && while true; do sleep 60; done'` @@ -135,7 +137,7 @@ uvx --from pyro-mcp pyro env list Expected output: ```bash -Catalog version: 2.7.0 +Catalog version: 2.8.0 debian:12 [installed|not installed] Debian 12 environment with Git preinstalled for common agent workflows. debian:12-base [installed|not installed] Minimal Debian 12 environment for shell and core Unix tooling. debian:12-build [installed|not installed] Debian 12 environment with Git and common build tools preinstalled. @@ -214,6 +216,9 @@ pyro workspace create debian:12 --seed-path ./repo pyro workspace sync push WORKSPACE_ID ./changes --dest src pyro workspace exec WORKSPACE_ID -- cat src/note.txt pyro workspace diff WORKSPACE_ID +pyro workspace snapshot create WORKSPACE_ID checkpoint +pyro workspace reset WORKSPACE_ID --snapshot checkpoint +pyro workspace reset WORKSPACE_ID pyro workspace export WORKSPACE_ID src/note.txt --output ./note.txt pyro workspace shell open WORKSPACE_ID pyro workspace shell write WORKSPACE_ID SHELL_ID --input 'pwd' @@ -234,10 +239,12 @@ Persistent workspaces start in `/workspace` and keep command history until you d machine consumption, add `--json` and read the returned `workspace_id`. Use `--seed-path` when you want the workspace to start from a host directory or a local `.tar` / `.tar.gz` / `.tgz` archive instead of an empty workspace. Use `pyro workspace sync push` when you want to import -later host-side changes into a started workspace. Sync is non-atomic in `2.7.0`; if it fails -partway through, delete and recreate the workspace from its seed. Use `pyro workspace diff` to -compare the live `/workspace` tree to its immutable create-time baseline, and `pyro workspace export` -to copy one changed file or directory back to the host. Use `pyro workspace exec` for one-shot +later host-side changes into a started workspace. Sync is non-atomic in `2.8.0`; if it fails +partway through, prefer `pyro workspace reset` to recover from `baseline` or one named snapshot. +Use `pyro workspace diff` to compare the live `/workspace` tree to its immutable create-time +baseline, and `pyro workspace export` to copy one changed file or directory back to the host. Use +`pyro workspace snapshot *` and `pyro workspace reset` when you want explicit checkpoints and +full-sandbox recovery. Use `pyro workspace exec` for one-shot non-interactive commands inside a live workspace, and `pyro workspace shell *` when you need a persistent PTY session that keeps interactive shell state between calls. Use `pyro workspace service *` when the workspace needs one or more long-running background processes. diff --git a/docs/first-run.md b/docs/first-run.md index c76c45f..0eb14da 100644 --- a/docs/first-run.md +++ b/docs/first-run.md @@ -22,7 +22,7 @@ Networking: tun=yes ip_forward=yes ```bash $ uvx --from pyro-mcp pyro env list -Catalog version: 2.7.0 +Catalog version: 2.8.0 debian:12 [installed|not installed] Debian 12 environment with Git preinstalled for common agent workflows. debian:12-base [installed|not installed] Minimal Debian 12 environment for shell and core Unix tooling. debian:12-build [installed|not installed] Debian 12 environment with Git and common build tools preinstalled. @@ -73,6 +73,8 @@ $ uvx --from pyro-mcp pyro demo $ uvx --from pyro-mcp pyro workspace create debian:12 --seed-path ./repo $ uvx --from pyro-mcp pyro workspace sync push WORKSPACE_ID ./changes $ uvx --from pyro-mcp pyro workspace diff WORKSPACE_ID +$ uvx --from pyro-mcp pyro workspace snapshot create WORKSPACE_ID checkpoint +$ uvx --from pyro-mcp pyro workspace reset WORKSPACE_ID --snapshot checkpoint $ uvx --from pyro-mcp pyro workspace export WORKSPACE_ID note.txt --output ./note.txt $ uvx --from pyro-mcp pyro workspace shell open WORKSPACE_ID $ uvx --from pyro-mcp pyro workspace service start WORKSPACE_ID app --ready-file .ready -- sh -lc 'touch .ready && while true; do sleep 60; done' @@ -107,6 +109,17 @@ $ uvx --from pyro-mcp pyro workspace diff WORKSPACE_ID +++ b/src/note.txt @@ ... +$ uvx --from pyro-mcp pyro workspace snapshot create WORKSPACE_ID checkpoint +[workspace-snapshot-create] snapshot_name=checkpoint kind=named entry_count=... bytes_written=... + +$ uvx --from pyro-mcp pyro workspace reset WORKSPACE_ID --snapshot checkpoint +Workspace reset from snapshot: checkpoint (named) +[workspace-reset] destination=/workspace entry_count=... bytes_written=... +Workspace ID: ... +State: started +Command count: 0 +Reset count: 1 + $ uvx --from pyro-mcp pyro workspace export WORKSPACE_ID src/note.txt --output ./note.txt [workspace-export] workspace_id=... workspace_path=/workspace/src/note.txt output_path=... artifact_type=file entry_count=... bytes_written=... execution_mode=guest_vsock @@ -157,10 +170,11 @@ $ uvx --from pyro-mcp pyro workspace service stop WORKSPACE_ID worker Use `--seed-path` when the workspace should start from a host directory or a local `.tar` / `.tar.gz` / `.tgz` archive instead of an empty `/workspace`. Use `pyro workspace sync push` when you need to import later host-side changes into a started -workspace. Sync is non-atomic in `2.7.0`; if it fails partway through, delete and recreate the -workspace. Use `pyro workspace diff` to compare the current `/workspace` tree to its immutable -create-time baseline, and `pyro workspace export` to copy one changed file or directory back to -the host. Use `pyro workspace exec` for one-shot commands and `pyro workspace shell *` when you +workspace. Sync is non-atomic in `2.8.0`; if it fails partway through, prefer `pyro workspace reset` +to recover from `baseline` or one named snapshot. Use `pyro workspace diff` to compare the current +`/workspace` tree to its immutable create-time baseline, `pyro workspace snapshot *` to create +named checkpoints, and `pyro workspace export` to copy one changed file or directory back to the +host. Use `pyro workspace exec` for one-shot commands and `pyro workspace shell *` when you need a persistent interactive PTY session in that same workspace. Use `pyro workspace service *` when the workspace needs long-running background processes with typed readiness checks. Internal service state and logs stay outside `/workspace`, so service runtime data does not appear in diff --git a/docs/install.md b/docs/install.md index fc4d5c7..89d41b6 100644 --- a/docs/install.md +++ b/docs/install.md @@ -83,7 +83,7 @@ uvx --from pyro-mcp pyro env list Expected output: ```bash -Catalog version: 2.7.0 +Catalog version: 2.8.0 debian:12 [installed|not installed] Debian 12 environment with Git preinstalled for common agent workflows. debian:12-base [installed|not installed] Minimal Debian 12 environment for shell and core Unix tooling. debian:12-build [installed|not installed] Debian 12 environment with Git and common build tools preinstalled. @@ -177,6 +177,7 @@ After the CLI path works, you can move on to: - persistent workspaces: `pyro workspace create debian:12 --seed-path ./repo` - live workspace updates: `pyro workspace sync push WORKSPACE_ID ./changes` - baseline diff: `pyro workspace diff WORKSPACE_ID` +- snapshots and reset: `pyro workspace snapshot create WORKSPACE_ID checkpoint` and `pyro workspace reset WORKSPACE_ID --snapshot checkpoint` - host export: `pyro workspace export WORKSPACE_ID note.txt --output ./note.txt` - interactive shells: `pyro workspace shell open WORKSPACE_ID` - long-running services: `pyro workspace service start WORKSPACE_ID app --ready-file .ready -- sh -lc 'touch .ready && while true; do sleep 60; done'` @@ -193,6 +194,9 @@ pyro workspace create debian:12 --seed-path ./repo pyro workspace sync push WORKSPACE_ID ./changes --dest src pyro workspace exec WORKSPACE_ID -- cat src/note.txt pyro workspace diff WORKSPACE_ID +pyro workspace snapshot create WORKSPACE_ID checkpoint +pyro workspace reset WORKSPACE_ID --snapshot checkpoint +pyro workspace reset WORKSPACE_ID pyro workspace export WORKSPACE_ID src/note.txt --output ./note.txt pyro workspace shell open WORKSPACE_ID pyro workspace shell write WORKSPACE_ID SHELL_ID --input 'pwd' @@ -213,9 +217,10 @@ Workspace commands default to the persistent `/workspace` directory inside the g the identifier programmatically, use `--json` and read the `workspace_id` field. Use `--seed-path` when the workspace should start from a host directory or a local `.tar` / `.tar.gz` / `.tgz` archive. Use `pyro workspace sync push` for later host-side changes to a started workspace. Sync -is non-atomic in `2.7.0`; if it fails partway through, delete and recreate the workspace from its -seed. Use `pyro workspace diff` to compare the current workspace tree to its immutable create-time -baseline, and `pyro workspace export` to copy one changed file or directory back to the host. Use +is non-atomic in `2.8.0`; if it fails partway through, prefer `pyro workspace reset` to recover +from `baseline` or one named snapshot. Use `pyro workspace diff` to compare the current workspace +tree to its immutable create-time baseline, `pyro workspace snapshot *` to capture named +checkpoints, and `pyro workspace export` to copy one changed file or directory back to the host. Use `pyro workspace exec` for one-shot commands and `pyro workspace shell *` when you need an interactive PTY that survives across separate calls. Use `pyro workspace service *` when the workspace needs long-running background processes with typed readiness probes. Service metadata and diff --git a/docs/public-contract.md b/docs/public-contract.md index 8f4d5af..4e06343 100644 --- a/docs/public-contract.md +++ b/docs/public-contract.md @@ -24,6 +24,10 @@ Top-level commands: - `pyro workspace exec` - `pyro workspace export` - `pyro workspace diff` +- `pyro workspace snapshot create` +- `pyro workspace snapshot list` +- `pyro workspace snapshot delete` +- `pyro workspace reset` - `pyro workspace service start` - `pyro workspace service list` - `pyro workspace service status` @@ -63,11 +67,14 @@ Behavioral guarantees: - `pyro workspace sync push WORKSPACE_ID SOURCE_PATH [--dest WORKSPACE_PATH]` imports later host-side directory or archive content into a started workspace. - `pyro workspace export WORKSPACE_ID PATH --output HOST_PATH` exports one file or directory from `/workspace` back to the host. - `pyro workspace diff WORKSPACE_ID` compares the current `/workspace` tree to the immutable create-time baseline. +- `pyro workspace snapshot *` manages explicit named snapshots in addition to the implicit `baseline`. +- `pyro workspace reset WORKSPACE_ID [--snapshot SNAPSHOT_NAME|baseline]` recreates the full sandbox and restores `/workspace` from the chosen snapshot. - `pyro workspace service *` manages long-running named services inside one started workspace with typed readiness probes. - `pyro workspace exec` runs in the persistent `/workspace` for that workspace and does not auto-clean. - `pyro workspace shell *` manages persistent PTY sessions inside a started workspace. - `pyro workspace logs` returns persisted command history for that workspace until `pyro workspace delete`. - Workspace create/status results expose `workspace_seed` metadata describing how `/workspace` was initialized. +- Workspace create/status/reset results expose `reset_count` and `last_reset_at`. - `pyro workspace status` includes aggregate `service_count` and `running_service_count` fields. ## Python SDK Contract @@ -89,6 +96,10 @@ Supported public entrypoints: - `Pyro.push_workspace_sync(workspace_id, source_path, *, dest="/workspace")` - `Pyro.export_workspace(workspace_id, path, *, output_path)` - `Pyro.diff_workspace(workspace_id)` +- `Pyro.create_snapshot(workspace_id, snapshot_name)` +- `Pyro.list_snapshots(workspace_id)` +- `Pyro.delete_snapshot(workspace_id, snapshot_name)` +- `Pyro.reset_workspace(workspace_id, *, snapshot="baseline")` - `Pyro.start_service(workspace_id, service_name, *, command, cwd="/workspace", readiness=None, ready_timeout_seconds=30, ready_interval_ms=500)` - `Pyro.list_services(workspace_id)` - `Pyro.status_service(workspace_id, service_name)` @@ -124,6 +135,10 @@ Stable public method names: - `push_workspace_sync(workspace_id, source_path, *, dest="/workspace")` - `export_workspace(workspace_id, path, *, output_path)` - `diff_workspace(workspace_id)` +- `create_snapshot(workspace_id, snapshot_name)` +- `list_snapshots(workspace_id)` +- `delete_snapshot(workspace_id, snapshot_name)` +- `reset_workspace(workspace_id, *, snapshot="baseline")` - `start_service(workspace_id, service_name, *, command, cwd="/workspace", readiness=None, ready_timeout_seconds=30, ready_interval_ms=500)` - `list_services(workspace_id)` - `status_service(workspace_id, service_name)` @@ -157,6 +172,10 @@ Behavioral defaults: - `Pyro.push_workspace_sync(...)` imports later host-side directory or archive content into a started workspace. - `Pyro.export_workspace(...)` exports one file or directory from `/workspace` to an explicit host path. - `Pyro.diff_workspace(...)` compares the current `/workspace` tree to the immutable create-time baseline. +- `Pyro.create_snapshot(...)` captures one named `/workspace` checkpoint. +- `Pyro.list_snapshots(...)` lists the implicit `baseline` plus any named snapshots. +- `Pyro.delete_snapshot(...)` deletes one named snapshot while leaving `baseline` intact. +- `Pyro.reset_workspace(...)` recreates the full sandbox from `baseline` or one named snapshot and clears command, shell, and service history. - `Pyro.start_service(...)` starts one named long-running process in a started workspace and waits for its typed readiness probe when configured. - `Pyro.list_services(...)`, `Pyro.status_service(...)`, `Pyro.logs_service(...)`, and `Pyro.stop_service(...)` manage those persisted workspace services. - `Pyro.exec_vm(...)` runs one command and auto-cleans that VM after the exec completes. @@ -190,6 +209,10 @@ Persistent workspace tools: - `workspace_exec` - `workspace_export` - `workspace_diff` +- `snapshot_create` +- `snapshot_list` +- `snapshot_delete` +- `workspace_reset` - `service_start` - `service_list` - `service_status` @@ -214,6 +237,8 @@ Behavioral defaults: - `workspace_sync_push` imports later host-side directory or archive content into a started workspace, with an optional `dest` under `/workspace`. - `workspace_export` exports one file or directory from `/workspace` to an explicit host path. - `workspace_diff` compares the current `/workspace` tree to the immutable create-time baseline. +- `snapshot_create`, `snapshot_list`, and `snapshot_delete` manage explicit named snapshots in addition to the implicit `baseline`. +- `workspace_reset` recreates the full sandbox and restores `/workspace` from `baseline` or one named snapshot. - `service_start`, `service_list`, `service_status`, `service_logs`, and `service_stop` manage persistent named services inside a started workspace. - `vm_exec` runs one command and auto-cleans that VM after the exec completes. - `workspace_exec` runs one command in a persistent `/workspace` and leaves the workspace alive. diff --git a/docs/roadmap/task-workspace-ga.md b/docs/roadmap/task-workspace-ga.md index 91702dc..a8461c8 100644 --- a/docs/roadmap/task-workspace-ga.md +++ b/docs/roadmap/task-workspace-ga.md @@ -2,14 +2,15 @@ This roadmap turns the agent-workspace vision into release-sized milestones. -Current baseline is `2.7.0`: +Current baseline is `2.8.0`: - workspace persistence exists and the public surface is now workspace-first - host crossing currently covers create-time seeding, later sync push, and explicit export - persistent PTY shell sessions exist alongside one-shot `workspace exec` - immutable create-time baselines now power whole-workspace diff - multi-service lifecycle exists with typed readiness and aggregate workspace status counts -- no snapshot, reset, or secrets contract exists yet +- named snapshots and full workspace reset now exist +- no secrets or explicit host port publication contract exists yet Locked roadmap decisions: @@ -32,7 +33,7 @@ also expected to update: 2. [`2.5.0` PTY Shell Sessions](task-workspace-ga/2.5.0-pty-shell-sessions.md) - Done 3. [`2.6.0` Structured Export And Baseline Diff](task-workspace-ga/2.6.0-structured-export-and-baseline-diff.md) - Done 4. [`2.7.0` Service Lifecycle And Typed Readiness](task-workspace-ga/2.7.0-service-lifecycle-and-typed-readiness.md) - Done -5. [`2.8.0` Named Snapshots And Reset](task-workspace-ga/2.8.0-named-snapshots-and-reset.md) +5. [`2.8.0` Named Snapshots And Reset](task-workspace-ga/2.8.0-named-snapshots-and-reset.md) - Done 6. [`2.9.0` Secrets](task-workspace-ga/2.9.0-secrets.md) 7. [`2.10.0` Network Policy And Host Port Publication](task-workspace-ga/2.10.0-network-policy-and-host-port-publication.md) 8. [`3.0.0` Stable Workspace Product](task-workspace-ga/3.0.0-stable-workspace-product.md) diff --git a/docs/roadmap/task-workspace-ga/2.8.0-named-snapshots-and-reset.md b/docs/roadmap/task-workspace-ga/2.8.0-named-snapshots-and-reset.md index 0a0ec16..21de48c 100644 --- a/docs/roadmap/task-workspace-ga/2.8.0-named-snapshots-and-reset.md +++ b/docs/roadmap/task-workspace-ga/2.8.0-named-snapshots-and-reset.md @@ -1,5 +1,7 @@ # `2.8.0` Named Snapshots And Reset +Status: Done + ## Goal Turn reset into a first-class workflow primitive and add explicit named diff --git a/examples/python_workspace.py b/examples/python_workspace.py index f6393a8..f26957a 100644 --- a/examples/python_workspace.py +++ b/examples/python_workspace.py @@ -23,6 +23,8 @@ def main() -> None: print(result["stdout"], end="") diff_result = pyro.diff_workspace(workspace_id) print(f"changed={diff_result['changed']} total={diff_result['summary']['total']}") + snapshot = pyro.create_snapshot(workspace_id, "checkpoint") + print(snapshot["snapshot"]["snapshot_name"]) exported_path = Path(export_dir, "note.txt") pyro.export_workspace(workspace_id, "note.txt", output_path=exported_path) print(exported_path.read_text(encoding="utf-8"), end="") @@ -39,6 +41,8 @@ def main() -> None: service_logs = pyro.logs_service(workspace_id, "web", tail_lines=20) print(f"service_stdout_len={len(service_logs['stdout'])}") pyro.stop_service(workspace_id, "web") + reset = pyro.reset_workspace(workspace_id, snapshot="checkpoint") + print(f"reset_count={reset['reset_count']}") logs = pyro.logs_workspace(workspace_id) print(f"workspace_id={workspace_id} command_count={logs['count']}") finally: diff --git a/pyproject.toml b/pyproject.toml index 6d5b25c..ed8a41d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pyro-mcp" -version = "2.7.0" +version = "2.8.0" description = "Ephemeral Firecracker sandboxes with curated environments, persistent workspaces, and MCP tools." readme = "README.md" license = { file = "LICENSE" } diff --git a/src/pyro_mcp/api.py b/src/pyro_mcp/api.py index a364c78..0e5ea1e 100644 --- a/src/pyro_mcp/api.py +++ b/src/pyro_mcp/api.py @@ -146,6 +146,23 @@ class Pyro: def diff_workspace(self, workspace_id: str) -> dict[str, Any]: return self._manager.diff_workspace(workspace_id) + def create_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + return self._manager.create_snapshot(workspace_id, snapshot_name) + + def list_snapshots(self, workspace_id: str) -> dict[str, Any]: + return self._manager.list_snapshots(workspace_id) + + def delete_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + return self._manager.delete_snapshot(workspace_id, snapshot_name) + + def reset_workspace( + self, + workspace_id: str, + *, + snapshot: str = "baseline", + ) -> dict[str, Any]: + return self._manager.reset_workspace(workspace_id, snapshot=snapshot) + def open_shell( self, workspace_id: str, @@ -444,6 +461,29 @@ class Pyro: """Compare `/workspace` to the immutable create-time baseline.""" return self.diff_workspace(workspace_id) + @server.tool() + async def snapshot_create(workspace_id: str, snapshot_name: str) -> dict[str, Any]: + """Create one named workspace snapshot from the current `/workspace` tree.""" + return self.create_snapshot(workspace_id, snapshot_name) + + @server.tool() + async def snapshot_list(workspace_id: str) -> dict[str, Any]: + """List the baseline plus named snapshots for one workspace.""" + return self.list_snapshots(workspace_id) + + @server.tool() + async def snapshot_delete(workspace_id: str, snapshot_name: str) -> dict[str, Any]: + """Delete one named snapshot from a workspace.""" + return self.delete_snapshot(workspace_id, snapshot_name) + + @server.tool() + async def workspace_reset( + workspace_id: str, + snapshot: str = "baseline", + ) -> dict[str, Any]: + """Recreate a workspace and restore `/workspace` from baseline or one named snapshot.""" + return self.reset_workspace(workspace_id, snapshot=snapshot) + @server.tool() async def shell_open( workspace_id: str, diff --git a/src/pyro_mcp/cli.py b/src/pyro_mcp/cli.py index 9e7a319..b80c153 100644 --- a/src/pyro_mcp/cli.py +++ b/src/pyro_mcp/cli.py @@ -174,6 +174,10 @@ def _print_workspace_summary_human(payload: dict[str, Any], *, action: str) -> N f"{int(payload.get('mem_mib', 0))} MiB" ) print(f"Command count: {int(payload.get('command_count', 0))}") + print(f"Reset count: {int(payload.get('reset_count', 0))}") + last_reset_at = payload.get("last_reset_at") + if last_reset_at is not None: + print(f"Last reset at: {last_reset_at}") print( "Services: " f"{int(payload.get('running_service_count', 0))}/" @@ -281,6 +285,55 @@ def _print_workspace_logs_human(payload: dict[str, Any]) -> None: print(stderr, end="" if stderr.endswith("\n") else "\n", file=sys.stderr) +def _print_workspace_snapshot_human(payload: dict[str, Any], *, prefix: str) -> None: + snapshot = payload.get("snapshot") + if not isinstance(snapshot, dict): + print(f"[{prefix}] workspace_id={str(payload.get('workspace_id', 'unknown'))}") + return + print( + f"[{prefix}] " + f"workspace_id={str(payload.get('workspace_id', 'unknown'))} " + f"snapshot_name={str(snapshot.get('snapshot_name', 'unknown'))} " + f"kind={str(snapshot.get('kind', 'unknown'))} " + f"entry_count={int(snapshot.get('entry_count', 0))} " + f"bytes_written={int(snapshot.get('bytes_written', 0))}" + ) + + +def _print_workspace_snapshot_list_human(payload: dict[str, Any]) -> None: + snapshots = payload.get("snapshots") + if not isinstance(snapshots, list) or not snapshots: + print("No workspace snapshots found.") + return + for snapshot in snapshots: + if not isinstance(snapshot, dict): + continue + print( + f"{str(snapshot.get('snapshot_name', 'unknown'))} " + f"[{str(snapshot.get('kind', 'unknown'))}] " + f"entry_count={int(snapshot.get('entry_count', 0))} " + f"bytes_written={int(snapshot.get('bytes_written', 0))} " + f"deletable={'yes' if bool(snapshot.get('deletable', False)) else 'no'}" + ) + + +def _print_workspace_reset_human(payload: dict[str, Any]) -> None: + _print_workspace_summary_human(payload, action="Reset workspace") + workspace_reset = payload.get("workspace_reset") + if isinstance(workspace_reset, dict): + print( + "Reset source: " + f"{str(workspace_reset.get('snapshot_name', 'unknown'))} " + f"({str(workspace_reset.get('kind', 'unknown'))})" + ) + print( + "Reset restore: " + f"destination={str(workspace_reset.get('destination', WORKSPACE_GUEST_PATH))} " + f"entry_count={int(workspace_reset.get('entry_count', 0))} " + f"bytes_written={int(workspace_reset.get('bytes_written', 0))}" + ) + + def _print_workspace_shell_summary_human(payload: dict[str, Any], *, prefix: str) -> None: print( f"[{prefix}] " @@ -592,6 +645,8 @@ def _build_parser() -> argparse.ArgumentParser: pyro workspace create debian:12 --seed-path ./repo pyro workspace sync push WORKSPACE_ID ./repo --dest src pyro workspace exec WORKSPACE_ID -- sh -lc 'printf "hello\\n" > note.txt' + pyro workspace snapshot create WORKSPACE_ID checkpoint + pyro workspace reset WORKSPACE_ID --snapshot checkpoint pyro workspace diff WORKSPACE_ID pyro workspace export WORKSPACE_ID src/note.txt --output ./note.txt pyro workspace shell open WORKSPACE_ID @@ -617,6 +672,8 @@ def _build_parser() -> argparse.ArgumentParser: pyro workspace create debian:12 pyro workspace create debian:12 --seed-path ./repo pyro workspace sync push WORKSPACE_ID ./changes + pyro workspace snapshot create WORKSPACE_ID checkpoint + pyro workspace reset WORKSPACE_ID --snapshot checkpoint pyro workspace diff WORKSPACE_ID pyro workspace service start WORKSPACE_ID app --ready-file .ready -- \ sh -lc 'touch .ready && while true; do sleep 60; done' @@ -720,7 +777,8 @@ def _build_parser() -> argparse.ArgumentParser: pyro workspace sync push WORKSPACE_ID ./repo pyro workspace sync push WORKSPACE_ID ./patches --dest src - Sync is non-atomic. If a sync fails partway through, delete and recreate the workspace. + Sync is non-atomic. If a sync fails partway through, prefer reset over repair with + `pyro workspace reset WORKSPACE_ID`. """ ), formatter_class=_HelpFormatter, @@ -808,6 +866,100 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Print structured JSON instead of human-readable output.", ) + workspace_snapshot_parser = workspace_subparsers.add_parser( + "snapshot", + help="Create, list, and delete workspace snapshots.", + description=( + "Manage explicit named snapshots in addition to the implicit create-time baseline." + ), + epilog=dedent( + """ + Examples: + pyro workspace snapshot create WORKSPACE_ID checkpoint + pyro workspace snapshot list WORKSPACE_ID + pyro workspace snapshot delete WORKSPACE_ID checkpoint + + Use `workspace reset` to restore `/workspace` from `baseline` or one named snapshot. + """ + ), + formatter_class=_HelpFormatter, + ) + workspace_snapshot_subparsers = workspace_snapshot_parser.add_subparsers( + dest="workspace_snapshot_command", + required=True, + metavar="SNAPSHOT", + ) + workspace_snapshot_create_parser = workspace_snapshot_subparsers.add_parser( + "create", + help="Create one named snapshot from the current workspace.", + description="Capture the current `/workspace` tree as one named snapshot.", + epilog="Example:\n pyro workspace snapshot create WORKSPACE_ID checkpoint", + formatter_class=_HelpFormatter, + ) + workspace_snapshot_create_parser.add_argument("workspace_id", metavar="WORKSPACE_ID") + workspace_snapshot_create_parser.add_argument("snapshot_name", metavar="SNAPSHOT_NAME") + workspace_snapshot_create_parser.add_argument( + "--json", + action="store_true", + help="Print structured JSON instead of human-readable output.", + ) + workspace_snapshot_list_parser = workspace_snapshot_subparsers.add_parser( + "list", + help="List the baseline plus named snapshots.", + description="List the implicit baseline snapshot plus any named snapshots for a workspace.", + epilog="Example:\n pyro workspace snapshot list WORKSPACE_ID", + formatter_class=_HelpFormatter, + ) + workspace_snapshot_list_parser.add_argument("workspace_id", metavar="WORKSPACE_ID") + workspace_snapshot_list_parser.add_argument( + "--json", + action="store_true", + help="Print structured JSON instead of human-readable output.", + ) + workspace_snapshot_delete_parser = workspace_snapshot_subparsers.add_parser( + "delete", + help="Delete one named snapshot.", + description="Delete one named snapshot while leaving the implicit baseline intact.", + epilog="Example:\n pyro workspace snapshot delete WORKSPACE_ID checkpoint", + formatter_class=_HelpFormatter, + ) + workspace_snapshot_delete_parser.add_argument("workspace_id", metavar="WORKSPACE_ID") + workspace_snapshot_delete_parser.add_argument("snapshot_name", metavar="SNAPSHOT_NAME") + workspace_snapshot_delete_parser.add_argument( + "--json", + action="store_true", + help="Print structured JSON instead of human-readable output.", + ) + workspace_reset_parser = workspace_subparsers.add_parser( + "reset", + help="Recreate a workspace from baseline or one named snapshot.", + description=( + "Recreate the full sandbox and restore `/workspace` from the baseline " + "or one named snapshot." + ), + epilog=dedent( + """ + Examples: + pyro workspace reset WORKSPACE_ID + pyro workspace reset WORKSPACE_ID --snapshot checkpoint + + Prefer reset over repair: reset clears command history, shells, and services so the + workspace comes back clean from `baseline` or one named snapshot. + """ + ), + formatter_class=_HelpFormatter, + ) + workspace_reset_parser.add_argument("workspace_id", metavar="WORKSPACE_ID") + workspace_reset_parser.add_argument( + "--snapshot", + default="baseline", + help="Snapshot name to restore. Defaults to the implicit `baseline` snapshot.", + ) + workspace_reset_parser.add_argument( + "--json", + action="store_true", + help="Print structured JSON instead of human-readable output.", + ) workspace_shell_parser = workspace_subparsers.add_parser( "shell", help="Open and manage persistent interactive shells.", @@ -1483,6 +1635,72 @@ def main() -> None: raise SystemExit(1) from exc _print_workspace_diff_human(payload) return + if args.workspace_command == "snapshot": + if args.workspace_snapshot_command == "create": + try: + payload = pyro.create_snapshot(args.workspace_id, args.snapshot_name) + except Exception as exc: # noqa: BLE001 + if bool(args.json): + _print_json({"ok": False, "error": str(exc)}) + else: + print(f"[error] {exc}", file=sys.stderr, flush=True) + raise SystemExit(1) from exc + if bool(args.json): + _print_json(payload) + else: + _print_workspace_snapshot_human( + payload, + prefix="workspace-snapshot-create", + ) + return + if args.workspace_snapshot_command == "list": + try: + payload = pyro.list_snapshots(args.workspace_id) + except Exception as exc: # noqa: BLE001 + if bool(args.json): + _print_json({"ok": False, "error": str(exc)}) + else: + print(f"[error] {exc}", file=sys.stderr, flush=True) + raise SystemExit(1) from exc + if bool(args.json): + _print_json(payload) + else: + _print_workspace_snapshot_list_human(payload) + return + if args.workspace_snapshot_command == "delete": + try: + payload = pyro.delete_snapshot(args.workspace_id, args.snapshot_name) + except Exception as exc: # noqa: BLE001 + if bool(args.json): + _print_json({"ok": False, "error": str(exc)}) + else: + print(f"[error] {exc}", file=sys.stderr, flush=True) + raise SystemExit(1) from exc + if bool(args.json): + _print_json(payload) + else: + print( + "Deleted workspace snapshot: " + f"{str(payload.get('snapshot_name', 'unknown'))}" + ) + return + if args.workspace_command == "reset": + try: + payload = pyro.reset_workspace( + args.workspace_id, + snapshot=args.snapshot, + ) + except Exception as exc: # noqa: BLE001 + if bool(args.json): + _print_json({"ok": False, "error": str(exc)}) + else: + print(f"[error] {exc}", file=sys.stderr, flush=True) + raise SystemExit(1) from exc + if bool(args.json): + _print_json(payload) + else: + _print_workspace_reset_human(payload) + return if args.workspace_command == "shell": if args.workspace_shell_command == "open": try: diff --git a/src/pyro_mcp/contract.py b/src/pyro_mcp/contract.py index 9a16cc3..cd92478 100644 --- a/src/pyro_mcp/contract.py +++ b/src/pyro_mcp/contract.py @@ -12,13 +12,16 @@ PUBLIC_CLI_WORKSPACE_SUBCOMMANDS = ( "exec", "export", "logs", + "reset", "service", "shell", + "snapshot", "status", "sync", ) PUBLIC_CLI_WORKSPACE_SERVICE_SUBCOMMANDS = ("list", "logs", "start", "status", "stop") PUBLIC_CLI_WORKSPACE_SHELL_SUBCOMMANDS = ("close", "open", "read", "signal", "write") +PUBLIC_CLI_WORKSPACE_SNAPSHOT_SUBCOMMANDS = ("create", "delete", "list") PUBLIC_CLI_WORKSPACE_SYNC_SUBCOMMANDS = ("push",) PUBLIC_CLI_WORKSPACE_CREATE_FLAGS = ( "--vcpu-count", @@ -31,6 +34,7 @@ PUBLIC_CLI_WORKSPACE_CREATE_FLAGS = ( ) PUBLIC_CLI_WORKSPACE_DIFF_FLAGS = ("--json",) PUBLIC_CLI_WORKSPACE_EXPORT_FLAGS = ("--output", "--json") +PUBLIC_CLI_WORKSPACE_RESET_FLAGS = ("--snapshot", "--json") PUBLIC_CLI_WORKSPACE_SERVICE_LIST_FLAGS = ("--json",) PUBLIC_CLI_WORKSPACE_SERVICE_LOGS_FLAGS = ("--tail-lines", "--all", "--json") PUBLIC_CLI_WORKSPACE_SERVICE_START_FLAGS = ( @@ -50,6 +54,9 @@ PUBLIC_CLI_WORKSPACE_SHELL_READ_FLAGS = ("--cursor", "--max-chars", "--json") PUBLIC_CLI_WORKSPACE_SHELL_WRITE_FLAGS = ("--input", "--no-newline", "--json") PUBLIC_CLI_WORKSPACE_SHELL_SIGNAL_FLAGS = ("--signal", "--json") PUBLIC_CLI_WORKSPACE_SHELL_CLOSE_FLAGS = ("--json",) +PUBLIC_CLI_WORKSPACE_SNAPSHOT_CREATE_FLAGS = ("--json",) +PUBLIC_CLI_WORKSPACE_SNAPSHOT_DELETE_FLAGS = ("--json",) +PUBLIC_CLI_WORKSPACE_SNAPSHOT_LIST_FLAGS = ("--json",) PUBLIC_CLI_WORKSPACE_SYNC_PUSH_FLAGS = ("--dest", "--json") PUBLIC_CLI_RUN_FLAGS = ( "--vcpu-count", @@ -64,8 +71,10 @@ PUBLIC_CLI_RUN_FLAGS = ( PUBLIC_SDK_METHODS = ( "close_shell", "create_server", + "create_snapshot", "create_vm", "create_workspace", + "delete_snapshot", "delete_vm", "delete_workspace", "diff_workspace", @@ -75,6 +84,7 @@ PUBLIC_SDK_METHODS = ( "inspect_environment", "list_environments", "list_services", + "list_snapshots", "logs_service", "logs_workspace", "network_info_vm", @@ -84,6 +94,7 @@ PUBLIC_SDK_METHODS = ( "push_workspace_sync", "read_shell", "reap_expired", + "reset_workspace", "run_in_vm", "signal_shell", "start_service", @@ -107,6 +118,9 @@ PUBLIC_MCP_TOOLS = ( "shell_read", "shell_signal", "shell_write", + "snapshot_create", + "snapshot_delete", + "snapshot_list", "vm_create", "vm_delete", "vm_exec", @@ -123,6 +137,7 @@ PUBLIC_MCP_TOOLS = ( "workspace_exec", "workspace_export", "workspace_logs", + "workspace_reset", "workspace_status", "workspace_sync_push", ) diff --git a/src/pyro_mcp/vm_environments.py b/src/pyro_mcp/vm_environments.py index 2a8ce4c..3611bf1 100644 --- a/src/pyro_mcp/vm_environments.py +++ b/src/pyro_mcp/vm_environments.py @@ -19,7 +19,7 @@ from typing import Any from pyro_mcp.runtime import DEFAULT_PLATFORM, RuntimePaths DEFAULT_ENVIRONMENT_VERSION = "1.0.0" -DEFAULT_CATALOG_VERSION = "2.7.0" +DEFAULT_CATALOG_VERSION = "2.8.0" OCI_MANIFEST_ACCEPT = ", ".join( ( "application/vnd.oci.image.index.v1+json", diff --git a/src/pyro_mcp/vm_manager.py b/src/pyro_mcp/vm_manager.py index 713769f..9f9cc51 100644 --- a/src/pyro_mcp/vm_manager.py +++ b/src/pyro_mcp/vm_manager.py @@ -49,9 +49,10 @@ DEFAULT_TIMEOUT_SECONDS = 30 DEFAULT_TTL_SECONDS = 600 DEFAULT_ALLOW_HOST_COMPAT = False -WORKSPACE_LAYOUT_VERSION = 4 +WORKSPACE_LAYOUT_VERSION = 5 WORKSPACE_BASELINE_DIRNAME = "baseline" WORKSPACE_BASELINE_ARCHIVE_NAME = "workspace.tar" +WORKSPACE_SNAPSHOTS_DIRNAME = "snapshots" WORKSPACE_DIRNAME = "workspace" WORKSPACE_COMMANDS_DIRNAME = "commands" WORKSPACE_SHELLS_DIRNAME = "shells" @@ -68,10 +69,12 @@ DEFAULT_SERVICE_READY_INTERVAL_MS = 500 DEFAULT_SERVICE_LOG_TAIL_LINES = 200 WORKSPACE_SHELL_SIGNAL_NAMES = shell_signal_names() WORKSPACE_SERVICE_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$") +WORKSPACE_SNAPSHOT_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$") WorkspaceSeedMode = Literal["empty", "directory", "tar_archive"] WorkspaceArtifactType = Literal["file", "directory", "symlink"] WorkspaceServiceReadinessType = Literal["file", "tcp", "http", "command"] +WorkspaceSnapshotKind = Literal["baseline", "named"] @dataclass @@ -116,6 +119,8 @@ class WorkspaceRecord: command_count: int = 0 last_command: dict[str, Any] | None = None workspace_seed: dict[str, Any] = field(default_factory=dict) + reset_count: int = 0 + last_reset_at: float | None = None @classmethod def from_instance( @@ -144,6 +149,8 @@ class WorkspaceRecord: command_count=command_count, last_command=last_command, workspace_seed=dict(workspace_seed or _empty_workspace_seed_payload()), + reset_count=0, + last_reset_at=None, ) def to_instance(self, *, workdir: Path) -> VmInstance: @@ -185,6 +192,8 @@ class WorkspaceRecord: "command_count": self.command_count, "last_command": self.last_command, "workspace_seed": self.workspace_seed, + "reset_count": self.reset_count, + "last_reset_at": self.last_reset_at, } @classmethod @@ -207,6 +216,46 @@ class WorkspaceRecord: command_count=int(payload.get("command_count", 0)), last_command=_optional_dict(payload.get("last_command")), workspace_seed=_workspace_seed_dict(payload.get("workspace_seed")), + reset_count=int(payload.get("reset_count", 0)), + last_reset_at=( + None + if payload.get("last_reset_at") is None + else float(payload.get("last_reset_at", 0.0)) + ), + ) + + +@dataclass +class WorkspaceSnapshotRecord: + """Persistent snapshot metadata stored on disk per workspace.""" + + workspace_id: str + snapshot_name: str + kind: WorkspaceSnapshotKind + created_at: float + entry_count: int + bytes_written: int + + def to_payload(self) -> dict[str, Any]: + return { + "layout_version": WORKSPACE_LAYOUT_VERSION, + "workspace_id": self.workspace_id, + "snapshot_name": self.snapshot_name, + "kind": self.kind, + "created_at": self.created_at, + "entry_count": self.entry_count, + "bytes_written": self.bytes_written, + } + + @classmethod + def from_payload(cls, payload: dict[str, Any]) -> WorkspaceSnapshotRecord: + return cls( + workspace_id=str(payload["workspace_id"]), + snapshot_name=str(payload["snapshot_name"]), + kind=cast(WorkspaceSnapshotKind, str(payload.get("kind", "named"))), + created_at=float(payload["created_at"]), + entry_count=int(payload.get("entry_count", 0)), + bytes_written=int(payload.get("bytes_written", 0)), ) @@ -864,6 +913,24 @@ def _normalize_workspace_service_name(service_name: str) -> str: return normalized +def _normalize_workspace_snapshot_name( + snapshot_name: str, + *, + allow_baseline: bool = False, +) -> str: + normalized = snapshot_name.strip() + if normalized == "": + raise ValueError("snapshot_name must not be empty") + if normalized == "baseline" and not allow_baseline: + raise ValueError("snapshot_name 'baseline' is reserved") + if WORKSPACE_SNAPSHOT_NAME_RE.fullmatch(normalized) is None: + raise ValueError( + "snapshot_name must match " + r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$" + ) + return normalized + + def _normalize_workspace_service_readiness( readiness: dict[str, Any] | None, ) -> dict[str, Any] | None: @@ -2646,12 +2713,14 @@ class VmManager: commands_dir = self._workspace_commands_dir(workspace_id) shells_dir = self._workspace_shells_dir(workspace_id) services_dir = self._workspace_services_dir(workspace_id) + snapshots_dir = self._workspace_snapshots_dir(workspace_id) baseline_archive_path = self._workspace_baseline_archive_path(workspace_id) workspace_dir.mkdir(parents=True, exist_ok=False) host_workspace_dir.mkdir(parents=True, exist_ok=True) commands_dir.mkdir(parents=True, exist_ok=True) shells_dir.mkdir(parents=True, exist_ok=True) services_dir.mkdir(parents=True, exist_ok=True) + snapshots_dir.mkdir(parents=True, exist_ok=True) _persist_workspace_baseline( prepared_seed, baseline_archive_path=baseline_archive_path, @@ -2859,6 +2928,192 @@ class VmManager: diff_payload["workspace_id"] = workspace_id return diff_payload + def create_snapshot( + self, + workspace_id: str, + snapshot_name: str, + ) -> dict[str, Any]: + normalized_snapshot_name = _normalize_workspace_snapshot_name(snapshot_name) + with self._lock: + workspace = self._load_workspace_locked(workspace_id) + self._ensure_workspace_not_expired_locked(workspace, time.time()) + self._workspace_baseline_snapshot_locked(workspace) + if ( + self._load_workspace_snapshot_locked_optional( + workspace_id, + normalized_snapshot_name, + ) + is not None + ): + raise ValueError( + f"snapshot {normalized_snapshot_name!r} already exists in workspace " + f"{workspace_id!r}" + ) + instance = self._workspace_instance_for_live_operation_locked( + workspace, + operation_name="workspace_snapshot_create", + ) + with tempfile.TemporaryDirectory(prefix="pyro-workspace-snapshot-") as temp_dir: + temp_archive_path = Path(temp_dir) / f"{normalized_snapshot_name}.tar" + exported = self._backend.export_archive( + instance, + workspace_path=WORKSPACE_GUEST_PATH, + archive_path=temp_archive_path, + ) + snapshot = WorkspaceSnapshotRecord( + workspace_id=workspace_id, + snapshot_name=normalized_snapshot_name, + kind="named", + created_at=time.time(), + entry_count=int(exported["entry_count"]), + bytes_written=int(exported["bytes_written"]), + ) + with self._lock: + workspace = self._load_workspace_locked(workspace_id) + self._ensure_workspace_not_expired_locked(workspace, time.time()) + if ( + self._load_workspace_snapshot_locked_optional( + workspace_id, + normalized_snapshot_name, + ) + is not None + ): + raise ValueError( + f"snapshot {normalized_snapshot_name!r} already exists in workspace " + f"{workspace_id!r}" + ) + archive_path = self._workspace_snapshot_archive_path( + workspace_id, + normalized_snapshot_name, + ) + archive_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(temp_archive_path, archive_path) + workspace.state = instance.state + workspace.firecracker_pid = instance.firecracker_pid + workspace.last_error = instance.last_error + workspace.metadata = dict(instance.metadata) + self._save_workspace_locked(workspace) + self._save_workspace_snapshot_locked(snapshot) + return { + "workspace_id": workspace_id, + "snapshot": self._serialize_workspace_snapshot(snapshot), + "execution_mode": instance.metadata.get("execution_mode", "pending"), + } + + def list_snapshots(self, workspace_id: str) -> dict[str, Any]: + with self._lock: + workspace = self._load_workspace_locked(workspace_id) + self._ensure_workspace_not_expired_locked(workspace, time.time()) + snapshots = self._list_workspace_snapshots_locked(workspace) + return { + "workspace_id": workspace_id, + "count": len(snapshots), + "snapshots": [ + self._serialize_workspace_snapshot(snapshot) for snapshot in snapshots + ], + } + + def delete_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + normalized_snapshot_name = _normalize_workspace_snapshot_name( + snapshot_name, + allow_baseline=True, + ) + if normalized_snapshot_name == "baseline": + raise ValueError("cannot delete the baseline snapshot") + with self._lock: + workspace = self._load_workspace_locked(workspace_id) + self._ensure_workspace_not_expired_locked(workspace, time.time()) + self._workspace_baseline_snapshot_locked(workspace) + self._load_workspace_snapshot_locked(workspace_id, normalized_snapshot_name) + self._delete_workspace_snapshot_locked(workspace_id, normalized_snapshot_name) + return { + "workspace_id": workspace_id, + "snapshot_name": normalized_snapshot_name, + "deleted": True, + } + + def reset_workspace( + self, + workspace_id: str, + *, + snapshot: str = "baseline", + ) -> dict[str, Any]: + with self._lock: + workspace = self._load_workspace_locked(workspace_id) + self._ensure_workspace_not_expired_locked(workspace, time.time()) + self._refresh_workspace_liveness_locked(workspace) + selected_snapshot, archive_path = self._resolve_workspace_snapshot_locked( + workspace, + snapshot, + ) + instance = workspace.to_instance( + workdir=self._workspace_runtime_dir(workspace.workspace_id) + ) + self._stop_workspace_services_locked(workspace, instance) + self._close_workspace_shells_locked(workspace, instance) + if workspace.state == "started": + self._backend.stop(instance) + workspace.state = "stopped" + self._backend.delete(instance) + workspace.state = "stopped" + workspace.firecracker_pid = None + workspace.last_error = None + self._reset_workspace_runtime_dirs(workspace_id) + self._save_workspace_locked(workspace) + recreated: VmInstance | None = None + try: + recreated = workspace.to_instance( + workdir=self._workspace_runtime_dir(workspace.workspace_id) + ) + self._backend.create(recreated) + if self._runtime_capabilities.supports_guest_exec: + self._ensure_workspace_guest_agent_support(recreated) + with self._lock: + self._start_instance_locked(recreated) + self._require_guest_exec_or_opt_in(recreated) + reset_summary = self._backend.import_archive( + recreated, + archive_path=archive_path, + destination=WORKSPACE_GUEST_PATH, + ) + workspace = self._load_workspace_locked(workspace_id) + workspace.state = recreated.state + workspace.firecracker_pid = recreated.firecracker_pid + workspace.last_error = recreated.last_error + workspace.metadata = dict(recreated.metadata) + workspace.command_count = 0 + workspace.last_command = None + workspace.reset_count += 1 + workspace.last_reset_at = time.time() + self._save_workspace_locked(workspace) + payload = self._serialize_workspace(workspace) + except Exception: + try: + if recreated is not None and recreated.state == "started": + self._backend.stop(recreated) + except Exception: + pass + try: + if recreated is not None: + self._backend.delete(recreated) + except Exception: + pass + with self._lock: + workspace = self._load_workspace_locked(workspace_id) + workspace.state = "stopped" + workspace.firecracker_pid = None + workspace.last_error = None + self._save_workspace_locked(workspace) + raise + payload["workspace_reset"] = { + "snapshot_name": selected_snapshot.snapshot_name, + "kind": selected_snapshot.kind, + "destination": str(reset_summary["destination"]), + "entry_count": int(reset_summary["entry_count"]), + "bytes_written": int(reset_summary["bytes_written"]), + } + return payload + def exec_workspace( self, workspace_id: str, @@ -3372,6 +3627,8 @@ class VmManager: "workspace_seed": _workspace_seed_dict(workspace.workspace_seed), "command_count": workspace.command_count, "last_command": workspace.last_command, + "reset_count": workspace.reset_count, + "last_reset_at": workspace.last_reset_at, "service_count": service_count, "running_service_count": running_service_count, "metadata": workspace.metadata, @@ -3408,6 +3665,17 @@ class VmManager: "stop_reason": service.stop_reason, } + def _serialize_workspace_snapshot(self, snapshot: WorkspaceSnapshotRecord) -> dict[str, Any]: + return { + "workspace_id": snapshot.workspace_id, + "snapshot_name": snapshot.snapshot_name, + "kind": snapshot.kind, + "created_at": snapshot.created_at, + "entry_count": snapshot.entry_count, + "bytes_written": snapshot.bytes_written, + "deletable": snapshot.kind != "baseline", + } + def _require_guest_boot_or_opt_in(self, instance: VmInstance) -> None: if self._runtime_capabilities.supports_vm_boot or instance.allow_host_compat: return @@ -3589,6 +3857,15 @@ class VmManager: def _workspace_baseline_archive_path(self, workspace_id: str) -> Path: return self._workspace_baseline_dir(workspace_id) / WORKSPACE_BASELINE_ARCHIVE_NAME + def _workspace_snapshots_dir(self, workspace_id: str) -> Path: + return self._workspace_dir(workspace_id) / WORKSPACE_SNAPSHOTS_DIRNAME + + def _workspace_snapshot_archive_path(self, workspace_id: str, snapshot_name: str) -> Path: + return self._workspace_snapshots_dir(workspace_id) / f"{snapshot_name}.tar" + + def _workspace_snapshot_metadata_path(self, workspace_id: str, snapshot_name: str) -> Path: + return self._workspace_snapshots_dir(workspace_id) / f"{snapshot_name}.json" + def _workspace_commands_dir(self, workspace_id: str) -> Path: return self._workspace_dir(workspace_id) / WORKSPACE_COMMANDS_DIRNAME @@ -3846,6 +4123,41 @@ class VmManager: services = self._list_workspace_services_locked(workspace_id) return len(services), sum(1 for service in services if service.state == "running") + def _workspace_baseline_snapshot_locked( + self, + workspace: WorkspaceRecord, + ) -> WorkspaceSnapshotRecord: + baseline_archive_path = self._workspace_baseline_archive_path(workspace.workspace_id) + if not baseline_archive_path.exists(): + raise RuntimeError( + "workspace snapshots and reset require a baseline snapshot. " + "Recreate the workspace to use snapshot/reset features." + ) + entry_count, bytes_written = _inspect_seed_archive(baseline_archive_path) + return WorkspaceSnapshotRecord( + workspace_id=workspace.workspace_id, + snapshot_name="baseline", + kind="baseline", + created_at=workspace.created_at, + entry_count=entry_count, + bytes_written=bytes_written, + ) + + def _resolve_workspace_snapshot_locked( + self, + workspace: WorkspaceRecord, + snapshot_name: str, + ) -> tuple[WorkspaceSnapshotRecord, Path]: + normalized_name = _normalize_workspace_snapshot_name(snapshot_name, allow_baseline=True) + if normalized_name == "baseline": + baseline = self._workspace_baseline_snapshot_locked(workspace) + return baseline, self._workspace_baseline_archive_path(workspace.workspace_id) + snapshot = self._load_workspace_snapshot_locked(workspace.workspace_id, normalized_name) + return ( + snapshot, + self._workspace_snapshot_archive_path(workspace.workspace_id, normalized_name), + ) + def _load_workspace_service_locked( self, workspace_id: str, @@ -3861,6 +4173,34 @@ class VmManager: raise RuntimeError(f"service record at {record_path} is invalid") return WorkspaceServiceRecord.from_payload(payload) + def _load_workspace_snapshot_locked( + self, + workspace_id: str, + snapshot_name: str, + ) -> WorkspaceSnapshotRecord: + record_path = self._workspace_snapshot_metadata_path(workspace_id, snapshot_name) + if not record_path.exists(): + raise ValueError( + f"snapshot {snapshot_name!r} does not exist in workspace {workspace_id!r}" + ) + payload = json.loads(record_path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise RuntimeError(f"snapshot record at {record_path} is invalid") + return WorkspaceSnapshotRecord.from_payload(payload) + + def _load_workspace_snapshot_locked_optional( + self, + workspace_id: str, + snapshot_name: str, + ) -> WorkspaceSnapshotRecord | None: + record_path = self._workspace_snapshot_metadata_path(workspace_id, snapshot_name) + if not record_path.exists(): + return None + payload = json.loads(record_path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise RuntimeError(f"snapshot record at {record_path} is invalid") + return WorkspaceSnapshotRecord.from_payload(payload) + def _load_workspace_service_locked_optional( self, workspace_id: str, @@ -3885,6 +4225,17 @@ class VmManager: encoding="utf-8", ) + def _save_workspace_snapshot_locked(self, snapshot: WorkspaceSnapshotRecord) -> None: + record_path = self._workspace_snapshot_metadata_path( + snapshot.workspace_id, + snapshot.snapshot_name, + ) + record_path.parent.mkdir(parents=True, exist_ok=True) + record_path.write_text( + json.dumps(snapshot.to_payload(), indent=2, sort_keys=True), + encoding="utf-8", + ) + def _delete_workspace_service_artifacts_locked( self, workspace_id: str, @@ -3897,6 +4248,10 @@ class VmManager: _workspace_service_status_path(services_dir, service_name).unlink(missing_ok=True) _workspace_service_runner_path(services_dir, service_name).unlink(missing_ok=True) + def _delete_workspace_snapshot_locked(self, workspace_id: str, snapshot_name: str) -> None: + self._workspace_snapshot_metadata_path(workspace_id, snapshot_name).unlink(missing_ok=True) + self._workspace_snapshot_archive_path(workspace_id, snapshot_name).unlink(missing_ok=True) + def _list_workspace_services_locked(self, workspace_id: str) -> list[WorkspaceServiceRecord]: services_dir = self._workspace_services_dir(workspace_id) if not services_dir.exists(): @@ -3909,6 +4264,26 @@ class VmManager: services.append(WorkspaceServiceRecord.from_payload(payload)) return services + def _list_workspace_snapshots_locked( + self, + workspace: WorkspaceRecord, + ) -> list[WorkspaceSnapshotRecord]: + snapshots_dir = self._workspace_snapshots_dir(workspace.workspace_id) + snapshots: list[WorkspaceSnapshotRecord] = [ + self._workspace_baseline_snapshot_locked(workspace) + ] + if not snapshots_dir.exists(): + return snapshots + named_snapshots: list[WorkspaceSnapshotRecord] = [] + for record_path in snapshots_dir.glob("*.json"): + payload = json.loads(record_path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + continue + named_snapshots.append(WorkspaceSnapshotRecord.from_payload(payload)) + named_snapshots.sort(key=lambda item: (-item.created_at, item.snapshot_name)) + snapshots.extend(named_snapshots) + return snapshots + def _save_workspace_shell_locked(self, shell: WorkspaceShellRecord) -> None: record_path = self._workspace_shell_record_path(shell.workspace_id, shell.shell_id) record_path.parent.mkdir(parents=True, exist_ok=True) @@ -3950,6 +4325,17 @@ class VmManager: pass self._delete_workspace_shell_locked(workspace.workspace_id, shell.shell_id) + def _reset_workspace_runtime_dirs(self, workspace_id: str) -> None: + shutil.rmtree(self._workspace_runtime_dir(workspace_id), ignore_errors=True) + shutil.rmtree(self._workspace_host_dir(workspace_id), ignore_errors=True) + shutil.rmtree(self._workspace_commands_dir(workspace_id), ignore_errors=True) + shutil.rmtree(self._workspace_shells_dir(workspace_id), ignore_errors=True) + shutil.rmtree(self._workspace_services_dir(workspace_id), ignore_errors=True) + self._workspace_host_dir(workspace_id).mkdir(parents=True, exist_ok=True) + self._workspace_commands_dir(workspace_id).mkdir(parents=True, exist_ok=True) + self._workspace_shells_dir(workspace_id).mkdir(parents=True, exist_ok=True) + self._workspace_services_dir(workspace_id).mkdir(parents=True, exist_ok=True) + def _refresh_workspace_service_locked( self, workspace: WorkspaceRecord, diff --git a/tests/test_api.py b/tests/test_api.py index 65762ec..db3dbc5 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -52,6 +52,10 @@ def test_pyro_create_server_registers_vm_run(tmp_path: Path) -> None: assert "workspace_diff" in tool_names assert "workspace_sync_push" in tool_names assert "workspace_export" in tool_names + assert "snapshot_create" in tool_names + assert "snapshot_list" in tool_names + assert "snapshot_delete" in tool_names + assert "workspace_reset" in tool_names assert "shell_open" in tool_names assert "shell_read" in tool_names assert "shell_write" in tool_names @@ -143,6 +147,8 @@ def test_pyro_workspace_methods_delegate_to_manager(tmp_path: Path) -> None: synced = pyro.push_workspace_sync(workspace_id, updated_dir, dest="subdir") executed = pyro.exec_workspace(workspace_id, command="cat note.txt") diff_payload = pyro.diff_workspace(workspace_id) + snapshot = pyro.create_snapshot(workspace_id, "checkpoint") + snapshots = pyro.list_snapshots(workspace_id) export_path = tmp_path / "exported-note.txt" exported = pyro.export_workspace(workspace_id, "note.txt", output_path=export_path) service = pyro.start_service( @@ -155,6 +161,8 @@ def test_pyro_workspace_methods_delegate_to_manager(tmp_path: Path) -> None: service_status = pyro.status_service(workspace_id, "app") service_logs = pyro.logs_service(workspace_id, "app", all=True) service_stopped = pyro.stop_service(workspace_id, "app") + reset = pyro.reset_workspace(workspace_id, snapshot="checkpoint") + deleted_snapshot = pyro.delete_snapshot(workspace_id, "checkpoint") status = pyro.status_workspace(workspace_id) logs = pyro.logs_workspace(workspace_id) deleted = pyro.delete_workspace(workspace_id) @@ -163,6 +171,8 @@ def test_pyro_workspace_methods_delegate_to_manager(tmp_path: Path) -> None: assert created["workspace_seed"]["mode"] == "directory" assert synced["workspace_sync"]["destination"] == "/workspace/subdir" assert diff_payload["changed"] is True + assert snapshot["snapshot"]["snapshot_name"] == "checkpoint" + assert snapshots["count"] == 2 assert exported["output_path"] == str(export_path) assert export_path.read_text(encoding="utf-8") == "ok\n" assert service["state"] == "running" @@ -170,7 +180,9 @@ def test_pyro_workspace_methods_delegate_to_manager(tmp_path: Path) -> None: assert service_status["state"] == "running" assert service_logs["tail_lines"] is None assert service_stopped["state"] == "stopped" - assert status["command_count"] == 1 - assert status["service_count"] == 1 - assert logs["count"] == 1 + assert reset["workspace_reset"]["snapshot_name"] == "checkpoint" + assert deleted_snapshot["deleted"] is True + assert status["command_count"] == 0 + assert status["service_count"] == 0 + assert logs["count"] == 0 assert deleted["deleted"] is True diff --git a/tests/test_cli.py b/tests/test_cli.py index 2fb8f35..f487007 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -66,6 +66,8 @@ def test_cli_subcommand_help_includes_examples_and_guidance() -> None: assert "pyro workspace exec WORKSPACE_ID" in workspace_help assert "pyro workspace diff WORKSPACE_ID" in workspace_help assert "pyro workspace export WORKSPACE_ID src/note.txt --output ./note.txt" in workspace_help + assert "pyro workspace snapshot create WORKSPACE_ID checkpoint" in workspace_help + assert "pyro workspace reset WORKSPACE_ID --snapshot checkpoint" in workspace_help assert "pyro workspace shell open WORKSPACE_ID" in workspace_help workspace_create_help = _subparser_choice( @@ -107,6 +109,34 @@ def test_cli_subcommand_help_includes_examples_and_guidance() -> None: assert "immutable workspace baseline" in workspace_diff_help assert "workspace export" in workspace_diff_help + workspace_snapshot_help = _subparser_choice( + _subparser_choice(parser, "workspace"), + "snapshot", + ).format_help() + assert "pyro workspace snapshot create WORKSPACE_ID checkpoint" in workspace_snapshot_help + assert "baseline" in workspace_snapshot_help + + workspace_snapshot_create_help = _subparser_choice( + _subparser_choice(_subparser_choice(parser, "workspace"), "snapshot"), "create" + ).format_help() + assert "Capture the current `/workspace` tree" in workspace_snapshot_create_help + + workspace_snapshot_list_help = _subparser_choice( + _subparser_choice(_subparser_choice(parser, "workspace"), "snapshot"), "list" + ).format_help() + assert "baseline snapshot plus any named snapshots" in workspace_snapshot_list_help + + workspace_snapshot_delete_help = _subparser_choice( + _subparser_choice(_subparser_choice(parser, "workspace"), "snapshot"), "delete" + ).format_help() + assert "leaving the implicit baseline intact" in workspace_snapshot_delete_help + + workspace_reset_help = _subparser_choice( + _subparser_choice(parser, "workspace"), "reset" + ).format_help() + assert "--snapshot" in workspace_reset_help + assert "reset over repair" in workspace_reset_help + workspace_shell_help = _subparser_choice( _subparser_choice(parser, "workspace"), "shell", @@ -651,6 +681,359 @@ def test_cli_workspace_diff_prints_human_output( assert "--- a/note.txt" in output +def test_cli_workspace_snapshot_create_list_delete_and_reset_prints_json( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + class StubPyro: + def create_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + assert workspace_id == "workspace-123" + assert snapshot_name == "checkpoint" + return { + "workspace_id": workspace_id, + "snapshot": { + "snapshot_name": snapshot_name, + "kind": "named", + "entry_count": 3, + "bytes_written": 42, + }, + } + + def list_snapshots(self, workspace_id: str) -> dict[str, Any]: + assert workspace_id == "workspace-123" + return { + "workspace_id": workspace_id, + "count": 2, + "snapshots": [ + { + "snapshot_name": "baseline", + "kind": "baseline", + "entry_count": 1, + "bytes_written": 10, + "deletable": False, + }, + { + "snapshot_name": "checkpoint", + "kind": "named", + "entry_count": 3, + "bytes_written": 42, + "deletable": True, + }, + ], + } + + def reset_workspace(self, workspace_id: str, *, snapshot: str) -> dict[str, Any]: + assert workspace_id == "workspace-123" + assert snapshot == "checkpoint" + return { + "workspace_id": workspace_id, + "state": "started", + "workspace_path": "/workspace", + "reset_count": 2, + "workspace_reset": { + "snapshot_name": snapshot, + "kind": "named", + "destination": "/workspace", + "entry_count": 3, + "bytes_written": 42, + }, + } + + def delete_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + assert workspace_id == "workspace-123" + assert snapshot_name == "checkpoint" + return { + "workspace_id": workspace_id, + "snapshot_name": snapshot_name, + "deleted": True, + } + + class SnapshotCreateParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="create", + workspace_id="workspace-123", + snapshot_name="checkpoint", + json=True, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: SnapshotCreateParser()) + monkeypatch.setattr(cli, "Pyro", StubPyro) + cli.main() + created = json.loads(capsys.readouterr().out) + assert created["snapshot"]["snapshot_name"] == "checkpoint" + + class SnapshotListParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="list", + workspace_id="workspace-123", + json=True, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: SnapshotListParser()) + cli.main() + listed = json.loads(capsys.readouterr().out) + assert listed["count"] == 2 + + class ResetParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="workspace", + workspace_command="reset", + workspace_id="workspace-123", + snapshot="checkpoint", + json=True, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: ResetParser()) + cli.main() + reset = json.loads(capsys.readouterr().out) + assert reset["workspace_reset"]["snapshot_name"] == "checkpoint" + + class SnapshotDeleteParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="delete", + workspace_id="workspace-123", + snapshot_name="checkpoint", + json=True, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: SnapshotDeleteParser()) + cli.main() + deleted = json.loads(capsys.readouterr().out) + assert deleted["deleted"] is True + + +def test_cli_workspace_reset_prints_human_output( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + class StubPyro: + def reset_workspace(self, workspace_id: str, *, snapshot: str) -> dict[str, Any]: + assert workspace_id == "workspace-123" + assert snapshot == "baseline" + return { + "workspace_id": workspace_id, + "state": "started", + "environment": "debian:12", + "workspace_path": "/workspace", + "workspace_seed": { + "mode": "directory", + "seed_path": "/tmp/repo", + "destination": "/workspace", + "entry_count": 1, + "bytes_written": 4, + }, + "execution_mode": "guest_vsock", + "command_count": 0, + "service_count": 0, + "running_service_count": 0, + "reset_count": 3, + "last_reset_at": 123.0, + "workspace_reset": { + "snapshot_name": "baseline", + "kind": "baseline", + "destination": "/workspace", + "entry_count": 1, + "bytes_written": 4, + }, + } + + class StubParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="workspace", + workspace_command="reset", + workspace_id="workspace-123", + snapshot="baseline", + json=False, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: StubParser()) + monkeypatch.setattr(cli, "Pyro", StubPyro) + cli.main() + output = capsys.readouterr().out + assert "Reset source: baseline (baseline)" in output + assert "Reset count: 3" in output + + +def test_cli_workspace_snapshot_prints_human_output( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + class StubPyro: + def create_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + assert workspace_id == "workspace-123" + assert snapshot_name == "checkpoint" + return { + "workspace_id": workspace_id, + "snapshot": { + "snapshot_name": snapshot_name, + "kind": "named", + "entry_count": 3, + "bytes_written": 42, + }, + } + + def list_snapshots(self, workspace_id: str) -> dict[str, Any]: + assert workspace_id == "workspace-123" + return { + "workspace_id": workspace_id, + "count": 2, + "snapshots": [ + { + "snapshot_name": "baseline", + "kind": "baseline", + "entry_count": 1, + "bytes_written": 10, + "deletable": False, + }, + { + "snapshot_name": "checkpoint", + "kind": "named", + "entry_count": 3, + "bytes_written": 42, + "deletable": True, + }, + ], + } + + def delete_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + assert workspace_id == "workspace-123" + assert snapshot_name == "checkpoint" + return { + "workspace_id": workspace_id, + "snapshot_name": snapshot_name, + "deleted": True, + } + + class SnapshotCreateParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="create", + workspace_id="workspace-123", + snapshot_name="checkpoint", + json=False, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: SnapshotCreateParser()) + monkeypatch.setattr(cli, "Pyro", StubPyro) + cli.main() + create_output = capsys.readouterr().out + assert "[workspace-snapshot-create] workspace_id=workspace-123" in create_output + assert "snapshot_name=checkpoint kind=named" in create_output + + class SnapshotListParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="list", + workspace_id="workspace-123", + json=False, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: SnapshotListParser()) + cli.main() + list_output = capsys.readouterr().out + assert "baseline [baseline]" in list_output + assert "checkpoint [named]" in list_output + + class SnapshotDeleteParser: + def parse_args(self) -> argparse.Namespace: + return argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="delete", + workspace_id="workspace-123", + snapshot_name="checkpoint", + json=False, + ) + + monkeypatch.setattr(cli, "_build_parser", lambda: SnapshotDeleteParser()) + cli.main() + delete_output = capsys.readouterr().out + assert "Deleted workspace snapshot: checkpoint" in delete_output + + +def test_cli_workspace_snapshot_error_paths( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + class StubPyro: + def create_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + del workspace_id, snapshot_name + raise RuntimeError("create boom") + + def list_snapshots(self, workspace_id: str) -> dict[str, Any]: + del workspace_id + raise RuntimeError("list boom") + + def delete_snapshot(self, workspace_id: str, snapshot_name: str) -> dict[str, Any]: + del workspace_id, snapshot_name + raise RuntimeError("delete boom") + + def _run(args: argparse.Namespace) -> tuple[str, str]: + class StubParser: + def parse_args(self) -> argparse.Namespace: + return args + + monkeypatch.setattr(cli, "_build_parser", lambda: StubParser()) + monkeypatch.setattr(cli, "Pyro", StubPyro) + with pytest.raises(SystemExit, match="1"): + cli.main() + captured = capsys.readouterr() + return captured.out, captured.err + + out, err = _run( + argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="create", + workspace_id="workspace-123", + snapshot_name="checkpoint", + json=True, + ) + ) + assert json.loads(out)["error"] == "create boom" + assert err == "" + + out, err = _run( + argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="list", + workspace_id="workspace-123", + json=False, + ) + ) + assert out == "" + assert "[error] list boom" in err + + out, err = _run( + argparse.Namespace( + command="workspace", + workspace_command="snapshot", + workspace_snapshot_command="delete", + workspace_id="workspace-123", + snapshot_name="checkpoint", + json=False, + ) + ) + assert out == "" + assert "[error] delete boom" in err + + def test_cli_workspace_sync_push_prints_json( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: diff --git a/tests/test_public_contract.py b/tests/test_public_contract.py index 39534f2..0dce0ed 100644 --- a/tests/test_public_contract.py +++ b/tests/test_public_contract.py @@ -20,6 +20,7 @@ from pyro_mcp.contract import ( PUBLIC_CLI_WORKSPACE_CREATE_FLAGS, PUBLIC_CLI_WORKSPACE_DIFF_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, @@ -32,6 +33,10 @@ from pyro_mcp.contract import ( 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, @@ -110,6 +115,35 @@ def test_public_cli_help_lists_commands_and_run_flags() -> None: ).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", diff --git a/tests/test_server.py b/tests/test_server.py index 87c3669..7c1cca7 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -46,6 +46,10 @@ def test_create_server_registers_vm_tools(tmp_path: Path) -> None: assert "service_status" in tool_names assert "service_logs" in tool_names assert "service_stop" in tool_names + assert "snapshot_create" in tool_names + assert "snapshot_delete" in tool_names + assert "snapshot_list" in tool_names + assert "workspace_reset" in tool_names def test_vm_run_round_trip(tmp_path: Path) -> None: @@ -234,6 +238,15 @@ def test_workspace_tools_round_trip(tmp_path: Path) -> None: diffed = _extract_structured( await server.call_tool("workspace_diff", {"workspace_id": workspace_id}) ) + snapshot = _extract_structured( + await server.call_tool( + "snapshot_create", + {"workspace_id": workspace_id, "snapshot_name": "checkpoint"}, + ) + ) + snapshots = _extract_structured( + await server.call_tool("snapshot_list", {"workspace_id": workspace_id}) + ) export_path = tmp_path / "exported-more.txt" exported = _extract_structured( await server.call_tool( @@ -287,6 +300,18 @@ def test_workspace_tools_round_trip(tmp_path: Path) -> None: }, ) ) + reset = _extract_structured( + await server.call_tool( + "workspace_reset", + {"workspace_id": workspace_id, "snapshot": "checkpoint"}, + ) + ) + deleted_snapshot = _extract_structured( + await server.call_tool( + "snapshot_delete", + {"workspace_id": workspace_id, "snapshot_name": "checkpoint"}, + ) + ) logs = _extract_structured( await server.call_tool("workspace_logs", {"workspace_id": workspace_id}) ) @@ -298,12 +323,16 @@ def test_workspace_tools_round_trip(tmp_path: Path) -> None: synced, executed, diffed, + snapshot, + snapshots, exported, service, services, service_status, service_logs, service_stopped, + reset, + deleted_snapshot, logs, deleted, ) @@ -313,12 +342,16 @@ def test_workspace_tools_round_trip(tmp_path: Path) -> None: synced, executed, diffed, + snapshot, + snapshots, exported, service, services, service_status, service_logs, service_stopped, + reset, + deleted_snapshot, logs, deleted, ) = asyncio.run(_run()) @@ -327,6 +360,11 @@ def test_workspace_tools_round_trip(tmp_path: Path) -> None: assert synced["workspace_sync"]["destination"] == "/workspace/subdir" assert executed["stdout"] == "more\n" assert diffed["changed"] is True + assert snapshot["snapshot"]["snapshot_name"] == "checkpoint" + assert [entry["snapshot_name"] for entry in snapshots["snapshots"]] == [ + "baseline", + "checkpoint", + ] assert exported["artifact_type"] == "file" assert Path(str(exported["output_path"])).read_text(encoding="utf-8") == "more\n" assert service["state"] == "running" @@ -334,5 +372,9 @@ def test_workspace_tools_round_trip(tmp_path: Path) -> None: assert service_status["state"] == "running" assert service_logs["tail_lines"] is None assert service_stopped["state"] == "stopped" - assert logs["count"] == 1 + assert reset["workspace_reset"]["snapshot_name"] == "checkpoint" + assert reset["command_count"] == 0 + assert reset["service_count"] == 0 + assert deleted_snapshot["deleted"] is True + assert logs["count"] == 0 assert deleted["deleted"] is True diff --git a/tests/test_vm_manager.py b/tests/test_vm_manager.py index 67b3bbd..146e6bd 100644 --- a/tests/test_vm_manager.py +++ b/tests/test_vm_manager.py @@ -545,10 +545,212 @@ def test_workspace_diff_requires_create_time_baseline(tmp_path: Path) -> None: baseline_path = tmp_path / "vms" / "workspaces" / workspace_id / "baseline" / "workspace.tar" baseline_path.unlink() - with pytest.raises(RuntimeError, match="requires a baseline snapshot"): + with pytest.raises(RuntimeError, match="require[s]? a baseline snapshot"): manager.diff_workspace(workspace_id) +def test_workspace_snapshots_and_reset_round_trip(tmp_path: Path) -> None: + seed_dir = tmp_path / "seed" + seed_dir.mkdir() + (seed_dir / "note.txt").write_text("seed\n", encoding="utf-8") + + manager = VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + + workspace_id = str( + manager.create_workspace( + environment="debian:12-base", + allow_host_compat=True, + seed_path=seed_dir, + )["workspace_id"] + ) + + manager.exec_workspace( + workspace_id, + command="printf 'checkpoint\\n' > note.txt", + timeout_seconds=30, + ) + created_snapshot = manager.create_snapshot(workspace_id, "checkpoint") + assert created_snapshot["snapshot"]["snapshot_name"] == "checkpoint" + + listed = manager.list_snapshots(workspace_id) + assert listed["count"] == 2 + assert [snapshot["snapshot_name"] for snapshot in listed["snapshots"]] == [ + "baseline", + "checkpoint", + ] + + manager.exec_workspace( + workspace_id, + command="printf 'after\\n' > note.txt", + timeout_seconds=30, + ) + manager.start_service( + workspace_id, + "app", + command="sh -lc 'touch .ready; while true; do sleep 60; done'", + readiness={"type": "file", "path": ".ready"}, + ) + + reset_to_snapshot = manager.reset_workspace(workspace_id, snapshot="checkpoint") + assert reset_to_snapshot["workspace_reset"]["snapshot_name"] == "checkpoint" + assert reset_to_snapshot["reset_count"] == 1 + assert reset_to_snapshot["last_command"] is None + assert reset_to_snapshot["command_count"] == 0 + assert reset_to_snapshot["service_count"] == 0 + assert reset_to_snapshot["running_service_count"] == 0 + + checkpoint_result = manager.exec_workspace( + workspace_id, + command="cat note.txt", + timeout_seconds=30, + ) + assert checkpoint_result["stdout"] == "checkpoint\n" + logs_after_snapshot_reset = manager.logs_workspace(workspace_id) + assert logs_after_snapshot_reset["count"] == 1 + + reset_to_baseline = manager.reset_workspace(workspace_id) + assert reset_to_baseline["workspace_reset"]["snapshot_name"] == "baseline" + assert reset_to_baseline["reset_count"] == 2 + assert reset_to_baseline["command_count"] == 0 + assert reset_to_baseline["service_count"] == 0 + assert manager.logs_workspace(workspace_id)["count"] == 0 + + baseline_result = manager.exec_workspace( + workspace_id, + command="cat note.txt", + timeout_seconds=30, + ) + assert baseline_result["stdout"] == "seed\n" + diff_payload = manager.diff_workspace(workspace_id) + assert diff_payload["changed"] is False + + deleted_snapshot = manager.delete_snapshot(workspace_id, "checkpoint") + assert deleted_snapshot["deleted"] is True + listed_after_delete = manager.list_snapshots(workspace_id) + assert [snapshot["snapshot_name"] for snapshot in listed_after_delete["snapshots"]] == [ + "baseline" + ] + + +def test_workspace_snapshot_and_reset_require_baseline(tmp_path: Path) -> None: + manager = VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + workspace_id = str( + manager.create_workspace( + environment="debian:12-base", + allow_host_compat=True, + )["workspace_id"] + ) + baseline_path = tmp_path / "vms" / "workspaces" / workspace_id / "baseline" / "workspace.tar" + baseline_path.unlink() + + with pytest.raises(RuntimeError, match="require[s]? a baseline snapshot"): + manager.list_snapshots(workspace_id) + with pytest.raises(RuntimeError, match="require[s]? a baseline snapshot"): + manager.create_snapshot(workspace_id, "checkpoint") + with pytest.raises(RuntimeError, match="require[s]? a baseline snapshot"): + manager.delete_snapshot(workspace_id, "checkpoint") + with pytest.raises(RuntimeError, match="require[s]? a baseline snapshot"): + manager.reset_workspace(workspace_id) + + +def test_workspace_delete_baseline_snapshot_is_rejected(tmp_path: Path) -> None: + manager = VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + workspace_id = str( + manager.create_workspace( + environment="debian:12-base", + allow_host_compat=True, + )["workspace_id"] + ) + + with pytest.raises(ValueError, match="cannot delete the baseline snapshot"): + manager.delete_snapshot(workspace_id, "baseline") + + +def test_workspace_reset_recreates_stopped_workspace(tmp_path: Path) -> None: + seed_dir = tmp_path / "seed" + seed_dir.mkdir() + (seed_dir / "note.txt").write_text("seed\n", encoding="utf-8") + + manager = VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + workspace_id = str( + manager.create_workspace( + environment="debian:12-base", + allow_host_compat=True, + seed_path=seed_dir, + )["workspace_id"] + ) + + with manager._lock: # noqa: SLF001 + workspace = manager._load_workspace_locked(workspace_id) # noqa: SLF001 + workspace.state = "stopped" + workspace.firecracker_pid = None + manager._save_workspace_locked(workspace) # noqa: SLF001 + + reset_payload = manager.reset_workspace(workspace_id) + + assert reset_payload["state"] == "started" + assert reset_payload["workspace_reset"]["snapshot_name"] == "baseline" + result = manager.exec_workspace(workspace_id, command="cat note.txt", timeout_seconds=30) + assert result["stdout"] == "seed\n" + + +def test_workspace_reset_failure_leaves_workspace_stopped(tmp_path: Path) -> None: + seed_dir = tmp_path / "seed" + seed_dir.mkdir() + (seed_dir / "note.txt").write_text("seed\n", encoding="utf-8") + + manager = VmManager( + backend_name="mock", + base_dir=tmp_path / "vms", + network_manager=TapNetworkManager(enabled=False), + ) + workspace_id = str( + manager.create_workspace( + environment="debian:12-base", + allow_host_compat=True, + seed_path=seed_dir, + )["workspace_id"] + ) + manager.create_snapshot(workspace_id, "checkpoint") + + def _failing_import_archive(*args: Any, **kwargs: Any) -> dict[str, Any]: + del args, kwargs + raise RuntimeError("boom") + + manager._backend.import_archive = _failing_import_archive # type: ignore[method-assign] # noqa: SLF001 + + with pytest.raises(RuntimeError, match="boom"): + manager.reset_workspace(workspace_id, snapshot="checkpoint") + + with manager._lock: # noqa: SLF001 + workspace = manager._load_workspace_locked(workspace_id) # noqa: SLF001 + assert workspace.state == "stopped" + assert workspace.firecracker_pid is None + assert workspace.reset_count == 0 + + listed = manager.list_snapshots(workspace_id) + assert [snapshot["snapshot_name"] for snapshot in listed["snapshots"]] == [ + "baseline", + "checkpoint", + ] + + def test_workspace_export_helpers_preserve_directory_symlinks(tmp_path: Path) -> None: workspace_dir = tmp_path / "workspace" workspace_dir.mkdir() diff --git a/uv.lock b/uv.lock index d7370c6..436eb8c 100644 --- a/uv.lock +++ b/uv.lock @@ -706,7 +706,7 @@ crypto = [ [[package]] name = "pyro-mcp" -version = "2.7.0" +version = "2.8.0" source = { editable = "." } dependencies = [ { name = "mcp" },