orchestrata_ultimate.gateway.session

Reference › orchestrata_ultimate › gateway › session

session holds ConnectionManager, the registry of live WebSocket connections (SPEC-014 US-007). It is not the same as the conversation-session feature served by the session handler. It is built from config and used by server, protocol, rpc and some handlers.

register creates a ConnectionInfo in the ACTIVE state under an asyncio lock, or re-activates an existing record, and updates the connection gauge and counter in metrics. unregister removes a record and returns it, or None. get and get_active_connections read records, get_user_connections filters them by user, and connection_count returns the total. update_heartbeat validates the id, stamps the record with the current UTC time and returns False for an unknown connection. See gateway.

At a glance

Class diagram

Dependency diagram

API

Classes

ConnectionManager

Line 18

Manages WebSocket connection lifecycle.

Source

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

"""WebSocket connection lifecycle management."""

import asyncio
import logging
from datetime import UTC, datetime

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.gateway.config import GatewayConfig
from orchestrata_ultimate.gateway.metrics import connection_closed, connection_opened
from orchestrata_ultimate.gateway.models import (
    ConnectionInfo,
    ConnectionState,
)

logger = logging.getLogger(__name__)


class ConnectionManager:
    """Manages WebSocket connection lifecycle."""

    def __init__(self, config: GatewayConfig) -> None:
        """Initialise with gateway configuration."""
        self._config = config
        self._connections: dict[str, ConnectionInfo] = {}
        self._lock = asyncio.Lock()

    @traces(
        spec="014",
        story="US-007",
        reqs=["REQ-014-112"],
    )
    async def register(self, connection_id: str, user_id: str) -> ConnectionInfo:
        """Register a new active connection.

        Args:
            connection_id: Unique connection identifier.
            user_id: Authenticated user identifier.

        Returns:
            The newly created ConnectionInfo.

        Raises:
            ConnectionLifecycleError: If already registered.
        """
        async with self._lock:
            if connection_id in self._connections:
                logger.debug("Re-registering existing connection %s", connection_id)
                self._connections[connection_id].state = ConnectionState.ACTIVE
                return self._connections[connection_id]
            info = ConnectionInfo(
                connection_id=connection_id,
                user_id=user_id,
                state=ConnectionState.ACTIVE,
            )
            self._connections[connection_id] = info
            logger.info(
                "Registered %s for %s",
                connection_id,
                user_id,
            )
            # REQ-014-142/143: gauge of live connections + running total
            # accepted. Recorded inside the lock so it stays consistent
            # with the dict mutation it mirrors.
            connection_opened()
            return info

    @traces(spec="014", story="US-011", reqs=["REQ-014-168"])
    async def unregister(self, connection_id: str) -> ConnectionInfo | None:
        """Unregister and mark connection closed.

        Args:
            connection_id: Connection to remove.

        Returns:
            The closed ConnectionInfo, or None if not found.

        Raises:
            TypeError: If connection_id is not a str.
            ValueError: If connection_id is empty.
        """
        if not isinstance(connection_id, str):
            raise TypeError("connection_id must be a str")
        if not connection_id:
            raise ValueError("connection_id must not be empty")
        async with self._lock:
            info = self._connections.pop(connection_id, None)
            if info:
                info.state = ConnectionState.CLOSED
                logger.info("Unregistered %s", connection_id)
                # REQ-014-142: drop the live-connections gauge.
                connection_closed()
            return info

    async def get(self, connection_id: str) -> ConnectionInfo | None:
        """Look up connection info by ID.

        Args:
            connection_id: Connection to look up.

        Returns:
            ConnectionInfo or None if not found.

        Raises:
            TypeError: If connection_id is not a str.
            ValueError: If connection_id is empty.
        """
        if not isinstance(connection_id, str):
            raise TypeError("connection_id must be a str")
        if not connection_id:
            raise ValueError("connection_id must not be empty")
        return self._connections.get(connection_id)

    @traces(
        spec="014",
        story="US-007",
        reqs=["REQ-014-109"],
    )
    async def update_heartbeat(self, connection_id: str) -> bool:
        """Update last-heartbeat timestamp.

        Args:
            connection_id: Connection to update.

        Returns:
            True if updated, False if connection not found.

        Raises:
            TypeError: If connection_id is not a str.
            ValueError: If connection_id is empty.
        """
        if not isinstance(connection_id, str):
            raise TypeError("connection_id must be a str")
        if not connection_id:
            raise ValueError("connection_id must not be empty")
        info = self._connections.get(connection_id)
        if not info:
            return False
        info.last_heartbeat = datetime.now(UTC).isoformat()
        return True

    @traces(
        spec="014",
        story="US-007",
        reqs=["REQ-014-111"],
    )
    async def get_active_connections(
        self,
    ) -> list[ConnectionInfo]:
        """Return all connections in ACTIVE state."""
        # Snapshot under the lock — register/unregister mutate
        # _connections; iterating .values() without the lock can raise
        # "dictionary changed size during iteration".
        async with self._lock:
            conns = list(self._connections.values())
        return [c for c in conns if c.state == ConnectionState.ACTIVE]

    async def get_user_connections(self, user_id: str) -> list[ConnectionInfo]:
        """Return active connections for a user.

        Args:
            user_id: User to filter by.

        Returns:
            List of active ConnectionInfo for the user.

        Raises:
            TypeError: If user_id is not a str.
            ValueError: If user_id is empty.
        """
        if not isinstance(user_id, str):
            raise TypeError("user_id must be a str")
        if not user_id:
            raise ValueError("user_id must not be empty")
        async with self._lock:
            conns = list(self._connections.values())
        return [
            c
            for c in conns
            if c.user_id == user_id and c.state == ConnectionState.ACTIVE
        ]

    @property
    def connection_count(self) -> int:
        """Total tracked connections."""
        return len(self._connections)