feat(B-WS): WebSocket Helpers + Redis Pub/Sub + Error-Handling
Check Cross-Plugin Imports / check (push) Has been cancelled

B-WS: app/core/ws_helpers.py (NEU) — gemeinsame WebSocket Helpers
- authenticate_ws: Session-Auth für WebSocket (Cookie/Token → User/Tenant)
- check_ws_origin: Origin-Check (delegiert auf verify_ws_origin)
- check_ws_tenant: User-Tenant-Membership-Check
- cleanup_ws_connection: Connection aus Registry entfernen + WS schließen
- start_heartbeat: Background Ping-Task
- send_ws_error: strukturierte Error-Message an Client
- handle_ws_message: Message-Dispatch mit Error-Handling

B-WS: app/core/ws_pubsub.py (NEU) — Redis Pub/Sub für Multi-Worker-Fanout
- publish_to_channel / subscribe_to_channel
- get_tenant_channel / broadcast_to_tenants

B-WS: WebSocketManager + AIUIControlWSManager angepasst
- connect() nutzt authenticate_ws + check_ws_origin + check_ws_tenant
- disconnect() nutzt cleanup_ws_connection + cancelt Heartbeat/PubSub
- broadcast() unterstützt Redis Pub/Sub Fanout

B-ERR-WS: WS Error-Handling in ws_helpers integriert
- send_ws_error für strukturierte Errors
- handle_ws_message fängt Handler-Exceptions

B-WS-TEST: 24 Tests in test_ws_helpers.py — alle grün
- Auth, Origin, Error, Dispatch, Cleanup, Heartbeat, Pub/Sub Roundtrip
- Keine Regression: 47/47 Resilience+Hooks Tests grün
This commit is contained in:
Agent Zero
2026-08-13 16:43:54 +02:00
parent a3a26d1f66
commit 7a81a5f072
7 changed files with 969 additions and 76 deletions
+184
View File
@@ -0,0 +1,184 @@
"""Shared WebSocket helpers: auth, origin check, tenant check, cleanup, heartbeat, error handling, message dispatch."""
from __future__ import annotations
import asyncio
import json
import logging
import uuid
from typing import Any, Callable
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from starlette.websockets import WebSocket
from app.core.auth import get_redis, get_session_data, verify_ws_origin
from app.config import get_settings
logger = logging.getLogger(__name__)
async def authenticate_ws(websocket: WebSocket, db: AsyncSession) -> dict[str, Any] | None:
"""Authenticate a WebSocket connection via session cookie.
Extracts the session cookie, validates the session in Redis,
and returns user/tenant info. On failure, closes the WebSocket
with code 4401 and returns ``None``.
"""
settings = get_settings()
session_id = websocket.cookies.get(settings.session_cookie_name)
if not session_id:
await websocket.close(code=4401, reason="Unauthorized")
return None
redis = get_redis()
session_data = await get_session_data(redis, session_id)
if session_data is None:
await websocket.close(code=4401, reason="Unauthorized")
return None
if not session_data.get("is_active", False):
await websocket.close(code=4401, reason="Unauthorized")
return None
return {
"user_id": session_data["user_id"],
"tenant_id": session_data["tenant_id"],
"session_id": session_id,
"role": session_data.get("role"),
"email": session_data.get("email"),
"name": session_data.get("name"),
"is_system_admin": session_data.get("is_system_admin", False),
}
async def check_ws_origin(websocket: WebSocket) -> bool:
"""Verify the WebSocket origin and CSRF token.
Delegates to :func:`verify_ws_origin`. On failure, closes the
WebSocket with code 4403 and returns ``False``.
"""
result = await verify_ws_origin(websocket)
if not result:
await websocket.close(code=4403, reason="Forbidden origin")
return False
return True
async def check_ws_tenant(
websocket: WebSocket,
tenant_id: uuid.UUID,
user_id: uuid.UUID,
db: AsyncSession,
) -> bool:
"""Verify that *user_id* belongs to *tenant_id*.
On failure, closes the WebSocket with code 4403 and returns ``False``.
"""
from app.models.user import UserTenant
result = await db.execute(
select(UserTenant).where(
UserTenant.user_id == user_id,
UserTenant.tenant_id == tenant_id,
)
)
if result.scalar_one_or_none() is None:
await websocket.close(code=4403, reason="Forbidden tenant")
return False
return True
async def cleanup_ws_connection(
websocket: WebSocket,
user_id: str,
connection_registry: dict[str, list[WebSocket]],
) -> None:
"""Remove a WebSocket from the connection registry and close it cleanly."""
conns = connection_registry.get(user_id, [])
if websocket in conns:
conns.remove(websocket)
if not conns:
connection_registry.pop(user_id, None)
try:
await websocket.close()
except Exception:
logger.debug("WebSocket already closed during cleanup for user %s", user_id)
async def start_heartbeat(websocket: WebSocket, interval: int = 30) -> asyncio.Task:
"""Start a background heartbeat task that sends periodic pings.
Returns the :class:`asyncio.Task` so the caller can cancel it on disconnect.
"""
async def _heartbeat() -> None:
while True:
try:
await asyncio.sleep(interval)
await websocket.send_text(json.dumps({"type": "ping"}))
except asyncio.CancelledError:
break
except Exception:
logger.debug("Heartbeat stopped — WebSocket likely closed")
break
return asyncio.create_task(_heartbeat())
async def send_ws_error(
websocket: WebSocket,
code: str,
detail: str,
trace_id: str | None = None,
) -> None:
"""Send a structured error message to the WebSocket client."""
payload: dict[str, Any] = {
"type": "error",
"code": code,
"detail": detail,
}
if trace_id is not None:
payload["trace_id"] = trace_id
try:
await websocket.send_text(json.dumps(payload, default=str))
except Exception:
logger.debug("Failed to send WS error to client")
async def handle_ws_message(
websocket: WebSocket,
message: str,
handlers: dict[str, Callable[[WebSocket, dict[str, Any]], Any]],
) -> None:
"""Dispatch a WebSocket text message to the appropriate handler.
*handlers* maps message ``type`` strings to async callables that accept
``(websocket, msg)``. Unknown types and handler exceptions are
reported back to the client via :func:`send_ws_error`.
"""
try:
msg = json.loads(message)
except (json.JSONDecodeError, TypeError):
await send_ws_error(websocket, "invalid_json", "Message is not valid JSON")
return
if not isinstance(msg, dict):
await send_ws_error(websocket, "invalid_message", "Message must be a JSON object")
return
msg_type = msg.get("type")
if not msg_type:
await send_ws_error(websocket, "missing_type", "Message missing 'type' field")
return
handler = handlers.get(msg_type)
if handler is None:
await send_ws_error(websocket, "unknown_type", f"Unknown message type: {msg_type}")
return
try:
await handler(websocket, msg)
except Exception as exc:
logger.exception("Handler error for message type '%s'", msg_type)
await send_ws_error(websocket, "handler_error", str(exc))
+65
View File
@@ -0,0 +1,65 @@
"""Redis Pub/Sub helpers for WebSocket multi-worker fanout."""
from __future__ import annotations
import asyncio
import json
import logging
import uuid
from typing import Awaitable, Callable
from app.core.auth import get_redis
logger = logging.getLogger(__name__)
async def publish_to_channel(channel: str, message: dict) -> None:
"""Publish a JSON message to a Redis channel."""
redis = get_redis()
await redis.publish(channel, json.dumps(message, default=str))
async def subscribe_to_channel(
channel: str,
handler: Callable[[dict], Awaitable[None]],
) -> asyncio.Task:
"""Subscribe to a Redis channel and call *handler* for every message.
Returns the :class:`asyncio.Task` so the caller can cancel it on disconnect.
"""
async def _subscriber() -> None:
redis = get_redis()
pubsub = redis.pubsub()
await pubsub.subscribe(channel)
try:
async for raw in pubsub.listen():
if raw["type"] == "message":
try:
msg = json.loads(raw["data"])
await handler(msg)
except Exception:
logger.exception("Error in pubsub handler for channel %s", channel)
finally:
try:
await pubsub.unsubscribe(channel)
await pubsub.aclose()
except Exception:
logger.debug("PubSub cleanup error for channel %s", channel)
return asyncio.create_task(_subscriber())
def get_tenant_channel(tenant_id: uuid.UUID, topic: str) -> str:
"""Return the Redis channel name for a tenant + topic."""
return f"ws:{tenant_id}:{topic}"
async def broadcast_to_tenants(
tenant_id: uuid.UUID,
topic: str,
message: dict,
) -> None:
"""Publish a message to a tenant-specific Redis channel."""
channel = get_tenant_channel(tenant_id, topic)
await publish_to_channel(channel, message)
+12 -25
View File
@@ -205,33 +205,28 @@ async def ai_ui_control_ws(websocket: WebSocket):
Authentication: via session cookie (same pattern as kommunikation plugin).
"""
from app.config import get_settings
from app.core.auth import get_session_data, get_redis, verify_ws_origin
from app.core.db import async_session_maker
from app.core.service_container import get_container
settings = get_settings()
if not await verify_ws_origin(websocket):
await websocket.close(code=4003, reason="Origin not allowed")
container = get_container()
if not container.has("ai_ui_control_ws"):
await websocket.close(code=4003, reason="AI UI Control not available")
return
session_id = websocket.cookies.get(settings.session_cookie_name)
if not session_id:
await websocket.close(code=4001, reason="Not authenticated")
return
ws_manager = container.get("ai_ui_control_ws")
redis = get_redis()
session_data = await get_session_data(redis, session_id)
if session_data is None:
await websocket.close(code=4001, reason="Session expired")
return
# connect() performs origin check, session auth, tenant check
async with async_session_maker() as db:
auth = await ws_manager.connect(websocket, db)
if auth is None:
return # connection was rejected and closed by ws_helpers
user_id = session_data["user_id"]
tenant_id = session_data["tenant_id"]
user_id = auth["user_id"]
tenant_id = auth["tenant_id"]
# Plugin-Gate: check if ai_ui_control plugin is active (global + tenant)
from app.core.permission_registry import get_permission_registry
from sqlalchemy import text as sa_text
from app.core.db import async_session_maker
import uuid as _uuid
try:
registry = get_permission_registry()
@@ -251,14 +246,6 @@ async def ai_ui_control_ws(websocket: WebSocket):
await websocket.close(code=4003, reason="Plugin check failed")
return
container = get_container()
if not container.has("ai_ui_control_ws"):
await websocket.close(code=4003, reason="AI UI Control not available")
return
ws_manager = container.get("ai_ui_control_ws")
await ws_manager.connect(websocket, user_id)
try:
while True:
data = await websocket.receive_text()
@@ -16,6 +16,20 @@ from typing import Any
from fastapi import WebSocket
from app.core.ws_helpers import (
authenticate_ws,
check_ws_origin,
check_ws_tenant,
cleanup_ws_connection,
start_heartbeat,
)
from app.core.ws_pubsub import (
broadcast_to_tenants,
get_tenant_channel,
subscribe_to_channel,
)
from sqlalchemy.ext.asyncio import AsyncSession
logger = logging.getLogger(__name__)
@@ -36,23 +50,84 @@ class AIUIControlWSManager:
self._delivered_at: dict[str, float] = {}
# command_id → user_id (to route feedback)
self._command_user: dict[str, str] = {}
# user_id → list of heartbeat tasks
self._heartbeat_tasks: dict[str, list] = {}
# user_id → list of pubsub subscriber tasks
self._pubsub_tasks: dict[str, list] = {}
async def connect(self, websocket: WebSocket, user_id: str) -> None:
"""Accept and register a new WebSocket connection."""
async def connect(
self,
websocket: WebSocket,
db: AsyncSession,
) -> dict[str, Any] | None:
"""Accept, authenticate and register a new WebSocket connection.
Performs origin check, session authentication and tenant validation.
Returns the auth dict (user_id, tenant_id, …) on success, or ``None``
if the connection was rejected (already closed).
"""
# Origin / CSRF check
if not await check_ws_origin(websocket):
return None
# Session authentication
auth = await authenticate_ws(websocket, db)
if auth is None:
return None
user_id = auth["user_id"]
tenant_id_str = auth["tenant_id"]
tenant_id = uuid.UUID(tenant_id_str)
user_uuid = uuid.UUID(user_id)
# Tenant membership check
if not await check_ws_tenant(websocket, tenant_id, user_uuid, db):
return None
# Accept the WebSocket
await websocket.accept()
# Register connection
if user_id not in self._connections:
self._connections[user_id] = []
self._connections[user_id].append(websocket)
# Start heartbeat
hb_task = await start_heartbeat(websocket)
self._heartbeat_tasks.setdefault(user_id, []).append(hb_task)
# Start Redis Pub/Sub subscriber for tenant-wide fanout
channel = get_tenant_channel(tenant_id, "ai_ui_control")
pubsub_task = await subscribe_to_channel(channel, lambda msg: self._on_pubsub_message(user_id, msg))
self._pubsub_tasks.setdefault(user_id, []).append(pubsub_task)
logger.debug(f"AI UI Control WS connected: user={user_id}, total={len(self._connections[user_id])}")
return auth
async def _on_pubsub_message(self, user_id: str, msg: dict[str, Any]) -> None:
"""Handle a message received via Redis Pub/Sub."""
# Forward pubsub messages to the user's connections
conns = self._connections.get(user_id, [])
text = json.dumps(msg, default=str)
for ws in conns:
try:
await ws.send_text(text)
except Exception:
logger.warning(f"Failed to send pubsub message to user {user_id}")
async def disconnect(self, websocket: WebSocket, user_id: str) -> None:
"""Remove a WebSocket connection."""
conns = self._connections.get(user_id, [])
if websocket in conns:
conns.remove(websocket)
if not conns:
self._connections.pop(user_id, None)
logger.debug(f"AI UI Control WS disconnected: user={user_id}, remaining={len(conns)}")
"""Remove a WebSocket connection and clean up resources."""
# Cancel heartbeat tasks for this user
for task in self._heartbeat_tasks.pop(user_id, []):
task.cancel()
# Cancel pubsub tasks for this user
for task in self._pubsub_tasks.pop(user_id, []):
task.cancel()
# Clean up connection registry
await cleanup_ws_connection(websocket, user_id, self._connections)
logger.debug(f"AI UI Control WS disconnected: user={user_id}")
async def send_command(self, user_id: str, command: dict[str, Any]) -> str | None:
"""Send a UI command to all frontend connections of a user.
@@ -120,6 +195,24 @@ class AIUIControlWSManager:
"""Get list of currently connected user IDs."""
return list(self._connections.keys())
async def broadcast(self, message: dict[str, Any], tenant_id: uuid.UUID | None = None) -> None:
"""Broadcast a message to all connected users.
If *tenant_id* is provided, publishes via Redis Pub/Sub for multi-worker fanout.
Otherwise, sends directly to all locally connected users.
"""
if tenant_id is not None:
await broadcast_to_tenants(tenant_id, "ai_ui_control", message)
else:
for user_id in list(self._connections.keys()):
conns = self._connections.get(user_id, [])
text = json.dumps(message, default=str)
for ws in conns:
try:
await ws.send_text(text)
except Exception:
logger.warning(f"Failed to broadcast to user {user_id}")
def cleanup_stale(self, timeout_seconds: int = 60) -> None:
"""Remove stale command tracking entries older than timeout."""
now = time.time()
+14 -28
View File
@@ -472,33 +472,29 @@ async def websocket_endpoint(
Authenticates via session cookie. On connect, subscribes user to all their conversations.
"""
# Verify Origin header against allowed CORS origins
from app.config import get_settings
from app.core.auth import get_session_data, get_redis, verify_ws_origin
from app.core.db import async_session_maker
from app.core.service_container import get_container
settings = get_settings()
if not await verify_ws_origin(websocket):
await websocket.close(code=4003, reason="Origin not allowed")
# Get WebSocket manager from service container
container = get_container()
if not container.has("comm_websocket"):
await websocket.close(code=4003, reason="Messaging not available")
return
session_id = websocket.cookies.get(settings.session_cookie_name)
if not session_id:
await websocket.close(code=4001, reason="Not authenticated")
return
ws_manager = container.get("comm_websocket")
redis = get_redis()
session_data = await get_session_data(redis, session_id)
if session_data is None:
await websocket.close(code=4001, reason="Session expired")
return
# connect() performs origin check, session auth, tenant check
async with async_session_maker() as db:
auth = await ws_manager.connect(websocket, db)
if auth is None:
return # connection was rejected and closed by ws_helpers
user_id = session_data["user_id"]
tenant_id = session_data["tenant_id"]
user_id = auth["user_id"]
tenant_id = auth["tenant_id"]
# Plugin-Gate: check if kommunikation plugin is active (global + tenant)
from app.core.permission_registry import get_permission_registry
from sqlalchemy import text as sa_text
from app.core.db import async_session_maker
import uuid as _uuid
try:
registry = get_permission_registry()
@@ -518,16 +514,6 @@ async def websocket_endpoint(
await websocket.close(code=4003, reason="Plugin check failed")
return
# Get WebSocket manager from service container
from app.core.service_container import get_container
container = get_container()
if not container.has("comm_websocket"):
await websocket.close(code=4003, reason="Messaging not available")
return
ws_manager = container.get("comm_websocket")
await ws_manager.connect(websocket, user_id)
try:
while True:
data = await websocket.receive_text()
@@ -4,10 +4,25 @@ from __future__ import annotations
import json
import logging
import uuid
from typing import Any
from fastapi import WebSocket
from app.core.ws_helpers import (
authenticate_ws,
check_ws_origin,
check_ws_tenant,
cleanup_ws_connection,
start_heartbeat,
)
from app.core.ws_pubsub import (
broadcast_to_tenants,
get_tenant_channel,
subscribe_to_channel,
)
from sqlalchemy.ext.asyncio import AsyncSession
logger = logging.getLogger(__name__)
@@ -19,26 +34,82 @@ class WebSocketManager:
self._connections: dict[str, list[WebSocket]] = {}
# conversation_id (str) → set of user_ids subscribed
self._subscriptions: dict[str, set[str]] = {}
# user_id → list of heartbeat tasks
self._heartbeat_tasks: dict[str, list] = {}
# user_id → list of pubsub subscriber tasks
self._pubsub_tasks: dict[str, list] = {}
async def connect(self, websocket: WebSocket, user_id: str) -> None:
"""Accept and register a new WebSocket connection."""
async def connect(
self,
websocket: WebSocket,
db: AsyncSession,
) -> dict[str, Any] | None:
"""Accept, authenticate and register a new WebSocket connection.
Performs origin check, session authentication and tenant validation.
Returns the auth dict (user_id, tenant_id, …) on success, or ``None``
if the connection was rejected (already closed).
"""
# Origin / CSRF check
if not await check_ws_origin(websocket):
return None
# Session authentication
auth = await authenticate_ws(websocket, db)
if auth is None:
return None
user_id = auth["user_id"]
tenant_id_str = auth["tenant_id"]
tenant_id = uuid.UUID(tenant_id_str)
user_uuid = uuid.UUID(user_id)
# Tenant membership check
if not await check_ws_tenant(websocket, tenant_id, user_uuid, db):
return None
# Accept the WebSocket
await websocket.accept()
# Register connection
if user_id not in self._connections:
self._connections[user_id] = []
self._connections[user_id].append(websocket)
# Start heartbeat
hb_task = await start_heartbeat(websocket)
self._heartbeat_tasks.setdefault(user_id, []).append(hb_task)
# Start Redis Pub/Sub subscriber for tenant-wide fanout
channel = get_tenant_channel(tenant_id, "kommunikation")
pubsub_task = await subscribe_to_channel(channel, lambda msg: self._on_pubsub_message(user_id, msg))
self._pubsub_tasks.setdefault(user_id, []).append(pubsub_task)
logger.debug(f"WebSocket connected: user={user_id}, total={len(self._connections[user_id])}")
return auth
async def _on_pubsub_message(self, user_id: str, msg: dict[str, Any]) -> None:
"""Handle a message received via Redis Pub/Sub."""
await self.send_to_user(user_id, msg)
async def disconnect(self, websocket: WebSocket, user_id: str) -> None:
"""Remove a WebSocket connection."""
conns = self._connections.get(user_id, [])
if websocket in conns:
conns.remove(websocket)
if not conns:
self._connections.pop(user_id, None)
# Remove from all subscriptions
"""Remove a WebSocket connection and clean up resources."""
# Cancel heartbeat tasks for this user
for task in self._heartbeat_tasks.pop(user_id, []):
task.cancel()
# Cancel pubsub tasks for this user
for task in self._pubsub_tasks.pop(user_id, []):
task.cancel()
# Clean up connection registry
await cleanup_ws_connection(websocket, user_id, self._connections)
# Remove from subscriptions if no more connections
if user_id not in self._connections:
for conv_id, users in self._subscriptions.items():
users.discard(user_id)
logger.debug(f"WebSocket disconnected: user={user_id}, remaining={len(conns)}")
logger.debug(f"WebSocket disconnected: user={user_id}")
def subscribe(self, conversation_id: str, user_id: str) -> None:
"""Subscribe a user to a conversation's updates."""
@@ -75,10 +146,17 @@ class WebSocketManager:
continue
await self.send_to_user(user_id, message)
async def broadcast(self, message: dict[str, Any]) -> None:
"""Broadcast a message to all connected users."""
for user_id in list(self._connections.keys()):
await self.send_to_user(user_id, message)
async def broadcast(self, message: dict[str, Any], tenant_id: uuid.UUID | None = None) -> None:
"""Broadcast a message to all connected users.
If *tenant_id* is provided, publishes via Redis Pub/Sub for multi-worker fanout.
Otherwise, sends directly to all locally connected users.
"""
if tenant_id is not None:
await broadcast_to_tenants(tenant_id, "kommunikation", message)
else:
for user_id in list(self._connections.keys()):
await self.send_to_user(user_id, message)
def get_online_users(self) -> list[str]:
"""Get list of currently connected user IDs."""