Merge origin/dev: لایهٔ realtime (WebSocket/RPC) + فلو دعوت عضو
ادغام ۱۲ کامیت ریموت (لایهٔ realtime روی WebSocket، حذف REST عمومی، رفع
باگ‌های صفحهٔ مدل‌ها) با ۸ کامیت محلی (ارائه‌دهندهٔ agy + رفع فلو دعوت).
تنها کانفلیکت واقعی در `app/main.py` بود و ماهوی بود، نه متنی:
ریموت اپ عمومی را به «فقط /health + /api/ws» کاهش داده و همهٔ هندلرهای HTTP
را به اپ داخلی `interfaces/realtime/rpc_app.py` منتقل کرده است. بنابراین
روتر `invites` من نباید در اپ عمومی mount شود.
راه‌حل:
- `app/main.py`: فقط `health` و `realtime_router` می‌مانند؛ ایمپورت‌های اضافی
حذف شدند (نسخهٔ ریموت).
- `interfaces/realtime/rpc_app.py`: روتر `invites` به اپ داخلی اضافه شد تا
`GET /api/invites/{token}` و `POST /api/invites/accept` از مسیر WS-RPC
قابل دسترس باشند.
`api.ts` و `contracts/index.ts` خودکار و بدون کانفلیکت ادغام شدند: کلاینت
`previewInvite`/`acceptInvite` حالا از `getRealtimeClient().rpc()` رد می‌شود،
پس هیچ تغییری در فراخوانی‌ها لازم نبود.
رفع تست شکستهٔ ریموت: `src/test/version.test.tsx` مقدار `AppShell` را بدون
`RealtimeProvider` رندر می‌کرد و `useRealtime` استثنا می‌داد. حالا اسکلت
همان‌طور که در `app/layout.tsx` سیم‌کشی شده رندر می‌شود و سوکت در تست stub
شده است.
اعتبارسنجی روی درخت ادغام‌شده:
- بک‌اند ۲۷۰ تست سبز (شامل ۶ تست دعوت روی معماری جدید)؛ mypy تمیز.
- فرانت ۶۱ تست سبز؛ tsc و eslint تمیز.
- تأیید شد روترهای invite روی اپ داخلی RPC حاضرند و سطح عمومی تغییر نکرده.
نکته: ۵ خطای ruff در `tests/conftest.py`، `tests/test_error_log.py` و
`tests/test_realtime.py` از خود `origin/dev` می‌آیند و به این ادغام ربطی
ندارند؛ دست‌نخورده رها شدند.
This commit is contained in:
mjalmoudy94
@@ -0,0 +1,20 @@
|
||||
"""Realtime push layer — the WebSocket *interface* adapter.
|
||||
|
||||
The adapter is transport and plumbing only: who is connected to which channel
|
||||
(:mod:`manager`), how an event travels between processes (:mod:`bus`), which
|
||||
channels a socket may subscribe to (:mod:`authz`), and the wire protocol
|
||||
(:mod:`router`).
|
||||
|
||||
The event contract and the typed publishers live one layer in, in
|
||||
``app.application.ports.realtime`` — services depend on the port, this package
|
||||
implements it. ``app.main`` is the composition root that installs the bus as
|
||||
the process publisher and starts the cross-process relay.
|
||||
"""
|
||||
|
||||
from app.interfaces.realtime.bus import get_bus, start_relay, stop_relay
|
||||
|
||||
__all__ = [
|
||||
"get_bus",
|
||||
"start_relay",
|
||||
"stop_relay",
|
||||
]
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Channel authorization — the only gate on what a socket may receive.
|
||||
|
||||
Authorization happens once, at subscribe time. After that the connection holds
|
||||
plain channel names and the manager routes on them, so there is no code path
|
||||
that can deliver an event to a connection that was never checked.
|
||||
|
||||
Every rule resolves the channel's tenant through the repositories and then the
|
||||
caller's membership in it — exactly like the HTTP layer's ``get_org_context``.
|
||||
A tenant that does not own the object gets a rejection, not an empty stream:
|
||||
the difference matters when debugging "why is my page not live?".
|
||||
|
||||
The caller's identity arrives as two primitives (id, platform-owner flag) —
|
||||
what an authenticated socket actually knows — rather than a ``User`` row, so a
|
||||
subscribe costs one fewer database read.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.infrastructure.repositories.identity import MembershipRepository
|
||||
from app.infrastructure.repositories.sandbox import (
|
||||
SandboxSessionRepository,
|
||||
WorkspaceRepository,
|
||||
)
|
||||
from app.infrastructure.repositories.workforce import (
|
||||
RunRepository,
|
||||
TaskRepository,
|
||||
ThreadRepository,
|
||||
)
|
||||
|
||||
|
||||
async def authorize_channel(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
user_id: str,
|
||||
is_platform_owner: bool,
|
||||
channel: str,
|
||||
) -> bool:
|
||||
"""May this identity subscribe to ``channel``? Deny by default."""
|
||||
kind, _, identifier = channel.partition(":")
|
||||
if not identifier and kind != "admin":
|
||||
return False
|
||||
|
||||
if kind == "admin":
|
||||
return is_platform_owner
|
||||
if kind == "user":
|
||||
return identifier == user_id
|
||||
if kind == "org":
|
||||
return await _is_member(
|
||||
session=session, user_id=user_id, organization_id=identifier
|
||||
)
|
||||
if kind == "task":
|
||||
return await _task_is_visible(
|
||||
session=session, user_id=user_id, task_id=identifier
|
||||
)
|
||||
if kind == "run":
|
||||
return await _run_is_visible(
|
||||
session=session, user_id=user_id, run_id=identifier
|
||||
)
|
||||
if kind == "thread":
|
||||
return await _thread_is_visible(
|
||||
session=session, user_id=user_id, thread_id=identifier
|
||||
)
|
||||
if kind == "session":
|
||||
return await _session_is_visible(
|
||||
session=session, user_id=user_id, session_id=identifier
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
async def _is_member(*, session: AsyncSession, user_id: str, organization_id: str) -> bool:
|
||||
membership = await MembershipRepository(session).get(organization_id, user_id)
|
||||
return membership is not None and membership.is_active
|
||||
|
||||
|
||||
async def _task_is_visible(*, session: AsyncSession, user_id: str, task_id: str) -> bool:
|
||||
task = await TaskRepository(session).get(task_id)
|
||||
if task is None:
|
||||
return False
|
||||
return await _is_member(
|
||||
session=session, user_id=user_id, organization_id=task.organization_id
|
||||
)
|
||||
|
||||
|
||||
async def _run_is_visible(*, session: AsyncSession, user_id: str, run_id: str) -> bool:
|
||||
run = await RunRepository(session).get(run_id)
|
||||
if run is None:
|
||||
return False
|
||||
return await _task_is_visible(session=session, user_id=user_id, task_id=run.task_id)
|
||||
|
||||
|
||||
async def _thread_is_visible(*, session: AsyncSession, user_id: str, thread_id: str) -> bool:
|
||||
thread = await ThreadRepository(session).get(thread_id)
|
||||
if thread is None:
|
||||
return False
|
||||
return await _is_member(
|
||||
session=session, user_id=user_id, organization_id=thread.organization_id
|
||||
)
|
||||
|
||||
|
||||
async def _session_is_visible(*, session: AsyncSession, user_id: str, session_id: str) -> bool:
|
||||
"""A sandbox session belongs to a workspace, which belongs to a tenant."""
|
||||
sandbox_session = await SandboxSessionRepository(session).get(session_id)
|
||||
if sandbox_session is None:
|
||||
return False
|
||||
workspace = await WorkspaceRepository(session).get(sandbox_session.workspace_id)
|
||||
if workspace is None:
|
||||
return False
|
||||
return await _is_member(
|
||||
session=session, user_id=user_id, organization_id=workspace.organization_id
|
||||
)
|
||||
@@ -0,0 +1,260 @@
|
||||
"""The event bus: how a published event reaches connected browsers.
|
||||
|
||||
Two transports, in this order:
|
||||
|
||||
1. **In-process.** Always on. When the API itself records the event (a task run
|
||||
driven by a button, a sandbox command, an error handler), the fan-out is a
|
||||
direct call into the ``ConnectionManager`` — microseconds, no broker.
|
||||
2. **Redis pub/sub.** Carries events published by *other* processes: the
|
||||
standalone workers (``pnpm worker:ai-runtime``) run tasks in their own
|
||||
process, and without a broker their events would die there. The API process
|
||||
relays whatever arrives on the shared channel into the same manager.
|
||||
|
||||
Redis is a bonus, never a dependency:
|
||||
|
||||
* The bus tries it once per cooldown window. If nothing is listening, the
|
||||
attempt fails quietly and in-process delivery continues — the product stays
|
||||
fully live for every API-driven action, only cross-process events are lost.
|
||||
* The relay task owns the subscription lifecycle; it starts with the app and
|
||||
stops on shutdown, and reconnects on its own after a broker restart.
|
||||
|
||||
Nothing here ever raises. Publishing is called from request handlers, from
|
||||
inside exception handlers and from worker passes; an exception in the feed
|
||||
must not be able to fail any of those.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import suppress
|
||||
|
||||
from app.application.ports.realtime import RealtimeEvent, set_publisher
|
||||
from app.core.config import settings
|
||||
|
||||
logger = logging.getLogger("agy.realtime")
|
||||
|
||||
#: Redis pub/sub channel shared by every process of one deployment.
|
||||
DEFAULT_REDIS_CHANNEL = "agy:realtime"
|
||||
|
||||
#: After a failed broker attempt, wait this long before trying again. Keeps a
|
||||
#: dev machine without Redis from adding a connection timeout to every event.
|
||||
REDIS_RETRY_COOLDOWN_SECONDS = 10.0
|
||||
|
||||
#: How long a single connection attempt may take before it is abandoned.
|
||||
REDIS_CONNECT_TIMEOUT_SECONDS = 1.0
|
||||
|
||||
#: Listener signature: ``await listener(event)``. The ConnectionManager
|
||||
#: registers one; tests register fakes. The return value is ignored so a
|
||||
#: listener may report how many connections it fanned out to.
|
||||
EventListener = Callable[[RealtimeEvent], Awaitable[object]]
|
||||
|
||||
|
||||
class RealtimeBus:
|
||||
"""Fan-out hub shared by every publisher in this process."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
redis_url: str | None = None,
|
||||
channel: str = DEFAULT_REDIS_CHANNEL,
|
||||
) -> None:
|
||||
self._listeners: set[EventListener] = set()
|
||||
self._redis_url = redis_url if redis_url is not None else settings.redis_url
|
||||
self._channel = channel
|
||||
self._redis = None
|
||||
self._pubsub = None
|
||||
self._relay_task: asyncio.Task | None = None
|
||||
self._next_redis_attempt = 0.0
|
||||
self._redis_healthy = False
|
||||
|
||||
# -- listeners (in-process fan-out) --------------------------------------
|
||||
|
||||
def attach(self, listener: EventListener) -> None:
|
||||
self._listeners.add(listener)
|
||||
|
||||
def detach(self, listener: EventListener) -> None:
|
||||
self._listeners.discard(listener)
|
||||
|
||||
# -- publishing -----------------------------------------------------------
|
||||
|
||||
async def publish(self, event: RealtimeEvent) -> None:
|
||||
"""Deliver locally, then mirror to other processes. Never raises."""
|
||||
try:
|
||||
await self._dispatch_local(event)
|
||||
except Exception: # noqa: BLE001 - the feed must never fail a caller
|
||||
logger.exception("realtime: local dispatch failed for %s", event.type)
|
||||
|
||||
if not self._redis_enabled():
|
||||
return
|
||||
try:
|
||||
await self._publish_redis(event)
|
||||
except Exception: # noqa: BLE001 - same rule
|
||||
logger.exception("realtime: redis publish failed for %s", event.type)
|
||||
|
||||
async def _dispatch_local(self, event: RealtimeEvent) -> None:
|
||||
for listener in list(self._listeners):
|
||||
try:
|
||||
await listener(event)
|
||||
except Exception: # noqa: BLE001 - one bad listener is not fatal
|
||||
logger.exception("realtime: listener failed for %s", event.type)
|
||||
|
||||
def _redis_enabled(self) -> bool:
|
||||
# "local" pins the bus to in-process delivery — used by tests and by
|
||||
# single-process deployments that have no broker on purpose.
|
||||
return getattr(settings, "realtime_transport", "auto").lower() != "local"
|
||||
|
||||
async def _publish_redis(self, event: RealtimeEvent) -> None:
|
||||
client = await self._redis_client()
|
||||
if client is None:
|
||||
return
|
||||
await client.publish(self._channel, event.to_json())
|
||||
|
||||
async def _redis_client(self):
|
||||
"""Lazily connect, with a cooldown after failure. ``None`` if down."""
|
||||
if self._redis is not None:
|
||||
return self._redis
|
||||
now = time.monotonic()
|
||||
if now < self._next_redis_attempt:
|
||||
return None
|
||||
self._next_redis_attempt = now + REDIS_RETRY_COOLDOWN_SECONDS
|
||||
try:
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
client = aioredis.from_url(
|
||||
self._redis_url,
|
||||
encoding="utf-8",
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=REDIS_CONNECT_TIMEOUT_SECONDS,
|
||||
socket_timeout=REDIS_CONNECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
await client.ping()
|
||||
self._redis = client
|
||||
self._redis_healthy = True
|
||||
logger.info("realtime: redis pub/sub connected (%s)", self._channel)
|
||||
except Exception as exc: # noqa: BLE001 - broker is optional
|
||||
self._redis = None
|
||||
self._redis_healthy = False
|
||||
logger.debug("realtime: redis unavailable (%s); in-process only", exc)
|
||||
return None
|
||||
return self._redis
|
||||
|
||||
# -- relay (other processes -> this one's browsers) ----------------------
|
||||
|
||||
async def start_relay(self) -> None:
|
||||
"""Subscribe to the shared channel and feed the local manager."""
|
||||
if self._relay_task is not None and not self._relay_task.done():
|
||||
return
|
||||
self._relay_task = asyncio.create_task(self._relay_loop(), name="realtime-relay")
|
||||
|
||||
async def stop_relay(self) -> None:
|
||||
task = self._relay_task
|
||||
self._relay_task = None
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
await self._close_redis()
|
||||
|
||||
async def _relay_loop(self) -> None:
|
||||
while True:
|
||||
client = await self._redis_client()
|
||||
if client is None:
|
||||
# Broker absent (or down mid-flight): wait out the cooldown and
|
||||
# try again — a Redis started after the API must still work.
|
||||
await asyncio.sleep(REDIS_RETRY_COOLDOWN_SECONDS)
|
||||
continue
|
||||
try:
|
||||
pubsub = client.pubsub()
|
||||
await pubsub.subscribe(self._channel)
|
||||
self._pubsub = pubsub
|
||||
logger.info("realtime: relay subscribed to %s", self._channel)
|
||||
async for message in pubsub.listen():
|
||||
if message.get("type") != "message":
|
||||
continue
|
||||
event = _decode_message(message.get("data"))
|
||||
if event is not None:
|
||||
# Dedupe by id happens in the manager, so an event this
|
||||
# process published itself is not delivered twice.
|
||||
await self._dispatch_local(event)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001 - relay must survive
|
||||
self._redis_healthy = False
|
||||
logger.warning("realtime: relay interrupted (%s); reconnecting", exc)
|
||||
finally:
|
||||
if self._pubsub is not None:
|
||||
with suppress(Exception):
|
||||
await self._pubsub.unsubscribe(self._channel)
|
||||
self._pubsub = None
|
||||
self._redis = None
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
async def _close_redis(self) -> None:
|
||||
client, self._redis = self._redis, None
|
||||
if client is not None:
|
||||
with suppress(Exception):
|
||||
await client.aclose()
|
||||
self._redis_healthy = False
|
||||
|
||||
# -- introspection -------------------------------------------------------
|
||||
|
||||
@property
|
||||
def redis_healthy(self) -> bool:
|
||||
return self._redis_healthy
|
||||
|
||||
@property
|
||||
def listener_count(self) -> int:
|
||||
return len(self._listeners)
|
||||
|
||||
|
||||
def _decode_message(raw: object) -> RealtimeEvent | None:
|
||||
if not isinstance(raw, (str, bytes)):
|
||||
return None
|
||||
try:
|
||||
return RealtimeEvent.from_json(raw)
|
||||
except Exception: # noqa: BLE001 - a bad frame is skipped, not fatal
|
||||
logger.warning("realtime: undecodable relay frame dropped")
|
||||
return None
|
||||
|
||||
|
||||
_bus: RealtimeBus | None = None
|
||||
|
||||
|
||||
def get_bus() -> RealtimeBus:
|
||||
"""Process-wide bus. Created lazily so importing never opens a socket."""
|
||||
global _bus
|
||||
if _bus is None:
|
||||
_bus = RealtimeBus()
|
||||
return _bus
|
||||
|
||||
|
||||
def install() -> RealtimeBus:
|
||||
"""Make this process's bus the application-layer publisher.
|
||||
|
||||
Called by the composition roots (``app.main`` for the API,
|
||||
``app.workers.__main__`` for the workers). Idempotent, so a worker that
|
||||
imports several modules pays nothing for calling it more than once.
|
||||
"""
|
||||
bus = get_bus()
|
||||
set_publisher(bus)
|
||||
return bus
|
||||
|
||||
|
||||
async def publish(event: RealtimeEvent) -> None:
|
||||
"""Publish through the shared bus. Never raises."""
|
||||
bus = get_bus()
|
||||
try:
|
||||
await bus.publish(event)
|
||||
except Exception: # noqa: BLE001 - belt and braces; publish is best-effort
|
||||
logger.exception("realtime: publish failed for %s", event.type)
|
||||
|
||||
|
||||
async def start_relay() -> None:
|
||||
await get_bus().start_relay()
|
||||
|
||||
|
||||
async def stop_relay() -> None:
|
||||
await get_bus().stop_relay()
|
||||
@@ -0,0 +1,220 @@
|
||||
"""Connection bookkeeping: who is connected, and to which channels.
|
||||
|
||||
The manager is deliberately transport-agnostic — it does not know it is talking
|
||||
to a WebSocket, only to an object with ``send_text``. That keeps it testable
|
||||
without a socket and lets a future transport (SSE, push bridge) reuse it.
|
||||
|
||||
Two invariants:
|
||||
|
||||
* **A connection only ever receives channels it was authorized for.** The
|
||||
authorizer runs at subscribe time; the routing table built afterwards can
|
||||
only contain pre-approved channel names.
|
||||
* **An event is delivered at most once.** ``dispatch`` de-duplicates by event
|
||||
id, because the same event reaches a process twice in the normal case: once
|
||||
in-process and once relayed from Redis.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections import deque
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from starlette.websockets import WebSocket, WebSocketState
|
||||
|
||||
from app.application.ports.realtime import RealtimeEvent
|
||||
|
||||
logger = logging.getLogger("agy.realtime")
|
||||
|
||||
#: Bounded dedupe memory. Large enough that a relayed duplicate is always
|
||||
#: caught, small enough that a chatty day costs a few kilobytes.
|
||||
MAX_SEEN_EVENTS = 2048
|
||||
|
||||
#: Per-connection outbox. If a browser stops reading (a frozen tab), a
|
||||
#: producer must not be able to grow this without bound. Overflow drops the
|
||||
#: oldest frame — ordering of what follows stays meaningful.
|
||||
MAX_QUEUE_DEPTH = 256
|
||||
|
||||
#: What a subscribe check looks like: ``(connection, channel) -> allowed``.
|
||||
ChannelAuthorizer = Callable[["Connection", str], Awaitable[bool]]
|
||||
|
||||
|
||||
#: ``eq=False`` on purpose: a connection is identified by object identity. The
|
||||
#: default dataclass ``__eq__`` would also make it unhashable (a ``set`` field
|
||||
#: cannot be compared safely), and the manager keys connections by identity in
|
||||
#: both its connection set and its channel routing table.
|
||||
@dataclass(slots=True, eq=False)
|
||||
class Connection:
|
||||
"""One browser socket. ``user_id`` is ``None`` for an anonymous (pre-login)
|
||||
connection: it may only call the bootstrap RPCs (login/signup) and may not
|
||||
subscribe to live channels."""
|
||||
|
||||
user_id: str | None
|
||||
is_platform_owner: bool
|
||||
websocket: WebSocket
|
||||
token: str | None = None
|
||||
channels: set[str] = field(default_factory=set)
|
||||
_queue: asyncio.Queue[str] = field(
|
||||
default_factory=lambda: asyncio.Queue(maxsize=MAX_QUEUE_DEPTH)
|
||||
)
|
||||
_sender: asyncio.Task | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Single writer task per connection.
|
||||
|
||||
One task owns the socket: concurrent ``send_text`` calls on one
|
||||
WebSocket interleave frames and corrupt the stream, so everything goes
|
||||
through this queue instead.
|
||||
"""
|
||||
if self._sender is None or self._sender.done():
|
||||
self._sender = asyncio.create_task(self._drain(), name="realtime-sender")
|
||||
|
||||
async def _drain(self) -> None:
|
||||
while True:
|
||||
payload = await self._queue.get()
|
||||
if payload is None: # shutdown sentinel
|
||||
return
|
||||
try:
|
||||
if self.websocket.client_state == WebSocketState.CONNECTED:
|
||||
await self.websocket.send_text(payload)
|
||||
except Exception: # noqa: BLE001 - a dead socket ends the connection
|
||||
logger.debug("realtime: send failed; closing connection", exc_info=True)
|
||||
return
|
||||
|
||||
def offer(self, payload: str) -> None:
|
||||
"""Queue a frame, never blocking the publisher.
|
||||
|
||||
A slow consumer loses the oldest frames rather than stalling every
|
||||
other tenant's stream behind it.
|
||||
"""
|
||||
if self._queue.full():
|
||||
with suppress(asyncio.QueueFull):
|
||||
# Drop the oldest frame to make room: ordering of what follows
|
||||
# stays meaningful, which is what the UI renders.
|
||||
self._queue.get_nowait()
|
||||
with suppress(asyncio.QueueFull): # pragma: no cover - defensive
|
||||
self._queue.put_nowait(payload)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._sender is not None:
|
||||
self._sender.cancel()
|
||||
self._sender = None
|
||||
|
||||
|
||||
class ConnectionManager:
|
||||
"""The routing table between events and open sockets."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._connections: set[Connection] = set()
|
||||
self._by_channel: dict[str, set[Connection]] = {}
|
||||
self._seen: deque[str] = deque(maxlen=MAX_SEEN_EVENTS)
|
||||
self._seen_set: set[str] = set()
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
|
||||
async def register(self, connection: Connection) -> None:
|
||||
self._connections.add(connection)
|
||||
await connection.start()
|
||||
|
||||
async def unregister(self, connection: Connection) -> None:
|
||||
self._connections.discard(connection)
|
||||
for channel in list(connection.channels):
|
||||
subscribers = self._by_channel.get(channel)
|
||||
if subscribers is not None:
|
||||
subscribers.discard(connection)
|
||||
if not subscribers:
|
||||
self._by_channel.pop(channel, None)
|
||||
connection.channels.clear()
|
||||
await connection.close()
|
||||
|
||||
# -- subscriptions ---------------------------------------------------------
|
||||
|
||||
async def subscribe(
|
||||
self,
|
||||
connection: Connection,
|
||||
channels: list[str],
|
||||
authorizer: ChannelAuthorizer,
|
||||
) -> tuple[list[str], list[str]]:
|
||||
"""Authorize and subscribe. Returns ``(accepted, rejected)``."""
|
||||
accepted: list[str] = []
|
||||
rejected: list[str] = []
|
||||
for channel in channels:
|
||||
if channel in connection.channels:
|
||||
accepted.append(channel)
|
||||
continue
|
||||
try:
|
||||
allowed = await authorizer(connection, channel)
|
||||
except Exception: # noqa: BLE001 - a broken check denies, never allows
|
||||
logger.exception("realtime: channel authorization crashed")
|
||||
allowed = False
|
||||
if not allowed:
|
||||
rejected.append(channel)
|
||||
continue
|
||||
connection.channels.add(channel)
|
||||
self._by_channel.setdefault(channel, set()).add(connection)
|
||||
accepted.append(channel)
|
||||
return accepted, rejected
|
||||
|
||||
async def unsubscribe(self, connection: Connection, channels: list[str]) -> None:
|
||||
for channel in channels:
|
||||
connection.channels.discard(channel)
|
||||
subscribers = self._by_channel.get(channel)
|
||||
if subscribers is not None:
|
||||
subscribers.discard(connection)
|
||||
if not subscribers:
|
||||
self._by_channel.pop(channel, None)
|
||||
|
||||
# -- fan-out ---------------------------------------------------------------
|
||||
|
||||
async def dispatch(self, event: RealtimeEvent) -> int:
|
||||
"""Deliver to every subscriber of every channel. Returns fan-out count.
|
||||
|
||||
The id check is what makes "publish locally, then receive the same event
|
||||
from Redis" deliver once instead of twice.
|
||||
"""
|
||||
if event.id in self._seen_set:
|
||||
return 0
|
||||
self._remember(event.id)
|
||||
|
||||
recipients: set[Connection] = set()
|
||||
for channel in event.channels:
|
||||
recipients.update(self._by_channel.get(channel, ()))
|
||||
|
||||
envelope = event.envelope()
|
||||
# Control frames and event frames share one top-level discriminator
|
||||
# (``type``), so the client switches on one field for everything.
|
||||
payload = _dumps({"type": "event", "event": envelope})
|
||||
delivered = 0
|
||||
for connection in recipients:
|
||||
connection.offer(payload)
|
||||
delivered += 1
|
||||
return delivered
|
||||
|
||||
def _remember(self, event_id: str) -> None:
|
||||
if len(self._seen) == self._seen.maxlen:
|
||||
evicted = self._seen[0]
|
||||
self._seen_set.discard(evicted)
|
||||
self._seen.append(event_id)
|
||||
self._seen_set.add(event_id)
|
||||
|
||||
# -- introspection -------------------------------------------------------
|
||||
|
||||
@property
|
||||
def connection_count(self) -> int:
|
||||
return len(self._connections)
|
||||
|
||||
@property
|
||||
def channel_count(self) -> int:
|
||||
return len(self._by_channel)
|
||||
|
||||
def channel_load(self) -> dict[str, int]:
|
||||
return {channel: len(subs) for channel, subs in self._by_channel.items()}
|
||||
|
||||
|
||||
def _dumps(envelope: dict) -> str:
|
||||
import json
|
||||
|
||||
return json.dumps(envelope, ensure_ascii=False, default=str)
|
||||
@@ -0,0 +1,300 @@
|
||||
"""The WebSocket endpoint — one socket per browser, every live feed and every
|
||||
RPC on it.
|
||||
|
||||
``GET /api/ws?token=<jwt>``
|
||||
|
||||
The token travels as a query parameter because the browser WebSocket API cannot
|
||||
set an ``Authorization`` header. That is a deliberate, contained exception to the
|
||||
header-only rule: the token is short-lived, the socket is TLS-terminated in any
|
||||
real deployment, and the alternative (a cookie) would drag CSRF into a protocol
|
||||
that has no CSRF story.
|
||||
|
||||
Protocol — small, symmetric, JSON in both directions. Server->client frames
|
||||
share the ``type`` discriminator; client->server frames use ``action``::
|
||||
|
||||
client -> {"action": "subscribe", "channels": ["org:…"]}
|
||||
client -> {"action": "unsubscribe", "channels": ["task:…"]}
|
||||
client -> {"action": "ping"}
|
||||
client -> {"action": "rpc", "id": "…", "method": "POST",
|
||||
"path": "/api/organizations", "body": {"name": "…"}}
|
||||
server -> {"type": "hello", "user_id": "…"|null}
|
||||
server -> {"type": "subscribed", "accepted": [...], "rejected": [...]}
|
||||
server -> {"type": "event", "event": {…envelope…}}
|
||||
server -> {"type": "rpc_result", "id": "…", "status": 200, "body": …}
|
||||
server -> {"type": "pong"}
|
||||
server -> {"type": "error", "message": "…"}
|
||||
|
||||
A connection may be **anonymous** (no token): that is the bootstrap state that
|
||||
lets a brand-new visitor call ``/api/auth/login`` and ``/api/auth/signup`` over
|
||||
RPC. Every other RPC and every ``subscribe`` requires an authenticated
|
||||
connection; the token presented at connect time is re-injected as a bearer
|
||||
header when the RPC is dispatched into the internal app, so the handlers' own
|
||||
auth dependencies still run unchanged.
|
||||
|
||||
Authorization queries open their own short-lived session: a socket can stay open
|
||||
for hours, and one SQLAlchemy session per open tab would quietly exhaust the
|
||||
connection pool.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
|
||||
|
||||
from app.application.ports.realtime import dedupe_channels
|
||||
from app.core.database import SessionFactory
|
||||
from app.core.errors import AuthenticationError
|
||||
from app.core.security import decode_access_token
|
||||
from app.infrastructure.repositories.identity import UserRepository
|
||||
from app.interfaces.realtime.authz import authorize_channel
|
||||
from app.interfaces.realtime.manager import Connection, ConnectionManager
|
||||
|
||||
logger = logging.getLogger("agy.realtime")
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
#: The manager instance this process fans out with. Created at import: the
|
||||
#: endpoint and the lifecycle hooks can share it without a wiring ceremony.
|
||||
_manager = ConnectionManager()
|
||||
|
||||
|
||||
def realtime_manager() -> ConnectionManager:
|
||||
"""The process-wide manager — attached to the bus by ``app.main``."""
|
||||
return _manager
|
||||
|
||||
|
||||
#: Paths an anonymous (not-yet-logged-in) connection is allowed to call. These are
|
||||
#: the only two calls a visitor can make; everything else needs a session.
|
||||
_ANONYMOUS_RPC_PATHS = frozenset({"/api/auth/login", "/api/auth/signup"})
|
||||
|
||||
#: One in-process HTTP client that dispatches RPCs into the internal app. Built
|
||||
#: lazily so importing this module never pulls the app into existence.
|
||||
_rpc_client: httpx.AsyncClient | None = None
|
||||
|
||||
|
||||
def _get_rpc_client() -> httpx.AsyncClient:
|
||||
global _rpc_client
|
||||
if _rpc_client is None:
|
||||
# The internal app holds every HTTP handler; it is never socket-bound, so
|
||||
# it is reached here via an in-process ASGI transport instead of the wire.
|
||||
from app.interfaces.realtime.rpc_app import get_rpc_app
|
||||
|
||||
transport = httpx.ASGITransport(app=get_rpc_app())
|
||||
_rpc_client = httpx.AsyncClient(transport=transport, base_url="http://internal")
|
||||
return _rpc_client
|
||||
|
||||
|
||||
@router.websocket("/ws")
|
||||
async def realtime_socket(
|
||||
websocket: WebSocket,
|
||||
token: str | None = Query(default=None),
|
||||
) -> None:
|
||||
connection = await _authenticate(websocket, token)
|
||||
if connection is None:
|
||||
return # _authenticate already closed with a policy-violation code
|
||||
|
||||
await _manager.register(connection)
|
||||
|
||||
_send_frame(
|
||||
connection,
|
||||
{"type": "hello", "user_id": connection.user_id, "server_time": _now_iso()},
|
||||
)
|
||||
|
||||
try:
|
||||
while True:
|
||||
raw = await websocket.receive_text()
|
||||
await _handle_message(connection, raw)
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
except Exception: # noqa: BLE001 - any read failure ends the socket cleanly
|
||||
logger.debug("realtime: socket loop ended", exc_info=True)
|
||||
finally:
|
||||
await _manager.unregister(connection)
|
||||
|
||||
|
||||
async def _authenticate(websocket: WebSocket, token: str | None) -> Connection | None:
|
||||
"""Accept an anonymous or an authenticated, active user.
|
||||
|
||||
* No token -> an anonymous connection (bootstrap for login/signup).
|
||||
* A token that decodes to a real, active user -> an authenticated connection
|
||||
whose raw token is kept for RPC re-injection.
|
||||
* A presented-but-invalid token -> 1008, so the client can tell "bad token"
|
||||
from "network dropped" and stop retrying.
|
||||
"""
|
||||
user = None
|
||||
raw_token: str | None = None
|
||||
if token:
|
||||
try:
|
||||
payload = decode_access_token(token)
|
||||
user_id = payload.get("sub")
|
||||
if user_id:
|
||||
# Own session: the socket must not depend on request scope.
|
||||
async with SessionFactory() as session:
|
||||
user = await UserRepository(session).get(str(user_id))
|
||||
if user is not None and user.is_active:
|
||||
raw_token = token
|
||||
except AuthenticationError:
|
||||
user = None
|
||||
|
||||
if token and user is None:
|
||||
await websocket.close(code=1008, reason="unauthorized")
|
||||
return None
|
||||
|
||||
await websocket.accept()
|
||||
return Connection(
|
||||
user_id=user.id if user else None,
|
||||
is_platform_owner=user.is_platform_owner if user else False,
|
||||
token=raw_token,
|
||||
websocket=websocket,
|
||||
)
|
||||
|
||||
|
||||
async def _handle_message(connection: Connection, raw: str) -> None:
|
||||
try:
|
||||
message = json.loads(raw)
|
||||
except (TypeError, ValueError):
|
||||
_send_frame(connection, {"type": "error", "message": "malformed JSON frame"})
|
||||
return
|
||||
if not isinstance(message, dict):
|
||||
_send_frame(connection, {"type": "error", "message": "frame must be an object"})
|
||||
return
|
||||
|
||||
action = message.get("action")
|
||||
|
||||
if action == "ping":
|
||||
_send_frame(connection, {"type": "pong"})
|
||||
return
|
||||
|
||||
if action == "rpc":
|
||||
# Run off the receive loop: a slow handler must not stall pings.
|
||||
asyncio.create_task(_handle_rpc(connection, message))
|
||||
return
|
||||
|
||||
# subscribe / unsubscribe require a logged-in user.
|
||||
if connection.user_id is None:
|
||||
_send_frame(connection, {"type": "error", "message": "authentication required"})
|
||||
return
|
||||
|
||||
channels = dedupe_channels(
|
||||
str(channel) for channel in (message.get("channels") or []) if isinstance(channel, str)
|
||||
)
|
||||
|
||||
if action == "subscribe":
|
||||
accepted, rejected = await _subscribe(connection, channels)
|
||||
_send_frame(
|
||||
connection, {"type": "subscribed", "accepted": accepted, "rejected": rejected}
|
||||
)
|
||||
return
|
||||
|
||||
if action == "unsubscribe":
|
||||
await _manager.unsubscribe(connection, list(channels))
|
||||
_send_frame(connection, {"type": "unsubscribed", "channels": channels})
|
||||
return
|
||||
|
||||
_send_frame(connection, {"type": "error", "message": f"unknown action: {action!r}"})
|
||||
|
||||
|
||||
async def _handle_rpc(connection: Connection, message: dict) -> None:
|
||||
"""Dispatch one client RPC into the internal app and return the result."""
|
||||
req_id = message.get("id")
|
||||
if not req_id:
|
||||
_send_frame(connection, {"type": "error", "message": "rpc frame needs an id"})
|
||||
return
|
||||
|
||||
method = str(message.get("method") or "GET").upper()
|
||||
path = message.get("path")
|
||||
if not isinstance(path, str) or not path:
|
||||
_send_rpc_result(
|
||||
connection,
|
||||
req_id,
|
||||
400,
|
||||
{"error": {"code": "bad_request", "message": "rpc frame needs a path"}},
|
||||
)
|
||||
return
|
||||
|
||||
# Gate unauthenticated callers to login/signup only.
|
||||
if connection.user_id is None and path not in _ANONYMOUS_RPC_PATHS:
|
||||
_send_rpc_result(
|
||||
connection,
|
||||
req_id,
|
||||
401,
|
||||
{"error": {"code": "unauthenticated", "message": "نیاز به ورود دارید."}},
|
||||
)
|
||||
return
|
||||
|
||||
body = message.get("body")
|
||||
headers = {"X-Request-ID": str(req_id), "User-Agent": "realtime-rpc"}
|
||||
if connection.token:
|
||||
headers["Authorization"] = f"Bearer {connection.token}"
|
||||
|
||||
try:
|
||||
client = _get_rpc_client()
|
||||
if isinstance(body, (dict, list)):
|
||||
response = await client.request(method, path, json=body, headers=headers)
|
||||
elif body is None:
|
||||
response = await client.request(method, path, headers=headers)
|
||||
else:
|
||||
response = await client.request(method, path, content=str(body), headers=headers)
|
||||
except Exception: # noqa: BLE001 - reported to the caller as a 500 envelope
|
||||
logger.exception("realtime: rpc dispatch to internal app failed")
|
||||
_send_rpc_result(
|
||||
connection,
|
||||
req_id,
|
||||
500,
|
||||
{"error": {"code": "internal_error", "message": "خطای داخلی سرور."}},
|
||||
)
|
||||
return
|
||||
|
||||
status = response.status_code
|
||||
try:
|
||||
payload = response.json()
|
||||
except Exception:
|
||||
payload = None
|
||||
_send_rpc_result(connection, req_id, status, payload)
|
||||
|
||||
|
||||
def _send_rpc_result(
|
||||
connection: Connection, req_id: object, status: int, body: object
|
||||
) -> None:
|
||||
_send_frame(
|
||||
connection,
|
||||
{"type": "rpc_result", "id": req_id, "status": status, "body": body},
|
||||
)
|
||||
|
||||
|
||||
async def _subscribe(
|
||||
connection: Connection, channels: tuple[str, ...]
|
||||
) -> tuple[list[str], list[str]]:
|
||||
if not channels:
|
||||
return [], []
|
||||
|
||||
async def _authorize(conn: Connection, channel: str) -> bool:
|
||||
async with SessionFactory() as session:
|
||||
return await authorize_channel(
|
||||
session,
|
||||
user_id=conn.user_id or "",
|
||||
is_platform_owner=conn.is_platform_owner,
|
||||
channel=channel,
|
||||
)
|
||||
|
||||
return await _manager.subscribe(connection, list(channels), _authorize)
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
from app.domain.base import utcnow
|
||||
|
||||
return utcnow().isoformat()
|
||||
|
||||
|
||||
def _send_frame(connection: Connection, payload: dict) -> None:
|
||||
"""Queue a frame on the connection's single-writer outbox.
|
||||
|
||||
Every frame — control, event and rpc_result — goes through the same queue so
|
||||
a slow RPC task can never interleave bytes with the receive loop's own
|
||||
frames and corrupt the stream.
|
||||
"""
|
||||
connection.offer(json.dumps(payload, ensure_ascii=False, default=str))
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Internal, never-socket-bound FastAPI app that holds every HTTP handler.
|
||||
|
||||
The public app (``app.main.create_app``) exposes only ``/health`` and the
|
||||
``/api/ws`` WebSocket. All request/response logic lives here and is reached in
|
||||
exactly one of two ways:
|
||||
|
||||
* from the WS-RPC proxy (``app.interfaces.realtime.router``), which dispatches
|
||||
client RPCs into this app over an in-process ASGI transport, and
|
||||
* from the test-suite's ``AsyncClient``/``TestClient``.
|
||||
|
||||
It is intentionally never bound to a socket, so including the routers here adds no
|
||||
public REST surface — the only public routes are the ones ``create_app``
|
||||
registers. Keeping the handlers in one FastAPI app also means they run through the
|
||||
real middleware, dependencies, validation and error handling, unchanged.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from app.core.config import settings
|
||||
from app.interfaces.http.errors import register_exception_handlers
|
||||
from app.interfaces.http.request_context import RequestContextMiddleware
|
||||
from app.interfaces.http.routers import (
|
||||
admin,
|
||||
agy,
|
||||
auth,
|
||||
billing,
|
||||
health,
|
||||
invites,
|
||||
knowledge,
|
||||
organizations,
|
||||
projects,
|
||||
sandbox,
|
||||
workforce,
|
||||
)
|
||||
from app.interfaces.realtime.router import router as realtime_router
|
||||
from app.main import lifespan # shared lifespan installs the realtime bus
|
||||
|
||||
#: Module-level singleton so the RPC proxy and the tests share one app instance.
|
||||
_rpc_app: FastAPI | None = None
|
||||
|
||||
|
||||
def build_rpc_app() -> FastAPI:
|
||||
"""Construct the full handler app. Called once per process."""
|
||||
application = FastAPI(
|
||||
title=settings.app_name,
|
||||
version="0.1.0",
|
||||
description="AI Workforce Platform — internal RPC dispatch surface.",
|
||||
lifespan=lifespan,
|
||||
# This app is reached in-process, never over the network, so the
|
||||
# interactive docs would only be a liability.
|
||||
docs_url=None,
|
||||
openapi_url=None,
|
||||
)
|
||||
application.add_middleware(RequestContextMiddleware)
|
||||
register_exception_handlers(application)
|
||||
|
||||
# Same mounting as the old public app: health at the root, everything else
|
||||
# under /api, and the live socket alongside the RPC surface.
|
||||
application.include_router(health.router)
|
||||
application.include_router(auth.router, prefix="/api")
|
||||
application.include_router(invites.router, prefix="/api")
|
||||
application.include_router(organizations.router, prefix="/api")
|
||||
application.include_router(projects.router, prefix="/api")
|
||||
application.include_router(workforce.router, prefix="/api")
|
||||
application.include_router(billing.router, prefix="/api")
|
||||
application.include_router(knowledge.router, prefix="/api")
|
||||
application.include_router(sandbox.router, prefix="/api")
|
||||
application.include_router(admin.router, prefix="/api")
|
||||
application.include_router(agy.router, prefix="/api")
|
||||
application.include_router(realtime_router, prefix="/api")
|
||||
return application
|
||||
|
||||
|
||||
def get_rpc_app() -> FastAPI:
|
||||
"""Return the process-wide internal app, building it on first use."""
|
||||
global _rpc_app
|
||||
if _rpc_app is None:
|
||||
_rpc_app = build_rpc_app()
|
||||
return _rpc_app
|
||||
Reference in New Issue
Block a user