orchestrata_ultimate.session.redis_store

Reference › orchestrata_ultimate › session › redis_store

RedisSessionStore is the hot-path layer of the dual-storage design (spec 006): low-latency state for sessions that are in progress. It wraps a RedisStateManager passed to the constructor and uses its synchronous client. The cold layer is GraphSessionStore, and SessionSync moves data from one to the other.

Per session it keeps separate keys, built by _session_key, _facts_key, _context_key, _checkpoints_key and _carryover_key. Session metadata is a hash (save_session, get_session, delete_session). Active sessions are tracked in a set (set_active, clear_active, get_active_sessions, is_active). Facts are appended to a list, singly or in a pipeline (save_fact, save_facts_batch), and read back through get_facts with an optional FactFilter and get_fact_count. Context is a key-value map (save_context, get_context, get_all_context, delete_context, clear_context). Checkpoints are sorted-set members scored by timestamp (save_checkpoint, get_latest_checkpoint, get_checkpoints, delete_checkpoints), and carryover items are a list (save_carryover, get_carryover, clear_carryover). cleanup_session removes everything for a session. The types come from types.

At a glance

Class diagram

Dependency diagram

API

Classes

RedisSessionStore

Decorators: traces(spec='006', story='US-001', reqs=['REQ-006-US-005-005']), traces(spec='006', story='US-003', reqs=['REQ-006-US-003-001', 'REQ-006-US-003-002', 'REQ-006-US-003-003', 'REQ-006-US-… · Line 64

Redis-backed session storage for the hot-path layer.

Source

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

# mypy: ignore-errors
"""Redis session store for hot storage — spec 006 compliant.

Provides low-latency session operations backed by Redis.
All operations target < 50 ms p95 latency.

Key schema::

    orchestrata_ultimate.session:{id}              -- HASH   (session fields)
    orchestrata_ultimate.session:active            -- SET    (active session IDs)
    orchestrata_ultimate.session:{id}:facts        -- LIST   (fact JSONs)
    orchestrata_ultimate.session:{id}:context      -- HASH   (key-value context)
    orchestrata_ultimate.session:{id}:checkpoints  -- SORTED SET (score=timestamp)
    orchestrata_ultimate.session:{id}:carryover    -- LIST   (work item JSONs)
"""

import json
import logging
from datetime import datetime
from typing import TYPE_CHECKING, Any, cast

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.infra.redis_state import RedisStateManager
from orchestrata_ultimate.session.constants import (
    CHECKPOINT_TTL,
    MAX_CHECKPOINTS_PER_SESSION,
    MAX_FACTS_PER_SESSION,
    SESSION_ACTIVE_SET,
    SESSION_DEFAULT_TTL,
    SESSION_REDIS_PREFIX,
)
from orchestrata_ultimate.session.types import (
    Checkpoint,
    Fact,
    FactFilter,
)

if TYPE_CHECKING:
    # `RedisStateManager.client` is the SYNC client; import the sync type
    # explicitly (this module formerly aliased `redis.asyncio as redis`, under
    # which `redis.Redis` wrongly resolved to the async client type).
    from redis import Redis as SyncRedis

logger = logging.getLogger(__name__)


@traces(
    spec="006",
    story="US-001",
    reqs=["REQ-006-US-005-005"],
)
@traces(
    spec="006",
    story="US-003",
    reqs=[
        "REQ-006-US-003-001",
        "REQ-006-US-003-002",
        "REQ-006-US-003-003",
        "REQ-006-US-003-004",
        "REQ-006-US-003-005",
        "REQ-006-US-003-006",
    ],
)
class RedisSessionStore:
    """Redis-backed session storage for the hot-path layer.

    Provides CRUD operations for sessions, facts, context key-values,
    checkpoints, and carryover items. All data is stored under the
    ``orchestrata_ultimate.session:`` key prefix with automatic TTL expiration.

    Args:
        redis: A ``RedisStateManager`` providing the underlying client.

    Example:
        >>> store = RedisSessionStore(redis_manager)
        >>> store.save_session({"id": "sess_001", "status": "active"})
        >>> store.get_session("sess_001")
        {'id': 'sess_001', 'status': 'active'}
    """

    def __init__(self, redis: RedisStateManager) -> None:
        """Initialize the store with a Redis state manager.

        Args:
            redis: ``RedisStateManager`` wrapping the Redis connection.
        """
        self._redis = redis

    @property
    def client(self) -> SyncRedis:
        """Get the SYNC Redis client (delegates to RedisStateManager.client)."""
        return self._redis.client

    @traces(spec="006", story="US-007", reqs=["REQ-006-FR4-004", "REQ-006-US-007-002"])
    def _session_key(self, session_id: str) -> str:
        return f"{SESSION_REDIS_PREFIX}:{session_id}"

    @traces(
        spec="006",
        story="US-005",
        reqs=[
            "REQ-006-FR4-001",
            "REQ-006-FR4-005",
            "REQ-006-US-005-004",
            "REQ-006-US-007-003",
        ],
    )
    def _facts_key(self, session_id: str) -> str:
        return f"{SESSION_REDIS_PREFIX}:{session_id}:facts"

    @traces(
        spec="006", story="US-005", reqs=["REQ-006-US-005-004", "REQ-006-US-007-003"]
    )
    def _context_key(self, session_id: str) -> str:
        return f"{SESSION_REDIS_PREFIX}:{session_id}:context"

    @traces(spec="006", story="US-002", reqs=["REQ-006-US-002-002"])
    def _checkpoints_key(self, session_id: str) -> str:
        return f"{SESSION_REDIS_PREFIX}:{session_id}:checkpoints"

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-013"])
    def _carryover_key(self, session_id: str) -> str:
        return f"{SESSION_REDIS_PREFIX}:{session_id}:carryover"

    # -- Session CRUD ----------------------------------------------

    @traces(
        spec="006",
        story="US-001",
        reqs=[
            "REQ-006-FR2-004",
            "REQ-006-FR7-001",
            "REQ-006-US-001-005",
        ],
    )
    def save_session(self, session: dict) -> None:
        """Persist session metadata to a Redis HASH.

        The dict must contain an ``"id"`` key. All values are JSON-serialized
        or stringified before storage. The key expires after
        ``SESSION_DEFAULT_TTL`` seconds (24 h).

        Args:
            session: Session dict. Must include ``"id"``.

        Example:
            >>> store.save_session({"id": "sess_001", "status": "active"})
        """
        session_id = session["id"]
        key = self._session_key(session_id)
        serialized = {}
        for k, v in session.items():
            if isinstance(v, datetime):
                serialized[k] = v.isoformat()
            elif isinstance(v, str):
                serialized[k] = v
            else:
                # JSON-serialize dict/list/bool/None/int/float so they
                # round-trip as their real types via json.loads on read.
                serialized[k] = json.dumps(v)
        self.client.hset(key, mapping=serialized)
        self.client.expire(key, SESSION_DEFAULT_TTL)

    @traces(
        spec="006",
        story="US-003",
        reqs=[
            "REQ-006-FR10-004",
            "REQ-006-FR2-005",
            "REQ-006-FR6-001",
            "REQ-006-FR7-003",
            "REQ-006-US-003-001",
            "REQ-006-US-003-003",
        ],
    )
    def get_session(self, session_id: str) -> dict | None:
        """Retrieve session metadata from Redis.

        JSON-encoded values are automatically deserialized.

        Args:
            session_id: The session identifier.

        Returns:
            Session dict, or ``None`` if the session does not exist.
        """
        key = self._session_key(session_id)
        # hgetall() is typed as redis-py's shared ResponseT (Awaitable[dict] |
        # dict) because the same command mixin backs both sync and async
        # clients; this sync client always returns a plain dict.
        data = cast("dict[bytes | str, bytes | str]", self.client.hgetall(key))
        if not data:
            return None
        result = {}
        for k, v in data.items():
            k_str = k if isinstance(k, str) else k.decode()
            v_str = v if isinstance(v, str) else v.decode()
            try:
                result[k_str] = json.loads(v_str)
            except json.JSONDecodeError, TypeError:
                result[k_str] = v_str
        return result

    @traces(spec="006", story="US-002", reqs=["REQ-006-US-002-002"])
    def delete_session(self, session_id: str) -> bool:
        """Delete a session and all associated Redis keys.

        Removes the session hash, facts list, context hash, checkpoints
        sorted set, carryover list, and active-set membership.

        Args:
            session_id: The session identifier to delete.

        Returns:
            ``True`` if at least one key was deleted, ``False`` otherwise.
        """
        keys = [
            self._session_key(session_id),
            self._facts_key(session_id),
            self._context_key(session_id),
            self._checkpoints_key(session_id),
            self._carryover_key(session_id),
        ]
        # delete() is typed as shared ResponseT (Awaitable[int] | int); this
        # sync client always returns a plain int (count of keys deleted).
        deleted = cast(int, self.client.delete(*keys))
        self.client.srem(SESSION_ACTIVE_SET, session_id)
        return deleted > 0

    # -- Active session tracking -----------------------------------

    @traces(spec="006", story="US-007", reqs=["REQ-006-US-007-001"])
    def set_active(self, session_id: str) -> None:
        """Mark session as active."""
        self.client.sadd(SESSION_ACTIVE_SET, session_id)

    @traces(spec="006", story="US-002", reqs=["REQ-006-US-002-002"])
    def clear_active(self, session_id: str) -> None:
        """Remove session from active set."""
        self.client.srem(SESSION_ACTIVE_SET, session_id)

    @traces(
        spec="006",
        story="US-007",
        reqs=["REQ-006-FR7-005", "REQ-006-US-007-005", "REQ-006-US-007-006"],
    )
    def get_active_sessions(self) -> list[str]:
        """Get all active session IDs."""
        # smembers() is typed as shared ResponseT (Awaitable[set] | set); this
        # sync client always returns a plain set of str/bytes ids.
        members = cast("set[bytes | str]", self.client.smembers(SESSION_ACTIVE_SET))
        return [m if isinstance(m, str) else m.decode() for m in members]

    @traces(spec="006", story="US-007", reqs=["REQ-006-FR6-002", "REQ-006-US-007-006"])
    def is_active(self, session_id: str) -> bool:
        """Check if session is active."""
        return bool(self.client.sismember(SESSION_ACTIVE_SET, session_id))

    # -- Fact storage ----------------------------------------------

    @traces(
        spec="006",
        story="US-004",
        reqs=["REQ-006-FR3-001", "REQ-006-US-004-002", "REQ-006-US-004-004"],
    )
    def save_fact(self, session_id: str, fact: Fact) -> None:
        """Append a fact to the session's fact list."""
        key = self._facts_key(session_id)
        self.client.rpush(key, json.dumps(fact.to_dict()))
        self.client.ltrim(key, -MAX_FACTS_PER_SESSION, -1)
        self.client.expire(key, SESSION_DEFAULT_TTL)

    @traces(spec="006", story="US-004", reqs=["REQ-006-US-004-005"])
    def save_facts_batch(self, session_id: str, facts: list[Fact]) -> int:
        """Save multiple facts in a pipeline."""
        key = self._facts_key(session_id)
        pipe = self.client.pipeline()
        for fact in facts:
            pipe.rpush(key, json.dumps(fact.to_dict()))
        pipe.ltrim(key, -MAX_FACTS_PER_SESSION, -1)
        pipe.expire(key, SESSION_DEFAULT_TTL)
        pipe.execute()
        return len(facts)

    @traces(
        spec="006",
        story="US-003",
        reqs=[
            "REQ-006-FR3-004",
            "REQ-006-FR7-003",
            "REQ-006-US-003-002",
            "REQ-006-US-003-004",
            "REQ-006-US-005-003",
        ],
    )
    def get_facts(
        self, session_id: str, filter: FactFilter | None = None
    ) -> list[Fact]:
        """Get facts for a session, optionally filtered."""
        key = self._facts_key(session_id)
        # lrange() is typed as shared ResponseT (Awaitable[list] | list); this
        # sync client always returns a plain list of str/bytes entries.
        raw_list = cast("list[bytes | str]", self.client.lrange(key, 0, -1))
        facts = []
        for raw in raw_list:
            raw_str = raw if isinstance(raw, str) else raw.decode()
            data = json.loads(raw_str)
            fact = Fact.from_dict(data)
            if filter is None or filter.matches(fact):
                facts.append(fact)
        return facts

    @traces(spec="006", story="FR-013", reqs=["REQ-006-FR13-003"])
    def get_fact_count(self, session_id: str) -> int:
        """Get number of facts for a session."""
        key = self._facts_key(session_id)
        # llen() is typed as shared ResponseT (Awaitable[int] | int); this
        # sync client always returns a plain int.
        return cast(int, self.client.llen(key))

    # -- Context storage -------------------------------------------

    @traces(spec="006", story="US-005", reqs=["REQ-006-US-005-001"])
    def save_context(self, session_id: str, key: str, value: Any) -> None:
        """Save a context key-value pair."""
        ctx_key = self._context_key(session_id)
        self.client.hset(ctx_key, key, json.dumps(value))
        self.client.expire(ctx_key, SESSION_DEFAULT_TTL)

    @traces(spec="006", story="US-005", reqs=["REQ-006-US-005-006"])
    def get_context(self, session_id: str, key: str) -> Any:
        """Get a context value by key."""
        ctx_key = self._context_key(session_id)
        # hget() is typed as shared ResponseT (Awaitable[str | None] | str |
        # None); this sync client always returns a plain str/bytes or None.
        raw = cast("bytes | str | None", self.client.hget(ctx_key, key))
        if raw is None:
            return None
        raw_str = raw if isinstance(raw, str) else raw.decode()
        return json.loads(raw_str)

    @traces(
        spec="006",
        story="US-005",
        reqs=["REQ-006-US-005-001", "REQ-006-US-005-002", "REQ-006-US-005-006"],
    )
    def get_all_context(self, session_id: str) -> dict:
        """Get all context key-value pairs."""
        ctx_key = self._context_key(session_id)
        # hgetall() is typed as shared ResponseT (Awaitable[dict] | dict);
        # this sync client always returns a plain dict.
        data = cast("dict[bytes | str, bytes | str]", self.client.hgetall(ctx_key))
        result = {}
        for k, v in data.items():
            k_str = k if isinstance(k, str) else k.decode()
            v_str = v if isinstance(v, str) else v.decode()
            result[k_str] = json.loads(v_str)
        return result

    def delete_context(self, session_id: str, key: str) -> None:
        """Delete a context key."""
        ctx_key = self._context_key(session_id)
        self.client.hdel(ctx_key, key)

    @traces(spec="006", story="US-002", reqs=["REQ-006-US-002-005"])
    def clear_context(self, session_id: str) -> None:
        """Clear all context for a session."""
        ctx_key = self._context_key(session_id)
        self.client.delete(ctx_key)

    # -- Checkpoint storage ----------------------------------------

    def save_checkpoint(self, session_id: str, checkpoint: Checkpoint) -> None:
        """Save a checkpoint as a sorted set member (score=timestamp)."""
        key = self._checkpoints_key(session_id)
        score = checkpoint.created_at.timestamp()
        self.client.zadd(key, {json.dumps(checkpoint.to_dict()): score})
        self.client.expire(key, CHECKPOINT_TTL)
        # Prune old checkpoints if over limit.
        # zcard() is typed as shared ResponseT (Awaitable[int] | int); this
        # sync client always returns a plain int.
        count = cast(int, self.client.zcard(key))
        if count > MAX_CHECKPOINTS_PER_SESSION:
            self.client.zremrangebyrank(key, 0, count - MAX_CHECKPOINTS_PER_SESSION - 1)

    def get_latest_checkpoint(self, session_id: str) -> Checkpoint | None:
        """Get the most recent checkpoint."""
        key = self._checkpoints_key(session_id)
        # zrevrange() is typed as shared ResponseT (Awaitable[list] | list);
        # this sync client always returns a plain list of str/bytes entries.
        entries = cast("list[bytes | str]", self.client.zrevrange(key, 0, 0))
        if not entries:
            return None
        raw = entries[0] if isinstance(entries[0], str) else entries[0].decode()
        return Checkpoint.from_dict(json.loads(raw))

    def get_checkpoints(self, session_id: str) -> list[Checkpoint]:
        """Get all checkpoints ordered by time."""
        key = self._checkpoints_key(session_id)
        # zrange() is typed as shared ResponseT (Awaitable[list] | list); this
        # sync client always returns a plain list of str/bytes entries.
        entries = cast("list[bytes | str]", self.client.zrange(key, 0, -1))
        result = []
        for entry in entries:
            raw = entry if isinstance(entry, str) else entry.decode()
            result.append(Checkpoint.from_dict(json.loads(raw)))
        return result

    @traces(spec="006", story="US-002", reqs=["REQ-006-FR5-005", "REQ-006-US-002-005"])
    def delete_checkpoints(self, session_id: str) -> int:
        """Delete all checkpoints for a session."""
        key = self._checkpoints_key(session_id)
        # zcard() is typed as shared ResponseT (Awaitable[int] | int); this
        # sync client always returns a plain int.
        count = cast(int, self.client.zcard(key))
        self.client.delete(key)
        return count

    # -- Carryover storage -----------------------------------------

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-013"])
    def save_carryover(self, session_id: str, item: dict) -> None:
        """Append a carryover item."""
        key = self._carryover_key(session_id)
        self.client.rpush(key, json.dumps(item))

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-013"])
    def get_carryover(self, session_id: str) -> list[dict]:
        """Get all carryover items for a session."""
        key = self._carryover_key(session_id)
        # lrange() is typed as shared ResponseT (Awaitable[list] | list); this
        # sync client always returns a plain list of str/bytes entries.
        raw_list = cast("list[bytes | str]", self.client.lrange(key, 0, -1))
        result = []
        for raw in raw_list:
            raw_str = raw if isinstance(raw, str) else raw.decode()
            result.append(json.loads(raw_str))
        return result

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-013"])
    def clear_carryover(self, session_id: str) -> None:
        """Clear all carryover items."""
        key = self._carryover_key(session_id)
        self.client.delete(key)

    # -- Cleanup ---------------------------------------------------

    @traces(
        spec="006",
        story="US-002",
        reqs=[
            "REQ-006-FR5-002",
            "REQ-006-FR5-003",
            "REQ-006-FR5-005",
            "REQ-006-FR7-004",
            "REQ-006-US-002-004",
            "REQ-006-US-002-005",
        ],
    )
    def cleanup_session(self, session_id: str) -> None:
        """Remove all data for a session."""
        self.delete_session(session_id)