Coverage for src/local_deep_research/web/services/socketio_asgi.py: 97%
421 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2ASGI Socket.IO integration for FastAPI.
4Replaces Flask-SocketIO's SocketIOService with python-socketio's AsyncServer
5mounted as an ASGI sub-application. This is the pattern used by Open WebUI
6(129k stars) and recommended by the python-socketio documentation.
8Usage:
9 from .socketio_asgi import sio, socket_app, emit_to_subscribers
11 # Mount in FastAPI app:
12 app.mount("/ws", socket_app)
14 # Emit from anywhere (owner is required — see _subscriptions):
15 emit_to_subscribers("research_progress", research_id, data, owner=username)
16"""
18import asyncio
19import contextvars
20import hashlib
21from typing import Any
23import socketio
24from loguru import logger
26from ...constants import ResearchStatus
27from ...utilities.resource_utils import safe_close
28from ..research_state import get_active_research_snapshot
30# Determine WebSocket CORS policy from env var
31from ...settings.env_registry import get_env_setting
32from ...security.log_sanitizer import redact_and_bound_for_log
35def _resolve_socketio_cors(env_value: str | None) -> str | list[str] | None:
36 """Resolve the WebSocket origin allowlist, never yielding ``[]``.
38 None (NOT []) is load-bearing: engine.io treats None as "derive the
39 same-origin whitelist from Host/X-Forwarded-Proto", whereas an empty
40 list DISABLES origin validation entirely (engineio checks
41 ``if self.cors_allowed_origins != []`` before validating), i.e. [] is
42 effectively allow-all for WebSocket handshakes. Every degenerate
43 input (unset, empty, nothing but separators) therefore resolves to
44 None; ``*`` requires the explicit env opt-in. A ``*`` entry anywhere
45 inside a comma-separated list is treated as that same opt-in: engine.io
46 has no "allow-all plus named origins" mode, so a list containing ``*``
47 already accepts every Origin -- returning it as a list would silently
48 behave as allow-all while the startup log claims a finite allowlist.
49 """
50 if env_value is None or not env_value.strip():
51 return None
52 if env_value.strip() == "*":
53 return "*"
54 origins = [o.strip() for o in env_value.split(",") if o.strip()]
55 if "*" in origins:
56 return "*"
57 return origins or None
60_ws_origins_env = get_env_setting("security.websocket.allowed_origins")
61# Plain assignment (not annotated): the startup-ordering freeze contract
62# detects module-scope ast.Assign targets.
63_socketio_cors = _resolve_socketio_cors(_ws_origins_env)
66def _log_socketio_cors_policy(
67 cors: str | list[str] | None, env_value: str | None
68) -> None:
69 """Log the resolved WebSocket origin policy at a default-visible level.
71 The allow-all posture must never be hidden below the default INFO sink.
72 A ``*`` mixed with specific origins is logged at WARNING because it
73 almost certainly is a misconfiguration (the named origins have no
74 effect). Only the operator's own configured entries (stripped) are
75 logged, never request-supplied values.
76 """
77 if cors is None:
78 logger.info(
79 "Socket.IO CORS: same-origin only (set LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS to configure)"
80 )
81 elif cors == "*": 81 ↛ 98line 81 didn't jump to line 98 because the condition on line 81 was always true
82 named = [
83 o
84 for o in (env_value or "").split(",")
85 if o.strip() and o.strip() != "*"
86 ]
87 if named:
88 logger.warning(
89 "Socket.IO CORS: all origins allowed -- "
90 "LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS contains '*' alongside "
91 f"specific origins {[o.strip() for o in named]}; engine.io has "
92 "no allow-all-plus-named mode, so the named origins have no "
93 "effect. Remove '*' to restrict to them."
94 )
95 else:
96 logger.info("Socket.IO CORS: all origins allowed ('*')")
97 else:
98 logger.info(f"Socket.IO CORS: restricted to {cors}")
101def log_socketio_cors_policy() -> None:
102 """Log the import-time-resolved WebSocket origin policy.
104 Deliberately NOT called at module import: this module is imported at
105 the top of ``utilities/log_utils.py``, i.e. while the package's
106 ``logger.disable("local_deep_research")`` (``__init__.py``) is still in
107 effect and before ``config_logger()`` re-enables it, so an import-time
108 log line -- including the allow-all INFO/WARNING -- would be silently
109 dropped. The FastAPI lifespan startup calls this instead, once the
110 server is starting and logging is configured.
111 """
112 _log_socketio_cors_policy(_socketio_cors, _ws_origins_env)
115def _research_id_preview(research_id: Any) -> str:
116 """Bounded preview of a client-sent ``research_id`` for log lines.
118 (#6938: the Socket.IO payload's id is any JSON value of up to the
119 engine.io buffer size, logged on the rejection paths before any
120 ownership match. Imported at call time so this module never imports
121 the routers package at import time.)
122 """
123 from ..routers.notes import _log_value_preview
125 return _log_value_preview(research_id)
128# What engine.io's ``_log_error_once`` receives for a repeat it will not
129# emit (see ``_install_origin_rejection_logging``).
130_ENGINEIO_SUPPRESSED_TEXT = "Rejected (repeat; text not rendered)."
131# ``logging.INFO``: engine.io logs through the standard library, whose
132# level numbers are fixed; this module itself logs only through loguru.
133_STDLIB_INFO_LEVEL = 20
136def _install_origin_rejection_logging(sio: "socketio.AsyncServer") -> bool:
137 """Re-emit engine.io's silenced WebSocket origin rejections via loguru.
139 engine.io validates the Origin at handshake and calls
140 ``_log_error_once('<origin> is not an accepted origin.', 'bad-origin')``,
141 but the server runs with ``logger=False`` so that message never surfaces —
142 the only symptom of a misconfigured WebSocket origin is a frozen progress
143 UI. Wrap that one call to log a WARNING (deduped per origin) pointing at the
144 fix. An Origin is a scheme+host, not PII. Best-effort: a no-op (returns
145 False) if engine.io internals change, so it can never break startup.
147 The dedup set is capped: the handshake is pre-auth and ``Origin`` is
148 attacker-controlled, so an unbounded set would be a memory-growth + log-
149 amplification vector. After ``cap`` distinct origins we stop tracking/warning
150 (an operator has more than enough signal by then).
151 """
152 try:
153 eio = sio.eio
154 original = eio._log_error_once
155 except AttributeError:
156 logger.debug(
157 "Socket.IO: origin-rejection logging not installed "
158 "(engine.io internals changed); handshake rejections stay silent"
159 )
160 return False
162 warned: set[str] = set()
163 cap = 100
165 def _log_error_once(message, message_key):
166 # (#6938: the message embeds the client's Origin header, which
167 # has no length limit on the polling transport, so this module's
168 # warning and the text handed on to engine.io's own logger are the
169 # redacted, bounded form. Redaction still costs a few ms on
170 # adversarial input (``redact_and_bound_for_log`` scans up to 1,200
171 # characters for a 300-character result) and this path is pre-auth
172 # and not rate limited, so it runs only when that text is actually
173 # logged: for each of the first ``cap`` distinct messages, and when
174 # engine.io's logger will emit the line (the first time a key is
175 # seen, unless its level was lowered). The dedup key is a fixed-size SHA-256
176 # digest of the raw message, so deciding whether to warn needs no
177 # redaction; once the cap is reached it is not computed at all.)
178 is_text = isinstance(message, str)
179 redacted: list[str] = []
181 def _bounded() -> str:
182 if not redacted:
183 redacted.append(
184 redact_and_bound_for_log(message, 300)
185 if is_text
186 else type(message).__name__
187 )
188 return redacted[0]
190 if message_key == "bad-origin" and len(warned) < cap:
191 dedup_key = (
192 hashlib.sha256(
193 message.encode("utf-8", "surrogatepass")
194 ).hexdigest()
195 if is_text
196 else type(message).__name__
197 )
198 if dedup_key not in warned:
199 warned.add(dedup_key)
200 logger.warning(
201 f"Socket.IO rejected a WebSocket handshake: {_bounded()} "
202 "Set LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS to this origin "
203 "if it is your front-end; behind a TLS-terminating proxy, "
204 "also forward X-Forwarded-Proto so the same-origin check "
205 "sees https."
206 )
207 if not is_text: 207 ↛ 208line 207 didn't jump to line 208 because the condition on line 207 was never true
208 return original(message, message_key)
209 if _engineio_will_log(message_key):
210 return original(_bounded(), message_key)
211 # engine.io drops this line (a repeat key, logged at INFO under its
212 # ERROR-level logger); hand it a constant so its bookkeeping runs
213 # unchanged without paying for a redaction nobody reads.
214 return original(_ENGINEIO_SUPPRESSED_TEXT, message_key)
216 def _engineio_will_log(message_key) -> bool:
217 """Whether engine.io's ``_log_error_once`` emits for this key.
219 It logs at ERROR the first time a key is seen and at INFO after
220 that (``engineio.base_server``). Anything this cannot determine
221 counts as "will log", so the redacted text is passed on.
222 """
223 try:
224 if message_key not in eio.log_message_keys:
225 return True
226 return bool(eio.logger.isEnabledFor(_STDLIB_INFO_LEVEL))
227 except Exception:
228 return True
230 eio._log_error_once = _log_error_once
231 return True
234# Create the async Socket.IO server
235sio = socketio.AsyncServer(
236 async_mode="asgi",
237 cors_allowed_origins=_socketio_cors,
238 logger=False,
239 engineio_logger=False,
240 ping_timeout=20,
241 ping_interval=5,
242)
244# Make a rejected WebSocket origin diagnosable (otherwise it is a silent
245# frozen progress UI). Skipped for the allow-all case, which rejects
246# nothing.
247if _socketio_cors != "*":
248 _install_origin_rejection_logging(sio)
250# Create the ASGI app for mounting
251# socketio_path must match how the client connects
252socket_app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
254# Subscription tracking: (username, research_id) -> set of sids.
255#
256# Keyed by the OWNER as well as the id, never by the id alone. Research ids
257# are UUIDs, but the benchmark page subscribes with a numeric
258# ``BenchmarkRun.id``, and that id autoincrements inside each user's own
259# encrypted database — so every user's first benchmark run is id 1.
260#
261# The ownership check in ``_owns_research_sync`` cannot catch this, and it is
262# not wrong: it asks "does THIS user own this id in THEIR database?", and for
263# two different users' run 1 the honest answer is yes for both. Keyed by the
264# bare id, both then landed in ``_subscriptions["1"]`` and one user's
265# benchmark progress was emitted to the other's browser. Reproduced:
266# ``_subscriptions[1] = {'sid-alice', 'sid-bob'}``, with Bob receiving
267# Alice's run data.
268#
269# So the id alone is not a key — a correct per-user check plus a global map
270# keyed by a per-user id is still a cross-user leak, and it is invisible in
271# either file on its own. See ADR-0009; ``routers/rag.py`` already keys
272# ``_active_sse_indexers`` this way.
273_subscriptions: dict[tuple[str, str], set[str]] = {}
274# Authenticated socket sessions: sid -> username (resolved from session cookie at connect)
275_sid_users: dict[str, str] = {}
277# sid -> the auth session_id that authorised the handshake. Tracked so
278# logout can sever only the sockets of the session being logged out,
279# leaving the user's other tabs and devices connected (#5535). Without
280# this, single-session logout would either leave every socket alive or
281# have to disconnect all of them. Populated in `connect` and popped in
282# `disconnect`, under the same `_lock` as `_sid_users`.
283_sid_sessions: dict[str, str] = {}
284# Eagerly initialised in init_lock() during lifespan startup. We can't
285# create asyncio.Lock at import time because that binds it to whatever
286# loop (or no loop) is running at import — typically wrong, and on
287# Python 3.12+ raises during pre-startup imports. Initialising once in
288# the lifespan hook (after the running loop exists) gives us a single
289# Lock instance bound to the right loop and avoids the previous lazy
290# double-checked-locking race in _get_lock(): two concurrent first-time
291# callers could each create a fresh Lock, with the second write
292# stomping the first instance and orphaning its awaiters.
293_lock: asyncio.Lock | None = None
296def init_lock() -> None:
297 """Initialise the module-level asyncio.Lock. Call from lifespan
298 startup, after the event loop is running. Idempotent."""
299 global _lock
300 if _lock is None:
301 _lock = asyncio.Lock()
304# Flag suppressing logging during an emit. Set by the log sink itself so an
305# emit's own log lines cannot recurse back into the sink that triggered them.
306#
307# This MUST be a ContextVar, not a threading.local(). The emit no longer runs
308# on the caller's thread: `emit_to_subscribers` schedules a coroutine onto the
309# event loop via `run_coroutine_threadsafe`, so a thread-local set on the
310# calling thread is invisible to the coroutine, which then reads the default
311# (True) and logs anyway. Worse, the caller's `finally` restores the previous
312# value before the coroutine has even run, so the window closed early too.
313#
314# Losing suppression is not cosmetic — it is unbounded amplification. The only
315# caller that disables logging is `frontend_progress_sink`, a loguru sink: a
316# failed emit logs, that record carries the inherited research_id, the sink
317# fires again, emits again, fails again. A single research log line was
318# measured producing 500+ emits before a hard cap tripped.
319#
320# A ContextVar is exactly right here: `run_coroutine_threadsafe` snapshots the
321# calling context, so the value at scheduling time propagates INTO the
322# coroutine, while the caller's own reset stays confined to the caller. It
323# also keeps the per-thread isolation the threading.local was introduced for,
324# since contexts are per-thread and per-task.
325_logging_enabled_var: contextvars.ContextVar[bool] = contextvars.ContextVar(
326 "socketio_logging_enabled", default=True
327)
330def _logging_is_enabled() -> bool:
331 return _logging_enabled_var.get()
334# uvicorn's main event loop, captured at app startup. Background threads
335# (research workers, log queue, scheduler) need this to schedule emits via
336# `asyncio.run_coroutine_threadsafe`. Without it, `asyncio.get_event_loop()`
337# from a worker thread either raises (Python 3.12+) or creates a fresh loop
338# that the AsyncServer doesn't know about, so emits silently no-op.
339_main_loop: asyncio.AbstractEventLoop | None = None
342def set_main_loop(loop: asyncio.AbstractEventLoop) -> None:
343 """Capture the running event loop so background threads can dispatch to it."""
344 global _main_loop
345 _main_loop = loop
348def _get_main_loop() -> asyncio.AbstractEventLoop | None:
349 """Return the captured main loop, or try to find it as a fallback."""
350 if _main_loop is not None and not _main_loop.is_closed():
351 return _main_loop
352 # Last-resort fallback: try the current running loop (works only if
353 # called from the main thread inside an async context).
354 try:
355 return asyncio.get_running_loop()
356 except RuntimeError:
357 return None
360def _schedule_coroutine_threadsafe(coro, loop) -> None:
361 """Hand *coro* to ``loop`` without leaking it on scheduling failure.
363 ``asyncio.run_coroutine_threadsafe`` owns and eventually closes the
364 coroutine after it successfully queues the callback. During shutdown,
365 however, ``loop.call_soon_threadsafe`` can raise synchronously after the
366 caller has already created the coroutine. In that case ownership never
367 transfers, so close it here before preserving the original exception for
368 the public wrapper's best-effort error handling.
369 """
370 try:
371 asyncio.run_coroutine_threadsafe(coro, loop)
372 except Exception:
373 safe_close(coro, "rejected Socket.IO coroutine")
374 raise
377def _decode_session_cookie(cookie_header: str) -> dict | None:
378 """Decode a Starlette SessionMiddleware cookie to extract session data.
380 Uses the same itsdangerous-based scheme Starlette uses, with the
381 project's SECRET_KEY. Returns None on any failure (no session, bad
382 signature, expired).
383 """
384 if not cookie_header:
385 return None
386 try:
387 from http.cookies import SimpleCookie
388 import base64
389 import json
390 import itsdangerous
392 # Look up the session cookie value (same name as SessionMiddleware uses)
393 jar = SimpleCookie()
394 jar.load(cookie_header)
395 morsel = jar.get("session")
396 if morsel is None:
397 return None
398 cookie_value = morsel.value
400 from ..fastapi_app import SECRET_KEY
401 from ...security import get_security_default
403 signer = itsdangerous.TimestampSigner(SECRET_KEY)
404 # Must mirror SessionMiddleware's max_age so an expired/revoked
405 # session cookie that the HTTP path would reject does not
406 # authenticate a WebSocket connection forever.
407 max_age_seconds = (
408 get_security_default("security.session_remember_me_days", 30)
409 * 24
410 * 3600
411 )
412 unsigned = signer.unsign(cookie_value, max_age=max_age_seconds)
413 return json.loads(base64.b64decode(unsigned))
414 except Exception:
415 return None
418@sio.event
419async def connect(sid, environ, auth=None):
420 """Handle client connection.
422 Parity with main's Flask-SocketIO handler: unauthenticated sockets
423 are REJECTED at the handshake (returning False refuses the
424 connection). Accepting them would hand every roomless broadcast —
425 e.g. parallel_search_started, which carries the searching user's
426 query text — to clients that never logged in. The JS client simply
427 reconnects after login, exactly as it did under Flask.
429 The user's DB engine is lazily opened here (race vs first XHR after
430 page load, server restart, idle eviction) so subscribe_to_research's
431 ownership check works immediately after connect.
432 """
433 cookie_header = environ.get("HTTP_COOKIE", "")
434 session_data = _decode_session_cookie(cookie_header) or {}
435 username = session_data.get("username")
437 if not username:
438 logger.info(f"Rejected unauthenticated WebSocket connection from {sid}")
439 return False
441 # The cookie's signature only proves it was issued by us, not that the
442 # session behind it is still alive. Validate the session id against the
443 # server-side store, exactly as require_auth does for HTTP.
444 #
445 # Without this a captured or stale cookie opens a BRAND NEW authenticated
446 # socket after logout / password change / session expiry, and can then
447 # subscribe to that user's live research events. The `is_user_connected`
448 # check below is username-scoped, so it is true whenever ANY of the
449 # user's sessions has the database open -- another device, or an
450 # in-flight research run (logout deliberately leaves the DB open in that
451 # case). That is a common condition, not an edge case.
452 #
453 # This is main's `__session_authorizes` connect-time gate (#5535), which
454 # had no successor here: DatabaseMiddleware's `_enforce_session_revocation`
455 # cannot cover it, because that middleware returns early for non-HTTP
456 # scopes and additionally skips the "/ws/" prefix.
457 from ..auth.session_manager import session_manager
459 handshake_session_id = session_data.get("session_id")
460 if (
461 not handshake_session_id
462 or session_manager.validate_session(handshake_session_id) != username
463 ):
464 logger.warning(
465 f"Rejected WebSocket connection for {username} from {sid}: "
466 "session is not valid server-side (logged out, expired, or "
467 "revoked)"
468 )
469 return False
471 from ...database.encrypted_db import db_manager
472 from ...database.rekey_lock import RekeyOperationalError
473 from ..dependencies.threadpool import run_db_sync
475 # The probe is not the microsecond dict lookup it once was: it routes
476 # through _acquire_current_connection_generation, which takes the
477 # per-user lifecycle lock and the cross-process rekey sidecar
478 # (filesystem I/O, and a wait of up to a few seconds for another
479 # holder). Run it through run_db_sync exactly like the
480 # open_user_database call below, so a rekey in progress holds up only
481 # this handshake's worker thread — not the event loop, which this
482 # single-worker deployment shares with every user's sockets and async
483 # routes.
484 try:
485 user_connected = await run_db_sync(
486 db_manager.is_user_connected, username
487 )
488 except RekeyOperationalError:
489 # The probe already reports a sidecar that cannot be opened or
490 # locked (fd exhaustion, EROFS, EACCES, ENOSPC), and an epoch read
491 # that fails, as a plain False. A failure to release the sidecar
492 # afterwards (unlock/close) raises RekeySidecarIOError instead:
493 # degrade that to "not connected" too — ensure_user_database's
494 # contract — so the handshake fails closed through the normal
495 # no-connection path instead of raising out of connect(). Either
496 # way, while an open/lock or epoch-read fault hides a cached
497 # engine, the lazy open below raises that fault before deriving a
498 # key, so the refused handshake costs no PBKDF2 derivation.
499 user_connected = False
501 if not user_connected:
502 from ...database.session_passwords import session_password_store
504 session_id = session_data.get("session_id")
505 password = (
506 session_password_store.get_session_password(username, session_id)
507 if session_id
508 else None
509 )
510 if not password:
511 logger.info(
512 f"Rejected WebSocket connection for {username}: "
513 "no active DB session and no stored password"
514 )
515 return False
516 try:
517 # SQLCipher open is sync I/O — keep it off the event loop.
518 #
519 # Use run_db_sync, not bare asyncio.to_thread: a cold
520 # open_user_database() call runs the phase-2 rekey path
521 # (_run_phase2_rekey -> rekey_user_indexes), which opens a
522 # get_user_db_session() unconditionally, before its own marker
523 # check (vector_stores/legacy_rekey.py). That leaves a thread-local
524 # session on whichever asyncio default-executor thread served
525 # this call, with no cleanup boundary. run_db_sync runs the sync
526 # call on that same default executor, then cleans up the
527 # thread-local DB session (and other ambient thread-local
528 # state) on that worker before returning -- same shape as the
529 # equivalent fix at fastapi_app.py's DatabaseMiddleware
530 # (ensure_user_database).
531 await run_db_sync(db_manager.open_user_database, username, password)
532 except Exception as e:
533 # Capture the type name only and log OUTSIDE the handler:
534 # logger.exception's diagnose=True traceback would render the
535 # `password` local, so the traceback must never be logged here.
536 open_failure = type(e).__name__
537 else:
538 open_failure = None
539 if open_failure is not None:
540 logger.error(
541 f"Lazy DB open failed for {username} at WebSocket "
542 f"connect: {open_failure}"
543 )
544 return False
546 async with _lock:
547 _sid_users[sid] = username
548 # Remember which auth session authorised this socket so logout can
549 # sever exactly this session's sockets (#5535). Always present here:
550 # the connect gate above rejects a cookie without a valid session id.
551 _sid_sessions[sid] = handshake_session_id
553 # Close the connect-vs-teardown race (main #5572, ported: main fixed this
554 # in web/services/socket_service.py, which this branch replaced, so the
555 # fix had no successor here and would otherwise have been lost).
556 #
557 # The gate above and this registration are not atomic with a concurrent
558 # logout or session expiry. `disconnect_session` severs a session's
559 # sockets by enumerating `_sid_sessions`; if that enumeration ran BEFORE
560 # the block above added this sid, the teardown never sees this socket. It
561 # would then sit registered with a dead session, still receiving that
562 # user's events -- including `settings_changed` -- until some other
563 # session ends.
564 #
565 # Re-validate after registering and fail closed. Returning False rejects
566 # the handshake, so python-socketio never fires our `disconnect` handler
567 # for this sid; the dict entries have to be removed here explicitly.
568 #
569 # This is correct only because teardown invalidates the session BEFORE
570 # enumerating sockets -- logout destroys the session first, then
571 # disconnects -- so a socket that registers after invalidation is
572 # guaranteed to observe the failed re-check.
573 if session_manager.validate_session(handshake_session_id) != username:
574 async with _lock:
575 _sid_users.pop(sid, None)
576 _sid_sessions.pop(sid, None)
577 logger.warning(
578 f"WebSocket connect/teardown race for {username}: session was "
579 f"invalidated while socket {sid} was registering; rejecting it"
580 )
581 return False
583 logger.info(f"Client connected: {sid} (user={username})")
584 return True
587@sio.event
588async def disconnect(sid):
589 """Handle client disconnection.
591 The whole body is guarded, as main's ``__handle_disconnect`` was.
592 ``AsyncServer._trigger_event`` does not swallow handler exceptions and
593 ``_handle_disconnect`` calls ``manager.disconnect(sid, ...)`` -- which
594 removes the sid from engine.io's own room bookkeeping -- only *after*
595 this handler returns. So a raise from here would leave engine.io still
596 holding the sid while our maps are half-swept, which is the opposite of
597 what a disconnect is for.
598 """
599 logger.info(f"Client {sid} disconnected")
600 try:
601 # Clean up subscriptions and identity for this client
602 async with _lock:
603 _sid_users.pop(sid, None)
604 _sid_sessions.pop(sid, None)
605 empty_keys = []
606 # Keys are (username, research_id) tuples; a disconnecting sid is
607 # dropped from every subscription set regardless of owner.
608 for sub_key, sids in _subscriptions.items():
609 sids.discard(sid)
610 if not sids:
611 empty_keys.append(sub_key)
612 for key in empty_keys:
613 del _subscriptions[key]
615 # Clean up thread-local database sessions
616 try:
617 from ...database.thread_local_session import (
618 cleanup_current_thread,
619 )
621 cleanup_current_thread()
622 except Exception:
623 logger.debug("Error cleaning up thread session on disconnect")
624 except Exception:
625 # logger.exception already attaches the traceback; interpolating the
626 # exception as well is redundant and the repo's custom-code check
627 # rejects it.
628 logger.exception(f"Error handling disconnect for {sid}")
631def _owns_research_sync(username: str, research_id: str) -> bool:
632 """Query the user's encrypted DB for research/benchmark ownership.
634 Synchronous (SQLAlchemy) — must be called via ``_user_owns_research``
635 (which offloads to a thread), never directly on the event loop.
636 Shared by on_subscribe and on_unsubscribe so both handlers apply the
637 exact same ownership rule.
638 """
639 from ...database.models import ResearchHistory
640 from ...database.session_context import get_user_db_session
642 with get_user_db_session(username) as db_session:
643 if (
644 db_session.query(ResearchHistory).filter_by(id=research_id).first()
645 is not None
646 ):
647 return True
648 # The benchmark page subscribes with a numeric BenchmarkRun.id,
649 # which has no matching ResearchHistory row. Without this branch
650 # benchmark live-progress subscriptions are silently rejected.
651 if str(research_id).isdigit():
652 from ...database.models.benchmark import BenchmarkRun
654 return (
655 db_session.query(BenchmarkRun.id)
656 .filter(BenchmarkRun.id == research_id)
657 .first()
658 is not None
659 )
660 return False
663def _subscription_key(owner: str, research_id) -> tuple:
664 """Key under which ``research_id``'s subscriptions are stored.
666 Ported from main's ``SocketIOService.__subscription_key`` (#5600) while
667 merging that security fix in. This branch already keyed subscriptions by
668 ``(owner, research_id)`` -- and more strictly than main, which composites
669 only numeric ids -- but it keyed on the RAW id, and main's version also
670 normalizes a numeric id to ``int``.
672 That normalization is load-bearing: subscribe-time ids arrive from a JSON
673 socket payload (so a benchmark run arrives as ``"1"``) while emit-time ids
674 come from the database (``1``). Keyed raw, those are two different dict
675 keys, and the event is silently never delivered to a subscriber who is
676 correctly subscribed. ``isdigit()`` guarantees the ``int()`` is safe.
678 UUID research ids are left exactly as they are.
679 """
680 if str(research_id).isdigit():
681 return (owner, int(research_id))
682 return (owner, research_id)
685async def _user_owns_research(username: str, research_id: str) -> bool:
686 """Return True if `username` owns `research_id` (research or benchmark).
688 Offloaded to a thread — SQLAlchemy is sync and must not block the
689 uvicorn event loop that's serving every other WebSocket + HTTP
690 request. Used by both on_subscribe and on_unsubscribe as the shared
691 authorization boundary for the per-research subscription map.
692 """
693 try:
694 # Use run_db_sync — _owns_research_sync opens get_user_db_session,
695 # so the worker's thread-local DB session needs cleanup before
696 # the next task lands on the same worker.
697 from ..dependencies.threadpool import run_db_sync
699 return await run_db_sync(_owns_research_sync, username, research_id)
700 except Exception:
701 logger.exception(
702 "Ownership check failed for user={} rid={}",
703 username,
704 _research_id_preview(research_id),
705 )
706 return False
709async def _socket_session_still_valid(
710 sid: str, username: str
711) -> tuple[str | None, bool]:
712 """Re-check that this socket's originating session is still alive.
714 Returns ``(session_id, is_valid)``. The id comes back so the caller can
715 tear down EVERY socket of a dead session, not just the one that asked --
716 see the rejection path in ``on_subscribe``.
718 The connect-time gate is not sufficient on its own: identity is captured
719 into ``_sid_users`` at handshake and then frozen for the socket's whole
720 lifetime. A session that expires, is logged out, or has its password
721 changed WHILE the socket is open leaves that socket able to keep acting --
722 including subscribing to research it had not subscribed to before.
724 Logout and password change disconnect sockets actively. Idle expiry has a
725 bounded revocation-latency window until the periodic sweep (5 minutes in
726 production). Revalidating on subscribe and unsubscribe closes that window
727 whenever the socket next performs either action; the sweep remains the
728 fallback for an otherwise idle socket.
730 Fails closed: no recorded session id means the socket predates the gate,
731 which is not a state a current client can reach.
732 """
733 async with _lock:
734 session_id = _sid_sessions.get(sid)
735 if not session_id:
736 return None, False
738 from ..auth.session_manager import session_manager
740 return session_id, session_manager.validate_session(session_id) == username
743@sio.on("subscribe_to_research")
744async def on_subscribe(sid, data):
745 """Handle client subscription to research updates.
747 Verifies the requesting socket is authenticated AND that the
748 requested research_id belongs to that user. Without this check
749 any logged-in user could spy on any other user's research progress
750 in real time by guessing/enumerating UUIDs.
751 """
752 research_id = data.get("research_id") if isinstance(data, dict) else None
753 if not research_id:
754 return
756 async with _lock:
757 username = _sid_users.get(sid)
759 if not username:
760 logger.warning(
761 "Rejected subscribe from unauthenticated sid {} for {}",
762 sid,
763 _research_id_preview(research_id),
764 )
765 await sio.emit(
766 "subscribe_error",
767 {"error": "Authentication required", "research_id": research_id},
768 room=sid,
769 )
770 return
772 stale_session_id, session_ok = await _socket_session_still_valid(
773 sid, username
774 )
775 if not session_ok:
776 logger.warning(
777 f"Rejected subscribe from {sid} ({username}): originating session "
778 "is no longer valid; disconnecting that session's sockets"
779 )
780 try:
781 await sio.emit(
782 "subscribe_error",
783 {"error": "Session expired", "research_id": research_id},
784 room=sid,
785 )
786 except Exception:
787 # Revocation is the security boundary; the explanatory frame is
788 # best-effort. A broken transport must not strand this socket (or
789 # its siblings) after session validation has already failed.
790 if _logging_is_enabled(): 790 ↛ 803line 790 didn't jump to line 803 because the condition on line 790 was always true
791 logger.debug(f"Error notifying stale socket {sid}")
792 # Sever EVERY socket of that session, not just this one.
793 #
794 # `validate_session` DELETES an expired session as a side effect of
795 # checking it, so by the time we get here the session is already gone
796 # from the store. Disconnecting only the calling sid would leave a
797 # sibling socket (a second tab on the same session) connected AND
798 # permanently un-revocable: the periodic idle sweep can never find
799 # that session again to tear it down, so the sibling keeps receiving
800 # the user's events -- including `settings_changed`, which carries
801 # plaintext secrets -- indefinitely. Main revoked per-session for
802 # exactly this reason.
803 if stale_session_id:
804 disconnect_session(stale_session_id)
805 else:
806 try:
807 await sio.disconnect(sid)
808 except Exception:
809 if _logging_is_enabled(): 809 ↛ 811line 809 didn't jump to line 811 because the condition on line 809 was always true
810 logger.debug(f"Error disconnecting stale socket {sid}")
811 return
813 # Verify ownership: the research must exist in this user's encrypted DB.
814 owns = await _user_owns_research(username, research_id)
816 if not owns:
817 logger.warning(
818 "Rejected subscribe: user {} does not own research {}",
819 username,
820 _research_id_preview(research_id),
821 )
822 await sio.emit(
823 "subscribe_error",
824 {"error": "Not authorized", "research_id": research_id},
825 room=sid,
826 )
827 return
829 async with _lock:
830 key = _subscription_key(username, research_id)
831 if key not in _subscriptions: 831 ↛ 833line 831 didn't jump to line 833 because the condition on line 831 was always true
832 _subscriptions[key] = set()
833 _subscriptions[key].add(sid)
835 logger.info(
836 f"Client {sid} (user={username}) subscribed to research {research_id}"
837 )
839 # Send current status immediately if available
840 snapshot = get_active_research_snapshot(research_id)
841 if snapshot is not None:
842 progress = snapshot["progress"]
843 latest_log = snapshot["log"][-1] if snapshot["log"] else None
845 if latest_log:
846 await sio.emit(
847 f"research_progress_{research_id}",
848 {
849 "progress": progress,
850 "message": latest_log.get("message", "Processing..."),
851 "status": ResearchStatus.IN_PROGRESS,
852 "log_entry": latest_log,
853 },
854 room=sid,
855 )
858@sio.on("unsubscribe_from_research")
859async def on_unsubscribe(sid, data):
860 """Handle client unsubscribe from research updates.
862 Applies the same ownership gate as on_subscribe before mutating the
863 per-research subscription set. This restores authorization-boundary
864 consistency with on_subscribe (matching main's own rationale for the
865 check) — it is not closing an active exploit: ``subs.discard(sid)``
866 can only ever remove the CALLER's own sid, so without this gate an
867 unauthorized unsubscribe attempt against another user's research_id
868 was already a silent no-op (that sid was never a member of the set),
869 and the victim's subscription was never affected.
871 Also re-validates the originating session, same as on_subscribe. This
872 reduces idle-expiry revocation latency whenever a socket next sends an
873 unsubscribe: an expired session is severed immediately instead of waiting
874 for the bounded periodic sweep. An otherwise idle socket is still handled
875 by that sweep.
876 """
877 research_id = data.get("research_id") if isinstance(data, dict) else None
878 if not research_id:
879 return
881 async with _lock:
882 username = _sid_users.get(sid)
884 if not username:
885 logger.info(
886 "Rejected unsubscribe from unauthenticated sid {} for {}",
887 sid,
888 _research_id_preview(research_id),
889 )
890 return
892 stale_session_id, session_ok = await _socket_session_still_valid(
893 sid, username
894 )
895 if not session_ok:
896 logger.warning(
897 f"Rejected unsubscribe from {sid} ({username}): originating "
898 "session is no longer valid; disconnecting that session's sockets"
899 )
900 # Same rationale as the rejection branch in on_subscribe: sever
901 # EVERY socket of that session, not just this one. validate_session
902 # already deleted the session from the store as a side effect of
903 # checking it, so a sibling socket left connected would be
904 # permanently un-revocable -- the idle sweep can never find that
905 # session again to tear it down.
906 if stale_session_id:
907 disconnect_session(stale_session_id)
908 else:
909 try:
910 await sio.disconnect(sid)
911 except Exception:
912 if _logging_is_enabled(): 912 ↛ 914line 912 didn't jump to line 914 because the condition on line 912 was always true
913 logger.debug(f"Error disconnecting stale socket {sid}")
914 return
916 if not await _user_owns_research(username, research_id):
917 logger.info(
918 "Rejected unsubscribe from sid {}: user does not own research {}",
919 sid,
920 _research_id_preview(research_id),
921 )
922 return
924 async with _lock:
925 key = _subscription_key(username, research_id)
926 subs = _subscriptions.get(key)
927 if subs:
928 subs.discard(sid)
929 if not subs:
930 _subscriptions.pop(key, None)
933def emit_socket_event(event: str, data: Any, room: str | None = None) -> bool:
934 """Emit a socket event (sync wrapper for use from background threads).
936 Schedules the emit on uvicorn's main event loop via
937 `asyncio.run_coroutine_threadsafe`. Without the captured loop, emits
938 from worker threads silently no-op because `asyncio.get_event_loop()`
939 in a non-main thread on Python 3.12+ either raises or creates a new
940 isolated loop.
941 """
942 loop = _get_main_loop()
943 if loop is None or not loop.is_running():
944 if _logging_is_enabled(): 944 ↛ 946line 944 didn't jump to line 946 because the condition on line 944 was always true
945 logger.debug(f"Cannot emit {event}: main event loop not available")
946 return False
947 try:
948 _schedule_coroutine_threadsafe(_async_emit(event, data, room), loop)
949 return True
950 except Exception:
951 if _logging_is_enabled(): 951 ↛ 953line 951 didn't jump to line 953 because the condition on line 951 was always true
952 logger.debug(f"Error emitting socket event {event}")
953 return False
956def emit_to_user(event: str, username: str, data: Any) -> bool:
957 """Emit a socket event to all sockets authenticated as `username`.
959 Prevents cross-user leaks for events like settings_changed — each
960 user's sockets get the event, nobody else's do.
961 """
962 loop = _get_main_loop()
963 if loop is None or not loop.is_running():
964 if _logging_is_enabled(): 964 ↛ 966line 964 didn't jump to line 966 because the condition on line 964 was always true
965 logger.debug(f"Cannot emit {event}: main event loop not available")
966 return False
968 async def _emit() -> None:
969 # Snapshot sids for this user under the lock — connect/disconnect
970 # mutate _sid_users on this same event loop and could otherwise
971 # interleave at any await point, raising
972 # ``RuntimeError: dictionary changed size during iteration``.
973 async with _lock:
974 sids = [sid for sid, u in _sid_users.items() if u == username]
975 for sid in sids:
976 try:
977 await sio.emit(event, data, room=sid)
978 except Exception:
979 if _logging_is_enabled(): 979 ↛ 975line 979 didn't jump to line 975 because the condition on line 979 was always true
980 logger.debug(f"Error emitting {event} to {sid}")
982 try:
983 _schedule_coroutine_threadsafe(_emit(), loop)
984 return True
985 except Exception:
986 if _logging_is_enabled(): 986 ↛ 988line 986 didn't jump to line 988 because the condition on line 986 was always true
987 logger.debug(f"Error emitting {event} to user {username}")
988 return False
991def disconnect_user(username: str) -> bool:
992 """Disconnect every socket authenticated as ``username``.
994 A socket is authorised once, at handshake, and never re-checked. So
995 when a user's session ends — logout, password change, or the idle
996 sweep closing their database because they have no session left — any
997 socket they still hold keeps receiving their events indefinitely. This
998 severs them.
1000 Port of ``SocketIOService.disconnect_user`` (#5535), which grouped
1001 sockets with Socket.IO rooms. This layer already tracks identity
1002 directly in ``_sid_users``, so it filters that instead of maintaining
1003 a parallel per-user room; the snapshot is taken under ``_lock`` for
1004 the same reason ``emit_to_user`` takes it — connect/disconnect mutate
1005 ``_sid_users`` on this event loop and would otherwise interleave at an
1006 await point.
1008 Dropping subscriptions is handled for free: ``sio.disconnect(sid)``
1009 fires the ``disconnect`` handler above, which removes the sid from
1010 ``_sid_users`` and from every entry in ``_subscriptions``.
1012 Best-effort and non-blocking, like the emit helpers: schedules onto
1013 the main loop and returns whether it could be scheduled, not whether
1014 the sockets are gone. Callers are teardown paths that must not block
1015 or raise.
1016 """
1017 return _disconnect_matching(
1018 lambda: [sid for sid, u in _sid_users.items() if u == username],
1019 f"user {username}",
1020 )
1023def disconnect_session(session_id: str) -> bool:
1024 """Disconnect only the sockets authorised by ``session_id``.
1026 Single-session teardown for logout: the tab being logged out loses its
1027 sockets while the user's other still-valid sessions keep theirs. Port
1028 of ``SocketIOService.disconnect_session`` (#5535), which used a
1029 per-session Socket.IO room; this layer matches on the handshake
1030 session_id recorded in ``_sid_sessions`` instead.
1032 Use ``disconnect_user`` for the all-sessions cases — password change
1033 and the idle sweep — where every session is destroyed anyway.
1034 """
1035 return _disconnect_matching(
1036 lambda: [sid for sid, s in _sid_sessions.items() if s == session_id],
1037 f"session {session_id[:8]}...",
1038 )
1041def _disconnect_matching(select_sids, description: str) -> bool:
1042 """Disconnect every sid returned by ``select_sids``.
1044 ``select_sids`` is called while holding ``_lock`` — connect/disconnect
1045 mutate the sid maps on this event loop, so the selection must be a
1046 snapshot taken under the lock rather than a live view, or iteration can
1047 raise ``RuntimeError: dictionary changed size during iteration``.
1048 """
1049 loop = _get_main_loop()
1050 if loop is None or not loop.is_running():
1051 if _logging_is_enabled(): 1051 ↛ 1056line 1051 didn't jump to line 1056 because the condition on line 1051 was always true
1052 logger.debug(
1053 f"Cannot disconnect sockets for {description}: "
1054 "main event loop not available"
1055 )
1056 return False
1058 async def _disconnect() -> None:
1059 async with _lock:
1060 sids = select_sids()
1061 for sid in sids:
1062 try:
1063 await sio.disconnect(sid)
1064 except Exception:
1065 if _logging_is_enabled(): 1065 ↛ 1061line 1065 didn't jump to line 1061 because the condition on line 1065 was always true
1066 logger.debug(f"Error disconnecting {sid}")
1067 if sids:
1068 logger.info(f"Disconnected {len(sids)} socket(s) for {description}")
1070 try:
1071 _schedule_coroutine_threadsafe(_disconnect(), loop)
1072 return True
1073 except Exception:
1074 if _logging_is_enabled(): 1074 ↛ 1076line 1074 didn't jump to line 1076 because the condition on line 1074 was always true
1075 logger.debug(f"Error disconnecting sockets for {description}")
1076 return False
1079def emit_to_subscribers(
1080 event_base: str,
1081 research_id: str,
1082 data: Any,
1083 *,
1084 owner: str,
1085 enable_logging: bool = True,
1086) -> bool:
1087 """Emit an event to all subscribers of a specific research.
1089 Sync wrapper for use from background threads.
1091 ``owner`` is the username whose research this is, and is REQUIRED and
1092 keyword-only on purpose: subscriptions are keyed by ``(owner,
1093 research_id)`` because a benchmark id is only unique within one user's
1094 database (see ``_subscriptions``). Making it required means a call site
1095 that forgets it fails loudly at the call rather than silently reverting
1096 to the id-only lookup that leaked one user's progress to another.
1098 `enable_logging` is carried by a ContextVar so concurrent emits from
1099 different worker threads don't corrupt each other's suppression windows
1100 (as a module global did), AND so the suppression reaches the coroutine
1101 scheduled onto the event loop below — a threading.local could not, since
1102 that coroutine runs on a different thread. See `_logging_enabled_var`.
1103 """
1104 token = None
1105 if not enable_logging:
1106 token = _logging_enabled_var.set(False)
1108 try:
1109 loop = _get_main_loop()
1110 if loop is None or not loop.is_running():
1111 if _logging_is_enabled():
1112 logger.debug(
1113 f"Cannot emit subscribers for {research_id}: loop unavailable"
1114 )
1115 return False
1116 try:
1117 _schedule_coroutine_threadsafe(
1118 _async_emit_to_subscribers(
1119 event_base, research_id, data, owner
1120 ),
1121 loop,
1122 )
1123 return True
1124 except Exception:
1125 if _logging_is_enabled():
1126 logger.debug(
1127 f"Error emitting to subscribers for research {research_id}"
1128 )
1129 return False
1130 finally:
1131 # Reset only this caller's context. The coroutine scheduled above got
1132 # its own snapshot taken while the flag was still False, so it keeps
1133 # the suppression it needs regardless of this reset.
1134 if token is not None:
1135 _logging_enabled_var.reset(token)
1138async def _async_emit(event: str, data: Any, room: str | None = None) -> None:
1139 """Async emit implementation."""
1140 try:
1141 if room:
1142 await sio.emit(event, data, room=room)
1143 else:
1144 await sio.emit(event, data)
1145 except Exception:
1146 # The concurrent Future returned by run_coroutine_threadsafe is
1147 # intentionally discarded by the synchronous wrapper, so an
1148 # exception escaping this coroutine would otherwise be silent.
1149 if _logging_is_enabled(): 1149 ↛ exitline 1149 didn't return from function '_async_emit' because the condition on line 1149 was always true
1150 logger.debug(f"Error emitting socket event {event}")
1153async def _async_emit_to_subscribers(
1154 event_base: str, research_id: str, data: Any, owner: str
1155) -> None:
1156 """Async emit to subscribers implementation.
1158 Looks up ``(owner, research_id)``. An unknown or wrong owner simply
1159 matches no entry and the event is dropped — the map cannot be reached
1160 with an id alone, so a caller that loses track of whose research it is
1161 emitting delivers nothing rather than delivering to everyone.
1162 """
1163 full_event = f"{event_base}_{research_id}"
1165 async with _lock:
1166 subscriptions = _subscriptions.get(
1167 _subscription_key(owner, research_id)
1168 )
1169 if subscriptions:
1170 subscriptions = subscriptions.copy()
1171 else:
1172 subscriptions = None
1174 if subscriptions is not None:
1175 from ..auth.session_manager import session_manager
1177 for sid in subscriptions:
1178 try:
1179 await sio.emit(full_event, data, room=sid)
1180 except Exception:
1181 if _logging_is_enabled():
1182 logger.debug(f"Error emitting to subscriber {sid}")
1184 # Watching a run over the socket IS activity. Without this, a
1185 # user who starts a long research and then just watches the
1186 # progress UI issues no HTTP requests, so nothing refreshes
1187 # their idle timer and the session is reaped mid-run -- the
1188 # idle sweep then disconnects the socket too, so the run
1189 # appears to die in front of them.
1190 #
1191 # This is the second half of main's #5535 (778688295). The
1192 # disconnect_user/disconnect_session half was ported; this half
1193 # was not, which is why touch_session had zero call sites here.
1194 # Own try/except, as main has, so a refresh failure cannot
1195 # strand delivery to the remaining subscribers.
1196 try:
1197 sid_session = _sid_sessions.get(sid)
1198 if sid_session:
1199 session_manager.touch_session(sid_session)
1200 except Exception:
1201 if _logging_is_enabled(): 1201 ↛ 1177line 1201 didn't jump to line 1177 because the condition on line 1201 was always true
1202 logger.debug(f"Error refreshing session for {sid}")
1203 # No subscribers: drop the event. We must NOT broadcast — that would
1204 # leak one user's research progress to every connected client.
1207def remove_subscriptions_for_research(research_id: str, owner: str) -> None:
1208 """Remove all socket subscriptions for a completed research.
1210 ``owner`` scopes the removal to the user whose research finished, for
1211 the same reason subscriptions are keyed that way: without it, one
1212 user's benchmark finishing would tear down every other user's
1213 subscription to their own run of the same numeric id.
1215 Sync function safe to call from any thread. Schedules the async
1216 cleanup on the main loop. If the loop is not yet available (very
1217 early startup) or has been closed (shutdown), the cleanup is
1218 silently dropped — disconnect handlers will reap stale entries on
1219 their own, so the worst case is a small temporary leak that does
1220 not race with the async path.
1221 """
1222 loop = _get_main_loop()
1223 if loop is None or not loop.is_running():
1224 logger.debug(
1225 f"Skipping subscription cleanup for {research_id}: loop unavailable"
1226 )
1227 return
1228 try:
1229 _schedule_coroutine_threadsafe(
1230 _async_remove_subscriptions(research_id, owner), loop
1231 )
1232 except Exception:
1233 logger.debug(
1234 f"Async cleanup scheduling failed for {research_id}", exc_info=True
1235 )
1238async def _async_remove_subscriptions(research_id: str, owner: str) -> None:
1239 """Async subscription removal, scoped to the research's owner."""
1240 async with _lock:
1241 removed = _subscriptions.pop(
1242 _subscription_key(owner, research_id), None
1243 )
1244 if removed is not None:
1245 logger.info(
1246 f"Removed {len(removed)} subscription(s) for research {research_id}"
1247 )