orchestrata_ultimate.kernel.session_manager

Reference › orchestrata_ultimate › kernel › session_manager

The persistence layer for user sessions. A Session is a dataclass entity (a BaseEntity from the core package) with a session id, user id, start and end times, metadata, a status and the list of attached volks; SessionStatus is active, closed or archived. SessionManager stores live sessions in Redis and archives finished ones to the graph. It serves SPEC-001 story US-004 (REQ-001-US-004-001 to -026). No other module imports it. It is distinct from the session controller that pipeline_session uses for work logging, and it works with session_context, which holds per-session context. The curated kernel page lists it among the kernel services.

The constructor builds a synchronous Redis client through core.redis_factory (or takes an injected one), pings it, and records the graph URI, user and password settings; any failure closes what was opened, logs the error and re-raises. create_session makes a UUID, writes the session as JSON under session:<id>, and adds the id to the set user_sessions:<user>; if the second write fails it deletes the first. resume_session reads and parses the record and raises ValueError when the session is missing or its data is invalid. close_session marks it closed with an end time and removes it from the user's set. archive_session is async: it merges a :Session node attached to the project node, links the attached volks that exist in the graph with worked-on edges, and marks the Redis record archived. get_active_sessions and query_session_history list sessions, and delete_expired_sessions scans session:* keys and deletes those older than a time-to-live in batches. attach_volk and detach_volk on Session edit its volk list.

At a glance

Class diagram

Dependency diagram

API

Classes

SessionStatus (enum)

Bases: Enum · Line 66

Session status enumeration.

Session (dataclass)

Bases: BaseEntity · Decorators: traces(spec='001', story='US-004', reqs=['REQ-001-US-004-001']), dataclass(eq=False) · Line 76

Session data model.

SessionManager

Line 198

Manages session lifecycle and persistence.

Functions

Source

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

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

This module implements session lifecycle management, persistence, and routing
as specified in US-004 of the Orchestration Kernel spec.

Requirements Implemented:
- REQ-001-US-004-001: Create session with unique UUID
- REQ-001-US-004-002: Resume session from Redis
- REQ-001-US-004-003: Close session and mark inactive
- REQ-001-US-004-004: Archive session to graph
- REQ-001-US-004-005: Get active sessions for user
- REQ-001-US-004-006: Delete expired sessions
- REQ-001-US-004-007: Query session history from graph
- REQ-001-US-004-008: Session TTL and expiration
- REQ-001-US-004-009: Session persistence across restarts
- REQ-001-US-004-010: Session metadata management
"""

import json
import logging
import os
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from enum import Enum
from typing import TYPE_CHECKING, Any, cast

if TYPE_CHECKING:
    from redis import Redis as SyncRedis

import redis.asyncio as redis

from orchestrata_ultimate.core.base import BaseEntity
from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.graph import (
    AsyncGraphRepository,
    GraphRepositoryConfig,
    create_async_graph_repository,
)
from orchestrata_ultimate.core.graph_factory import resolve_uri
from orchestrata_ultimate.knowledge_graph.canonical_schema import (
    LABEL_PROJECT,
    LABEL_SESSION,
    REL_HAS_SESSION,
)

logger = logging.getLogger(__name__)

# Default session TTL: 24 hours in seconds
DEFAULT_SESSION_TTL_SECONDS = 86400

# The root Project slug the archive-session hub edge below MERGEs onto —
# same default every other KG writer in this codebase uses (e.g.
# factory/promotion/store.py's `_project_slug()`). Read at call time via
# `_project_slug()` (RUL-003), never cached at import time.
_DEFAULT_PROJECT_SLUG = "orchestrata-ultimate"


@traces(spec="001", story="US-004", reqs=["REQ-001-US-004-022"])
def _project_slug() -> str:
    """Resolve the root ``:Project`` slug the archive hub edge MERGEs onto."""
    return os.environ.get("ORCHESTRATA_PROJECT_ID", _DEFAULT_PROJECT_SLUG)


class SessionStatus(Enum):
    """Session status enumeration."""

    ACTIVE = "active"
    CLOSED = "closed"
    ARCHIVED = "archived"


@traces(spec="001", story="US-004", reqs=["REQ-001-US-004-001"])
@dataclass(eq=False)
class Session(BaseEntity):
    """Session data model.

    Inherits from BaseEntity for UUID identity, timestamps,
    serialization, and identity-based equality.

    The session_id field is the public-facing identifier; the
    inherited id field is set to session_id for BaseEntity compat.
    """

    session_id: str = ""
    user_id: str = ""
    start_time: datetime = field(default_factory=lambda: datetime.now(UTC))
    end_time: datetime | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    status: SessionStatus = SessionStatus.ACTIVE
    volks: list[str] = field(default_factory=list)

    def __post_init__(self) -> None:
        """Sync id with session_id for BaseEntity compatibility."""
        if self.session_id:
            self.id = self.session_id
        else:
            self.session_id = self.id

    def attach_volk(self, volk_alias: str) -> None:
        """Attach a Volk to this session.

        Args:
            volk_alias: Volk alias (e.g., "U:Ox", "U:PS")
        """
        if volk_alias not in self.volks:
            self.volks.append(volk_alias)
            logger.info(f"Attached {volk_alias} to session {self.session_id}")

    def detach_volk(self, volk_alias: str) -> None:
        """Detach a Volk from this session.

        Args:
            volk_alias: Volk alias to detach
        """
        if volk_alias in self.volks:
            self.volks.remove(volk_alias)
            logger.info(f"Detached {volk_alias} from session {self.session_id}")

    @traces(spec="006", story="US-007", reqs=["REQ-006-US-007-005"])
    def is_active(self) -> bool:
        """Check if session is active.

        Returns:
            True if session status is ACTIVE
        """
        return self.status == SessionStatus.ACTIVE

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-005"])
    def to_dict(self) -> dict[str, Any]:
        """Convert session to dictionary for serialization.

        Merges BaseEntity fields (id, created_at, updated_at) with
        Session-specific fields.

        Returns:
            Dictionary representation.
        """
        base = super().to_dict()
        base.update(
            {
                "session_id": self.session_id,
                "user_id": self.user_id,
                "start_time": self.start_time.isoformat(),
                "end_time": (self.end_time.isoformat() if self.end_time else None),
                "metadata": self.metadata,
                "status": self.status.value,
                "volks": self.volks,
            }
        )
        return base

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Session:
        """Create Session from dictionary.

        Args:
            data: Dictionary with session data

        Returns:
            Session instance

        Raises:
            ValueError: If required fields missing or datetime parsing fails
        """
        if not isinstance(data, dict):
            raise ValueError(f"Expected dict, got {type(data).__name__}")

        try:
            start_time = datetime.fromisoformat(data["start_time"])
        except (ValueError, KeyError) as e:
            raise ValueError(f"Invalid start_time format: {e}")

        end_time = None
        if data.get("end_time"):
            try:
                end_time = datetime.fromisoformat(data["end_time"])
            except ValueError as e:
                raise ValueError(f"Invalid end_time format: {e}")

        try:
            return cls(
                session_id=data["session_id"],
                user_id=data["user_id"],
                start_time=start_time,
                end_time=end_time,
                metadata=data.get("metadata", {}),
                status=SessionStatus(data.get("status", "active")),
                volks=data.get("volks", []),
            )
        except KeyError as e:
            raise ValueError(f"Missing required field in session data: {e}")
        except ValueError as e:
            raise ValueError(f"Invalid session data: {e}")


class SessionManager:
    """Manages session lifecycle and persistence.

    Implements REQ-001-US-004-001 through REQ-001-US-004-010 for session management.
    Uses Redis for hot storage and graph for archival.
    """

    def __init__(
        self,
        redis_host: str | None = None,
        redis_port: int | None = None,
        redis_db: int = 0,
        graph_uri: str | None = None,
        graph_user: str | None = None,
        graph_password: str | None = None,
        redis_client: SyncRedis | None = None,
    ) -> None:
        """Initialize SessionManager.

        Args:
            redis_host: Redis server hostname
            redis_port: Redis server port
            redis_db: Redis database number
            graph_uri: Graph connection URI
            graph_user: Graph username
            graph_password: Graph password
            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.

        Raises:
            Exception: If Redis or graph connection fails
        """
        self.redis_client = None
        self.graph_repo: AsyncGraphRepository | None = None

        try:
            # Redis connection for hot storage
            from orchestrata_ultimate.core.redis_factory import create_sync_redis

            self.redis_client = (
                redis_client
                if redis_client is not None
                else create_sync_redis(
                    host=redis_host,
                    port=redis_port,
                    db=redis_db,
                )
            )

            # Test Redis connection
            self.redis_client.ping()
            logger.info(
                "Connected to Redis at %s:%s",
                redis_host or os.environ.get("REDIS_HOST"),
                redis_port or os.environ.get("REDIS_PORT"),
            )

            # Graph connection parameters (driver created in initialize()).
            # Resolved at CALL time with no fallback values: localhost/
            # "password" defaults only mask misconfiguration with dead
            # endpoints in this infra — unset env raises here instead.
            # resolve_uri() reads KG_ROUTING_URI/KG_URI and raises with the key
            # name when unset — the same fail-loud contract, on the keys that
            # still exist. NEO4J_URI/USERNAME/PASSWORD were retired with the
            # Neo4j containers on 2026-07-27.
            self._graph_uri = graph_uri or resolve_uri()
            # No credentials: every remaining endpoint negotiates no auth.
            self._graph_user = graph_user or ""
            self._graph_password = graph_password or ""
            logger.info(f"Graph connection configured for {self._graph_uri}")

        except Exception as e:
            # Clean up any partially initialized resources
            if self.redis_client is not None:
                try:
                    self.redis_client.close()
                except Exception as cleanup_error:
                    logger.warning(f"Error during Redis cleanup: {cleanup_error}")

            if self.graph_repo is not None:  # pragma: no cover
                try:  # pragma: no cover
                    self.graph_repo.close()  # pragma: no cover
                except Exception as cleanup_error:  # pragma: no cover
                    logger.warning(
                        f"Error during graph cleanup: {cleanup_error}"
                    )  # pragma: no cover

            logger.error(f"Failed to initialize SessionManager: {e}")
            raise

    @traces(spec="001", reqs=["REQ-001-US-004-009"])
    async def initialize(self) -> None:
        """Initialize async graph driver and verify connectivity.

        Must be called after __init__ to set up the async driver.

        Raises:
            Exception: If graph connection fails
        """
        try:
            self.graph_repo = create_async_graph_repository(
                GraphRepositoryConfig(
                    uri=self._graph_uri,
                    auth=(self._graph_user, self._graph_password),
                )
            )
            await self.graph_repo.verify_connectivity()
            logger.info(f"Connected to graph at {self._graph_uri}")
        except Exception as e:
            logger.error(f"Failed to initialize graph driver: {e}")
            raise

    @property
    def _redis(self) -> SyncRedis:
        """Non-Optional accessor for the sync Redis client.

        ``redis_client`` is declared ``Redis | None`` because ``__init__``
        sets it to ``None`` up front so the except-block can null-check it
        during cleanup on a failed connection. Every other method in this
        class runs only after ``__init__`` has returned successfully, at
        which point the client is always connected.
        """
        if self.redis_client is None:
            raise RuntimeError("SessionManager Redis client not initialized")
        return cast("SyncRedis", self.redis_client)

    @property
    def _graph(self) -> AsyncGraphRepository:
        """Non-Optional accessor for the owned async graph repository.

        ``graph_repo`` stays ``None`` until ``initialize()`` completes;
        the methods that use this (``archive_session``,
        ``query_session_history``) are only meant to be called afterward.
        """
        if self.graph_repo is None:
            raise RuntimeError(
                "SessionManager graph driver not initialized; call initialize() first"
            )
        return self.graph_repo

    @traces(
        spec="001",
        story="US-004",
        reqs=[
            "REQ-001-US-004-001",
            "REQ-001-US-004-005",
        ],
    )
    def create_session(
        self, user_id: str, metadata: dict[str, Any] | None = None
    ) -> str:
        """Create a new session (REQ-001-US-004-001).

        Args:
            user_id: User identifier
            metadata: Optional session metadata

        Returns:
            Unique session ID (UUID)
        """
        session_id = str(uuid.uuid4())
        session = Session(
            session_id=session_id,
            user_id=user_id,
            start_time=datetime.now(UTC),
            metadata=metadata or {},
            status=SessionStatus.ACTIVE,
        )

        # Store in Redis
        key = f"session:{session_id}"
        try:
            self._redis.set(key, json.dumps(session.to_dict()))
        except redis.RedisError as e:
            logger.error(f"Failed to store session {session_id} in Redis: {e}")
            raise

        # Add to user's active sessions set
        user_sessions_key = f"user_sessions:{user_id}"
        try:
            self._redis.sadd(user_sessions_key, session_id)
        except redis.RedisError as e:
            logger.error(f"Failed to add session {session_id} to user sessions: {e}")
            # Rollback session creation
            self._redis.delete(key)
            raise

        logger.info(f"Created session {session_id} for user {user_id}")
        return session_id

    @traces(spec="001", reqs=["REQ-001-US-004-009"])
    def resume_session(self, session_id: str) -> Session:
        """Resume an existing session from Redis (REQ-001-US-004-002).

        Args:
            session_id: Session ID to resume

        Returns:
            Session object

        Raises:
            ValueError: If session not found
        """
        key = f"session:{session_id}"
        try:
            # get() is typed as redis-py's shared ResponseT (Awaitable[Any] |
            # Any); this sync client with decode_responses=True always
            # returns str or None.
            data = cast("str | None", self._redis.get(key))
        except redis.RedisError as e:
            logger.error(f"Failed to retrieve session {session_id} from Redis: {e}")
            raise

        if not data:
            raise ValueError(f"Session {session_id} not found")

        try:
            session = Session.from_dict(json.loads(data))
        except (json.JSONDecodeError, ValueError) as e:
            logger.error(f"Failed to parse session {session_id} data: {e}")
            raise ValueError(f"Invalid session data for {session_id}: {e}")

        logger.info(f"Resumed session {session_id}")
        return session

    @traces(
        spec="001",
        story="US-004",
        reqs=[
            "REQ-001-US-004-021",
            "REQ-001-US-004-023",
        ],
    )
    def close_session(self, session_id: str) -> None:
        """Close a session and mark it inactive (REQ-001-US-004-003).

        Args:
            session_id: Session ID to close
        """
        session = self.resume_session(session_id)
        session.status = SessionStatus.CLOSED
        session.end_time = datetime.now(UTC)

        # Update in Redis
        key = f"session:{session_id}"
        self._redis.set(key, json.dumps(session.to_dict()))

        # Remove from user's active sessions
        user_sessions_key = f"user_sessions:{session.user_id}"
        self._redis.srem(user_sessions_key, session_id)

        logger.info(f"Closed session {session_id}")

    @traces(
        spec="001",
        story="US-004",
        reqs=[
            "REQ-001-US-004-022",
            "REQ-001-US-004-026",
        ],
    )
    async def archive_session(self, session_id: str) -> None:
        """Archive a session to graph (REQ-001-US-004-004).

        Args:
            session_id: Session ID to archive
        """
        session = self.resume_session(session_id)

        # Archive to graph. Session-254 root-connectivity ruling: this
        # writer previously MERGEd `:Session` with no edge back to the root
        # `:Project` node (13 of the 25 orphaned Sessions
        # scripts/check_kg_connectivity.py found) -- the Project hub MERGE
        # below closes that gap.
        await self._graph.execute_write(
            f"""
            MERGE (s:{LABEL_SESSION} {{session_id: $session_id}})
            SET s.user_id = $user_id,
                s.start_time = $start_time,
                s.end_time = $end_time,
                s.metadata = $metadata,
                s.status = $status,
                s.volks = $volks
            WITH s
            MERGE (p:{LABEL_PROJECT} {{slug: $project}})
            MERGE (p)-[:{REL_HAS_SESSION}]->(s)
            """,
            {
                "session_id": session.session_id,
                "user_id": session.user_id,
                "start_time": session.start_time.isoformat(),
                "end_time": session.end_time.isoformat() if session.end_time else None,
                "metadata": json.dumps(session.metadata),
                "status": session.status.value,
                "volks": session.volks,
                "project": _project_slug(),
            },
        )

        # SPEC-008 gap #2: ``volks`` above is a flat string-list PROPERTY
        # (kept for backward-compat reads); also materialize real
        # WORKED_ON / WORKED_ON_BY edges to existing Volk nodes so the
        # provenance is graph-traversable. Only link Volks that already
        # exist -- never MERGE-create a Volk from a list entry.
        if session.volks:
            await self._graph.execute_write(
                """
                MATCH (s:Session {session_id: $session_id})
                UNWIND $volks AS volk_alias
                OPTIONAL MATCH (v:Volk {alias: volk_alias})
                FOREACH (_ IN CASE WHEN v IS NOT NULL THEN [1] ELSE [] END |
                    MERGE (v)-[:WORKED_ON]->(s)
                    MERGE (s)-[:WORKED_ON_BY]->(v)
                )
                """,
                {
                    "session_id": session.session_id,
                    "volks": session.volks,
                },
            )

        # Update status in Redis
        session.status = SessionStatus.ARCHIVED
        key = f"session:{session_id}"
        self._redis.set(key, json.dumps(session.to_dict()))

        logger.info(f"Archived session {session_id} to graph")

    @traces(spec="006", story="US-007", reqs=["REQ-006-US-007-005"])
    def get_active_sessions(self, user_id: str) -> list[Session]:
        """Get all active sessions for a user (REQ-001-US-004-005).

        Args:
            user_id: User identifier

        Returns:
            List of active Session objects
        """
        user_sessions_key = f"user_sessions:{user_id}"
        # smembers() is typed as Awaitable[Set] | Set (shared sync/async
        # command mixin); this sync client always returns a plain set of
        # decoded session-id strings.
        session_ids = cast("set[str]", self._redis.smembers(user_sessions_key))

        sessions = []
        for session_id in session_ids:
            try:
                session = self.resume_session(session_id)
                if session.is_active():
                    sessions.append(session)
            except ValueError:
                # Session not found, remove from set
                self._redis.srem(user_sessions_key, session_id)

        logger.info(f"Retrieved {len(sessions)} active sessions for user {user_id}")
        return sessions

    async def query_session_history(
        self,
        user_id: str,
        limit: int = 100,
        status: SessionStatus | None = None,
    ) -> list[dict[str, Any]]:
        """Query session history from graph (REQ-001-US-004-007).

        Args:
            user_id: User identifier
            limit: Maximum number of sessions to return
            status: Optional status filter

        Returns:
            List of session dictionaries
        """
        if status:
            query = """
            MATCH (s:Session {user_id: $user_id, status: $status})
            RETURN s
            ORDER BY s.start_time DESC
            LIMIT $limit
            """
            params = {"user_id": user_id, "status": status.value, "limit": limit}
        else:
            query = """
            MATCH (s:Session {user_id: $user_id})
            RETURN s
            ORDER BY s.start_time DESC
            LIMIT $limit
            """
            params = {"user_id": user_id, "limit": limit}

        rows = await self._graph.execute_read(query, params)

        sessions = []
        for row in rows:
            session_data = dict(row["s"])
            sessions.append(session_data)

        logger.info(f"Retrieved {len(sessions)} historical sessions for user {user_id}")
        return sessions

    def _fetch_all_session_data(self, keys: list) -> list:
        """Batch-fetch raw session data from Redis for the given keys.

        Args:
            keys: List of Redis key strings to fetch.

        Returns:
            List of raw data values in the same order as keys.

        Raises:
            redis.RedisError: If Redis pipeline execution fails.
        """
        pipeline = self._redis.pipeline()
        for key in keys:
            pipeline.get(key)
        try:
            return pipeline.execute()
        except redis.RedisError as e:
            logger.error(f"Failed to batch retrieve sessions: {e}")
            raise

    def _collect_expired_sessions(
        self, keys: list, results: list, cutoff_time: datetime
    ) -> list:
        """Identify expired sessions from batch-fetched Redis data.

        Args:
            keys: Redis key strings corresponding to results.
            results: Raw data values from Redis pipeline.
            cutoff_time: Sessions started before this time are expired.

        Returns:
            List of (key, session_id, user_id) tuples for expired sessions.
        """
        expired_sessions = []
        for key, data in zip(keys, results):
            if data:
                try:
                    session_data = json.loads(data)
                    start_time = datetime.fromisoformat(session_data["start_time"])
                    if start_time < cutoff_time:
                        expired_sessions.append(
                            (key, session_data["session_id"], session_data["user_id"])
                        )
                except (json.JSONDecodeError, ValueError, KeyError) as e:
                    logger.warning(f"Skipping invalid session data in {key}: {e}")
        return expired_sessions

    def _batch_delete_expired_sessions(self, expired_sessions: list) -> int:
        """Delete expired sessions from Redis in a single pipeline call.

        Args:
            expired_sessions: List of (key, session_id, user_id) tuples.

        Returns:
            Number of sessions deleted.

        Raises:
            redis.RedisError: If Redis pipeline execution fails.
        """
        if not expired_sessions:
            return 0
        delete_pipeline = self._redis.pipeline()
        for key, session_id, user_id in expired_sessions:
            delete_pipeline.delete(key)
            delete_pipeline.srem(f"user_sessions:{user_id}", session_id)
        try:
            delete_pipeline.execute()
            return len(expired_sessions)
        except redis.RedisError as e:
            logger.error(f"Failed to batch delete sessions: {e}")
            raise

    def delete_expired_sessions(
        self, ttl_seconds: int = DEFAULT_SESSION_TTL_SECONDS
    ) -> int:
        """Delete sessions older than TTL (REQ-001-US-004-006, REQ-001-US-004-008).

        Args:
            ttl_seconds: Time-to-live in seconds (default: DEFAULT_SESSION_TTL_SECONDS)

        Returns:
            Number of sessions deleted
        """
        cutoff_time = datetime.now(UTC) - timedelta(seconds=ttl_seconds)

        # Scan all session keys — collect once to avoid double-scan mismatch
        keys = list(self._redis.scan_iter("session:*"))
        results = self._fetch_all_session_data(keys)
        expired_sessions = self._collect_expired_sessions(keys, results, cutoff_time)
        deleted_count = self._batch_delete_expired_sessions(expired_sessions)

        logger.info(f"Deleted {deleted_count} expired sessions (TTL: {ttl_seconds}s)")
        return deleted_count

    @traces(spec="001", reqs=["REQ-001-US-004-009"])
    async def close(self) -> None:
        """Close all connections."""
        self._redis.close()
        if self.graph_repo:
            await self.graph_repo.close()
        logger.info("Closed SessionManager connections")