feat: realtime WebSocket layer for live events and streaming logs
- Backend: ConnectionManager, Redis-backed bus with in-process fallback, channel authorization, and typed emit helpers wired into workers (ai_runtime, sandbox_gateway), workforce/sandbox/billing/audit/error services. - Frontend: single shared WS client with backoff + heartbeat, RealtimeProvider, useLiveChannel/useLiveEvent hooks, live indicators on task detail, manager chat, dashboard, traces, sandbox terminal, admin errors, and billing. - Contracts: WsEvent/WsServerFrame/WsClientFrame types mirroring the port. All 255 backend tests pass; web typecheck and build pass.
This commit is contained in:
mjalmoudy94
@@ -0,0 +1,416 @@
|
||||
"""Realtime push layer: authorization, fan-out, and the live socket end-to-end.
|
||||
|
||||
Three levels, mirroring the package's structure:
|
||||
|
||||
* **Authorization** — a channel is only ever granted through
|
||||
``authorize_channel``; every rule is tested for both the allow and the deny
|
||||
side, because a rule that always denies is indistinguishable from a bug.
|
||||
* **Fan-out** — the manager delivers to the subscribers of a channel, keeps
|
||||
ordering, de-duplicates relayed events and cleans up on unsubscribe.
|
||||
* **The socket** — a real WebSocket against the real app: subscribe, act over
|
||||
HTTP, and assert the event arrives without a reload.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
from sqlalchemy.pool import NullPool
|
||||
|
||||
import app.infrastructure.db.models as _orm_models # noqa: F401 - registers metadata
|
||||
from app.application.ports.realtime import RealtimeEvent, set_publisher
|
||||
from app.core.config import settings
|
||||
from app.core.database import Base, get_session
|
||||
from app.core.security import create_access_token
|
||||
from app.domain.base import new_id
|
||||
from app.infrastructure.db.models import UserRow
|
||||
from app.interfaces.realtime.authz import authorize_channel
|
||||
from app.interfaces.realtime.bus import get_bus
|
||||
from app.interfaces.realtime.manager import Connection, ConnectionManager
|
||||
from app.main import app
|
||||
|
||||
API_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FakeSocket:
|
||||
"""The manager only needs ``send_text`` and a connected ``client_state``.
|
||||
|
||||
Both matter: the sender task checks ``client_state`` before writing, so a
|
||||
fake without it would silently swallow every frame and make the manager
|
||||
look broken.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.sent: list[str] = []
|
||||
from starlette.websockets import WebSocketState
|
||||
|
||||
self.client_state = WebSocketState.CONNECTED
|
||||
|
||||
async def send_text(self, payload: str) -> None:
|
||||
self.sent.append(payload)
|
||||
|
||||
|
||||
def _event(event_id: str, *channels: str) -> RealtimeEvent:
|
||||
return RealtimeEvent(type="test.event", channels=channels, data={"id": event_id}, id=event_id)
|
||||
|
||||
|
||||
async def _allow(_: Connection, __: str) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Authorization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _user_with_org(session: AsyncSession, *, platform_owner: bool = False):
|
||||
"""A user, and optionally an organization they belong to."""
|
||||
user_id = new_id()
|
||||
session.add(
|
||||
UserRow(
|
||||
id=user_id,
|
||||
email=f"{user_id}@example.com",
|
||||
password_hash="x",
|
||||
is_platform_owner=platform_owner,
|
||||
)
|
||||
)
|
||||
await session.flush()
|
||||
return user_id
|
||||
|
||||
|
||||
async def _org_with_member(session: AsyncSession, user_id: str) -> str:
|
||||
from app.domain.identity.models import Membership, Organization, Role
|
||||
from app.infrastructure.repositories.identity import (
|
||||
MembershipRepository,
|
||||
OrganizationRepository,
|
||||
)
|
||||
|
||||
organization = await OrganizationRepository(session).create(
|
||||
Organization(name="سازمان تست", slug=new_id()[:8], owner_user_id=user_id)
|
||||
)
|
||||
await MembershipRepository(session).add(
|
||||
Membership(organization_id=organization.id, user_id=user_id, role=Role.ADMIN)
|
||||
)
|
||||
return organization.id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_member_may_subscribe_to_own_org(session_factory):
|
||||
async with session_factory() as session:
|
||||
user_id = await _user_with_org(session)
|
||||
org_id = await _org_with_member(session, user_id)
|
||||
allowed = await authorize_channel(
|
||||
session, user_id=user_id, is_platform_owner=False, channel=f"org:{org_id}"
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_outsider_may_not_subscribe_to_foreign_org(session_factory):
|
||||
async with session_factory() as session:
|
||||
insider = await _user_with_org(session)
|
||||
outsider = await _user_with_org(session)
|
||||
org_id = await _org_with_member(session, insider)
|
||||
allowed = await authorize_channel(
|
||||
session, user_id=outsider, is_platform_owner=False, channel=f"org:{org_id}"
|
||||
)
|
||||
assert allowed is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_channel_kinds_are_denied(session_factory):
|
||||
async with session_factory() as session:
|
||||
user_id = await _user_with_org(session)
|
||||
allowed = await authorize_channel(
|
||||
session, user_id=user_id, is_platform_owner=False, channel="nonsense:whatever"
|
||||
)
|
||||
assert allowed is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_channel_needs_platform_owner(session_factory):
|
||||
async with session_factory() as session:
|
||||
owner = await _user_with_org(session, platform_owner=True)
|
||||
member = await _user_with_org(session)
|
||||
owner_ok = await authorize_channel(
|
||||
session, user_id=owner, is_platform_owner=True, channel="admin"
|
||||
)
|
||||
member_ok = await authorize_channel(
|
||||
session, user_id=member, is_platform_owner=False, channel="admin"
|
||||
)
|
||||
assert owner_ok is True
|
||||
assert member_ok is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_channel_is_private(session_factory):
|
||||
async with session_factory() as session:
|
||||
user_id = await _user_with_org(session)
|
||||
other = await _user_with_org(session)
|
||||
own = await authorize_channel(
|
||||
session, user_id=user_id, is_platform_owner=False, channel=f"user:{user_id}"
|
||||
)
|
||||
foreign = await authorize_channel(
|
||||
session, user_id=other, is_platform_owner=False, channel=f"user:{user_id}"
|
||||
)
|
||||
assert own is True
|
||||
assert foreign is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fan-out
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_delivers_to_subscribers_only():
|
||||
manager = ConnectionManager()
|
||||
listening = FakeSocket()
|
||||
other = FakeSocket()
|
||||
|
||||
subscribed = Connection(user_id="u1", is_platform_owner=False, websocket=listening)
|
||||
stranger = Connection(user_id="u2", is_platform_owner=False, websocket=other)
|
||||
for connection in (subscribed, stranger):
|
||||
await manager.register(connection)
|
||||
await manager.subscribe(subscribed, ["org:1"], _allow)
|
||||
|
||||
delivered = await manager.dispatch(_event("e1", "org:1"))
|
||||
assert delivered == 1
|
||||
|
||||
# The write happens on the connection's own sender task; give the loop the
|
||||
# turn it needs to flush the queue.
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(listening.sent) == 1
|
||||
assert other.sent == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_wraps_events_in_the_wire_envelope():
|
||||
"""A browser sees one discriminator: ``{"type": "event", "event": {...}}``."""
|
||||
import json
|
||||
|
||||
manager = ConnectionManager()
|
||||
socket = FakeSocket()
|
||||
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
|
||||
await manager.register(connection)
|
||||
await manager.subscribe(connection, ["org:1"], _allow)
|
||||
await manager.dispatch(_event("e1", "org:1"))
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(socket.sent) == 1
|
||||
frame = json.loads(socket.sent[0])
|
||||
assert frame["type"] == "event"
|
||||
assert frame["event"]["id"] == "e1"
|
||||
assert frame["event"]["type"] == "test.event"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_drops_duplicate_events():
|
||||
"""In-process publish plus a relayed copy deliver exactly once."""
|
||||
manager = ConnectionManager()
|
||||
socket = FakeSocket()
|
||||
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
|
||||
await manager.register(connection)
|
||||
await manager.subscribe(connection, ["org:1"], _allow)
|
||||
|
||||
event = _event("dup", "org:1")
|
||||
first = await manager.dispatch(event)
|
||||
second = await manager.dispatch(event)
|
||||
assert (first, second) == (1, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unsubscribe_stops_delivery():
|
||||
manager = ConnectionManager()
|
||||
socket = FakeSocket()
|
||||
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
|
||||
await manager.register(connection)
|
||||
await manager.subscribe(connection, ["org:1", "org:2"], _allow)
|
||||
await manager.unsubscribe(connection, ["org:1"])
|
||||
|
||||
assert await manager.dispatch(_event("e1", "org:1")) == 0
|
||||
assert await manager.dispatch(_event("e2", "org:2")) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publisher_never_raises_on_bad_listener():
|
||||
"""A listener that explodes must not break the publisher or its peers."""
|
||||
|
||||
async def exploding(event: RealtimeEvent) -> None:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
bus = get_bus()
|
||||
bus.attach(exploding)
|
||||
try:
|
||||
set_publisher(bus)
|
||||
# Must not raise despite the exploding listener.
|
||||
from app.application.ports.realtime import publish
|
||||
|
||||
await publish(_event("e1", "org:1"))
|
||||
finally:
|
||||
bus.detach(exploding)
|
||||
set_publisher(_NullPublisher())
|
||||
|
||||
|
||||
class _NullPublisher:
|
||||
async def publish(self, event: RealtimeEvent) -> None:
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# The socket, end to end
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _create_schema(engine) -> None:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def ws_client(monkeypatch):
|
||||
"""A TestClient whose API and socket both talk to one throwaway database.
|
||||
|
||||
File-based SQLite with ``NullPool`` on purpose: the app runs on the test
|
||||
portal's event loop while the schema is created from the test's own, and a
|
||||
pooled connection cannot cross that boundary. The database file lives in
|
||||
the package directory — ``*.db`` is gitignored — so the fixture needs no
|
||||
writable temp directory outside the workspace.
|
||||
"""
|
||||
db_path = API_ROOT / f"realtime-test-{new_id()}.db"
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}", poolclass=NullPool)
|
||||
asyncio.run(_create_schema(engine))
|
||||
factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
||||
|
||||
monkeypatch.setattr("app.main.engine", engine)
|
||||
monkeypatch.setattr("app.interfaces.realtime.router.SessionFactory", factory)
|
||||
monkeypatch.setattr(settings, "realtime_transport", "local")
|
||||
|
||||
async def _override_session():
|
||||
async with factory() as session:
|
||||
try:
|
||||
yield session
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
app.dependency_overrides[get_session] = _override_session
|
||||
# The manager is module state; start each run from a clean slate so a
|
||||
# connection from another test cannot colour this one's fan-out.
|
||||
from app.interfaces.realtime.router import realtime_manager
|
||||
|
||||
realtime_manager()._connections.clear()
|
||||
realtime_manager()._by_channel.clear()
|
||||
try:
|
||||
with TestClient(app) as client:
|
||||
yield client
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
asyncio.run(engine.dispose())
|
||||
with suppress(OSError):
|
||||
db_path.unlink()
|
||||
|
||||
|
||||
def _frames_until(websocket, wanted_type: str, *, max_frames: int = 20) -> dict:
|
||||
"""Read frames until ``wanted_type`` arrives — hello/subscribed come first."""
|
||||
for _ in range(max_frames):
|
||||
frame = websocket.receive_json()
|
||||
if frame.get("type") == "event" and frame.get("event", {}).get("type") == wanted_type:
|
||||
return frame["event"]
|
||||
raise AssertionError(f"no {wanted_type} event arrived")
|
||||
|
||||
|
||||
def test_socket_rejects_a_missing_token(ws_client):
|
||||
with pytest.raises(Exception), ws_client.websocket_connect("/api/ws"): # noqa: B017, SIM117
|
||||
# No token: the endpoint closes with 1008 before accepting.
|
||||
pass
|
||||
|
||||
|
||||
def test_task_creation_is_pushed_live(ws_client):
|
||||
# Sign up, create an organization — the tenant the socket subscribes to.
|
||||
signup = ws_client.post(
|
||||
"/api/auth/signup",
|
||||
json={"email": "live@example.com", "password": "password123", "name": "زنده"},
|
||||
)
|
||||
assert signup.status_code == 201, signup.text
|
||||
token = signup.json()["access_token"]
|
||||
org = ws_client.post(
|
||||
"/api/organizations",
|
||||
json={"name": "سازمان زنده"},
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
assert org.status_code == 201, org.text
|
||||
org_id = org.json()["id"]
|
||||
|
||||
with ws_client.websocket_connect(f"/api/ws?token={token}") as socket:
|
||||
hello = socket.receive_json()
|
||||
assert hello["type"] == "hello"
|
||||
|
||||
socket.send_json({"action": "subscribe", "channels": [f"org:{org_id}"]})
|
||||
subscribed = socket.receive_json()
|
||||
assert subscribed["type"] == "subscribed"
|
||||
assert subscribed["accepted"] == [f"org:{org_id}"]
|
||||
|
||||
# A foreign channel is refused, not silently ignored.
|
||||
socket.send_json({"action": "subscribe", "channels": ["org:does-not-exist"]})
|
||||
refused = socket.receive_json()
|
||||
assert refused["type"] == "subscribed"
|
||||
assert refused["rejected"] == ["org:does-not-exist"]
|
||||
|
||||
created = ws_client.post(
|
||||
f"/api/organizations/{org_id}/tasks",
|
||||
json={"title": "کار زنده"},
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
assert created.status_code == 201, created.text
|
||||
|
||||
event = _frames_until(socket, "task.created")
|
||||
assert event["data"]["title"] == "کار زنده"
|
||||
assert event["organization_id"] == org_id
|
||||
|
||||
|
||||
def test_second_user_cannot_read_another_org(ws_client):
|
||||
"""Two tenants, one socket each: an outsider's org is out of reach."""
|
||||
tokens = []
|
||||
org_ids = []
|
||||
for index in range(2):
|
||||
signup = ws_client.post(
|
||||
"/api/auth/signup",
|
||||
json={"email": f"user{index}@example.com", "password": "password123"},
|
||||
)
|
||||
assert signup.status_code == 201, signup.text
|
||||
token = signup.json()["access_token"]
|
||||
org = ws_client.post(
|
||||
"/api/organizations",
|
||||
json={"name": f"org {index}"},
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
assert org.status_code == 201, org.text
|
||||
tokens.append(token)
|
||||
org_ids.append(org.json()["id"])
|
||||
|
||||
with ws_client.websocket_connect(f"/api/ws?token={tokens[0]}") as socket:
|
||||
socket.receive_json() # hello
|
||||
socket.send_json({"action": "subscribe", "channels": [f"org:{org_ids[1]}"]})
|
||||
refused = socket.receive_json()
|
||||
assert refused["rejected"] == [f"org:{org_ids[1]}"]
|
||||
|
||||
|
||||
def test_socket_token_identifies_a_real_user(ws_client):
|
||||
"""A well-formed token for a deleted user must not open a socket."""
|
||||
user_id = new_id()
|
||||
token = create_access_token(user_id)
|
||||
with pytest.raises(Exception), ws_client.websocket_connect(f"/api/ws?token={token}"): # noqa: B017, SIM117
|
||||
# Closed with 1008: the token is well-formed but names no user.
|
||||
pass
|
||||
Reference in New Issue
Block a user