Coverage for src/local_deep_research/web/services/socketio_asgi.py: 97%

1""" 

2ASGI Socket.IO integration for FastAPI. 

3 

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. 

7 

8Usage: 

9 from .socketio_asgi import sio, socket_app, emit_to_subscribers 

10 

11 # Mount in FastAPI app: 

12 app.mount("/ws", socket_app) 

13 

14 # Emit from anywhere (owner is required — see _subscriptions): 

15 emit_to_subscribers("research_progress", research_id, data, owner=username) 

16""" 

17 

18import asyncio 

19import contextvars 

20from typing import Any 

21 

22import socketio 

23from loguru import logger 

24 

25from ...constants import ResearchStatus 

26from ...utilities.resource_utils import safe_close 

27from ..research_state import get_active_research_snapshot 

28 

29# Determine WebSocket CORS policy from env var 

30from ...settings.env_registry import get_env_setting 

31 

32_ws_origins_env = get_env_setting("security.websocket.allowed_origins") 

33_socketio_cors: str | list[str] | None 

34if _ws_origins_env is not None: 

35 if _ws_origins_env == "*": 

36 _socketio_cors = "*" 

37 elif _ws_origins_env: 

38 _socketio_cors = [o.strip() for o in _ws_origins_env.split(",")] 

39 else: 

40 _socketio_cors = None 

41else: 

42 # No env var set — fail closed to same-origin only, matching HTTP CORS 

43 # default. Operators who need cross-origin WS access must set 

44 # LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS explicitly. 

45 # 

46 # None (NOT []) is load-bearing: engine.io treats None as "derive the 

47 # same-origin whitelist from Host/X-Forwarded-Proto", whereas an empty 

48 # list DISABLES origin validation entirely (engineio checks 

49 # `if self.cors_allowed_origins != []` before validating), i.e. [] is 

50 # effectively allow-all for WebSocket handshakes. 

51 _socketio_cors = None 

52 

53if _socketio_cors is None: 

54 logger.info( 

55 "Socket.IO CORS: same-origin only (set LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS to configure)" 

56 ) 

57elif _socketio_cors == "*": 

58 logger.debug("Socket.IO CORS: all origins allowed") 

59else: 

60 logger.info(f"Socket.IO CORS: restricted to {_socketio_cors}") 

61 

62 

63def _install_origin_rejection_logging(sio: "socketio.AsyncServer") -> bool: 

64 """Re-emit engine.io's silenced WebSocket origin rejections via loguru. 

65 

66 engine.io validates the Origin at handshake and calls 

67 ``_log_error_once('<origin> is not an accepted origin.', 'bad-origin')``, 

68 but the server runs with ``logger=False`` so that message never surfaces — 

69 the only symptom of a misconfigured WebSocket origin is a frozen progress 

70 UI. Wrap that one call to log a WARNING (deduped per origin) pointing at the 

71 fix. An Origin is a scheme+host, not PII. Best-effort: a no-op (returns 

72 False) if engine.io internals change, so it can never break startup. 

73 

74 The dedup set is capped: the handshake is pre-auth and ``Origin`` is 

75 attacker-controlled, so an unbounded set would be a memory-growth + log- 

76 amplification vector. After ``cap`` distinct origins we stop tracking/warning 

77 (an operator has more than enough signal by then). 

78 """ 

79 try: 

80 eio = sio.eio 

81 original = eio._log_error_once 

82 except AttributeError: 

83 logger.debug( 

84 "Socket.IO: origin-rejection logging not installed " 

85 "(engine.io internals changed); handshake rejections stay silent" 

86 ) 

87 return False 

88 

89 warned: set[str] = set() 

90 cap = 100 

91 

92 def _log_error_once(message, message_key): 

93 if ( 

94 message_key == "bad-origin" 

95 and len(warned) < cap 

96 and message not in warned 

97 ): 

98 warned.add(message) 

99 logger.warning( 

100 f"Socket.IO rejected a WebSocket handshake: {message} Set " 

101 "LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS to this origin if it is " 

102 "your front-end; behind a TLS-terminating proxy, also forward " 

103 "X-Forwarded-Proto so the same-origin check sees https." 

104 ) 

105 return original(message, message_key) 

106 

107 eio._log_error_once = _log_error_once 

108 return True 

109 

110 

111# Create the async Socket.IO server 

112sio = socketio.AsyncServer( 

113 async_mode="asgi", 

114 cors_allowed_origins=_socketio_cors, 

115 logger=False, 

116 engineio_logger=False, 

117 ping_timeout=20, 

118 ping_interval=5, 

119) 

120 

121# Make a rejected WebSocket origin diagnosable (otherwise it is a silent 

122# frozen progress UI). Skipped for the allow-all case, which rejects 

123# nothing. 

124if _socketio_cors != "*": 

125 _install_origin_rejection_logging(sio) 

126 

127# Create the ASGI app for mounting 

128# socketio_path must match how the client connects 

129socket_app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io") 

130 

131# Subscription tracking: (username, research_id) -> set of sids. 

132# 

133# Keyed by the OWNER as well as the id, never by the id alone. Research ids 

134# are UUIDs, but the benchmark page subscribes with a numeric 

135# ``BenchmarkRun.id``, and that id autoincrements inside each user's own 

136# encrypted database — so every user's first benchmark run is id 1. 

137# 

138# The ownership check in ``_owns_research_sync`` cannot catch this, and it is 

139# not wrong: it asks "does THIS user own this id in THEIR database?", and for 

140# two different users' run 1 the honest answer is yes for both. Keyed by the 

141# bare id, both then landed in ``_subscriptions["1"]`` and one user's 

142# benchmark progress was emitted to the other's browser. Reproduced: 

143# ``_subscriptions[1] = {'sid-alice', 'sid-bob'}``, with Bob receiving 

144# Alice's run data. 

145# 

146# So the id alone is not a key — a correct per-user check plus a global map 

147# keyed by a per-user id is still a cross-user leak, and it is invisible in 

148# either file on its own. See ADR-0009; ``routers/rag.py`` already keys 

149# ``_active_sse_indexers`` this way. 

150_subscriptions: dict[tuple[str, str], set[str]] = {} 

151# Authenticated socket sessions: sid -> username (resolved from session cookie at connect) 

152_sid_users: dict[str, str] = {} 

153 

154# sid -> the auth session_id that authorised the handshake. Tracked so 

155# logout can sever only the sockets of the session being logged out, 

156# leaving the user's other tabs and devices connected (#5535). Without 

157# this, single-session logout would either leave every socket alive or 

158# have to disconnect all of them. Populated in `connect` and popped in 

159# `disconnect`, under the same `_lock` as `_sid_users`. 

160_sid_sessions: dict[str, str] = {} 

161# Eagerly initialised in init_lock() during lifespan startup. We can't 

162# create asyncio.Lock at import time because that binds it to whatever 

163# loop (or no loop) is running at import — typically wrong, and on 

164# Python 3.12+ raises during pre-startup imports. Initialising once in 

165# the lifespan hook (after the running loop exists) gives us a single 

166# Lock instance bound to the right loop and avoids the previous lazy 

167# double-checked-locking race in _get_lock(): two concurrent first-time 

168# callers could each create a fresh Lock, with the second write 

169# stomping the first instance and orphaning its awaiters. 

170_lock: asyncio.Lock | None = None 

171 

172 

173def init_lock() -> None: 

174 """Initialise the module-level asyncio.Lock. Call from lifespan 

175 startup, after the event loop is running. Idempotent.""" 

176 global _lock 

177 if _lock is None: 

178 _lock = asyncio.Lock() 

179 

180 

181# Flag suppressing logging during an emit. Set by the log sink itself so an 

182# emit's own log lines cannot recurse back into the sink that triggered them. 

183# 

184# This MUST be a ContextVar, not a threading.local(). The emit no longer runs 

185# on the caller's thread: `emit_to_subscribers` schedules a coroutine onto the 

186# event loop via `run_coroutine_threadsafe`, so a thread-local set on the 

187# calling thread is invisible to the coroutine, which then reads the default 

188# (True) and logs anyway. Worse, the caller's `finally` restores the previous 

189# value before the coroutine has even run, so the window closed early too. 

190# 

191# Losing suppression is not cosmetic — it is unbounded amplification. The only 

192# caller that disables logging is `frontend_progress_sink`, a loguru sink: a 

193# failed emit logs, that record carries the inherited research_id, the sink 

194# fires again, emits again, fails again. A single research log line was 

195# measured producing 500+ emits before a hard cap tripped. 

196# 

197# A ContextVar is exactly right here: `run_coroutine_threadsafe` snapshots the 

198# calling context, so the value at scheduling time propagates INTO the 

199# coroutine, while the caller's own reset stays confined to the caller. It 

200# also keeps the per-thread isolation the threading.local was introduced for, 

201# since contexts are per-thread and per-task. 

202_logging_enabled_var: contextvars.ContextVar[bool] = contextvars.ContextVar( 

203 "socketio_logging_enabled", default=True 

204) 

205 

206 

207def _logging_is_enabled() -> bool: 

208 return _logging_enabled_var.get() 

209 

210 

211# uvicorn's main event loop, captured at app startup. Background threads 

212# (research workers, log queue, scheduler) need this to schedule emits via 

213# `asyncio.run_coroutine_threadsafe`. Without it, `asyncio.get_event_loop()` 

214# from a worker thread either raises (Python 3.12+) or creates a fresh loop 

215# that the AsyncServer doesn't know about, so emits silently no-op. 

216_main_loop: asyncio.AbstractEventLoop | None = None 

217 

218 

219def set_main_loop(loop: asyncio.AbstractEventLoop) -> None: 

220 """Capture the running event loop so background threads can dispatch to it.""" 

221 global _main_loop 

222 _main_loop = loop 

223 

224 

225def _get_main_loop() -> asyncio.AbstractEventLoop | None: 

226 """Return the captured main loop, or try to find it as a fallback.""" 

227 if _main_loop is not None and not _main_loop.is_closed(): 

228 return _main_loop 

229 # Last-resort fallback: try the current running loop (works only if 

230 # called from the main thread inside an async context). 

231 try: 

232 return asyncio.get_running_loop() 

233 except RuntimeError: 

234 return None 

235 

236 

237def _schedule_coroutine_threadsafe(coro, loop) -> None: 

238 """Hand *coro* to ``loop`` without leaking it on scheduling failure. 

239 

240 ``asyncio.run_coroutine_threadsafe`` owns and eventually closes the 

241 coroutine after it successfully queues the callback. During shutdown, 

242 however, ``loop.call_soon_threadsafe`` can raise synchronously after the 

243 caller has already created the coroutine. In that case ownership never 

244 transfers, so close it here before preserving the original exception for 

245 the public wrapper's best-effort error handling. 

246 """ 

247 try: 

248 asyncio.run_coroutine_threadsafe(coro, loop) 

249 except Exception: 

250 safe_close(coro, "rejected Socket.IO coroutine") 

251 raise 

252 

253 

254def _decode_session_cookie(cookie_header: str) -> dict | None: 

255 """Decode a Starlette SessionMiddleware cookie to extract session data. 

256 

257 Uses the same itsdangerous-based scheme Starlette uses, with the 

258 project's SECRET_KEY. Returns None on any failure (no session, bad 

259 signature, expired). 

260 """ 

261 if not cookie_header: 

262 return None 

263 try: 

264 from http.cookies import SimpleCookie 

265 import base64 

266 import json 

267 import itsdangerous 

268 

269 # Look up the session cookie value (same name as SessionMiddleware uses) 

270 jar = SimpleCookie() 

271 jar.load(cookie_header) 

272 morsel = jar.get("session") 

273 if morsel is None: 

274 return None 

275 cookie_value = morsel.value 

276 

277 from ..fastapi_app import SECRET_KEY 

278 from ...security import get_security_default 

279 

280 signer = itsdangerous.TimestampSigner(SECRET_KEY) 

281 # Must mirror SessionMiddleware's max_age so an expired/revoked 

282 # session cookie that the HTTP path would reject does not 

283 # authenticate a WebSocket connection forever. 

284 max_age_seconds = ( 

285 get_security_default("security.session_remember_me_days", 30) 

286 * 24 

287 * 3600 

288 ) 

289 unsigned = signer.unsign(cookie_value, max_age=max_age_seconds) 

290 return json.loads(base64.b64decode(unsigned)) 

291 except Exception: 

292 return None 

293 

294 

295@sio.event 

296async def connect(sid, environ, auth=None): 

297 """Handle client connection. 

298 

299 Parity with main's Flask-SocketIO handler: unauthenticated sockets 

300 are REJECTED at the handshake (returning False refuses the 

301 connection). Accepting them would hand every roomless broadcast — 

302 e.g. parallel_search_started, which carries the searching user's 

303 query text — to clients that never logged in. The JS client simply 

304 reconnects after login, exactly as it did under Flask. 

305 

306 The user's DB engine is lazily opened here (race vs first XHR after 

307 page load, server restart, idle eviction) so subscribe_to_research's 

308 ownership check works immediately after connect. 

309 """ 

310 cookie_header = environ.get("HTTP_COOKIE", "") 

311 session_data = _decode_session_cookie(cookie_header) or {} 

312 username = session_data.get("username") 

313 

314 if not username: 

315 logger.info(f"Rejected unauthenticated WebSocket connection from {sid}") 

316 return False 

317 

318 # The cookie's signature only proves it was issued by us, not that the 

319 # session behind it is still alive. Validate the session id against the 

320 # server-side store, exactly as require_auth does for HTTP. 

321 # 

322 # Without this a captured or stale cookie opens a BRAND NEW authenticated 

323 # socket after logout / password change / session expiry, and can then 

324 # subscribe to that user's live research events. The `is_user_connected` 

325 # check below is username-scoped, so it is true whenever ANY of the 

326 # user's sessions has the database open -- another device, or an 

327 # in-flight research run (logout deliberately leaves the DB open in that 

328 # case). That is a common condition, not an edge case. 

329 # 

330 # This is main's `__session_authorizes` connect-time gate (#5535), which 

331 # had no successor here: DatabaseMiddleware's `_enforce_session_revocation` 

332 # cannot cover it, because that middleware returns early for non-HTTP 

333 # scopes and additionally skips the "/ws/" prefix. 

334 from ..auth.session_manager import session_manager 

335 

336 handshake_session_id = session_data.get("session_id") 

337 if ( 

338 not handshake_session_id 

339 or session_manager.validate_session(handshake_session_id) != username 

340 ): 

341 logger.warning( 

342 f"Rejected WebSocket connection for {username} from {sid}: " 

343 "session is not valid server-side (logged out, expired, or " 

344 "revoked)" 

345 ) 

346 return False 

347 

348 from ...database.encrypted_db import db_manager 

349 

350 if not db_manager.is_user_connected(username): 

351 from ...database.session_passwords import session_password_store 

352 

353 session_id = session_data.get("session_id") 

354 password = ( 

355 session_password_store.get_session_password(username, session_id) 

356 if session_id 

357 else None 

358 ) 

359 if not password: 

360 logger.info( 

361 f"Rejected WebSocket connection for {username}: " 

362 "no active DB session and no stored password" 

363 ) 

364 return False 

365 try: 

366 # SQLCipher open is sync I/O — keep it off the event loop. 

367 # 

368 # Use run_db_sync, not bare asyncio.to_thread: a cold 

369 # open_user_database() call runs the phase-2 rekey path 

370 # (_run_phase2_rekey -> rekey_user_indexes), which opens a 

371 # get_user_db_session() unconditionally, before its own marker 

372 # check (vector_stores/legacy_rekey.py). That leaves a thread-local 

373 # session on whichever asyncio default-executor thread served 

374 # this call, with no cleanup boundary. run_db_sync runs the sync 

375 # call on that same default executor, then cleans up the 

376 # thread-local DB session (and other ambient thread-local 

377 # state) on that worker before returning -- same shape as the 

378 # equivalent fix at fastapi_app.py's DatabaseMiddleware 

379 # (ensure_user_database). 

380 from ..dependencies.threadpool import run_db_sync 

381 

382 await run_db_sync(db_manager.open_user_database, username, password) 

383 except Exception as e: 

384 # Capture the type name only and log OUTSIDE the handler: 

385 # logger.exception's diagnose=True traceback would render the 

386 # `password` local, so the traceback must never be logged here. 

387 open_failure = type(e).__name__ 

388 else: 

389 open_failure = None 

390 if open_failure is not None: 

391 logger.error( 

392 f"Lazy DB open failed for {username} at WebSocket " 

393 f"connect: {open_failure}" 

394 ) 

395 return False 

396 

397 async with _lock: 

398 _sid_users[sid] = username 

399 # Remember which auth session authorised this socket so logout can 

400 # sever exactly this session's sockets (#5535). Always present here: 

401 # the connect gate above rejects a cookie without a valid session id. 

402 _sid_sessions[sid] = handshake_session_id 

403 

404 # Close the connect-vs-teardown race (main #5572, ported: main fixed this 

405 # in web/services/socket_service.py, which this branch replaced, so the 

406 # fix had no successor here and would otherwise have been lost). 

407 # 

408 # The gate above and this registration are not atomic with a concurrent 

409 # logout or session expiry. `disconnect_session` severs a session's 

410 # sockets by enumerating `_sid_sessions`; if that enumeration ran BEFORE 

411 # the block above added this sid, the teardown never sees this socket. It 

412 # would then sit registered with a dead session, still receiving that 

413 # user's events -- including `settings_changed`, which carries plaintext 

414 # secrets -- until some other session ends. 

415 # 

416 # Re-validate after registering and fail closed. Returning False rejects 

417 # the handshake, so python-socketio never fires our `disconnect` handler 

418 # for this sid; the dict entries have to be removed here explicitly. 

419 # 

420 # This is correct only because teardown invalidates the session BEFORE 

421 # enumerating sockets -- logout destroys the session first, then 

422 # disconnects -- so a socket that registers after invalidation is 

423 # guaranteed to observe the failed re-check. 

424 if session_manager.validate_session(handshake_session_id) != username: 

425 async with _lock: 

426 _sid_users.pop(sid, None) 

427 _sid_sessions.pop(sid, None) 

428 logger.warning( 

429 f"WebSocket connect/teardown race for {username}: session was " 

430 f"invalidated while socket {sid} was registering; rejecting it" 

431 ) 

432 return False 

433 

434 logger.info(f"Client connected: {sid} (user={username})") 

435 return True 

436 

437 

438@sio.event 

439async def disconnect(sid): 

440 """Handle client disconnection. 

441 

442 The whole body is guarded, as main's ``__handle_disconnect`` was. 

443 ``AsyncServer._trigger_event`` does not swallow handler exceptions and 

444 ``_handle_disconnect`` calls ``manager.disconnect(sid, ...)`` -- which 

445 removes the sid from engine.io's own room bookkeeping -- only *after* 

446 this handler returns. So a raise from here would leave engine.io still 

447 holding the sid while our maps are half-swept, which is the opposite of 

448 what a disconnect is for. 

449 """ 

450 logger.info(f"Client {sid} disconnected") 

451 try: 

452 # Clean up subscriptions and identity for this client 

453 async with _lock: 

454 _sid_users.pop(sid, None) 

455 _sid_sessions.pop(sid, None) 

456 empty_keys = [] 

457 # Keys are (username, research_id) tuples; a disconnecting sid is 

458 # dropped from every subscription set regardless of owner. 

459 for sub_key, sids in _subscriptions.items(): 

460 sids.discard(sid) 

461 if not sids: 

462 empty_keys.append(sub_key) 

463 for key in empty_keys: 

464 del _subscriptions[key] 

465 

466 # Clean up thread-local database sessions 

467 try: 

468 from ...database.thread_local_session import ( 

469 cleanup_current_thread, 

470 ) 

471 

472 cleanup_current_thread() 

473 except Exception: 

474 logger.debug("Error cleaning up thread session on disconnect") 

475 except Exception: 

476 # logger.exception already attaches the traceback; interpolating the 

477 # exception as well is redundant and the repo's custom-code check 

478 # rejects it. 

479 logger.exception(f"Error handling disconnect for {sid}") 

480 

481 

482def _owns_research_sync(username: str, research_id: str) -> bool: 

483 """Query the user's encrypted DB for research/benchmark ownership. 

484 

485 Synchronous (SQLAlchemy) — must be called via ``_user_owns_research`` 

486 (which offloads to a thread), never directly on the event loop. 

487 Shared by on_subscribe and on_unsubscribe so both handlers apply the 

488 exact same ownership rule. 

489 """ 

490 from ...database.models import ResearchHistory 

491 from ...database.session_context import get_user_db_session 

492 

493 with get_user_db_session(username) as db_session: 

494 if ( 

495 db_session.query(ResearchHistory).filter_by(id=research_id).first() 

496 is not None 

497 ): 

498 return True 

499 # The benchmark page subscribes with a numeric BenchmarkRun.id, 

500 # which has no matching ResearchHistory row. Without this branch 

501 # benchmark live-progress subscriptions are silently rejected. 

502 if str(research_id).isdigit(): 

503 from ...database.models.benchmark import BenchmarkRun 

504 

505 return ( 

506 db_session.query(BenchmarkRun.id) 

507 .filter(BenchmarkRun.id == research_id) 

508 .first() 

509 is not None 

510 ) 

511 return False 

512 

513 

514def _subscription_key(owner: str, research_id) -> tuple: 

515 """Key under which ``research_id``'s subscriptions are stored. 

516 

517 Ported from main's ``SocketIOService.__subscription_key`` (#5600) while 

518 merging that security fix in. This branch already keyed subscriptions by 

519 ``(owner, research_id)`` -- and more strictly than main, which composites 

520 only numeric ids -- but it keyed on the RAW id, and main's version also 

521 normalizes a numeric id to ``int``. 

522 

523 That normalization is load-bearing: subscribe-time ids arrive from a JSON 

524 socket payload (so a benchmark run arrives as ``"1"``) while emit-time ids 

525 come from the database (``1``). Keyed raw, those are two different dict 

526 keys, and the event is silently never delivered to a subscriber who is 

527 correctly subscribed. ``isdigit()`` guarantees the ``int()`` is safe. 

528 

529 UUID research ids are left exactly as they are. 

530 """ 

531 if str(research_id).isdigit(): 

532 return (owner, int(research_id)) 

533 return (owner, research_id) 

534 

535 

536async def _user_owns_research(username: str, research_id: str) -> bool: 

537 """Return True if `username` owns `research_id` (research or benchmark). 

538 

539 Offloaded to a thread — SQLAlchemy is sync and must not block the 

540 uvicorn event loop that's serving every other WebSocket + HTTP 

541 request. Used by both on_subscribe and on_unsubscribe as the shared 

542 authorization boundary for the per-research subscription map. 

543 """ 

544 try: 

545 # Use run_db_sync — _owns_research_sync opens get_user_db_session, 

546 # so the worker's thread-local DB session needs cleanup before 

547 # the next task lands on the same worker. 

548 from ..dependencies.threadpool import run_db_sync 

549 

550 return await run_db_sync(_owns_research_sync, username, research_id) 

551 except Exception: 

552 logger.exception( 

553 f"Ownership check failed for user={username} rid={research_id}" 

554 ) 

555 return False 

556 

557 

558async def _socket_session_still_valid( 

559 sid: str, username: str 

560) -> tuple[str | None, bool]: 

561 """Re-check that this socket's originating session is still alive. 

562 

563 Returns ``(session_id, is_valid)``. The id comes back so the caller can 

564 tear down EVERY socket of a dead session, not just the one that asked -- 

565 see the rejection path in ``on_subscribe``. 

566 

567 The connect-time gate is not sufficient on its own: identity is captured 

568 into ``_sid_users`` at handshake and then frozen for the socket's whole 

569 lifetime. A session that expires, is logged out, or has its password 

570 changed WHILE the socket is open leaves that socket able to keep acting -- 

571 including subscribing to research it had not subscribed to before. 

572 

573 Logout and password change disconnect sockets actively. Idle expiry has a 

574 bounded revocation-latency window until the periodic sweep (5 minutes in 

575 production). Revalidating on subscribe and unsubscribe closes that window 

576 whenever the socket next performs either action; the sweep remains the 

577 fallback for an otherwise idle socket. 

578 

579 Fails closed: no recorded session id means the socket predates the gate, 

580 which is not a state a current client can reach. 

581 """ 

582 async with _lock: 

583 session_id = _sid_sessions.get(sid) 

584 if not session_id: 

585 return None, False 

586 

587 from ..auth.session_manager import session_manager 

588 

589 return session_id, session_manager.validate_session(session_id) == username 

590 

591 

592@sio.on("subscribe_to_research") 

593async def on_subscribe(sid, data): 

594 """Handle client subscription to research updates. 

595 

596 Verifies the requesting socket is authenticated AND that the 

597 requested research_id belongs to that user. Without this check 

598 any logged-in user could spy on any other user's research progress 

599 in real time by guessing/enumerating UUIDs. 

600 """ 

601 research_id = data.get("research_id") if isinstance(data, dict) else None 

602 if not research_id: 

603 return 

604 

605 async with _lock: 

606 username = _sid_users.get(sid) 

607 

608 if not username: 

609 logger.warning( 

610 f"Rejected subscribe from unauthenticated sid {sid} for {research_id}" 

611 ) 

612 await sio.emit( 

613 "subscribe_error", 

614 {"error": "Authentication required", "research_id": research_id}, 

615 room=sid, 

616 ) 

617 return 

618 

619 stale_session_id, session_ok = await _socket_session_still_valid( 

620 sid, username 

621 ) 

622 if not session_ok: 

623 logger.warning( 

624 f"Rejected subscribe from {sid} ({username}): originating session " 

625 "is no longer valid; disconnecting that session's sockets" 

626 ) 

627 try: 

628 await sio.emit( 

629 "subscribe_error", 

630 {"error": "Session expired", "research_id": research_id}, 

631 room=sid, 

632 ) 

633 except Exception: 

634 # Revocation is the security boundary; the explanatory frame is 

635 # best-effort. A broken transport must not strand this socket (or 

636 # its siblings) after session validation has already failed. 

637 if _logging_is_enabled(): 637 ↛ 650line 637 didn't jump to line 650 because the condition on line 637 was always true

638 logger.debug(f"Error notifying stale socket {sid}") 

639 # Sever EVERY socket of that session, not just this one. 

640 # 

641 # `validate_session` DELETES an expired session as a side effect of 

642 # checking it, so by the time we get here the session is already gone 

643 # from the store. Disconnecting only the calling sid would leave a 

644 # sibling socket (a second tab on the same session) connected AND 

645 # permanently un-revocable: the periodic idle sweep can never find 

646 # that session again to tear it down, so the sibling keeps receiving 

647 # the user's events -- including `settings_changed`, which carries 

648 # plaintext secrets -- indefinitely. Main revoked per-session for 

649 # exactly this reason. 

650 if stale_session_id: 

651 disconnect_session(stale_session_id) 

652 else: 

653 try: 

654 await sio.disconnect(sid) 

655 except Exception: 

656 if _logging_is_enabled(): 656 ↛ 658line 656 didn't jump to line 658 because the condition on line 656 was always true

657 logger.debug(f"Error disconnecting stale socket {sid}") 

658 return 

659 

660 # Verify ownership: the research must exist in this user's encrypted DB. 

661 owns = await _user_owns_research(username, research_id) 

662 

663 if not owns: 

664 logger.warning( 

665 f"Rejected subscribe: user {username} does not own research {research_id}" 

666 ) 

667 await sio.emit( 

668 "subscribe_error", 

669 {"error": "Not authorized", "research_id": research_id}, 

670 room=sid, 

671 ) 

672 return 

673 

674 async with _lock: 

675 key = _subscription_key(username, research_id) 

676 if key not in _subscriptions: 676 ↛ 678line 676 didn't jump to line 678 because the condition on line 676 was always true

677 _subscriptions[key] = set() 

678 _subscriptions[key].add(sid) 

679 

680 logger.info( 

681 f"Client {sid} (user={username}) subscribed to research {research_id}" 

682 ) 

683 

684 # Send current status immediately if available 

685 snapshot = get_active_research_snapshot(research_id) 

686 if snapshot is not None: 

687 progress = snapshot["progress"] 

688 latest_log = snapshot["log"][-1] if snapshot["log"] else None 

689 

690 if latest_log: 

691 await sio.emit( 

692 f"research_progress_{research_id}", 

693 { 

694 "progress": progress, 

695 "message": latest_log.get("message", "Processing..."), 

696 "status": ResearchStatus.IN_PROGRESS, 

697 "log_entry": latest_log, 

698 }, 

699 room=sid, 

700 ) 

701 

702 

703@sio.on("unsubscribe_from_research") 

704async def on_unsubscribe(sid, data): 

705 """Handle client unsubscribe from research updates. 

706 

707 Applies the same ownership gate as on_subscribe before mutating the 

708 per-research subscription set. This restores authorization-boundary 

709 consistency with on_subscribe (matching main's own rationale for the 

710 check) — it is not closing an active exploit: ``subs.discard(sid)`` 

711 can only ever remove the CALLER's own sid, so without this gate an 

712 unauthorized unsubscribe attempt against another user's research_id 

713 was already a silent no-op (that sid was never a member of the set), 

714 and the victim's subscription was never affected. 

715 

716 Also re-validates the originating session, same as on_subscribe. This 

717 reduces idle-expiry revocation latency whenever a socket next sends an 

718 unsubscribe: an expired session is severed immediately instead of waiting 

719 for the bounded periodic sweep. An otherwise idle socket is still handled 

720 by that sweep. 

721 """ 

722 research_id = data.get("research_id") if isinstance(data, dict) else None 

723 if not research_id: 

724 return 

725 

726 async with _lock: 

727 username = _sid_users.get(sid) 

728 

729 if not username: 

730 logger.info( 

731 f"Rejected unsubscribe from unauthenticated sid {sid} for {research_id}" 

732 ) 

733 return 

734 

735 stale_session_id, session_ok = await _socket_session_still_valid( 

736 sid, username 

737 ) 

738 if not session_ok: 

739 logger.warning( 

740 f"Rejected unsubscribe from {sid} ({username}): originating " 

741 "session is no longer valid; disconnecting that session's sockets" 

742 ) 

743 # Same rationale as the rejection branch in on_subscribe: sever 

744 # EVERY socket of that session, not just this one. validate_session 

745 # already deleted the session from the store as a side effect of 

746 # checking it, so a sibling socket left connected would be 

747 # permanently un-revocable -- the idle sweep can never find that 

748 # session again to tear it down. 

749 if stale_session_id: 

750 disconnect_session(stale_session_id) 

751 else: 

752 try: 

753 await sio.disconnect(sid) 

754 except Exception: 

755 if _logging_is_enabled(): 755 ↛ 757line 755 didn't jump to line 757 because the condition on line 755 was always true

756 logger.debug(f"Error disconnecting stale socket {sid}") 

757 return 

758 

759 if not await _user_owns_research(username, research_id): 

760 logger.info( 

761 f"Rejected unsubscribe from sid {sid}: user does not own research {research_id}" 

762 ) 

763 return 

764 

765 async with _lock: 

766 key = _subscription_key(username, research_id) 

767 subs = _subscriptions.get(key) 

768 if subs: 

769 subs.discard(sid) 

770 if not subs: 

771 _subscriptions.pop(key, None) 

772 

773 

774def emit_socket_event(event: str, data: Any, room: str | None = None) -> bool: 

775 """Emit a socket event (sync wrapper for use from background threads). 

776 

777 Schedules the emit on uvicorn's main event loop via 

778 `asyncio.run_coroutine_threadsafe`. Without the captured loop, emits 

779 from worker threads silently no-op because `asyncio.get_event_loop()` 

780 in a non-main thread on Python 3.12+ either raises or creates a new 

781 isolated loop. 

782 """ 

783 loop = _get_main_loop() 

784 if loop is None or not loop.is_running(): 

785 if _logging_is_enabled(): 785 ↛ 787line 785 didn't jump to line 787 because the condition on line 785 was always true

786 logger.debug(f"Cannot emit {event}: main event loop not available") 

787 return False 

788 try: 

789 _schedule_coroutine_threadsafe(_async_emit(event, data, room), loop) 

790 return True 

791 except Exception: 

792 if _logging_is_enabled(): 792 ↛ 794line 792 didn't jump to line 794 because the condition on line 792 was always true

793 logger.debug(f"Error emitting socket event {event}") 

794 return False 

795 

796 

797def emit_to_user(event: str, username: str, data: Any) -> bool: 

798 """Emit a socket event to all sockets authenticated as `username`. 

799 

800 Prevents cross-user leaks for events like settings_changed — each 

801 user's sockets get the event, nobody else's do. 

802 """ 

803 loop = _get_main_loop() 

804 if loop is None or not loop.is_running(): 

805 if _logging_is_enabled(): 805 ↛ 807line 805 didn't jump to line 807 because the condition on line 805 was always true

806 logger.debug(f"Cannot emit {event}: main event loop not available") 

807 return False 

808 

809 async def _emit() -> None: 

810 # Snapshot sids for this user under the lock — connect/disconnect 

811 # mutate _sid_users on this same event loop and could otherwise 

812 # interleave at any await point, raising 

813 # ``RuntimeError: dictionary changed size during iteration``. 

814 async with _lock: 

815 sids = [sid for sid, u in _sid_users.items() if u == username] 

816 for sid in sids: 

817 try: 

818 await sio.emit(event, data, room=sid) 

819 except Exception: 

820 if _logging_is_enabled(): 820 ↛ 816line 820 didn't jump to line 816 because the condition on line 820 was always true

821 logger.debug(f"Error emitting {event} to {sid}") 

822 

823 try: 

824 _schedule_coroutine_threadsafe(_emit(), loop) 

825 return True 

826 except Exception: 

827 if _logging_is_enabled(): 827 ↛ 829line 827 didn't jump to line 829 because the condition on line 827 was always true

828 logger.debug(f"Error emitting {event} to user {username}") 

829 return False 

830 

831 

832def disconnect_user(username: str) -> bool: 

833 """Disconnect every socket authenticated as ``username``. 

834 

835 A socket is authorised once, at handshake, and never re-checked. So 

836 when a user's session ends — logout, password change, or the idle 

837 sweep closing their database because they have no session left — any 

838 socket they still hold keeps receiving their events indefinitely. This 

839 severs them. 

840 

841 Port of ``SocketIOService.disconnect_user`` (#5535), which grouped 

842 sockets with Socket.IO rooms. This layer already tracks identity 

843 directly in ``_sid_users``, so it filters that instead of maintaining 

844 a parallel per-user room; the snapshot is taken under ``_lock`` for 

845 the same reason ``emit_to_user`` takes it — connect/disconnect mutate 

846 ``_sid_users`` on this event loop and would otherwise interleave at an 

847 await point. 

848 

849 Dropping subscriptions is handled for free: ``sio.disconnect(sid)`` 

850 fires the ``disconnect`` handler above, which removes the sid from 

851 ``_sid_users`` and from every entry in ``_subscriptions``. 

852 

853 Best-effort and non-blocking, like the emit helpers: schedules onto 

854 the main loop and returns whether it could be scheduled, not whether 

855 the sockets are gone. Callers are teardown paths that must not block 

856 or raise. 

857 """ 

858 return _disconnect_matching( 

859 lambda: [sid for sid, u in _sid_users.items() if u == username], 

860 f"user {username}", 

861 ) 

862 

863 

864def disconnect_session(session_id: str) -> bool: 

865 """Disconnect only the sockets authorised by ``session_id``. 

866 

867 Single-session teardown for logout: the tab being logged out loses its 

868 sockets while the user's other still-valid sessions keep theirs. Port 

869 of ``SocketIOService.disconnect_session`` (#5535), which used a 

870 per-session Socket.IO room; this layer matches on the handshake 

871 session_id recorded in ``_sid_sessions`` instead. 

872 

873 Use ``disconnect_user`` for the all-sessions cases — password change 

874 and the idle sweep — where every session is destroyed anyway. 

875 """ 

876 return _disconnect_matching( 

877 lambda: [sid for sid, s in _sid_sessions.items() if s == session_id], 

878 f"session {session_id[:8]}...", 

879 ) 

880 

881 

882def _disconnect_matching(select_sids, description: str) -> bool: 

883 """Disconnect every sid returned by ``select_sids``. 

884 

885 ``select_sids`` is called while holding ``_lock`` — connect/disconnect 

886 mutate the sid maps on this event loop, so the selection must be a 

887 snapshot taken under the lock rather than a live view, or iteration can 

888 raise ``RuntimeError: dictionary changed size during iteration``. 

889 """ 

890 loop = _get_main_loop() 

891 if loop is None or not loop.is_running(): 

892 if _logging_is_enabled(): 892 ↛ 897line 892 didn't jump to line 897 because the condition on line 892 was always true

893 logger.debug( 

894 f"Cannot disconnect sockets for {description}: " 

895 "main event loop not available" 

896 ) 

897 return False 

898 

899 async def _disconnect() -> None: 

900 async with _lock: 

901 sids = select_sids() 

902 for sid in sids: 

903 try: 

904 await sio.disconnect(sid) 

905 except Exception: 

906 if _logging_is_enabled(): 906 ↛ 902line 906 didn't jump to line 902 because the condition on line 906 was always true

907 logger.debug(f"Error disconnecting {sid}") 

908 if sids: 908 ↛ exitline 908 didn't return from function '_disconnect' because the condition on line 908 was always true

909 logger.info(f"Disconnected {len(sids)} socket(s) for {description}") 

910 

911 try: 

912 _schedule_coroutine_threadsafe(_disconnect(), loop) 

913 return True 

914 except Exception: 

915 if _logging_is_enabled(): 915 ↛ 917line 915 didn't jump to line 917 because the condition on line 915 was always true

916 logger.debug(f"Error disconnecting sockets for {description}") 

917 return False 

918 

919 

920def emit_to_subscribers( 

921 event_base: str, 

922 research_id: str, 

923 data: Any, 

924 *, 

925 owner: str, 

926 enable_logging: bool = True, 

927) -> bool: 

928 """Emit an event to all subscribers of a specific research. 

929 

930 Sync wrapper for use from background threads. 

931 

932 ``owner`` is the username whose research this is, and is REQUIRED and 

933 keyword-only on purpose: subscriptions are keyed by ``(owner, 

934 research_id)`` because a benchmark id is only unique within one user's 

935 database (see ``_subscriptions``). Making it required means a call site 

936 that forgets it fails loudly at the call rather than silently reverting 

937 to the id-only lookup that leaked one user's progress to another. 

938 

939 `enable_logging` is carried by a ContextVar so concurrent emits from 

940 different worker threads don't corrupt each other's suppression windows 

941 (as a module global did), AND so the suppression reaches the coroutine 

942 scheduled onto the event loop below — a threading.local could not, since 

943 that coroutine runs on a different thread. See `_logging_enabled_var`. 

944 """ 

945 token = None 

946 if not enable_logging: 

947 token = _logging_enabled_var.set(False) 

948 

949 try: 

950 loop = _get_main_loop() 

951 if loop is None or not loop.is_running(): 

952 if _logging_is_enabled(): 

953 logger.debug( 

954 f"Cannot emit subscribers for {research_id}: loop unavailable" 

955 ) 

956 return False 

957 try: 

958 _schedule_coroutine_threadsafe( 

959 _async_emit_to_subscribers( 

960 event_base, research_id, data, owner 

961 ), 

962 loop, 

963 ) 

964 return True 

965 except Exception: 

966 if _logging_is_enabled(): 

967 logger.debug( 

968 f"Error emitting to subscribers for research {research_id}" 

969 ) 

970 return False 

971 finally: 

972 # Reset only this caller's context. The coroutine scheduled above got 

973 # its own snapshot taken while the flag was still False, so it keeps 

974 # the suppression it needs regardless of this reset. 

975 if token is not None: 

976 _logging_enabled_var.reset(token) 

977 

978 

979async def _async_emit(event: str, data: Any, room: str | None = None) -> None: 

980 """Async emit implementation.""" 

981 try: 

982 if room: 

983 await sio.emit(event, data, room=room) 

984 else: 

985 await sio.emit(event, data) 

986 except Exception: 

987 # The concurrent Future returned by run_coroutine_threadsafe is 

988 # intentionally discarded by the synchronous wrapper, so an 

989 # exception escaping this coroutine would otherwise be silent. 

990 if _logging_is_enabled(): 990 ↛ exitline 990 didn't return from function '_async_emit' because the condition on line 990 was always true

991 logger.debug(f"Error emitting socket event {event}")