orchestrata_ultimate.session.checkpoint

Reference › orchestrata_ultimate › session › checkpoint

CheckpointManager creates and keeps session checkpoints. A checkpoint (the Checkpoint type in types) captures the current session state, plus optional extra state, and is stored by the RedisSessionStore as a sorted-set member scored by time. The manager offers create, restore (the state from a given checkpoint, or nothing), get_latest, list, count, and prune, which keeps the most recent checkpoints (default CHECKPOINT_PRUNE_KEEP) and returns the number removed. RecoveryManager builds its recovery and rollback on it.

At a glance

Class diagram

Dependency diagram

API

Classes

CheckpointManager

Decorators: traces(spec='006', story='US-005', reqs=['REQ-006-US-005-005']) · Line 21

Manages session checkpoints.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/session/checkpoint.py

"""Session checkpointing.

Creates and manages session state checkpoints for recovery.
"""

import logging
import uuid

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.session.constants import CHECKPOINT_PRUNE_KEEP
from orchestrata_ultimate.session.redis_store import RedisSessionStore
from orchestrata_ultimate.session.types import Checkpoint

# Spec: 006 - Session Management


logger = logging.getLogger(__name__)


@traces(spec="006", story="US-005", reqs=["REQ-006-US-005-005"])
class CheckpointManager:
    """Manages session checkpoints."""

    def __init__(self, redis_store: RedisSessionStore) -> None:
        self._redis = redis_store

    def create(self, session_id: str, extra_state: dict | None = None) -> Checkpoint:
        """Create a checkpoint capturing current session state."""
        session_data = self._redis.get_session(session_id)
        context_data = self._redis.get_all_context(session_id)
        fact_count = self._redis.get_fact_count(session_id)

        state_snapshot = {
            "session": session_data or {},
            "context": context_data,
        }
        if extra_state:
            state_snapshot["extra"] = extra_state

        checkpoint = Checkpoint(
            id=f"cp-{uuid.uuid4().hex[:12]}",
            session_id=session_id,
            state_snapshot=state_snapshot,
            fact_count=fact_count,
        )
        self._redis.save_checkpoint(session_id, checkpoint)
        return checkpoint

    def restore(self, checkpoint_id: str, session_id: str) -> dict | None:
        """Restore session state from a specific checkpoint."""
        checkpoints = self._redis.get_checkpoints(session_id)
        target = None
        for cp in checkpoints:
            if cp.id == checkpoint_id:
                target = cp
                break
        if target is None:
            return None

        # Restore session data
        session_data = target.state_snapshot.get("session", {})
        if session_data:
            self._redis.save_session(session_data)

        # Restore context
        context_data = target.state_snapshot.get("context", {})
        if context_data:
            self._redis.clear_context(session_id)
            for key, value in context_data.items():
                self._redis.save_context(session_id, key, value)

        return target.state_snapshot

    def get_latest(self, session_id: str) -> Checkpoint | None:
        """Get the most recent checkpoint."""
        return self._redis.get_latest_checkpoint(session_id)

    def list(self, session_id: str) -> list[Checkpoint]:
        """List all checkpoints for a session."""
        return self._redis.get_checkpoints(session_id)

    def prune(self, session_id: str, keep_last: int = CHECKPOINT_PRUNE_KEEP) -> int:
        """Prune old checkpoints, keeping the most recent."""
        checkpoints = self._redis.get_checkpoints(session_id)
        if len(checkpoints) <= keep_last:
            return 0

        # Delete all and re-save the latest
        self._redis.delete_checkpoints(session_id)
        to_keep = checkpoints[-keep_last:]
        for cp in to_keep:
            self._redis.save_checkpoint(session_id, cp)

        return len(checkpoints) - keep_last

    def count(self, session_id: str) -> int:
        """Count checkpoints for a session."""
        return len(self._redis.get_checkpoints(session_id))