orchestrata_ultimate.kernel.session_context

Reference › orchestrata_ultimate › kernel › session_context

Session-scoped context storage and routing. Where session_manager stores the sessions themselves, this module stores key-value context per session in Redis and uses it to keep packets inside the session they belong to. It serves SPEC-001 story US-004 (REQ-001-US-004-014 to -028) and story US-002 (REQ-001-US-002-001). It has no importers in the package; the three classes are used through their own tests.

• ContextValidator checks context schema, keys and values. validate_key raises ValueError for a non-string, empty, overlong or badly formed key.
• SessionContext keeps each session's variables in a Redis hash named session_context:<id> and its attached volks under a second key. set, get, update, delete and clear validate the key and value, and JSON-encode values; a Redis error is logged and re-raised. serialize and deserialize move a whole context in and out. attach_volk, detach_volk and get_session_volks manage the volk membership. inject_context writes the session id and a serialised context into a packet's payload context, and extract_context reads it back. validate_session_isolation raises ValueError when a packet has no session id or the wrong one.
• SessionRouter uses a SessionContext to route inside a session. route_within_session first checks isolation, then returns the packet's recipient only if that volk is attached to the session, otherwise raising ValueError. broadcast_to_session injects the context into the packet and returns every volk in the session.

The Redis client is created through core.redis_factory or injected, and close releases it. Compare this recipient-by-session routing with the global routing of router.

At a glance

Class diagram

Dependency diagram

API

Classes

ContextValidator

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

Validates session context data.

SessionContext

Decorators: traces(spec='001', story='US-004', reqs=['REQ-001-US-004-014', 'REQ-001-US-004-015', 'REQ-001-US-004-016', 'REQ-001-US-… · Line 160

Manages session-scoped context storage and routing.

SessionRouter

Decorators: traces(spec='001', story='US-002', reqs=['REQ-001-US-002-001']) · Line 612

Routes messages within session scope.

Source

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

# mypy: ignore-errors
"""Session Context Management for Orchestration Kernel.

This module implements session context storage, validation, and routing
as specified in US-004 of the Orchestration Kernel spec.

Requirements Implemented:
- REQ-001-US-004-011: Set context variable
- REQ-001-US-004-012: Get context with default
- REQ-001-US-004-013: Update context atomically
- REQ-001-US-004-014: Delete context key
- REQ-001-US-004-015: Clear all context
- REQ-001-US-004-016: Inject context into CXP packet
- REQ-001-US-004-017: Extract context from CXP packet
- REQ-001-US-004-018: Validate context schema
- REQ-001-US-004-019: Serialize context for storage
- REQ-001-US-004-020: Deserialize context from storage
- REQ-001-US-004-021: Route messages within session scope
- REQ-001-US-004-022: Prevent cross-session leakage
- REQ-001-US-004-023: Session-specific Volk instances
- REQ-001-US-004-024: Broadcast to session Volks
- REQ-001-US-004-025: Query active Volks for session
- REQ-001-US-004-026: Session-scoped multicast
- REQ-001-US-004-027: Session isolation validation
- REQ-001-US-004-028: Session routing metrics
"""

import json
import logging
import re
import threading
from typing import Any, cast

import redis.asyncio as redis
from redis import Redis as SyncRedis

from orchestrata_ultimate.core.base.traces import traces

logger = logging.getLogger(__name__)

# hget/hgetall/smembers are typed via redis-py's shared sync/async command
# mixin (``ResponseT`` = ``Awaitable[Any] | Any``); the ``cast()`` calls at
# each call site below document that ``self.redis_client`` (SyncRedis)
# always returns the plain, non-awaitable value at runtime.


@traces(spec="001", story="US-004", reqs=["REQ-001-US-004-017"])
class ContextValidator:
    """Validates session context data.

    Ensures context keys and values conform to schema requirements.
    """

    # Valid key pattern: alphanumeric, underscore, hyphen, dot
    KEY_PATTERN = re.compile(r"^[a-zA-Z0-9_\-\.]+$")

    # Maximum key length
    MAX_KEY_LENGTH = 256

    # Maximum value size (1MB in JSON)
    MAX_VALUE_SIZE = 1024 * 1024

    @traces(
        spec="001",
        story="US-004",
        reqs=["REQ-001-US-004-017", "REQ-001-US-004-018"],
    )
    @classmethod
    def validate_schema(cls, context: dict[str, Any]) -> bool:
        """Validate entire context schema (REQ-001-US-004-018).

        Args:
            context: Context dictionary to validate

        Returns:
            True if valid

        Raises:
            ValueError: If validation fails
        """
        if not isinstance(context, dict):
            raise ValueError("Context must be a dictionary")

        for key, value in context.items():
            cls.validate_key(key)
            cls.validate_value(value)

        return True

    @traces(
        spec="001", story="US-004", reqs=["REQ-001-US-004-017", "REQ-001-US-004-018"]
    )
    @classmethod
    def validate_key(cls, key: str) -> bool:
        """Validate a context key (REQ-001-US-004-018).

        Args:
            key: Context key to validate

        Returns:
            True if valid

        Raises:
            ValueError: If key is invalid
        """
        if not isinstance(key, str):
            raise ValueError(f"Key must be string, got {type(key)}")

        if not key:
            raise ValueError("Key cannot be empty")

        if len(key) > cls.MAX_KEY_LENGTH:
            raise ValueError(f"Key exceeds maximum length of {cls.MAX_KEY_LENGTH}")

        if not cls.KEY_PATTERN.match(key):
            raise ValueError(f"Key contains invalid characters: {key}")

        return True

    @traces(
        spec="001", story="US-004", reqs=["REQ-001-US-004-017", "REQ-001-US-004-018"]
    )
    @classmethod
    def validate_value(cls, value: Any) -> bool:
        """Validate a context value (REQ-001-US-004-018).

        Args:
            value: Value to validate

        Returns:
            True if valid

        Raises:
            ValueError: If value is invalid
        """
        # Serialize to check size
        try:
            serialized = json.dumps(value)
            if len(serialized) > cls.MAX_VALUE_SIZE:
                raise ValueError(
                    f"Value exceeds maximum size of {cls.MAX_VALUE_SIZE} bytes"
                )
        except (TypeError, ValueError) as e:
            raise ValueError(f"Value cannot be serialized to JSON: {e}")

        return True


@traces(
    spec="001",
    story="US-004",
    reqs=[
        "REQ-001-US-004-014",
        "REQ-001-US-004-015",
        "REQ-001-US-004-016",
        "REQ-001-US-004-018",
        "REQ-001-US-004-023",
    ],
)
class SessionContext:
    """Manages session-scoped context storage and routing.

    Provides context variable management with validation, serialization,
    and integration with CXP packets.
    """

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-018"])
    def __init__(
        self,
        redis_host: str | None = None,
        redis_port: int | None = None,
        redis_db: int = 0,
        redis_client: SyncRedis | None = None,
    ) -> None:
        """Initialize SessionContext.

        Args:
            redis_host: Redis server hostname
            redis_port: Redis server port
            redis_db: Redis database number
            redis_client: Pre-built sync client injected directly
                (the host/port/db args are then ignored for connection
                purposes; the connectivity ping still runs). Seam for
                callers needing a non-default retry policy — the
                init-failure test injects a fail-fast client so the ping
                against a dead port raises in ms instead of redis-py 7's
                4-13s default backoff loop.
        """
        from orchestrata_ultimate.core.redis_factory import (
            create_sync_redis,
            resolve_redis_host,
            resolve_redis_port,
        )

        try:
            if redis_client is None:
                redis_host = (
                    redis_host if redis_host is not None else resolve_redis_host()
                )
                redis_port = (
                    redis_port if redis_port is not None else resolve_redis_port()
                )
                redis_client = create_sync_redis(
                    host=redis_host,
                    port=redis_port,
                    db=redis_db,
                )

            # Test connection before assigning to instance
            redis_client.ping()
            self.redis_client = redis_client
            logger.info(
                f"Connected to Redis for session context at {redis_host}:{redis_port}"
            )
        except Exception as e:
            # Cleanup on failure
            if redis_client is not None:
                try:
                    redis_client.close()
                except Exception as exc:
                    logger.warning(
                        "Failed to close Redis client during cleanup: %s",
                        exc,
                        exc_info=True,
                    )
            logger.error(f"Failed to connect to Redis: {e}")
            raise

        self.validator = ContextValidator()
        self._lock = threading.Lock()

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-014"])
    def _context_key(self, session_id: str) -> str:
        """Build Redis key for session context.

        Args:
            session_id: Session identifier

        Returns:
            Redis key string
        """
        return f"session_context:{session_id}"

    def _volks_key(self, session_id: str) -> str:
        """Build Redis key for session Volks.

        Args:
            session_id: Session identifier

        Returns:
            Redis key string
        """
        return f"session_volks:{session_id}"

    @traces(
        spec="001",
        story="US-004",
        reqs=[
            "REQ-001-US-004-016",
            "REQ-001-US-004-018",
        ],
    )
    def set(self, session_id: str, key: str, value: Any) -> None:
        """Set a context variable (REQ-001-US-004-011).

        Args:
            session_id: Session identifier
            key: Context variable name
            value: Value to store
        """
        # Validate key and value
        self.validator.validate_key(key)
        self.validator.validate_value(value)

        # Store in Redis hash — lock for Python 3.14t thread safety
        context_key = self._context_key(session_id)
        serialized_value = json.dumps(value)

        try:
            with self._lock:
                self.redis_client.hset(context_key, key, serialized_value)
            logger.debug(f"Set context variable '{key}' for session {session_id}")
        except redis.RedisError as e:
            logger.error(
                f"Redis error setting context variable '{key}' for session {session_id}: {e}"
            )
            raise
        except TypeError as e:  # pragma: no cover
            logger.error(
                f"JSON encoding error for context variable '{key}' for session {session_id}: {e}"
            )  # pragma: no cover
            raise  # pragma: no cover
        except Exception as e:
            logger.error(
                f"Unexpected error setting context variable '{key}' for session {session_id}: {e}"
            )
            raise

    @traces(
        spec="001",
        story="US-004",
        reqs=["REQ-001-US-004-015"],
    )
    def get(self, session_id: str, key: str, default: Any = None) -> Any:
        """Get a context variable with default (REQ-001-US-004-012).

        Args:
            session_id: Session identifier
            key: Context variable name
            default: Default value if key not found

        Returns:
            Context value or default
        """
        context_key = self._context_key(session_id)

        try:
            serialized_value = cast(
                "str | None", self.redis_client.hget(context_key, key)
            )

            if serialized_value is None:
                return default

            try:
                return json.loads(serialized_value)
            except json.JSONDecodeError as e:
                logger.error(
                    f"Failed to decode context variable '{key}' for session {session_id}: {e}"
                )
                return default
        except redis.RedisError as e:
            logger.error(
                f"Redis error getting context variable '{key}' for session {session_id}: {e}"
            )
            raise
        except Exception as e:
            logger.error(
                f"Unexpected error getting context variable '{key}' for session {session_id}: {e}"
            )
            raise

    @traces(
        spec="001",
        story="US-004",
        reqs=["REQ-001-US-004-016"],
    )
    def update(self, session_id: str, updates: dict[str, Any]) -> None:
        """Update context atomically (REQ-001-US-004-013).

        Args:
            session_id: Session identifier
            updates: Dictionary of key-value pairs to update
        """
        # Validate all updates first
        self.validator.validate_schema(updates)

        # Serialize all values
        try:
            serialized_updates = {
                key: json.dumps(value) for key, value in updates.items()
            }
        except TypeError as e:  # pragma: no cover
            logger.error(
                f"JSON encoding error during update for session {session_id}: {e}"
            )  # pragma: no cover
            raise  # pragma: no cover

        # Atomic update using HMSET
        context_key = self._context_key(session_id)
        if serialized_updates:
            try:
                self.redis_client.hset(context_key, mapping=serialized_updates)
                logger.debug(
                    f"Updated {len(updates)} context variables for session {session_id}"
                )
            except redis.RedisError as e:
                logger.error(
                    f"Redis error updating {len(updates)} context variables for session {session_id}: {e}"
                )
                raise
            except Exception as e:
                logger.error(
                    f"Unexpected error updating {len(updates)} context variables for session {session_id}: {e}"
                )
                raise

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-016"])
    def delete(self, session_id: str, key: str) -> None:
        """Delete a context key (REQ-001-US-004-014).

        Args:
            session_id: Session identifier
            key: Context variable name to delete
        """
        context_key = self._context_key(session_id)
        try:
            self.redis_client.hdel(context_key, key)
            logger.debug(f"Deleted context variable '{key}' for session {session_id}")
        except redis.RedisError as e:
            logger.error(
                f"Redis error deleting context variable '{key}' for session {session_id}: {e}"
            )
            raise
        except Exception as e:
            logger.error(
                f"Unexpected error deleting context variable '{key}' for session {session_id}: {e}"
            )
            raise

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-016"])
    def clear(self, session_id: str) -> None:
        """Clear all context for a session (REQ-001-US-004-015).

        Args:
            session_id: Session identifier
        """
        context_key = self._context_key(session_id)
        try:
            self.redis_client.delete(context_key)
            logger.info(f"Cleared all context for session {session_id}")
        except redis.RedisError as e:
            logger.error(f"Redis error clearing context for session {session_id}: {e}")
            raise
        except Exception as e:
            logger.error(
                f"Unexpected error clearing context for session {session_id}: {e}"
            )
            raise

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-016"])
    def inject_context(self, packet: Any, session_id: str) -> Any:
        """Inject session context into CXP packet (REQ-001-US-004-016).

        Note: This method calls serialize() which reads from Redis.
        The serialize operation is synchronous and may be slow for large contexts.

        Args:
            packet: CXP packet object
            session_id: Session identifier

        Returns:
            Packet with injected context
        """
        # Get all context for session (reads from Redis)
        context_data = self.serialize(session_id)

        # Inject into packet payload context
        if hasattr(packet, "payload") and hasattr(packet.payload, "context"):
            packet.payload.context["session_id"] = session_id
            packet.payload.context["session_context"] = context_data

        logger.debug(f"Injected context into packet for session {session_id}")
        return packet

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-017"])
    def extract_context(self, packet: Any) -> dict[str, Any]:
        """Extract context from CXP packet (REQ-001-US-004-017).

        Args:
            packet: CXP packet object

        Returns:
            Extracted context dictionary (full payload context)
        """
        if hasattr(packet, "payload") and hasattr(packet.payload, "context"):
            return packet.payload.context

        return {}

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-015"])
    def serialize(self, session_id: str) -> dict[str, Any]:
        """Serialize session context (REQ-001-US-004-019).

        Note: This method performs hgetall() followed by N json.loads() operations.
        For large contexts, this could be slow. Consider pagination if needed.

        Args:
            session_id: Session identifier

        Returns:
            Dictionary of all context variables
        """
        context_key = self._context_key(session_id)

        try:
            serialized_data = cast(
                "dict[str, str]", self.redis_client.hgetall(context_key)
            )

            # Deserialize all values (could be slow for large contexts)
            context = {}
            for key, value in serialized_data.items():
                try:
                    context[key] = json.loads(value)
                except json.JSONDecodeError as e:
                    logger.warning(
                        f"Failed to decode context key '{key}' for session {session_id}: {e}"
                    )
                    # Skip malformed values
                    continue

            return context
        except redis.RedisError as e:
            logger.error(
                f"Redis error serializing context for session {session_id}: {e}"
            )
            raise
        except Exception as e:
            logger.error(
                f"Unexpected error serializing context for session {session_id}: {e}"
            )
            raise

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-016"])
    def deserialize(self, session_id: str, data: dict[str, Any]) -> None:
        """Deserialize and store session context (REQ-001-US-004-020).

        Args:
            session_id: Session identifier
            data: Context data to deserialize and store
        """
        # Validate and update context
        self.update(session_id, data)

        logger.info(
            f"Deserialized {len(data)} context variables for session {session_id}"
        )

    def attach_volk(self, session_id: str, volk_alias: str) -> None:
        """Attach a Volk to session for routing (REQ-001-US-004-023).

        Args:
            session_id: Session identifier
            volk_alias: Volk alias to attach
        """
        volks_key = self._volks_key(session_id)
        self.redis_client.sadd(volks_key, volk_alias)

        logger.info(f"Attached {volk_alias} to session {session_id}")

    def detach_volk(self, session_id: str, volk_alias: str) -> None:
        """Detach a Volk from session (REQ-001-US-004-023).

        Args:
            session_id: Session identifier
            volk_alias: Volk alias to detach
        """
        volks_key = self._volks_key(session_id)
        self.redis_client.srem(volks_key, volk_alias)

        logger.info(f"Detached {volk_alias} from session {session_id}")

    def get_session_volks(self, session_id: str) -> list[str]:
        """Get all Volks attached to session (REQ-001-US-004-025).

        Args:
            session_id: Session identifier

        Returns:
            List of Volk aliases
        """
        volks_key = self._volks_key(session_id)
        volks = list(cast("set[str]", self.redis_client.smembers(volks_key)))

        return volks

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-014"])
    def validate_session_isolation(
        self,
        packet: Any,
        expected_session_id: str,
    ) -> bool:
        """Validate packet belongs to expected session (REQ-001-US-004-027).

        Args:
            packet: CXP packet to validate
            expected_session_id: Expected session ID

        Returns:
            True if packet belongs to session

        Raises:
            ValueError: If session isolation is violated
        """
        extracted_context = self.extract_context(packet)
        packet_session_id = extracted_context.get("session_id")

        # Check for None explicitly to avoid comparing None with string
        if packet_session_id is None:
            raise ValueError(
                f"Session isolation violation: packet has no session_id, "
                f"expected {expected_session_id}"
            )

        if packet_session_id != expected_session_id:
            raise ValueError(
                f"Session isolation violation: packet from {packet_session_id}, "
                f"expected {expected_session_id}"
            )

        return True

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-023"])
    def close(self) -> None:
        """Close Redis connection."""
        self.redis_client.close()
        logger.info("Closed SessionContext connections")


@traces(spec="001", story="US-002", reqs=["REQ-001-US-002-001"])
class SessionRouter:
    """Routes messages within session scope.

    Provides session-scoped message routing with isolation guarantees.
    """

    def __init__(self, session_context: SessionContext) -> None:
        """Initialize SessionRouter.

        Args:
            session_context: SessionContext instance
        """
        self.context: SessionContext = session_context

    @traces(
        spec="001",
        story="US-004",
        reqs=[
            "REQ-001-US-004-014",
            "REQ-001-US-002-003",
            "REQ-001-US-004-021",
            "REQ-001-US-004-022",
            "REQ-001-US-004-027",
        ],
    )
    def route_within_session(
        self,
        packet: Any,
        session_id: str,
    ) -> list[str]:
        """Route message within session scope (REQ-001-US-004-021, REQ-001-US-004-022).

        Args:
            packet: CXP packet to route
            session_id: Session identifier

        Returns:
            List of Volk aliases to deliver to

        Raises:
            ValueError: If session isolation is violated
        """
        # Validate session isolation
        self.context.validate_session_isolation(packet, session_id)

        # Get recipient from packet
        if hasattr(packet, "header") and hasattr(packet.header, "recipient"):
            recipient = packet.header.recipient

            # Verify recipient is in session
            session_volks = self.context.get_session_volks(session_id)

            if recipient not in session_volks:
                raise ValueError(f"Recipient {recipient} not in session {session_id}")

            return [recipient]

        return []

    @traces(
        spec="001",
        story="US-002",
        reqs=[
            "REQ-001-US-002-005",
            "REQ-001-US-004-024",
            "REQ-001-US-004-026",
            "REQ-001-US-004-028",
        ],
    )
    def broadcast_to_session(
        self,
        packet: Any,
        session_id: str,
    ) -> list[str]:
        """Broadcast to all Volks in session (REQ-001-US-004-024, REQ-001-US-004-026).

        Args:
            packet: CXP packet to broadcast
            session_id: Session identifier

        Returns:
            List of Volk aliases to deliver to
        """
        # Inject session context
        packet = self.context.inject_context(packet, session_id)

        # Get all session Volks
        session_volks = self.context.get_session_volks(session_id)

        logger.info(
            f"Broadcasting to {len(session_volks)} Volks in session {session_id}"
        )
        return session_volks