State Management API

from marsys.coordination.state import StateManager, FileStorageBackend

StateManager(
    storage_backend: StorageBackend,
    enable_compression: bool = True,
    enable_checksum: bool = True,
    max_checkpoints_per_session: int = 10
)

async def save_state(
    session_id: str,
    snapshot: StateSnapshot
) -> None

async def load_state(
    session_id: str
) -> Optional[StateSnapshot]

async def pause_execution(
    session_id: str,
    state: Dict[str, Any]
) -> None

async def resume_execution(
    session_id: str
) -> Optional[Dict[str, Any]]

async def create_checkpoint(
    session_id: str,
    checkpoint_name: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None
) -> str

async def restore_checkpoint(
    checkpoint_id: str
) -> Optional[StateSnapshot]

from pathlib import Path

# Initialize with file storage
storage = FileStorageBackend(Path("./state"))
state_manager = StateManager(storage)

# Save state
snapshot = StateSnapshot(
    session_id="session_123",
    timestamp=time.time(),
    branches=serialized_branches,
    active_branches={"branch_1", "branch_2"},
    completed_branches={"branch_0"}
)
await state_manager.save_state("session_123", snapshot)

# Create checkpoint
checkpoint_id = await state_manager.create_checkpoint(
    "session_123",
    checkpoint_name="before_critical_section"
)

# Resume later
state = await state_manager.resume_execution("session_123")

# Calculate integrity checksum
snapshot.checksum = snapshot.calculate_checksum()

# Validate later
if not snapshot.validate_checksum():
    raise StateError("State corruption detected")

FileStorageBackend(
    base_path: Path,
    compression: Optional[str] = None,
    encryption_key: Optional[str] = None
)

base_path/
├── sessions/       # Active session states
├── checkpoints/    # Named checkpoints
└── metadata/       # Session metadata

storage = FileStorageBackend(
    base_path=Path("./state"),
    compression="gzip"
)

# Save data
await storage.save("session_123", state_data)

# Load data
data = await storage.load("session_123")

# List sessions
sessions = await storage.list_keys(prefix="session_")

# Check existence
if await storage.exists("session_123"):
    print("Session exists")

class RedisStorageBackend(StorageBackend):
    """Redis-based storage backend."""

    def __init__(self, redis_client):
        self.client = redis_client

    async def save(self, key: str, data: Dict[str, Any]) -> None:
        serialized = json.dumps(data)
        await self.client.set(key, serialized)

    async def load(self, key: str) -> Optional[Dict[str, Any]]:
        data = await self.client.get(key)
        return json.loads(data) if data else None

    async def delete(self, key: str) -> None:
        await self.client.delete(key)

    async def list_keys(self, prefix: str = "") -> List[str]:
        pattern = f"{prefix}*" if prefix else "*"
        return await self.client.keys(pattern)

    async def exists(self, key: str) -> bool:
        return await self.client.exists(key)

# Manual checkpoint
checkpoint_id = await state_manager.create_checkpoint(
    session_id="session_123",
    checkpoint_name="milestone_1",
    metadata={
        "progress": 0.5,
        "stage": "data_processing",
        "timestamp": datetime.now().isoformat()
    }
)

print(f"Checkpoint created: {checkpoint_id}")

class AutoCheckpointManager:
    """Automatic checkpoint creation."""

    def __init__(self, state_manager, interval_seconds=300):
        self.state_manager = state_manager
        self.interval = interval_seconds
        self.last_checkpoint = time.time()

    async def maybe_checkpoint(self, session_id: str, state: StateSnapshot):
        """Create checkpoint if interval elapsed."""
        if time.time() - self.last_checkpoint > self.interval:
            await self.state_manager.create_checkpoint(
                session_id,
                checkpoint_name=f"auto_{int(time.time())}"
            )
            self.last_checkpoint = time.time()

# List available checkpoints
checkpoints = await state_manager.list_checkpoints("session_123")
for cp in checkpoints:
    print(f"{cp['id']}: {cp['name']} - {cp['created_at']}")

# Restore specific checkpoint
state = await state_manager.restore_checkpoint("checkpoint_abc123")
if state:
    print(f"Restored from checkpoint: {state.session_id}")

# 1. Start session with state management
state_manager = StateManager(storage)

result = await Orchestra.run(
    task="Long running task",
    topology=topology,
    state_manager=state_manager,
    context={"session_id": "session_123"}
)

# 2. Session can be paused (manually or on error)
await state_manager.pause_execution("session_123", current_state)

# 3. Resume later
state = await state_manager.resume_execution("session_123")
if state:
    result = await Orchestra.resume(
        state=state,
        topology=topology
    )

# 4. Clean up completed session
await state_manager.delete_session("session_123")