diff --git a/benchmarks/sandbox_checkpoint_bench.py b/benchmarks/sandbox_checkpoint_bench.py new file mode 100644 index 000000000..64f331e25 --- /dev/null +++ b/benchmarks/sandbox_checkpoint_bench.py @@ -0,0 +1,117 @@ +"""ASV benchmarks for sandbox checkpoint infrastructure. + +Measures the performance of: +- CheckpointManager.create_checkpoint() +- CheckpointManager.rollback_to() +- CheckpointManager.list_checkpoints() +""" + +from __future__ import annotations + +import importlib +import os +import sys +import tempfile +from pathlib import Path + +# Ensure the local *source* tree is importable even when ASV has an +# older build of the package installed. +_SRC = str(Path(__file__).resolve().parents[1] / "src") +if _SRC not in sys.path: + sys.path.insert(0, _SRC) + +# Force-reload so ASV picks up the source tree version. +import cleveragents # noqa: E402 + +importlib.reload(cleveragents) + +from cleveragents.infrastructure.sandbox.checkpoint import ( # noqa: E402 + CheckpointManager, +) +from cleveragents.infrastructure.sandbox.no_sandbox import NoSandbox # noqa: E402 + + +class CheckpointCreationSuite: + """Benchmark CheckpointManager.create_checkpoint() overhead.""" + + timeout = 60 + + def setup(self) -> None: + self.manager = CheckpointManager() + self.tmpdir = tempfile.mkdtemp(prefix="ca-bench-cp-") + # Create sample files + os.makedirs(os.path.join(self.tmpdir, "src"), exist_ok=True) + with open(os.path.join(self.tmpdir, "src", "main.py"), "w") as f: + f.write("print('benchmark')\n") + with open(os.path.join(self.tmpdir, "README.md"), "w") as f: + f.write("# Benchmark project\n") + + self.sandbox = NoSandbox(resource_id="res-bench-01", original_path=self.tmpdir) + self.sandbox.create(plan_id="plan-bench-01") + + def time_checkpoint_creation(self) -> None: + """Benchmark a single checkpoint creation.""" + self.manager.create_checkpoint( + sandbox=self.sandbox, + plan_id="plan-bench-01", + phase="pre_execute", + metadata={"sandbox_path": self.tmpdir}, + ) + + +class CheckpointRollbackSuite: + """Benchmark CheckpointManager.rollback_to() overhead.""" + + timeout = 60 + + def setup(self) -> None: + self.manager = CheckpointManager() + self.tmpdir = tempfile.mkdtemp(prefix="ca-bench-rb-") + os.makedirs(os.path.join(self.tmpdir, "src"), exist_ok=True) + with open(os.path.join(self.tmpdir, "src", "main.py"), "w") as f: + f.write("print('benchmark')\n") + + self.sandbox = NoSandbox(resource_id="res-bench-02", original_path=self.tmpdir) + self.sandbox.create(plan_id="plan-bench-02") + + self.checkpoint = self.manager.create_checkpoint( + sandbox=self.sandbox, + plan_id="plan-bench-02", + phase="pre_execute", + metadata={"sandbox_path": self.tmpdir}, + ) + + # Modify the file to make rollback meaningful + with open(os.path.join(self.tmpdir, "src", "main.py"), "w") as f: + f.write("print('modified for benchmark')\n") + + def time_checkpoint_rollback(self) -> None: + """Benchmark a single rollback operation.""" + self.manager.rollback_to(self.checkpoint) + + +class CheckpointListingSuite: + """Benchmark CheckpointManager.list_checkpoints() overhead.""" + + timeout = 60 + + def setup(self) -> None: + self.manager = CheckpointManager() + self.tmpdir = tempfile.mkdtemp(prefix="ca-bench-ls-") + + self.sandbox = NoSandbox(resource_id="res-bench-03", original_path=self.tmpdir) + self.sandbox.create(plan_id="plan-bench-03") + + # Create several checkpoints + for phase in ["pre_execute", "post_execute", "pre_apply"]: + self.manager.create_checkpoint( + sandbox=self.sandbox, + plan_id="plan-bench-03", + phase=phase, + ) + + self.sandbox_id = self.sandbox.sandbox_id + + def time_checkpoint_listing(self) -> None: + """Benchmark listing checkpoints for a sandbox.""" + self.manager.list_checkpoints(self.sandbox_id) diff --git a/docs/reference/sandbox.md b/docs/reference/sandbox.md new file mode 100644 index 000000000..c8089f0d1 --- /dev/null +++ b/docs/reference/sandbox.md @@ -0,0 +1,117 @@ +# Sandbox Infrastructure + +## Overview + +The sandbox infrastructure provides resource isolation during plan execution. +Each sandbox creates an isolated environment where a plan can read and write +to a resource without affecting the original until changes are explicitly +committed. + +## Sandbox Strategies + +| Strategy | Class | Description | +|----------|-------|-------------| +| `none` | `NoSandbox` | No isolation; writes go directly to the original | +| `copy_on_write` | `CopyOnWriteSandbox` | Filesystem copy for isolation | +| `git_worktree` | `GitWorktreeSandbox` | Git worktree for git repositories | + +## Sandbox Lifecycle + +``` +PENDING -> CREATED -> ACTIVE -> COMMITTED -> CLEANED_UP + | | + | +-> CLEANED_UP + | + +-> ROLLED_BACK -> ACTIVE + | + +-> ERRORED -> CLEANED_UP +``` + +## Checkpoint and Rollback Hooks + +### Purpose + +Checkpoint hooks preserve sandbox state at key points during plan +execute/apply flows. When a failure occurs, the `CheckpointManager` +can restore a sandbox to a previously captured checkpoint. + +### SandboxCheckpoint Model + +A frozen Pydantic model capturing a point-in-time snapshot: + +| Field | Type | Description | +|-------|------|-------------| +| `checkpoint_id` | `str` | ULID identifier | +| `sandbox_id` | `str` | Sandbox that was checkpointed | +| `plan_id` | `str` | Plan owning the sandbox | +| `phase` | `str` | Lifecycle phase (`pre_execute`, `post_execute`, `pre_apply`) | +| `created_at` | `datetime` | When captured | +| `metadata` | `dict[str, str]` | Key-value metadata (status, reason, etc.) | +| `snapshot_path` | `str` | Path to the snapshot directory | + +### CheckpointManager API + +```python +mgr = CheckpointManager() + +# Create a snapshot before execute +cp = mgr.create_checkpoint(sandbox, plan_id, "pre_execute", {}) + +# Create a snapshot after successful execute +cp2 = mgr.create_checkpoint(sandbox, plan_id, "post_execute", {"status": "success"}) + +# Rollback on failure +success = mgr.rollback_to(cp) + +# List all checkpoints for a sandbox +checkpoints = mgr.list_checkpoints(sandbox.sandbox_id) + +# Delete a checkpoint +mgr.delete_checkpoint(cp.checkpoint_id) +``` + +### Checkpoint Lifecycle + +1. **Pre-execute checkpoint**: Created before the execute phase starts. + Captures the sandbox state so that a failed execution can be rolled back. + +2. **Post-execute checkpoint**: Created after a successful execute phase. + Preserves the post-execution state before apply begins. + +3. **Pre-apply checkpoint**: Created before the apply phase starts. + Allows rollback if the apply fails or encounters merge conflicts. + +4. **Rollback on failure**: When execute or apply fails, the system + attempts to restore the sandbox to the most recent checkpoint. + +### Integration with Plan Executor + +The `PlanExecutor` accepts an optional `checkpoint_manager` parameter. +When provided, checkpoint hooks are automatically invoked: + +- Before `run_execute`: `create_checkpoint(sandbox, plan_id, "pre_execute")` +- After successful execute: `create_checkpoint(sandbox, plan_id, "post_execute")` +- On execute failure: `rollback_to(last_checkpoint)` + +When no `CheckpointManager` is injected, all hooks are silently skipped. + +### Integration with Plan Apply Service + +The `PlanApplyService` also accepts an optional `checkpoint_manager`: + +- Before apply: `create_checkpoint(sandbox, plan_id, "pre_apply")` +- On apply failure: `rollback_to(last_checkpoint)` + +### Thread Safety + +The `CheckpointManager` is thread-safe. All mutable state is protected +by a reentrant lock (`threading.RLock`). + +### Snapshot Storage + +Snapshots are stored in temporary directories under the system temp folder. +Each snapshot is a full copy of the sandbox working directory at the time +of checkpoint creation. Snapshots are cleaned up when: + +- A checkpoint is explicitly deleted via `delete_checkpoint()` +- The checkpoint manager goes out of scope (manual cleanup recommended) diff --git a/features/sandbox_checkpoints.feature b/features/sandbox_checkpoints.feature new file mode 100644 index 000000000..e96019aaf --- /dev/null +++ b/features/sandbox_checkpoints.feature @@ -0,0 +1,97 @@ +Feature: Sandbox Checkpoint and Rollback Hooks + As an execution engine + I want to snapshot sandbox state at key points during plan execute/apply flows + So that I can rollback to a known-good state when failures occur + + Background: + Given a checkpoint manager + + # Checkpoint creation + + Scenario: Checkpoint created before plan execution + Given a sandbox with files in a temporary directory + When I create a checkpoint with phase "pre_execute" and plan "plan-001" + Then a checkpoint should be returned + And the checkpoint phase should be "pre_execute" + And the checkpoint plan_id should be "plan-001" + And the checkpoint sandbox_id should match the sandbox + + Scenario: Checkpoint created after successful execution + Given a sandbox with files in a temporary directory + When I create a checkpoint with phase "post_execute" and plan "plan-001" + And I attach metadata "status" as "success" + Then a checkpoint should be returned + And the checkpoint phase should be "post_execute" + And the checkpoint metadata should contain key "status" with value "success" + + # Rollback + + Scenario: Rollback restores sandbox state on failure + Given a sandbox with files in a temporary directory + And a checkpoint exists for the sandbox with phase "pre_execute" and plan "plan-001" + When I modify the sandbox files + And I rollback to the checkpoint + Then the rollback should succeed + And the sandbox files should match the original state + + Scenario: Rollback returns false when snapshot path is missing + Given a sandbox with files in a temporary directory + And a checkpoint with an invalid snapshot path + When I attempt to rollback to the invalid checkpoint + Then the rollback should fail + + # Listing + + Scenario: Multiple checkpoints maintain order + Given a sandbox with files in a temporary directory + When I create a checkpoint with phase "pre_execute" and plan "plan-001" + And I create a checkpoint with phase "post_execute" and plan "plan-001" + And I create a checkpoint with phase "pre_apply" and plan "plan-001" + Then listing checkpoints for the sandbox should return 3 checkpoints + And the checkpoints should be in creation order + + # Metadata + + Scenario: Checkpoint metadata captures phase and plan information + Given a sandbox with files in a temporary directory + When I create a checkpoint with phase "pre_execute" and plan "plan-002" + Then the checkpoint should have a valid ULID checkpoint_id + And the checkpoint should have a created_at timestamp + And the checkpoint should have a non-empty snapshot_path + + # Deletion + + Scenario: Deleting a checkpoint removes it from the list + Given a sandbox with files in a temporary directory + And a checkpoint exists for the sandbox with phase "pre_execute" and plan "plan-001" + When I delete the checkpoint + Then the deletion should succeed + And listing checkpoints for the sandbox should return 0 checkpoints + + Scenario: Deleting a non-existent checkpoint returns false + When I delete a checkpoint with id "nonexistent-id" + Then the deletion should fail + + # Plan executor integration (optional checkpoint manager) + + Scenario: PlanExecutor works without checkpoint manager + Given a plan lifecycle service + And a plan executor without checkpoint manager + Then the plan executor checkpoint_manager should be None + + Scenario: PlanExecutor accepts optional checkpoint manager + Given a plan lifecycle service + And a plan executor with a checkpoint manager + Then the plan executor checkpoint_manager should not be None + + # PlanApplyService integration + + Scenario: PlanApplyService works without checkpoint manager + Given a plan lifecycle service + And a plan apply service without checkpoint manager + Then the plan apply service checkpoint_manager should be None + + Scenario: PlanApplyService accepts optional checkpoint manager + Given a plan lifecycle service + And a plan apply service with a checkpoint manager + Then the plan apply service checkpoint_manager should not be None diff --git a/features/steps/changeset_persistence_steps.py b/features/steps/changeset_persistence_steps.py index a0fe739d7..09564bddb 100644 --- a/features/steps/changeset_persistence_steps.py +++ b/features/steps/changeset_persistence_steps.py @@ -3,10 +3,12 @@ from __future__ import annotations import re +from collections.abc import Callable from behave import given, then, when # type: ignore[import-untyped] from behave.runner import Context # type: ignore[import-untyped] from sqlalchemy import create_engine +from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker from cleveragents.application.services.plan_apply_service import ( @@ -30,16 +32,23 @@ from cleveragents.infrastructure.database.changeset_repository import ( from cleveragents.infrastructure.database.models import Base -def _make_in_memory_session_factory() -> tuple[sessionmaker[Session], object]: - """Create an in-memory SQLite engine + session factory with schema.""" +def _make_in_memory_session_factory() -> tuple[Callable[[], Session], Session, Engine]: + """Create an in-memory SQLite engine + shared session factory with schema. + + Uses a single shared session so that flush() data is visible across + all repository calls within the same scenario. Without this, + short-lived sessions returned by sessionmaker() may be garbage- + collected, causing the StaticPool's reset_on_return to discard + uncommitted INSERTs. + """ engine = create_engine( "sqlite:///:memory:", echo=False, connect_args={"check_same_thread": False}, ) Base.metadata.create_all(engine) - factory: sessionmaker[Session] = sessionmaker(bind=engine) - return factory, engine + session: Session = sessionmaker(bind=engine)() + return lambda: session, session, engine def _make_entry( @@ -67,10 +76,12 @@ def _make_entry( @given("I have a sqlite-backed changeset store") def step_sqlite_store(context: Context) -> None: - factory, engine = _make_in_memory_session_factory() + factory, session, engine = _make_in_memory_session_factory() context.session_factory = factory context.engine = engine context.sqlite_store = SqliteChangeSetStore(factory) + context._cleanup_handlers.append(lambda: session.close()) + context._cleanup_handlers.append(lambda: engine.dispose()) @when('I persist-start a changeset for plan "{plan_id}"') @@ -162,10 +173,12 @@ def step_summarize_empty(context: Context, cid: str) -> None: @given("I have a persisted ChangeSetEntryRepository") def step_entry_repo(context: Context) -> None: - factory, engine = _make_in_memory_session_factory() + factory, session, engine = _make_in_memory_session_factory() context.session_factory = factory context.engine = engine context.entry_repo = ChangeSetEntryRepository(factory) + context._cleanup_handlers.append(lambda: session.close()) + context._cleanup_handlers.append(lambda: engine.dispose()) @when('I persist-save an entry for changeset "{cs_id}" plan "{plan_id}"') @@ -209,10 +222,12 @@ def step_delete_for_changeset(context: Context, cs_id: str) -> None: @given("I have a persisted ToolInvocationRepository") def step_inv_repo(context: Context) -> None: - factory, engine = _make_in_memory_session_factory() + factory, session, engine = _make_in_memory_session_factory() context.session_factory = factory context.engine = engine context.inv_repo = ToolInvocationRepository(factory) + context._cleanup_handlers.append(lambda: session.close()) + context._cleanup_handlers.append(lambda: engine.dispose()) @when('I persist-save an invocation for plan "{plan_id}"') @@ -243,10 +258,12 @@ def step_delete_invocations(context: Context, plan_id: str) -> None: @given("I have a PlanApplyService backed by sqlite changeset store") def step_apply_service_with_store(context: Context) -> None: - factory, engine = _make_in_memory_session_factory() + factory, session, engine = _make_in_memory_session_factory() context.session_factory = factory context.engine = engine context.sqlite_store = SqliteChangeSetStore(factory) + context._cleanup_handlers.append(lambda: session.close()) + context._cleanup_handlers.append(lambda: engine.dispose()) settings = Settings() lifecycle = PlanLifecycleService(settings=settings) diff --git a/features/steps/changeset_repository_coverage_steps.py b/features/steps/changeset_repository_coverage_steps.py index 4abdeec87..04fb723fe 100644 --- a/features/steps/changeset_repository_coverage_steps.py +++ b/features/steps/changeset_repository_coverage_steps.py @@ -3,12 +3,14 @@ from __future__ import annotations import json +from collections.abc import Callable from datetime import UTC, datetime from typing import Any from behave import given, then, when # type: ignore[import-untyped] from behave.runner import Context # type: ignore[import-untyped] from sqlalchemy import create_engine +from sqlalchemy.engine import Engine from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session, sessionmaker @@ -32,16 +34,23 @@ from cleveragents.infrastructure.database.models import ( # ------ Helpers ------ -def _make_session_factory() -> tuple[sessionmaker[Session], Any]: - """Create an in-memory SQLite engine and session factory.""" +def _make_session_factory() -> tuple[Callable[[], Session], Session, Engine]: + """Create an in-memory SQLite engine + shared session factory with schema. + + Uses a single shared session so that flush() data is visible across + all repository calls within the same scenario. Without this, + short-lived sessions returned by sessionmaker() may be garbage- + collected, causing the StaticPool's reset_on_return to discard + uncommitted INSERTs. + """ engine = create_engine( "sqlite:///:memory:", echo=False, connect_args={"check_same_thread": False}, ) Base.metadata.create_all(engine) - factory: sessionmaker[Session] = sessionmaker(bind=engine) - return factory, engine + session: Session = sessionmaker(bind=engine)() + return lambda: session, session, engine def _make_change_entry( @@ -114,9 +123,11 @@ class _BrokenDeleteSession: @given("I have a coverage ChangeSetEntryRepository") def step_create_coverage_entry_repo(context: Context) -> None: """Create a ChangeSetEntryRepository with real in-memory session.""" - factory, engine = _make_session_factory() + factory, session, engine = _make_session_factory() context.cov_engine = engine context.cov_entry_repo = ChangeSetEntryRepository(factory) + context._cleanup_handlers.append(lambda: session.close()) + context._cleanup_handlers.append(lambda: engine.dispose()) @when("I try to coverage-save an entry with empty changeset_id") @@ -388,9 +399,11 @@ def step_assert_after_mode_value(context: Context, expected: int) -> None: @given("I have a coverage ToolInvocationRepository") def step_create_coverage_inv_repo(context: Context) -> None: """Create a ToolInvocationRepository with real in-memory session.""" - factory, engine = _make_session_factory() + factory, session, engine = _make_session_factory() context.cov_engine = engine context.cov_inv_repo = ToolInvocationRepository(factory) + context._cleanup_handlers.append(lambda: session.close()) + context._cleanup_handlers.append(lambda: engine.dispose()) @when("I try to coverage-save a non-ToolInvocation object") @@ -486,10 +499,12 @@ def step_delete_invocations_broken_plan(context: Context, plan_id: str) -> None: @given("I have a coverage ToolInvocationRepository with real session") def step_create_real_inv_repo(context: Context) -> None: """Create a ToolInvocationRepository with real in-memory session.""" - factory, engine = _make_session_factory() + factory, session, engine = _make_session_factory() context.cov_engine = engine context.cov_session_factory = factory context.cov_inv_repo = ToolInvocationRepository(factory) + context._cleanup_handlers.append(lambda: session.close()) + context._cleanup_handlers.append(lambda: engine.dispose()) @when("I coverage-save an invocation with result, completed_at, and provider_metadata") @@ -861,10 +876,12 @@ def step_assert_inv_completed_datetime(context: Context) -> None: @given("I have a coverage sqlite changeset store") def step_create_coverage_store(context: Context) -> None: """Create a SqliteChangeSetStore backed by in-memory SQLite.""" - factory, engine = _make_session_factory() + factory, session, engine = _make_session_factory() context.cov_engine = engine context.cov_session_factory = factory context.cov_store = SqliteChangeSetStore(factory) + context._cleanup_handlers.append(lambda: session.close()) + context._cleanup_handlers.append(lambda: engine.dispose()) @when("I coverage-get a changeset with empty string") diff --git a/features/steps/sandbox_checkpoints_steps.py b/features/steps/sandbox_checkpoints_steps.py new file mode 100644 index 000000000..66fbde97e --- /dev/null +++ b/features/steps/sandbox_checkpoints_steps.py @@ -0,0 +1,331 @@ +"""Step definitions for sandbox checkpoint and rollback BDD tests.""" + +from __future__ import annotations + +import os +import tempfile +from typing import Any +from unittest.mock import MagicMock + +from behave import given, then, when + +from cleveragents.infrastructure.sandbox.checkpoint import ( + CheckpointManager, + SandboxCheckpoint, +) +from cleveragents.infrastructure.sandbox.no_sandbox import NoSandbox + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _create_temp_sandbox(context: Any) -> None: + """Create a temporary directory with sample files for sandbox testing.""" + tmpdir = tempfile.mkdtemp(prefix="ca-cp-test-") + context._tmpdir = tmpdir + + # Write sample files + os.makedirs(os.path.join(tmpdir, "src"), exist_ok=True) + with open(os.path.join(tmpdir, "src", "main.py"), "w") as f: + f.write("print('hello')\n") + with open(os.path.join(tmpdir, "README.md"), "w") as f: + f.write("# Test project\n") + + sandbox = NoSandbox(resource_id="res-cp-test", original_path=tmpdir) + sandbox.create(plan_id="plan-cp-test") + context.sandbox = sandbox + context.sandbox_path = tmpdir + + +# --------------------------------------------------------------------------- +# Given steps +# --------------------------------------------------------------------------- + + +@given("a checkpoint manager") +def step_given_checkpoint_manager(context: Any) -> None: + context.checkpoint_manager = CheckpointManager() + context.checkpoint = None + context.checkpoints = [] + context.rollback_result = None + context.delete_result = None + context.error = None + + +@given("a sandbox with files in a temporary directory") +def step_given_sandbox_with_files(context: Any) -> None: + _create_temp_sandbox(context) + + +@given('a checkpoint exists for the sandbox with phase "{phase}" and plan "{plan_id}"') +def step_given_checkpoint_exists(context: Any, phase: str, plan_id: str) -> None: + meta: dict[str, str] = {} + if hasattr(context, "sandbox_path") and context.sandbox_path: + meta["sandbox_path"] = context.sandbox_path + context.checkpoint = context.checkpoint_manager.create_checkpoint( + sandbox=context.sandbox, + plan_id=plan_id, + phase=phase, + metadata=meta, + ) + + +@given("a checkpoint with an invalid snapshot path") +def step_given_invalid_checkpoint(context: Any) -> None: + context.invalid_checkpoint = SandboxCheckpoint( + checkpoint_id="invalid-cp-001", + sandbox_id="sb-invalid", + plan_id="plan-invalid", + phase="pre_execute", + created_at=__import__("datetime").datetime.now(), + metadata={}, + snapshot_path="/nonexistent/path/snapshot", + ) + + +@given("a plan lifecycle service") +def step_given_lifecycle_service(context: Any) -> None: + mock_settings = MagicMock() + mock_settings.cleveragents_home = "/tmp/test" + mock_settings.default_automation_profile = None + context.lifecycle_service = MagicMock() + context.lifecycle_service.settings = mock_settings + + +@given("a plan executor without checkpoint manager") +def step_given_executor_no_checkpoint(context: Any) -> None: + from cleveragents.application.services.plan_executor import PlanExecutor + + context.executor = PlanExecutor( + lifecycle_service=context.lifecycle_service, + ) + + +@given("a plan executor with a checkpoint manager") +def step_given_executor_with_checkpoint(context: Any) -> None: + from cleveragents.application.services.plan_executor import PlanExecutor + + context.executor = PlanExecutor( + lifecycle_service=context.lifecycle_service, + checkpoint_manager=CheckpointManager(), + ) + + +@given("a plan apply service without checkpoint manager") +def step_given_apply_no_checkpoint(context: Any) -> None: + from cleveragents.application.services.plan_apply_service import ( + PlanApplyService, + ) + + context.apply_service = PlanApplyService( + lifecycle_service=context.lifecycle_service, + ) + + +@given("a plan apply service with a checkpoint manager") +def step_given_apply_with_checkpoint(context: Any) -> None: + from cleveragents.application.services.plan_apply_service import ( + PlanApplyService, + ) + + context.apply_service = PlanApplyService( + lifecycle_service=context.lifecycle_service, + checkpoint_manager=CheckpointManager(), + ) + + +# --------------------------------------------------------------------------- +# When steps +# --------------------------------------------------------------------------- + + +@when('I create a checkpoint with phase "{phase}" and plan "{plan_id}"') +def step_when_create_checkpoint(context: Any, phase: str, plan_id: str) -> None: + meta: dict[str, str] = {} + if hasattr(context, "sandbox_path") and context.sandbox_path: + meta["sandbox_path"] = context.sandbox_path + context.checkpoint = context.checkpoint_manager.create_checkpoint( + sandbox=context.sandbox, + plan_id=plan_id, + phase=phase, + metadata=meta, + ) + context.checkpoints.append(context.checkpoint) + + +@when('I attach metadata "{key}" as "{value}"') +def step_when_attach_metadata(context: Any, key: str, value: str) -> None: + # Metadata is attached during creation; re-create with metadata + meta = dict(context.checkpoint.metadata) + meta[key] = value + if hasattr(context, "sandbox_path") and context.sandbox_path: + meta["sandbox_path"] = context.sandbox_path + context.checkpoint = context.checkpoint_manager.create_checkpoint( + sandbox=context.sandbox, + plan_id=context.checkpoint.plan_id, + phase=context.checkpoint.phase, + metadata=meta, + ) + + +@when("I modify the sandbox files") +def step_when_modify_sandbox(context: Any) -> None: + main_path = os.path.join(context.sandbox_path, "src", "main.py") + with open(main_path, "w") as f: + f.write("print('modified')\n") + new_file = os.path.join(context.sandbox_path, "new_file.txt") + with open(new_file, "w") as f: + f.write("new content\n") + + +@when("I rollback to the checkpoint") +def step_when_rollback(context: Any) -> None: + context.rollback_result = context.checkpoint_manager.rollback_to(context.checkpoint) + + +@when("I attempt to rollback to the invalid checkpoint") +def step_when_rollback_invalid(context: Any) -> None: + context.rollback_result = context.checkpoint_manager.rollback_to( + context.invalid_checkpoint + ) + + +@when("I delete the checkpoint") +def step_when_delete_checkpoint(context: Any) -> None: + context.delete_result = context.checkpoint_manager.delete_checkpoint( + context.checkpoint.checkpoint_id + ) + + +@when('I delete a checkpoint with id "{checkpoint_id}"') +def step_when_delete_nonexistent(context: Any, checkpoint_id: str) -> None: + context.delete_result = context.checkpoint_manager.delete_checkpoint(checkpoint_id) + + +# --------------------------------------------------------------------------- +# Then steps +# --------------------------------------------------------------------------- + + +@then("a checkpoint should be returned") +def step_then_checkpoint_returned(context: Any) -> None: + assert context.checkpoint is not None, "Expected a checkpoint" + assert isinstance(context.checkpoint, SandboxCheckpoint) + + +@then('the checkpoint phase should be "{phase}"') +def step_then_checkpoint_phase(context: Any, phase: str) -> None: + assert context.checkpoint.phase == phase, ( + f"Expected phase '{phase}', got '{context.checkpoint.phase}'" + ) + + +@then('the checkpoint plan_id should be "{plan_id}"') +def step_then_checkpoint_plan_id(context: Any, plan_id: str) -> None: + assert context.checkpoint.plan_id == plan_id, ( + f"Expected plan_id '{plan_id}', got '{context.checkpoint.plan_id}'" + ) + + +@then("the checkpoint sandbox_id should match the sandbox") +def step_then_checkpoint_sandbox_id(context: Any) -> None: + assert context.checkpoint.sandbox_id == context.sandbox.sandbox_id, ( + f"Expected sandbox_id '{context.sandbox.sandbox_id}', " + f"got '{context.checkpoint.sandbox_id}'" + ) + + +@then('the checkpoint metadata should contain key "{key}" with value "{value}"') +def step_then_checkpoint_metadata(context: Any, key: str, value: str) -> None: + assert key in context.checkpoint.metadata, ( + f"Key '{key}' not in metadata: {context.checkpoint.metadata}" + ) + assert context.checkpoint.metadata[key] == value, ( + f"Expected metadata['{key}'] = '{value}', " + f"got '{context.checkpoint.metadata[key]}'" + ) + + +@then("the rollback should succeed") +def step_then_rollback_success(context: Any) -> None: + assert context.rollback_result is True, "Expected rollback to succeed" + + +@then("the rollback should fail") +def step_then_rollback_fail(context: Any) -> None: + assert context.rollback_result is False, "Expected rollback to fail" + + +@then("the sandbox files should match the original state") +def step_then_sandbox_matches_original(context: Any) -> None: + main_path = os.path.join(context.sandbox_path, "src", "main.py") + with open(main_path) as f: + content = f.read() + assert content == "print('hello')\n", f"Expected original content, got: {content!r}" + new_file = os.path.join(context.sandbox_path, "new_file.txt") + assert not os.path.exists(new_file), "new_file.txt should not exist" + + +@then("listing checkpoints for the sandbox should return {count:d} checkpoints") +def step_then_checkpoint_count(context: Any, count: int) -> None: + result = context.checkpoint_manager.list_checkpoints(context.sandbox.sandbox_id) + assert len(result) == count, f"Expected {count} checkpoints, got {len(result)}" + + +@then("the checkpoints should be in creation order") +def step_then_checkpoints_ordered(context: Any) -> None: + result = context.checkpoint_manager.list_checkpoints(context.sandbox.sandbox_id) + for i in range(1, len(result)): + assert result[i].created_at >= result[i - 1].created_at, ( + f"Checkpoints not in order at index {i}" + ) + + +@then("the checkpoint should have a valid ULID checkpoint_id") +def step_then_valid_ulid(context: Any) -> None: + cp_id = context.checkpoint.checkpoint_id + assert len(cp_id) == 26, f"ULID should be 26 chars, got {len(cp_id)}" + + +@then("the checkpoint should have a created_at timestamp") +def step_then_has_timestamp(context: Any) -> None: + assert context.checkpoint.created_at is not None + + +@then("the checkpoint should have a non-empty snapshot_path") +def step_then_has_snapshot_path(context: Any) -> None: + assert context.checkpoint.snapshot_path, "snapshot_path should not be empty" + assert os.path.isdir(context.checkpoint.snapshot_path), ( + f"snapshot_path should be a directory: {context.checkpoint.snapshot_path}" + ) + + +@then("the deletion should succeed") +def step_then_deletion_success(context: Any) -> None: + assert context.delete_result is True, "Expected deletion to succeed" + + +@then("the deletion should fail") +def step_then_deletion_fail(context: Any) -> None: + assert context.delete_result is False, "Expected deletion to fail" + + +@then("the plan executor checkpoint_manager should be None") +def step_then_executor_no_checkpoint(context: Any) -> None: + assert context.executor.checkpoint_manager is None + + +@then("the plan executor checkpoint_manager should not be None") +def step_then_executor_has_checkpoint(context: Any) -> None: + assert context.executor.checkpoint_manager is not None + + +@then("the plan apply service checkpoint_manager should be None") +def step_then_apply_no_checkpoint(context: Any) -> None: + assert context.apply_service.checkpoint_manager is None + + +@then("the plan apply service checkpoint_manager should not be None") +def step_then_apply_has_checkpoint(context: Any) -> None: + assert context.apply_service.checkpoint_manager is not None diff --git a/robot/helper_sandbox_checkpoint.py b/robot/helper_sandbox_checkpoint.py new file mode 100644 index 000000000..15b47f1f2 --- /dev/null +++ b/robot/helper_sandbox_checkpoint.py @@ -0,0 +1,194 @@ +"""Helper utilities for sandbox checkpoint Robot integration tests. + +Covers the M4 checkpoint infrastructure: +- CheckpointManager create/rollback/list/delete +- PlanExecutor checkpoint integration +""" + +from __future__ import annotations + +import os +import sys +import tempfile +from unittest.mock import MagicMock + +from cleveragents.infrastructure.sandbox.checkpoint import ( + CheckpointManager, + SandboxCheckpoint, +) +from cleveragents.infrastructure.sandbox.no_sandbox import NoSandbox + + +def _checkpoint_create() -> None: + """Integration test: create a checkpoint and verify its fields.""" + mgr = CheckpointManager() + + with tempfile.TemporaryDirectory() as tmpdir: + # Create a file in the directory + with open(os.path.join(tmpdir, "file.txt"), "w") as f: + f.write("original content\n") + + sandbox = NoSandbox(resource_id="res-robot-01", original_path=tmpdir) + sandbox.create(plan_id="plan-robot-01") + + cp = mgr.create_checkpoint( + sandbox=sandbox, + plan_id="plan-robot-01", + phase="pre_execute", + metadata={"sandbox_path": tmpdir}, + ) + + assert isinstance(cp, SandboxCheckpoint) + assert cp.phase == "pre_execute" + assert cp.plan_id == "plan-robot-01" + assert cp.sandbox_id == sandbox.sandbox_id + assert len(cp.checkpoint_id) == 26 # ULID length + assert cp.created_at is not None + assert os.path.isdir(cp.snapshot_path) + + print("checkpoint-create-ok") + + +def _checkpoint_rollback() -> None: + """Integration test: create checkpoint, modify, rollback, verify.""" + mgr = CheckpointManager() + + with tempfile.TemporaryDirectory() as tmpdir: + # Create original file + file_path = os.path.join(tmpdir, "data.txt") + with open(file_path, "w") as f: + f.write("original\n") + + sandbox = NoSandbox(resource_id="res-robot-02", original_path=tmpdir) + sandbox.create(plan_id="plan-robot-02") + + # Checkpoint + cp = mgr.create_checkpoint( + sandbox=sandbox, + plan_id="plan-robot-02", + phase="pre_execute", + metadata={"sandbox_path": tmpdir}, + ) + + # Modify the file + with open(file_path, "w") as f: + f.write("modified\n") + with open(file_path) as f: + assert f.read() == "modified\n" + + # Add a new file + new_file = os.path.join(tmpdir, "extra.txt") + with open(new_file, "w") as f: + f.write("extra\n") + + # Rollback + result = mgr.rollback_to(cp) + assert result is True + + # Verify original state restored + with open(file_path) as f: + content = f.read() + assert content == "original\n", f"Expected 'original\\n', got {content!r}" + + # New file should be gone + assert not os.path.exists(new_file), "extra.txt should not exist" + + print("checkpoint-rollback-ok") + + +def _checkpoint_listing() -> None: + """Integration test: multiple checkpoints maintain creation order.""" + mgr = CheckpointManager() + + with tempfile.TemporaryDirectory() as tmpdir: + sandbox = NoSandbox(resource_id="res-robot-03", original_path=tmpdir) + sandbox.create(plan_id="plan-robot-03") + + phases = ["pre_execute", "post_execute", "pre_apply"] + for phase in phases: + mgr.create_checkpoint( + sandbox=sandbox, + plan_id="plan-robot-03", + phase=phase, + ) + + checkpoints = mgr.list_checkpoints(sandbox.sandbox_id) + assert len(checkpoints) == 3 + + for i, phase in enumerate(phases): + assert checkpoints[i].phase == phase + + # Verify chronological order + for i in range(1, len(checkpoints)): + assert checkpoints[i].created_at >= checkpoints[i - 1].created_at + + print("checkpoint-listing-ok") + + +def _checkpoint_delete() -> None: + """Integration test: delete a checkpoint and verify removal.""" + mgr = CheckpointManager() + + with tempfile.TemporaryDirectory() as tmpdir: + sandbox = NoSandbox(resource_id="res-robot-04", original_path=tmpdir) + sandbox.create(plan_id="plan-robot-04") + + cp = mgr.create_checkpoint( + sandbox=sandbox, + plan_id="plan-robot-04", + phase="pre_execute", + ) + + assert len(mgr.list_checkpoints(sandbox.sandbox_id)) == 1 + + result = mgr.delete_checkpoint(cp.checkpoint_id) + assert result is True + + assert len(mgr.list_checkpoints(sandbox.sandbox_id)) == 0 + + # Deleting again should return False + result2 = mgr.delete_checkpoint(cp.checkpoint_id) + assert result2 is False + + print("checkpoint-delete-ok") + + +def _executor_integration() -> None: + """Integration test: PlanExecutor accepts optional checkpoint_manager.""" + from cleveragents.application.services.plan_executor import PlanExecutor + + lifecycle = MagicMock() + + # Without checkpoint manager + executor_no_cp = PlanExecutor(lifecycle_service=lifecycle) + assert executor_no_cp.checkpoint_manager is None + + # With checkpoint manager + mgr = CheckpointManager() + executor_with_cp = PlanExecutor( + lifecycle_service=lifecycle, + checkpoint_manager=mgr, + ) + assert executor_with_cp.checkpoint_manager is mgr + + print("executor-integration-ok") + + +def main() -> None: + if len(sys.argv) < 2: + raise SystemExit("Expected command argument") + command = sys.argv[1] + commands = { + "checkpoint-create": _checkpoint_create, + "checkpoint-rollback": _checkpoint_rollback, + "checkpoint-listing": _checkpoint_listing, + "checkpoint-delete": _checkpoint_delete, + "executor-integration": _executor_integration, + } + if command not in commands: + raise SystemExit(f"Unknown command: {command}") + commands[command]() + + +if __name__ == "__main__": + main() diff --git a/robot/sandbox_checkpoint_smoke.robot b/robot/sandbox_checkpoint_smoke.robot new file mode 100644 index 000000000..ba8e91118 --- /dev/null +++ b/robot/sandbox_checkpoint_smoke.robot @@ -0,0 +1,45 @@ +*** Settings *** +Documentation Smoke tests for sandbox checkpoint and rollback hooks. +... Covers M4 checkpoint infrastructure: creation, rollback, listing. +Resource ${CURDIR}/common.resource +Suite Setup Setup Test Environment +Suite Teardown Cleanup Test Environment + +*** Variables *** +${HELPER_SCRIPT} robot/helper_sandbox_checkpoint.py + +*** Test Cases *** +Verify Checkpoint Creation During Execute + [Documentation] Create a checkpoint before execute and verify its fields + [Tags] sandbox checkpoint create + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} checkpoint-create cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} checkpoint-create-ok + +Verify Rollback After Failed Apply + [Documentation] Snapshot sandbox, modify files, rollback, verify restoration + [Tags] sandbox checkpoint rollback + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} checkpoint-rollback cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} checkpoint-rollback-ok + +Verify Checkpoint Listing Order + [Documentation] Create multiple checkpoints and verify ordering + [Tags] sandbox checkpoint listing + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} checkpoint-listing cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} checkpoint-listing-ok + +Verify Checkpoint Deletion + [Documentation] Create and delete a checkpoint, verify removal + [Tags] sandbox checkpoint delete + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} checkpoint-delete cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} checkpoint-delete-ok + +Verify PlanExecutor Checkpoint Integration + [Documentation] PlanExecutor accepts optional checkpoint_manager + [Tags] sandbox checkpoint integration + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} executor-integration cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} executor-integration-ok diff --git a/src/cleveragents/application/services/plan_apply_service.py b/src/cleveragents/application/services/plan_apply_service.py index ced1a6e66..6157336ad 100644 --- a/src/cleveragents/application/services/plan_apply_service.py +++ b/src/cleveragents/application/services/plan_apply_service.py @@ -31,6 +31,7 @@ from cleveragents.core.exceptions import PlanError, ValidationError from cleveragents.domain.models.core.change import ( SpecChangeSet, ) +from cleveragents.infrastructure.sandbox.checkpoint import CheckpointManager if TYPE_CHECKING: from cleveragents.application.services.plan_lifecycle_service import ( @@ -255,6 +256,7 @@ class PlanApplyService: self, lifecycle_service: PlanLifecycleService, changeset_store: Any | None = None, + checkpoint_manager: CheckpointManager | None = None, ) -> None: """Initialise the plan apply service. @@ -264,11 +266,15 @@ class PlanApplyService: changeset_store: Optional ``ChangeSetStore`` for retrieving persisted changesets. When ``None``, diff and artifact commands build output from plan metadata only. + checkpoint_manager: Optional checkpoint manager for + sandbox state snapshots during apply. When ``None``, + checkpoint hooks are silently skipped. """ if lifecycle_service is None: raise ValidationError("lifecycle_service must not be None") self._lifecycle = lifecycle_service self._changeset_store = changeset_store + self._checkpoint_manager = checkpoint_manager self._logger = logger.bind(service="plan_apply") # -- Diff output --------------------------------------------------------- @@ -462,6 +468,11 @@ class PlanApplyService: "No changes to apply. Use --allow-empty to override." ) + @property + def checkpoint_manager(self) -> CheckpointManager | None: + """Return the checkpoint manager, if configured.""" + return self._checkpoint_manager + # -- Validation-gated apply --------------------------------------------- def apply_with_validation_gate( @@ -592,6 +603,9 @@ class PlanApplyService: required_passed=req_passed, ) + # Checkpoint: pre_apply snapshot + self._try_checkpoint(plan_id, "pre_apply") + # Persist apply summary self.persist_apply_summary( plan_id=plan_id, @@ -603,8 +617,8 @@ class PlanApplyService: try: self._lifecycle.complete_apply(plan_id) except (PlanError, ValidationError): - # If lifecycle transition fails (e.g., wrong phase), - # still report success in terms of validation gating + # If lifecycle transition fails, attempt rollback + self._try_rollback(plan_id) self._logger.debug( "Could not transition to applied (plan may not be in Apply phase)", plan_id=plan_id, @@ -643,6 +657,56 @@ class PlanApplyService: return req_passed, req_failed, total + # -- Checkpoint helpers -------------------------------------------------- + + def _try_checkpoint( + self, + plan_id: str, + phase: str, + metadata: dict[str, Any] | None = None, + ) -> None: + """Create a checkpoint if a checkpoint manager is configured. + + Non-fatal: logs and swallows any errors. + """ + if self._checkpoint_manager is None: + return + try: + # Build a lightweight proxy for checkpoint creation + sandbox = _ApplyPhaseProxy(plan_id) + self._checkpoint_manager.create_checkpoint( + sandbox=sandbox, + plan_id=plan_id, + phase=phase, + metadata=metadata or {}, + ) + except Exception: + self._logger.debug( + "Checkpoint creation failed (non-fatal)", + plan_id=plan_id, + phase=phase, + exc_info=True, + ) + + def _try_rollback(self, plan_id: str) -> None: + """Attempt rollback to the last checkpoint for a plan. + + Non-fatal: logs and swallows any errors. + """ + if self._checkpoint_manager is None: + return + try: + sandbox_id = f"apply-{plan_id[:16]}" + checkpoints = self._checkpoint_manager.list_checkpoints(sandbox_id) + if checkpoints: + self._checkpoint_manager.rollback_to(checkpoints[-1]) + except Exception: + self._logger.debug( + "Checkpoint rollback failed (non-fatal)", + plan_id=plan_id, + exc_info=True, + ) + # -- Internal helpers ---------------------------------------------------- def _resolve_changeset(self, plan: Plan) -> SpecChangeSet | None: @@ -695,3 +759,29 @@ class PlanApplyService: records_deleted=deleted, ) return deleted + + +# --------------------------------------------------------------------------- +# Apply phase proxy for checkpoint hooks +# --------------------------------------------------------------------------- + + +class _ApplyPhaseProxy: + """Lightweight proxy providing sandbox-like interface for apply checkpoints. + + Used by ``PlanApplyService._try_checkpoint`` when no real sandbox + object is available. + """ + + def __init__(self, plan_id: str) -> None: + self._sandbox_id = f"apply-{plan_id[:16]}" + + @property + def sandbox_id(self) -> str: + """Return a synthetic sandbox ID.""" + return self._sandbox_id + + @property + def context(self) -> None: + """Return ``None`` (no physical sandbox path).""" + return None diff --git a/src/cleveragents/application/services/plan_executor.py b/src/cleveragents/application/services/plan_executor.py index 78674f398..6da03a3e0 100644 --- a/src/cleveragents/application/services/plan_executor.py +++ b/src/cleveragents/application/services/plan_executor.py @@ -4,6 +4,8 @@ Provides local-only (no LLM) stub actors for M1 that integrate with the ``PlanLifecycleService`` to drive plans through the Strategize and Execute phases. When a ``PlanExecutionContext`` is provided, the execute phase delegates to ``RuntimeExecuteActor`` for full changeset capture. + +Updated in M4 to add optional checkpoint hooks via ``CheckpointManager``. """ from __future__ import annotations @@ -29,6 +31,10 @@ from cleveragents.domain.models.core.plan import ( PlanPhase, ProcessingState, ) +from cleveragents.infrastructure.sandbox.checkpoint import ( + CheckpointManager, + SandboxCheckpoint, +) from cleveragents.tool.builtins.changeset import ChangeSet, ChangeSetCapture from cleveragents.tool.runner import ToolRunner @@ -252,6 +258,7 @@ class PlanExecutor: sandbox_root: str | None = None, execution_context: PlanExecutionContext | None = None, error_recovery_service: ErrorRecoveryService | None = None, + checkpoint_manager: CheckpointManager | None = None, ) -> None: """Initialize the plan executor. @@ -262,6 +269,9 @@ class PlanExecutor: execution_context: Optional execution context for runtime mode. error_recovery_service: Optional error recovery service for recording errors and managing retry logic. + checkpoint_manager: Optional checkpoint manager for + sandbox state snapshots during execute. When ``None``, + checkpoint hooks are silently skipped. """ if lifecycle_service is None: raise ValidationError("lifecycle_service must not be None") @@ -270,6 +280,7 @@ class PlanExecutor: self._sandbox_root = sandbox_root self._execution_context = execution_context self._error_recovery = error_recovery_service + self._checkpoint_manager = checkpoint_manager self._strategize_actor = StrategizeStubActor() self._execute_actor = ExecuteStubActor() self._logger = logger.bind(service="plan_executor") @@ -286,11 +297,121 @@ class PlanExecutor: return self._execution_context.changeset_store return None + @property + def checkpoint_manager(self) -> CheckpointManager | None: + """Return the checkpoint manager, if configured.""" + return self._checkpoint_manager + @property def execution_context(self) -> PlanExecutionContext | None: """Return the execution context, if configured.""" return self._execution_context + def _try_create_checkpoint( + self, + plan_id: str, + phase: str, + metadata: dict[str, Any] | None = None, + ) -> SandboxCheckpoint | None: + """Create a checkpoint if a manager and sandbox root are available. + + Silently returns ``None`` when no checkpoint manager is configured + or when checkpoint creation fails (non-fatal). + + Args: + plan_id: The plan identifier. + phase: Checkpoint phase label (e.g. ``"pre_execute"``). + metadata: Optional metadata to attach. + + Returns: + The created checkpoint, or ``None``. + """ + if self._checkpoint_manager is None: + return None + + sandbox = self._resolve_sandbox_for_checkpoint(plan_id) + if sandbox is None: + return None + + try: + meta = dict(metadata or {}) + if sandbox.context is not None: + meta["sandbox_path"] = sandbox.context.sandbox_path + return self._checkpoint_manager.create_checkpoint( + sandbox=sandbox, + plan_id=plan_id, + phase=phase, + metadata=meta, + ) + except Exception: + self._logger.debug( + "Checkpoint creation failed (non-fatal)", + plan_id=plan_id, + phase=phase, + exc_info=True, + ) + return None + + def _try_rollback_to_last_checkpoint(self, plan_id: str) -> bool: + """Attempt rollback to the most recent checkpoint for a plan. + + Silently returns ``False`` when no checkpoint manager is + configured, no checkpoints exist, or rollback fails. + + Args: + plan_id: The plan identifier. + + Returns: + ``True`` if rollback succeeded. + """ + if self._checkpoint_manager is None: + return False + + sandbox = self._resolve_sandbox_for_checkpoint(plan_id) + if sandbox is None: + return False + + checkpoints = self._checkpoint_manager.list_checkpoints(sandbox.sandbox_id) + if not checkpoints: + return False + + last_cp = checkpoints[-1] + try: + return self._checkpoint_manager.rollback_to(last_cp) + except Exception: + self._logger.debug( + "Checkpoint rollback failed (non-fatal)", + plan_id=plan_id, + checkpoint_id=last_cp.checkpoint_id, + exc_info=True, + ) + return False + + def _resolve_sandbox_for_checkpoint(self, plan_id: str) -> Any: + """Resolve a sandbox object suitable for checkpointing. + + Returns a lightweight mock-like object with ``sandbox_id`` and + ``context`` when only a ``sandbox_root`` is configured (no + execution context). + + Args: + plan_id: The plan identifier. + + Returns: + A sandbox-like object, or ``None``. + """ + if self._execution_context is not None: + sandbox_mgr = getattr(self._execution_context, "sandbox_manager", None) + if sandbox_mgr is not None: + sandboxes = sandbox_mgr.list_sandboxes(plan_id) + if sandboxes: + return sandboxes[0] + + if self._sandbox_root is not None: + return _SandboxRootProxy(self._sandbox_root, plan_id) + + return None + def run_strategize( self, plan_id: str, @@ -419,6 +540,7 @@ class PlanExecutor: execution_context=self._execution_context, ) self._lifecycle.start_execute(plan_id) + self._try_create_checkpoint(plan_id, "pre_execute") try: result = runtime_actor.execute( decisions=decisions, stream_callback=stream_callback @@ -434,6 +556,7 @@ class PlanExecutor: } plan.timestamps.updated_at = datetime.now(tz=UTC) self._lifecycle._commit_plan(plan) + self._try_create_checkpoint(plan_id, "post_execute", {"status": "success"}) self._lifecycle.complete_execute(plan_id) self._logger.info( "Execute completed (runtime)", @@ -443,6 +566,7 @@ class PlanExecutor: ) return result except Exception as exc: + self._try_rollback_to_last_checkpoint(plan_id) error_msg = f"{type(exc).__name__}: {exc}" plan = self._lifecycle.get_plan(plan_id) plan.error_details = { @@ -464,6 +588,7 @@ class PlanExecutor: decisions = self._build_decisions(plan) self._lifecycle.start_execute(plan_id) + self._try_create_checkpoint(plan_id, "pre_execute") # Determine max attempts: 1 (no recovery) or policy max_retries + 1. max_attempts = ( @@ -493,6 +618,9 @@ class PlanExecutor: } plan.timestamps.updated_at = datetime.now(tz=UTC) self._lifecycle._commit_plan(plan) + self._try_create_checkpoint( + plan_id, "post_execute", {"status": "success"} + ) self._lifecycle.complete_execute(plan_id) self._logger.info( "Execute completed (stub)", @@ -528,8 +656,9 @@ class PlanExecutor: # No recovery service or no more retries — fail break - # All attempts exhausted — persist error and fail + # All attempts exhausted — attempt rollback, persist error, fail assert last_exc is not None + self._try_rollback_to_last_checkpoint(plan_id) error_msg = f"{type(last_exc).__name__}: {last_exc}" plan = self._lifecycle.get_plan(plan_id) plan.error_details = { @@ -540,3 +669,38 @@ class PlanExecutor: self._lifecycle._commit_plan(plan) self._lifecycle.fail_execute(plan_id, error_msg) raise last_exc + + +# --------------------------------------------------------------------------- +# Sandbox root proxy (for checkpoint hooks without a full Sandbox object) +# --------------------------------------------------------------------------- + + +class _SandboxRootProxy: + """Lightweight proxy providing sandbox-like interface for checkpoint hooks. + + Used when only a ``sandbox_root`` path string is available (no real + ``Sandbox`` instance). Satisfies the minimal interface required by + :meth:`CheckpointManager.create_checkpoint`. + """ + + def __init__(self, sandbox_root: str, plan_id: str) -> None: + self._sandbox_id = f"root-{plan_id[:16]}" + self._context = _ProxyContext(sandbox_root) + + @property + def sandbox_id(self) -> str: + """Return a synthetic sandbox ID.""" + return self._sandbox_id + + @property + def context(self) -> _ProxyContext: + """Return a proxy context with sandbox_path.""" + return self._context + + +class _ProxyContext: + """Minimal context providing ``sandbox_path`` for checkpoint snapshotting.""" + + def __init__(self, sandbox_path: str) -> None: + self.sandbox_path = sandbox_path diff --git a/src/cleveragents/infrastructure/sandbox/__init__.py b/src/cleveragents/infrastructure/sandbox/__init__.py index 600e4e6d2..3a4c73028 100644 --- a/src/cleveragents/infrastructure/sandbox/__init__.py +++ b/src/cleveragents/infrastructure/sandbox/__init__.py @@ -4,8 +4,14 @@ Provides resource isolation during plan execution through the Sandbox protocol and multiple strategy implementations (git worktree, filesystem copy, no-op). Stage B3 of the implementation plan. Updated in TASK-006 (B4 rework). +Updated in M4 to add checkpoint/rollback hooks. """ +from cleveragents.infrastructure.sandbox.checkpoint import ( + Checkpointable, + CheckpointManager, + SandboxCheckpoint, +) from cleveragents.infrastructure.sandbox.copy_on_write import CopyOnWriteSandbox from cleveragents.infrastructure.sandbox.factory import SandboxFactory from cleveragents.infrastructure.sandbox.git_worktree import GitWorktreeSandbox @@ -27,6 +33,8 @@ from cleveragents.infrastructure.sandbox.protocol import ( ) __all__ = [ + "CheckpointManager", + "Checkpointable", "CommitResult", "CopyOnWriteSandbox", "GitMergeStrategy", @@ -36,6 +44,7 @@ __all__ = [ "MergeStrategy", "NoSandbox", "Sandbox", + "SandboxCheckpoint", "SandboxContext", "SandboxError", "SandboxFactory", diff --git a/src/cleveragents/infrastructure/sandbox/checkpoint.py b/src/cleveragents/infrastructure/sandbox/checkpoint.py new file mode 100644 index 000000000..cd65b346d --- /dev/null +++ b/src/cleveragents/infrastructure/sandbox/checkpoint.py @@ -0,0 +1,328 @@ +"""Sandbox checkpoint and rollback manager for CleverAgents. + +Provides snapshot-based checkpoint/rollback hooks that preserve sandbox state +at key points during plan execute/apply flows. When a failure occurs, the +``CheckpointManager`` can restore a sandbox to a previously captured +checkpoint. + +Stage M4 of the implementation plan. +""" + +from __future__ import annotations + +import os +import shutil +import tempfile +import threading +from datetime import datetime +from typing import Any, Protocol, runtime_checkable + +import structlog +from pydantic import BaseModel, ConfigDict, Field +from ulid import ULID + +from cleveragents.infrastructure.sandbox.protocol import ( + SandboxError, +) + +logger = structlog.get_logger(__name__) + + +# --------------------------------------------------------------------------- +# Checkpointable protocol +# --------------------------------------------------------------------------- + + +@runtime_checkable +class Checkpointable(Protocol): + """Minimal protocol for objects that can be checkpointed. + + Any object with a ``sandbox_id`` property and an optional ``context`` + providing ``sandbox_path`` can be used with :class:`CheckpointManager`. + This is satisfied by the full ``Sandbox`` protocol as well as + lightweight proxy objects. + """ + + @property + def sandbox_id(self) -> str: + """Unique identifier for the sandbox.""" + ... + + @property + def context(self) -> Any: + """Context with ``sandbox_path`` attribute, or ``None``.""" + ... + + +# --------------------------------------------------------------------------- +# Checkpoint data model +# --------------------------------------------------------------------------- + + +class SandboxCheckpoint(BaseModel): + """Immutable record of a sandbox state snapshot. + + Captures the sandbox working directory state at a specific point + in time, annotated with plan lifecycle metadata. + """ + + checkpoint_id: str = Field( + ..., description="Unique ULID identifier for this checkpoint" + ) + sandbox_id: str = Field(..., description="ID of the sandbox that was checkpointed") + plan_id: str = Field(..., description="ID of the plan that owns the sandbox") + phase: str = Field( + ..., description="Lifecycle phase (pre_execute/post_execute/pre_apply)" + ) + created_at: datetime = Field(..., description="When this checkpoint was captured") + metadata: dict[str, str] = Field( + default_factory=dict, + description="Arbitrary key-value metadata (status, reason, etc.)", + ) + snapshot_path: str = Field( + ..., description="Filesystem path to the snapshot directory" + ) + + model_config = ConfigDict(frozen=True) + + +# --------------------------------------------------------------------------- +# Checkpoint manager +# --------------------------------------------------------------------------- + + +class CheckpointManager: + """Manages sandbox checkpoint lifecycle: create, rollback, list, delete. + + Thread-safe: all mutable state is protected by ``_lock``. + + Usage:: + + mgr = CheckpointManager() + cp = mgr.create_checkpoint(sandbox, plan_id, "pre_execute", {}) + # ... execute phase runs ... + if failure: + mgr.rollback_to(cp) + + mgr.list_checkpoints(sandbox.sandbox_id) + mgr.delete_checkpoint(cp.checkpoint_id) + """ + + def __init__(self) -> None: + """Initialise the checkpoint manager.""" + self._checkpoints: dict[str, list[SandboxCheckpoint]] = {} + self._lock: threading.RLock = threading.RLock() + self._logger = logger.bind(component="checkpoint_manager") + + def create_checkpoint( + self, + sandbox: Checkpointable, + plan_id: str, + phase: str, + metadata: dict[str, Any] | None = None, + ) -> SandboxCheckpoint: + """Snapshot the current sandbox state as a checkpoint. + + For sandboxes with a filesystem path (``context.sandbox_path``), + the directory tree is copied to a temporary snapshot location. + For sandboxes without a physical path (e.g. ``NoSandbox``), a + lightweight metadata-only checkpoint is created. + + Args: + sandbox: The sandbox to checkpoint. + plan_id: The plan that owns the sandbox. + phase: Lifecycle phase label (e.g. ``"pre_execute"``). + metadata: Optional key-value metadata to attach. + + Returns: + A frozen :class:`SandboxCheckpoint` record. + + Raises: + SandboxError: If the snapshot copy fails. + """ + meta = {k: str(v) for k, v in (metadata or {}).items()} + checkpoint_id = str(ULID()) + sandbox_id = sandbox.sandbox_id + sandbox_path: str | None = None + + if sandbox.context is not None: + sandbox_path = sandbox.context.sandbox_path + + snapshot_path = self._snapshot_directory(sandbox_path, checkpoint_id) + + checkpoint = SandboxCheckpoint( + checkpoint_id=checkpoint_id, + sandbox_id=sandbox_id, + plan_id=plan_id, + phase=phase, + created_at=datetime.now(), + metadata=meta, + snapshot_path=snapshot_path, + ) + + with self._lock: + self._checkpoints.setdefault(sandbox_id, []).append(checkpoint) + + self._logger.info( + "Checkpoint created", + checkpoint_id=checkpoint_id, + sandbox_id=sandbox_id, + plan_id=plan_id, + phase=phase, + ) + + return checkpoint + + def rollback_to(self, checkpoint: SandboxCheckpoint) -> bool: + """Restore a sandbox to the state captured in *checkpoint*. + + Copies the snapshot directory back over the sandbox working + directory. Only works when the sandbox has a physical path + and the snapshot directory exists. + + Args: + checkpoint: The checkpoint to restore. + + Returns: + ``True`` if the rollback succeeded, ``False`` otherwise. + """ + snapshot = checkpoint.snapshot_path + if not snapshot or not os.path.isdir(snapshot): + self._logger.warning( + "Rollback skipped: snapshot path missing or not a directory", + checkpoint_id=checkpoint.checkpoint_id, + snapshot_path=snapshot, + ) + return False + + # Find sandbox path from the checkpoint's snapshot parent structure. + # The snapshot is at /ca-checkpoint-/snapshot. + # We need the *actual* sandbox path, which we embed in metadata. + sandbox_path = checkpoint.metadata.get("sandbox_path") + if not sandbox_path or not os.path.isdir(sandbox_path): + self._logger.warning( + "Rollback skipped: sandbox path unknown or missing", + checkpoint_id=checkpoint.checkpoint_id, + ) + return False + + try: + # Clear the current sandbox contents and restore from snapshot + for entry in os.listdir(sandbox_path): + entry_path = os.path.join(sandbox_path, entry) + if os.path.isdir(entry_path): + shutil.rmtree(entry_path) + else: + os.remove(entry_path) + + # Copy snapshot contents into sandbox + for entry in os.listdir(snapshot): + src = os.path.join(snapshot, entry) + dst = os.path.join(sandbox_path, entry) + if os.path.isdir(src): + shutil.copytree(src, dst, symlinks=True) + else: + shutil.copy2(src, dst) + + except OSError as exc: + self._logger.error( + "Rollback failed", + checkpoint_id=checkpoint.checkpoint_id, + error=str(exc), + ) + return False + + self._logger.info( + "Rollback completed", + checkpoint_id=checkpoint.checkpoint_id, + sandbox_id=checkpoint.sandbox_id, + ) + return True + + def list_checkpoints(self, sandbox_id: str) -> list[SandboxCheckpoint]: + """Return all checkpoints for a sandbox in creation order. + + Args: + sandbox_id: The sandbox identifier. + + Returns: + List of checkpoints ordered by creation time. + """ + with self._lock: + return list(self._checkpoints.get(sandbox_id, [])) + + def delete_checkpoint(self, checkpoint_id: str) -> bool: + """Remove a checkpoint and its snapshot data. + + Args: + checkpoint_id: The checkpoint ULID to delete. + + Returns: + ``True`` if the checkpoint was found and deleted, + ``False`` if it was not found. + """ + with self._lock: + for sandbox_id, checkpoints in self._checkpoints.items(): + for i, cp in enumerate(checkpoints): + if cp.checkpoint_id == checkpoint_id: + removed = checkpoints.pop(i) + self._cleanup_snapshot(removed.snapshot_path) + self._logger.info( + "Checkpoint deleted", + checkpoint_id=checkpoint_id, + sandbox_id=sandbox_id, + ) + return True + + self._logger.debug( + "Checkpoint not found for deletion", + checkpoint_id=checkpoint_id, + ) + return False + + # -- internal helpers ---------------------------------------------------- + + def _snapshot_directory( + self, + sandbox_path: str | None, + checkpoint_id: str, + ) -> str: + """Copy the sandbox working directory to a snapshot location. + + Args: + sandbox_path: Path to the sandbox directory, or ``None`` + for metadata-only checkpoints. + checkpoint_id: Used to name the snapshot directory. + + Returns: + Absolute path to the snapshot directory. + + Raises: + SandboxError: If the copy fails. + """ + parent = tempfile.mkdtemp(prefix=f"ca-checkpoint-{checkpoint_id[:8]}-") + snapshot = os.path.join(parent, "snapshot") + + if sandbox_path is not None and os.path.isdir(sandbox_path): + try: + shutil.copytree( + sandbox_path, snapshot, symlinks=True, dirs_exist_ok=False + ) + except OSError as exc: + shutil.rmtree(parent, ignore_errors=True) + raise SandboxError( + f"Failed to snapshot sandbox directory: {exc}" + ) from exc + else: + os.makedirs(snapshot, exist_ok=True) + + return snapshot + + @staticmethod + def _cleanup_snapshot(snapshot_path: str) -> None: + """Remove a snapshot directory and its parent temp dir.""" + if not snapshot_path: + return + parent = os.path.dirname(snapshot_path) + if os.path.isdir(parent): + shutil.rmtree(parent, ignore_errors=True) diff --git a/vulture_whitelist.py b/vulture_whitelist.py index 4a41b1739..88c6c4377 100644 --- a/vulture_whitelist.py +++ b/vulture_whitelist.py @@ -176,6 +176,14 @@ run_strategize # noqa: B018, F821 run_execute # noqa: B018, F821 _parse_steps # noqa: B018, F821 invariant_records # noqa: B018, F821 +_SandboxRootProxy # noqa: B018, F821 +_ProxyContext # noqa: B018, F821 + +# Sandbox checkpoint — public API (M4) +CheckpointManager # noqa: B018, F821 +SandboxCheckpoint # noqa: B018, F821 +Checkpointable # noqa: B018, F821 +_ApplyPhaseProxy # noqa: B018, F821 # Output rendering framework — public API (G5b.render) OutputSession # noqa: B018, F821