feat: snapshot system for docker provider

.gitignore

@@ -251,6 +251,7 @@ trajectories/
 
 # Lumier Storage
 storage/
+!libs/python/agent/agent/callbacks/snapshots/storage/
 
 # Trashes
 .Trashes

.vscode/launch.json

@@ -7,10 +7,10 @@
             "program": "examples/agent_ui_examples.py",
             "console": "integratedTerminal",
             "justMyCode": false,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer:${workspaceFolder:cua-root}/libs/python/agent:${workspaceFolder:cua-root}/libs/python/som:${workspaceFolder:cua-root}/libs/python/pylume"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
             }
         },
         {
@@ -20,10 +20,10 @@
             "program": "examples/computer_ui_examples.py",
             "console": "integratedTerminal",
             "justMyCode": false,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer:${workspaceFolder:cua-root}/libs/python/agent:${workspaceFolder:cua-root}/libs/python/som:${workspaceFolder:cua-root}/libs/python/pylume"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
             }
         },
         {
@@ -33,10 +33,10 @@
             "program": "examples/computer_examples.py",
             "console": "integratedTerminal",
             "justMyCode": true,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer:${workspaceFolder:cua-root}/libs/python/agent:${workspaceFolder:cua-root}/libs/python/som:${workspaceFolder:cua-root}/libs/python/pylume"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
             }
         },
         {
@@ -46,10 +46,23 @@
             "program": "examples/agent_examples.py",
             "console": "integratedTerminal",
             "justMyCode": false,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer:${workspaceFolder:cua-root}/libs/python/agent:${workspaceFolder:cua-root}/libs/python/som:${workspaceFolder:cua-root}/libs/python/pylume"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
+            }
+        },
+        {
+            "name": "Run Snapshot Restoration Examples",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "examples/restore_snapshot_example.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
+            "env": {
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
             }
         },
         {
@@ -59,10 +72,10 @@
             "program": "examples/pylume_examples.py",
             "console": "integratedTerminal",
             "justMyCode": true,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer:${workspaceFolder:cua-root}/libs/python/agent:${workspaceFolder:cua-root}/libs/python/som:${workspaceFolder:cua-root}/libs/python/pylume"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
             }
         },
         {
@@ -81,10 +94,10 @@
             ],
             "console": "integratedTerminal",
             "justMyCode": false,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer:${workspaceFolder:cua-root}/libs/python/agent:${workspaceFolder:cua-root}/libs/python/som:${workspaceFolder:cua-root}/libs/python/pylume"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
             }
         },
         {
@@ -103,10 +116,10 @@
             ],
             "console": "integratedTerminal",
             "justMyCode": false,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer:${workspaceFolder:cua-root}/libs/python/agent:${workspaceFolder:cua-root}/libs/python/som:${workspaceFolder:cua-root}/libs/python/pylume"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
             }
         },
         {
@@ -116,10 +129,10 @@
             "program": "${workspaceFolder}/libs/python/computer-server/run_server.py",
             "console": "integratedTerminal",
             "justMyCode": true,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer:${workspaceFolder:cua-root}/libs/python/agent:${workspaceFolder:cua-root}/libs/python/som:${workspaceFolder:cua-root}/libs/python/pylume"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer:${workspaceFolder:cua-snapshot}/libs/python/agent:${workspaceFolder:cua-snapshot}/libs/python/som:${workspaceFolder:cua-snapshot}/libs/python/pylume"
             }
         },
         {
@@ -137,28 +150,28 @@
             ],
             "console": "integratedTerminal",
             "justMyCode": false,
-            "python": "${workspaceFolder:cua-root}/.venv/bin/python",
-            "cwd": "${workspaceFolder:cua-root}",
+            "python": "${workspaceFolder:cua-snapshot}/.venv/bin/python",
+            "cwd": "${workspaceFolder:cua-snapshot}",
             "env": {
-                "PYTHONPATH": "${workspaceFolder:cua-root}/libs/python/core:${workspaceFolder:cua-root}/libs/python/computer-server"
+                "PYTHONPATH": "${workspaceFolder:cua-snapshot}/libs/python/core:${workspaceFolder:cua-snapshot}/libs/python/computer-server"
             }
         },
         {
             "type": "lldb",
             "request": "launch",
             "args": [],
-            "cwd": "${workspaceFolder:cua-root}/libs/lume",
+            "cwd": "${workspaceFolder:cua-snapshot}/libs/lume",
             "name": "Debug lume (libs/lume)",
-            "program": "${workspaceFolder:cua-root}/libs/lume/.build/debug/lume",
+            "program": "${workspaceFolder:cua-snapshot}/libs/lume/.build/debug/lume",
             "preLaunchTask": "swift: Build Debug lume (libs/lume)"
         },
         {
             "type": "lldb",
             "request": "launch",
             "args": [],
-            "cwd": "${workspaceFolder:cua-root}/libs/lume",
+            "cwd": "${workspaceFolder:cua-snapshot}/libs/lume",
             "name": "Release lume (libs/lume)",
-            "program": "${workspaceFolder:cua-root}/libs/lume/.build/release/lume",
+            "program": "${workspaceFolder:cua-snapshot}/libs/lume/.build/release/lume",
             "preLaunchTask": "swift: Build Release lume (libs/lume)"
         }
     ]

docs/SNAPSHOT_SYSTEM.md

@@ -0,0 +1,166 @@
+# CUA Snapshot System
+
+## Overview
+
+The CUA Snapshot System enables creating, managing, and restoring filesystem snapshots of Docker containers during agent execution by extending the Docker provider.
+
+## Architecture Diagram
+
+```mermaid
+flowchart TD
+    A[Agent / ComputerAgent] --> B[SnapshotManagerCallback]
+
+    B --> C{Trigger}
+    C -->|manual| D[create_manual_snapshot]
+    C -->|run_start / run_end / every_action| E[auto snapshot]
+
+    D --> F[SnapshotCreator]
+    E --> F
+    F --> G[ProviderAdapter]
+    G --> H{{Docker}}
+    H --> I[docker commit to image]
+    I --> J[Write metadata: labels and files]
+    J --> K[RetentionEnforcer]
+    K -->|count/age| L[delete old snapshots]
+
+    B --> M[restore snapshot by id]
+    M --> G
+    G --> N[restore via provider]
+    N --> O[Container filesystem replaced]
+    O --> P[Container restarted]
+    P --> Q[Agent continues from restored state]
+
+    classDef comp fill:#eaf7ff,stroke:#2b90d9,color:#0b3d62;
+    class A,B,D,E,F,G,M comp;
+```
+
+## Core Functionality
+
+### Snapshot Operations
+
+- **Create**: Capture current container state as a Docker image
+- **Restore**: Roll back container to a previous snapshot state  
+- **List**: View all available snapshots with metadata
+- **Delete**: Remove snapshots to free storage space
+
+### Scheduling Options
+
+- `manual`: Create snapshots only when explicitly requested
+- `every_action`: Snapshot after each computer action (debugging)
+- `run_start`: Snapshot at the beginning of each agent run
+- `run_end`: Snapshot at the end of each agent run  
+- `run_boundaries`: Snapshot at both start and end (recommended)
+
+### Retention Management
+
+- **Count-based**: Keep only the N most recent snapshots
+- **Age-based**: Delete snapshots older than specified days
+- **Automatic cleanup**: Configurable background cleanup
+
+## Basic Usage
+
+```python
+from computer import Computer
+from agent import ComputerAgent
+from libs.python.agent.agent.callbacks.snapshot_manager import SnapshotManagerCallback  # TODO This will eventually be a part of the pip package
+
+# Setup computer with Docker provider
+computer = Computer(
+    os_type="linux",
+    provider_type="docker", 
+    image="trycua/cua-ubuntu:latest",
+    name='my-container'
+)
+
+# Configure snapshot callback
+snapshot_callback = SnapshotManagerCallback(
+    computer=computer,
+    snapshot_interval="run_boundaries",  # When to create snapshots
+    max_snapshots=10,                    # Keep latest 10 snapshots
+    retention_days=7,                    # Delete snapshots older than 7 days
+    auto_cleanup=True                    # Enable automatic cleanup
+)
+
+# Create agent with snapshot support
+agent = ComputerAgent(
+    model="claude-3-5-sonnet-20241022",
+    callbacks=[snapshot_callback]
+)
+```
+
+## Manual Operations
+
+```python
+# Create manual snapshot
+result = await snapshot_callback.create_manual_snapshot("before-risky-operation")
+
+# List all snapshots  
+snapshots = await snapshot_callback.list_snapshots()
+
+# Restore to specific snapshot
+restore_result = await snapshot_callback.restore_snapshot(snapshot_id)
+
+# Delete snapshot
+delete_result = await snapshot_callback.delete_snapshot(snapshot_id)
+```
+
+## Configuration
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `snapshot_interval` | "manual" | When to create snapshots automatically |
+| `max_snapshots` | 10 | Maximum snapshots to retain |
+| `retention_days` | 7 | Delete snapshots older than N days |
+| `auto_cleanup` | True | Enable automatic cleanup |
+| `metadata_dir` | "/tmp/cua_snapshots" | Metadata storage location |
+
+## Implementation Details
+
+### Docker Integration
+
+Snapshots are implemented using Docker's native `docker commit` functionality:
+
+- Container state is captured as a Docker image
+- Metadata stored as image labels (prefixed with `cua.snapshot.`)
+- Container configuration (memory, CPU, ports) preserved during restore
+- Unique snapshot IDs generated with timestamps
+
+### Provider Interface
+
+VM providers must implement these methods:
+
+```python
+async def create_snapshot(name: str, snapshot_name: str = None, metadata: dict = None) -> dict
+async def restore_snapshot(name: str, snapshot_id: str) -> dict  
+async def list_snapshots(name: str) -> List[dict]
+async def delete_snapshot(snapshot_id: str) -> dict
+```
+
+### Error Handling
+
+All operations return status information:
+
+```python
+# Success response
+{
+    "id": "uuid-v4-string",
+    "status": "created", 
+    "timestamp": "2024-01-01T12:00:00",
+    "size": "150MB"
+}
+
+# Error response  
+{
+    "status": "error",
+    "error": "Detailed error message"
+}
+```
+
+## Best Practices
+
+1. Use `run_boundaries` interval for most cases
+2. Set reasonable retention limits to manage storage
+3. Include descriptive metadata for easier identification
+4. Test restore procedures regularly
+5. Monitor disk space usage
+6. Handle operation errors gracefully