orchestrata_ultimate.infra.redis_state

Reference › orchestrata_ultimate › infra › redis_state

Hot-state storage in Redis for sessions, caches, locks, publish and subscribe, and queues (Spec 006, session management; based on INFRASTRUCTURE.md). Agents runtime and task executor, governance audit and other modules import it, and it satisfies RedisStateManagerProtocol from core.protocols.decoupling.

RedisConfig takes its host and port from resolve_redis_host and resolve_redis_port in redis_factory, so an unset key raises instead of falling back to a default. RedisNamespace defines the key prefixes (orchestrata:session, volk, queue, cache, lock, pubsub and governance), and _key builds namespaced keys. RedisStateManager creates its clients lazily: client is a synchronous client and async_client an asyncio one, and both can be injected. It stores SessionState records with create_session (default TTL one hour), get_session, update_session and delete_session; a general cache with cache_set, cache_get and cache_delete, with cache_governance_decision and get_cached_decision keyed by an intent hash; distributed locks with acquire_lock, release_lock and extend_lock, where release and extension are single atomic Lua scripts that check the holder, which avoids a check-then-delete race; publish and subscribe; and list-based queues with queue_push, a blocking queue_pop and queue_length. JSON values are limited to 10 MB (MAX_JSON_SIZE_BYTES). health_check tests connectivity, and close and aclose release connections.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

RedisNamespace (enum)

Bases: Enum · Line 69

Redis key namespaces for data organization.

RedisConfig (dataclass)

Decorators: dataclass · Line 82

Redis connection configuration.

SessionState (dataclass)

Decorators: dataclass, traces(spec='001', story='US-004', reqs=['REQ-001-US-004-005']) · Line 101

Session state stored in Redis.

RedisStateManager

Line 114

Manages hot state in Redis for Orchestrata.

Functions

Source

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

# mypy: ignore-errors
"""Redis state management for Orchestrata.

Handles hot state storage including sessions, queues, cache, and pub/sub.
Based on INFRASTRUCTURE.md specifications.
"""

from __future__ import annotations

# Spec: 006 - Session Management
import json
import logging
import threading
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum
from typing import TYPE_CHECKING, Any, cast

import redis.asyncio as redis

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.redis_factory import (
    resolve_redis_host,
    resolve_redis_port,
)

if TYPE_CHECKING:
    # The `client` property returns a SYNC client (core.redis_factory.create_sync_redis);
    # only `async_client` is a redis.asyncio client. Import the sync types
    # explicitly so they are not shadowed by `import redis.asyncio as redis` above
    # (under which a bare `redis.Redis` would resolve to the *async* client type).
    from redis import Redis as SyncRedis
    from redis.client import PubSub as SyncPubSub
    from redis.client import PubSubWorkerThread
    from redis.commands.core import Script as SyncScript

# Module-level logger
_logger = logging.getLogger(__name__)

# Configuration constants
MAX_JSON_SIZE_BYTES: int = 10 * 1024 * 1024  # 10MB max JSON size


def _utc_now() -> datetime:
    """Get current UTC time with timezone info."""
    return datetime.now(UTC)


# Lua script for atomic lock release (prevents TOCTOU race condition)
RELEASE_LOCK_SCRIPT = """
if redis.call('get', KEYS[1]) == ARGV[1] then
    return redis.call('del', KEYS[1])
else
    return 0
end
"""

# Lua script for atomic lock extension
EXTEND_LOCK_SCRIPT = """
if redis.call('get', KEYS[1]) == ARGV[1] then
    return redis.call('expire', KEYS[1], ARGV[2])
else
    return 0
end
"""


class RedisNamespace(Enum):
    """Redis key namespaces for data organization."""

    SESSION = "orchestrata:session"
    VOLK = "orchestrata:volk"
    QUEUE = "orchestrata:queue"
    CACHE = "orchestrata:cache"
    LOCK = "orchestrata:lock"
    PUBSUB = "orchestrata:pubsub"
    GOVERNANCE = "orchestrata:governance"


@dataclass
class RedisConfig:
    """Redis connection configuration."""

    # No literal endpoint default (RUL-003/RUL-011): "localhost:6379" would
    # connect to nothing when REDIS_HOST/PORT aren't the target and the
    # resulting refusal reads as an outage instead of the config gap it is.
    host: str = field(default_factory=resolve_redis_host)
    port: int = field(default_factory=resolve_redis_port)
    db: int = 0
    password: str | None = None
    decode_responses: bool = True
    socket_timeout: float = 5.0
    socket_connect_timeout: float = 5.0
    retry_on_timeout: bool = True
    max_connections: int = 10


@dataclass
@traces(spec="001", story="US-004", reqs=["REQ-001-US-004-005"])
class SessionState:
    """Session state stored in Redis."""

    session_id: str
    volk_id: str
    status: str
    goals: list[str]
    context: dict[str, Any]
    started_at: datetime
    last_activity: datetime
    metadata: dict[str, Any] = field(default_factory=dict)


class RedisStateManager:
    """Manages hot state in Redis for Orchestrata.

    Provides session management, caching, distributed locks,
    and pub/sub messaging.
    """

    def __init__(
        self,
        config: RedisConfig | None = None,
        redis_client: SyncRedis | None = None,
        async_redis_client: redis.Redis | None = None,
    ) -> None:
        """Initialize Redis connection.

        Args:
            config: Redis configuration. Uses defaults if not provided.
            redis_client: Pre-built sync client injected directly (the
                production path still lazily builds one from ``config`` on
                first access to ``client`` when this is omitted). Exposed so
                tests can point ``RedisStateManager`` at a fail-fast client
                without paying redis-py's default connection-retry backoff.
            async_redis_client: Pre-built async client injected directly,
                mirroring ``redis_client`` for the ``async_client`` property.
        """
        self.config = config or RedisConfig()
        self._client: SyncRedis | None = redis_client
        self._pubsub: SyncPubSub | None = None
        # `PubSubWorkerThread` (not the bare `threading.Thread` base class):
        # `PubSub.run_in_thread()` always returns this subclass, which is
        # what exposes `.stop()` (used in `close()` below).
        self._pubsub_thread: PubSubWorkerThread | None = None
        self._subscribers: dict[str, list[Callable]] = {}
        self._lock = threading.RLock()
        self._release_lock_script: SyncScript | None = None
        self._extend_lock_script: SyncScript | None = None
        self._async_client: redis.Redis | None = async_redis_client
        if self._client is not None:
            self._release_lock_script = self._client.register_script(
                RELEASE_LOCK_SCRIPT
            )
            self._extend_lock_script = self._client.register_script(EXTEND_LOCK_SCRIPT)

    @property
    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-012"])
    def client(self) -> SyncRedis:
        """Get or create the SYNC Redis client (lazy initialization).

        Returns a synchronous ``redis.Redis`` (BlockingConnectionPool); its calls
        are NOT awaited. Use ``async_client`` for the asyncio path.
        """
        if self._client is None:
            with self._lock:
                if self._client is None:
                    from orchestrata_ultimate.core.redis_factory import (
                        create_sync_redis,
                    )

                    c = create_sync_redis(
                        host=self.config.host,
                        port=self.config.port,
                        db=self.config.db,
                        password=self.config.password,
                    )
                    self._release_lock_script = c.register_script(RELEASE_LOCK_SCRIPT)
                    self._extend_lock_script = c.register_script(EXTEND_LOCK_SCRIPT)
                    self._client = c  # publish LAST
        return self._client

    @property
    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-012"])
    def async_client(self) -> redis.Redis:
        """Get or create async Redis client (lazy initialization)."""
        if self._async_client is None:
            from orchestrata_ultimate.core.redis_factory import create_async_redis

            self._async_client = create_async_redis(
                host=self.config.host,
                port=self.config.port,
                db=self.config.db,
                password=self.config.password,
            )
        return self._async_client

    def _key(self, namespace: RedisNamespace, *parts: str) -> str:
        """Build namespaced Redis key."""
        return ":".join([namespace.value, *parts])

    # ===== Session Management =====

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-005"])
    def create_session(self, session: SessionState, ttl_seconds: int = 3600) -> bool:
        """Create a new session in Redis.

        Args:
            session: Session state to store.
            ttl_seconds: Time-to-live in seconds (default 1 hour).

        Returns:
            True if session was created, False if already exists.
        """
        key = self._key(RedisNamespace.SESSION, session.session_id)
        data = {
            "session_id": session.session_id,
            "volk_id": session.volk_id,
            "status": session.status,
            "goals": json.dumps(session.goals),
            "context": json.dumps(session.context),
            "started_at": session.started_at.isoformat(),
            "last_activity": session.last_activity.isoformat(),
            "metadata": json.dumps(session.metadata),
        }
        # Use NX to only set if not exists
        result = self.client.hset(key, mapping=data)
        if result:
            self.client.expire(key, ttl_seconds)
        return bool(result)

    def get_session(self, session_id: str) -> SessionState | None:
        """Retrieve session state from Redis.

        Args:
            session_id: Session identifier.

        Returns:
            SessionState if found, None otherwise.
        """
        key = self._key(RedisNamespace.SESSION, session_id)
        # redis-py's command mixins are shared between the sync `Redis` and
        # `redis.asyncio.Redis` classes, so `hgetall()` is stub-typed
        # `Awaitable[dict] | dict`; `self.client` is always the sync client
        # here, so the awaitable branch is never actually returned.
        data = cast("dict[str, str]", self.client.hgetall(key))
        if not data:
            return None
        return SessionState(
            session_id=data["session_id"],
            volk_id=data["volk_id"],
            status=data["status"],
            goals=json.loads(data["goals"]),
            context=json.loads(data["context"]),
            started_at=datetime.fromisoformat(data["started_at"]),
            last_activity=datetime.fromisoformat(data["last_activity"]),
            metadata=json.loads(data.get("metadata", "{}")),
        )

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-018"])
    def update_session(self, session_id: str, updates: dict[str, Any]) -> bool:
        """Update session fields in Redis.

        Args:
            session_id: Session identifier.
            updates: Dictionary of fields to update.

        Returns:
            True if session was updated.
        """
        key = self._key(RedisNamespace.SESSION, session_id)
        if not self.client.exists(key):
            return False

        # Serialize complex types
        for k, v in updates.items():
            if isinstance(v, (list, dict)):
                updates[k] = json.dumps(v)
            elif isinstance(v, datetime):
                updates[k] = v.isoformat()

        updates["last_activity"] = _utc_now().isoformat()
        self.client.hset(key, mapping=updates)
        return True

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-023"])
    def delete_session(self, session_id: str) -> bool:
        """Delete session from Redis.

        Args:
            session_id: Session identifier.

        Returns:
            True if session was deleted.
        """
        key = self._key(RedisNamespace.SESSION, session_id)
        return bool(self.client.delete(key))

    # ===== Caching =====

    def cache_set(self, key: str, value: Any, ttl_seconds: float = 300) -> None:
        """Set a cached value.

        Args:
            key: Cache key.
            value: Value to cache (will be JSON serialized).
            ttl_seconds: Time-to-live in seconds (default 5 minutes). Supports
                sub-second values (e.g. 0.1 for 100ms).
        """
        full_key = self._key(RedisNamespace.CACHE, key)
        ttl_ms = max(1, int(ttl_seconds * 1000))
        self.client.psetex(full_key, ttl_ms, json.dumps(value))

    def cache_get(self, key: str) -> Any | None:
        """Get a cached value.

        Args:
            key: Cache key.

        Returns:
            Cached value or None if not found/expired.
        """
        full_key = self._key(RedisNamespace.CACHE, key)
        # Stub deficiency (see get_session): `get()` is typed
        # `Awaitable[Any] | Any` on the shared sync/async command mixin.
        data = cast("str | None", self.client.get(full_key))
        return json.loads(data) if data else None

    def cache_delete(self, key: str) -> bool:
        """Delete a cached value.

        Args:
            key: Cache key.

        Returns:
            True if key was deleted.
        """
        full_key = self._key(RedisNamespace.CACHE, key)
        return bool(self.client.delete(full_key))

    # ===== Governance Cache =====

    @traces(
        spec="003",
        story="US-003",
        reqs=["REQ-003-US-003-009", "REQ-003-US-003-019"],
    )
    def cache_governance_decision(
        self, intent_hash: str, decision: dict[str, Any], ttl_seconds: float = 60
    ) -> None:
        """Cache a governance decision for fast lookup.

        Args:
            intent_hash: Hash of the intent for lookup.
            decision: Decision result to cache.
            ttl_seconds: Cache TTL (default 60 seconds).
        """
        key = self._key(RedisNamespace.GOVERNANCE, "decision", intent_hash)
        ttl_ms = max(1, int(ttl_seconds * 1000))
        self.client.psetex(key, ttl_ms, json.dumps(decision))

    @traces(
        spec="003",
        story="US-003",
        reqs=["REQ-003-US-003-009", "REQ-003-US-003-019"],
    )
    def get_cached_decision(self, intent_hash: str) -> dict[str, Any] | None:
        """Get a cached governance decision.

        Args:
            intent_hash: Hash of the intent.

        Returns:
            Cached decision or None.
        """
        key = self._key(RedisNamespace.GOVERNANCE, "decision", intent_hash)
        # Stub deficiency (see get_session).
        data = cast("str | None", self.client.get(key))
        return json.loads(data) if data else None

    # ===== Distributed Locks =====

    def acquire_lock(
        self, lock_name: str, holder_id: str, ttl_seconds: int = 30
    ) -> bool:
        """Acquire a distributed lock.

        Args:
            lock_name: Name of the lock.
            holder_id: Identifier of the lock holder.
            ttl_seconds: Lock TTL to prevent deadlocks.

        Returns:
            True if lock was acquired.
        """
        key = self._key(RedisNamespace.LOCK, lock_name)
        # Use SET NX EX for atomic lock acquisition
        return bool(self.client.set(key, holder_id, nx=True, ex=ttl_seconds))

    def release_lock(self, lock_name: str, holder_id: str) -> bool:
        """Release a distributed lock atomically.

        Uses Lua script to prevent TOCTOU race condition.

        Args:
            lock_name: Name of the lock.
            holder_id: Identifier of the lock holder.

        Returns:
            True if lock was released.
        """
        key = self._key(RedisNamespace.LOCK, lock_name)
        # Ensure client is initialized and scripts are registered
        _ = self.client
        # Use Lua script for atomic check-and-delete
        if self._release_lock_script is not None:
            result = self._release_lock_script(keys=[key], args=[holder_id])
            return bool(result)
        return False

    def extend_lock(
        self, lock_name: str, holder_id: str, ttl_seconds: int = 30
    ) -> bool:
        """Extend a lock's TTL atomically.

        Uses Lua script to prevent TOCTOU race condition.

        Args:
            lock_name: Name of the lock.
            holder_id: Identifier of the lock holder.
            ttl_seconds: New TTL.

        Returns:
            True if lock was extended.
        """
        key = self._key(RedisNamespace.LOCK, lock_name)
        # Ensure client is initialized and scripts are registered
        _ = self.client
        # Use Lua script for atomic check-and-extend
        if self._extend_lock_script is not None:
            result = self._extend_lock_script(keys=[key], args=[holder_id, ttl_seconds])
            return bool(result)
        return False

    # ===== Pub/Sub =====

    def publish(self, channel: str, message: dict[str, Any]) -> int:
        """Publish a message to a channel.

        Args:
            channel: Channel name.
            message: Message to publish.

        Returns:
            Number of subscribers that received the message.
        """
        full_channel = self._key(RedisNamespace.PUBSUB, channel)
        # Stub deficiency (see get_session).
        return cast(int, self.client.publish(full_channel, json.dumps(message)))

    def subscribe(
        self, channel: str, callback: Callable[[dict[str, Any]], None]
    ) -> None:
        """Subscribe to a channel.

        Args:
            channel: Channel name.
            callback: Function to call when message received.
        """
        full_channel = self._key(RedisNamespace.PUBSUB, channel)

        with self._lock:
            if full_channel not in self._subscribers:
                self._subscribers[full_channel] = []
            self._subscribers[full_channel].append(callback)

            if self._pubsub is None:
                self._pubsub = self.client.pubsub()
                # Start listener thread to process messages
                self._pubsub_thread = self._pubsub.run_in_thread(
                    sleep_time=0.001, daemon=True
                )

            self._pubsub.subscribe(**{full_channel: self._handle_message})

    def _handle_message(self, message: dict) -> None:
        """Handle incoming pub/sub message."""
        if message["type"] != "message":
            return

        channel = message["channel"]
        raw_data = message["data"]

        # Validate JSON size to prevent DoS
        if isinstance(raw_data, str) and len(raw_data) > MAX_JSON_SIZE_BYTES:
            _logger.warning(
                f"Dropping oversized message on channel {channel}: "
                f"{len(raw_data)} bytes > {MAX_JSON_SIZE_BYTES} limit"
            )
            return

        try:
            data = json.loads(raw_data)
        except json.JSONDecodeError as e:
            _logger.error(f"Invalid JSON in message on channel {channel}: {e}")
            return

        with self._lock:
            callbacks = self._subscribers.get(channel, [])
            for callback in callbacks:
                try:
                    callback(data)
                except Exception as e:
                    _logger.error(f"Error in callback for channel {channel}: {e}")

    # ===== Queue Operations =====

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-022"])
    def queue_push(self, queue_name: str, item: dict[str, Any]) -> int:
        """Push an item to a queue.

        Args:
            queue_name: Queue name.
            item: Item to push.

        Returns:
            Queue length after push.
        """
        key = self._key(RedisNamespace.QUEUE, queue_name)
        # Stub deficiency (see get_session).
        return cast(int, self.client.rpush(key, json.dumps(item)))

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-022"])
    def queue_pop(self, queue_name: str, timeout: int = 0) -> dict[str, Any] | None:
        """Pop an item from a queue (blocking).

        Args:
            queue_name: Queue name.
            timeout: Blocking timeout in seconds (0 = block forever).

        Returns:
            Item or None if timeout.
        """
        key = self._key(RedisNamespace.QUEUE, queue_name)
        # Stub deficiency (see get_session).
        result = cast("list[Any] | None", self.client.blpop(key, timeout=timeout))
        if result:
            return json.loads(result[1])
        return None

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-025"])
    def queue_length(self, queue_name: str) -> int:
        """Get queue length.

        Args:
            queue_name: Queue name.

        Returns:
            Number of items in queue.
        """
        key = self._key(RedisNamespace.QUEUE, queue_name)
        # Stub deficiency (see get_session).
        return cast(int, self.client.llen(key))

    # ===== Health Check =====

    def health_check(self) -> bool:
        """Check Redis connectivity.

        Returns:
            True if Redis is healthy.
        """
        # Check if client exists before trying to use it
        if self._client is None:
            return False

        try:
            # Stub deficiency (see get_session).
            return cast(bool, self._client.ping())
        except redis.RedisError:
            return False

    def close(self) -> None:
        """Close Redis connections."""
        if self._pubsub_thread:
            self._pubsub_thread.stop()
            self._pubsub_thread.join(timeout=2.0)
            self._pubsub_thread = None
        if self._pubsub:
            self._pubsub.close()
            self._pubsub = None
        if self._client:
            self._client.close()
            self._client = None

    async def aclose(self) -> None:
        """Close all connections including async client."""
        self.close()  # Sync cleanup
        if self._async_client:
            await self._async_client.aclose()
            self._async_client = None