# Copyright 2025 Softwell S.r.l.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Registry for active WSX WebSocket connections.
Tracks connected clients, supports lookup and broadcast.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Iterator
from .protocol import build_wsx_response
if TYPE_CHECKING:
from ..websocket import WebSocket
from ..types import Scope
__all__ = ["WsxRegistry", "WsxConnectionInfo"]
logger = logging.getLogger("genro_asgi.wsx")
[docs]
class WsxConnectionInfo:
"""Metadata for a registered WSX connection."""
__slots__ = ("connection_id", "websocket", "scope", "auth")
[docs]
def __init__(
self,
connection_id: str,
websocket: WebSocket,
scope: Scope,
auth: dict[str, Any] | None = None,
) -> None:
"""Args:
connection_id: Unique identifier for this connection.
websocket: Active WebSocket instance.
scope: ASGI scope dict for this connection.
auth: Authentication result dict (identity, tags), or None.
"""
self.connection_id = connection_id
self.websocket = websocket
self.scope = scope
self.auth = auth
@property
def identity(self) -> str | None:
"""Authenticated user identity, if available."""
if self.auth is None:
return None
return self.auth.get("identity")
@property
def client(self) -> tuple[str, int] | None:
"""Client address as (host, port)."""
return self.scope.get("client")
[docs]
class WsxRegistry:
"""Registry of active WSX WebSocket connections."""
__slots__ = ("_connections",)
def __init__(self) -> None:
self._connections: dict[str, WsxConnectionInfo] = {}
[docs]
def register(
self,
connection_id: str,
websocket: WebSocket,
scope: Scope,
auth: dict[str, Any] | None = None,
) -> WsxConnectionInfo:
"""Register a new WSX connection."""
info = WsxConnectionInfo(connection_id, websocket, scope, auth)
self._connections[connection_id] = info
logger.debug("WSX connection registered: %s", connection_id)
return info
[docs]
def unregister(self, connection_id: str) -> WsxConnectionInfo | None:
"""Unregister a WSX connection. Returns info if found."""
info = self._connections.pop(connection_id, None)
if info is not None:
logger.debug("WSX connection unregistered: %s", connection_id)
return info
[docs]
def get(self, connection_id: str) -> WsxConnectionInfo | None:
"""Lookup connection by ID."""
return self._connections.get(connection_id)
[docs]
def find_by_identity(self, identity: str) -> list[WsxConnectionInfo]:
"""Find all connections for a given authenticated identity."""
return [c for c in self._connections.values() if c.identity == identity]
[docs]
async def broadcast(self, data: Any, *, exclude: str | None = None) -> int:
"""Send a WSX response to all connected clients.
Args:
data: Payload to send.
exclude: Connection ID to skip.
Returns:
Number of clients the message was sent to.
"""
msg = build_wsx_response(id="_broadcast", status=200, data=data)
sent = 0
for cid, info in list(self._connections.items()):
if cid == exclude:
continue
try:
await info.websocket.send_text(msg)
sent += 1
except Exception:
logger.warning("Broadcast failed for connection %s", cid)
return sent
[docs]
async def send_to(self, identity: str, data: Any) -> int:
"""Send a WSX response to all connections of a given identity.
Returns:
Number of clients the message was sent to.
"""
msg = build_wsx_response(id="_directed", status=200, data=data)
sent = 0
for info in self.find_by_identity(identity):
try:
await info.websocket.send_text(msg)
sent += 1
except Exception:
logger.warning(
"Send to %s failed for connection %s",
identity,
info.connection_id,
)
return sent
def __len__(self) -> int:
return len(self._connections)
def __iter__(self) -> Iterator[WsxConnectionInfo]:
return iter(self._connections.values())
def __contains__(self, connection_id: str) -> bool:
return connection_id in self._connections
if __name__ == "__main__":
registry = WsxRegistry()
print(f"Empty registry: {len(registry)} connections")