orchestrata_ultimate.gateway.models

Reference › orchestrata_ultimate › gateway › models

models holds the data types that cross the gateway's wire protocol; the VizInt front end depends on their values, so changing them changes the protocol. It is imported by protocol, pubsub, session, auth_handler and mcp_transport, and re-exported by the package.

• MessageType has six values: RPC_REQUEST, RPC_RESPONSE, EVENT, ERROR, HEARTBEAT and AUTH. Its docstring maps the SPEC-098 FR-020 message kinds onto them: ping and pong share HEARTBEAT, and subscribe and unsubscribe are RPC methods, not message types.
• WSMessage is the frame envelope with a type, payload, message id, correlation id and timestamp; to_dict and from_dict convert it to and from JSON.
• ConnectionState lists CONNECTING, AUTHENTICATING, ACTIVE, CLOSING and CLOSED.
• ConnectionInfo is the per-socket record the connection manager keeps: connection id, user id, state and last heartbeat.
• Subscription is one pub/sub interest: connection id, topic pattern, key-value filters and a subscription id.

The private helpers generate UTC ISO timestamps and ids. See gateway.

At a glance

Class diagram

Dependency diagram

API

Classes

MessageType (enum)

Bases: Enum · Line 27

WebSocket message types.

WSMessage (dataclass)

Decorators: dataclass · Line 77

WebSocket protocol message.

ConnectionState (enum)

Bases: Enum · Line 121

WebSocket connection states.

ConnectionInfo (dataclass)

Decorators: traces(spec='014', story='US-007', reqs=['REQ-014-112']), dataclass · Line 133

Tracks a single WebSocket connection.

Subscription (dataclass)

Decorators: dataclass · Line 159

Tracks a single pub/sub subscription.

Functions

Source

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

"""Data models for the WebSocket gateway protocol."""

import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum
from typing import Any

from orchestrata_ultimate.core.base.traces import traces


def _utc_iso() -> str:
    """Current UTC time as ISO string."""
    return datetime.now(UTC).isoformat()


def _new_id() -> str:
    """Generate a new UUID string."""
    return str(uuid.uuid4())


def _new_sub_id() -> str:
    """Generate a subscription ID."""
    return f"sub-{uuid.uuid4()}"


class MessageType(Enum):
    """WebSocket message types.

    SPEC-098 FR-020 names six message types — subscribe, unsubscribe, event,
    error, ping, pong. This gateway carries all six semantics, but not all of
    them as message TYPES, because it is JSON-RPC over WebSocket rather than a
    bare message bus. The mapping, so the difference is recorded rather than
    rediscovered:

        FR-020 type   how this gateway provides it
        -----------   ---------------------------------------------------
        event         EVENT — identical
        error         ERROR — identical
        ping / pong   HEARTBEAT, one type for both directions; the JSON-RPC
                      envelope already distinguishes request from response,
                      so a second type would encode the same bit twice
        subscribe     a METHOD, not a type: RPC_REQUEST carrying the bare
                      "subscribe" — subscription is a call with parameters
                      (topic, filters) and a result, which a bare message
                      type cannot express.
                      MEASURED 2026-08-12: the registered key is "subscribe";
                      "gateway.subscribe" is NOT registered and dispatch
                      raises RPCError -32601 "Method not found". This
                      docstring named the namespaced form for months and is
                      the only place in the repo that string appears — it
                      seeded a wrong client contract. gateway/rpc.py
                      _build_method_registry lists subscribe/unsubscribe
                      ONLY under "Legacy short names (backward compat)",
                      unlike ping, which carries both "gateway.ping" and
                      "ping". Adding the missing namespaced aliases is owed
                      work, not something this docstring may assume.
        unsubscribe   likewise, the bare "unsubscribe"

    AUTH and RPC_REQUEST/RPC_RESPONSE have no FR-020 counterpart because the
    spec assumes REST for the request/response half and a WebSocket only for
    events; here both travel the same socket.

    Changing these values changes the wire protocol the vizint-ui front end
    speaks — they are a published contract, not an internal enum.
    """

    RPC_REQUEST = "rpc_request"
    RPC_RESPONSE = "rpc_response"
    EVENT = "event"
    ERROR = "error"
    HEARTBEAT = "heartbeat"
    AUTH = "auth"


@dataclass
class WSMessage:
    """WebSocket protocol message."""

    msg_type: MessageType
    payload: dict[str, Any] = field(default_factory=dict)
    msg_id: str = field(default_factory=_new_id)
    correlation_id: str = ""
    timestamp: str = field(default_factory=_utc_iso)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "msg_type": self.msg_type.value,
            "payload": self.payload,
            "msg_id": self.msg_id,
            "correlation_id": self.correlation_id,
            "timestamp": self.timestamp,
        }

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

        Args:
            data: Dictionary with msg_type and optional
                payload, msg_id, correlation_id, timestamp.

        Returns:
            Reconstructed WSMessage instance.

        Raises:
            TypeError: If data is not a dict.
        """
        if not isinstance(data, dict):
            raise TypeError("data must be a dict")
        return cls(
            msg_type=MessageType(data["msg_type"]),
            payload=data.get("payload", {}),
            msg_id=data.get("msg_id", _new_id()),
            correlation_id=data.get("correlation_id", ""),
            timestamp=data.get("timestamp", _utc_iso()),
        )


class ConnectionState(Enum):
    """WebSocket connection states."""

    CONNECTING = "connecting"
    AUTHENTICATING = "authenticating"
    ACTIVE = "active"
    CLOSING = "closing"
    CLOSED = "closed"


@traces(spec="014", story="US-007", reqs=["REQ-014-112"])
@dataclass
class ConnectionInfo:
    """Tracks a single WebSocket connection."""

    connection_id: str = field(default_factory=_new_id)
    user_id: str = ""
    session_id: str = ""
    state: ConnectionState = ConnectionState.CONNECTING
    connected_at: str = field(default_factory=_utc_iso)
    last_heartbeat: str = field(default_factory=_utc_iso)
    subscriptions: list[str] = field(default_factory=list)

    @traces(spec="014", story="US-007", reqs=["REQ-014-111"])
    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "connection_id": self.connection_id,
            "user_id": self.user_id,
            "session_id": self.session_id,
            "state": self.state.value,
            "connected_at": self.connected_at,
            "last_heartbeat": self.last_heartbeat,
            "subscriptions": list(self.subscriptions),
        }


@dataclass
class Subscription:
    """Tracks a single pub/sub subscription."""

    subscription_id: str = field(default_factory=_new_sub_id)
    connection_id: str = ""
    topic: str = ""
    filters: dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=_utc_iso)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "subscription_id": self.subscription_id,
            "connection_id": self.connection_id,
            "topic": self.topic,
            "filters": dict(self.filters),
            "created_at": self.created_at,
        }