mdb-engine 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
mdb_engine/core/engine.py CHANGED
@@ -149,6 +149,7 @@ class MongoDBEngine:
149
149
  self._encryption_service: EnvelopeEncryptionService | None = None
150
150
  self._app_secrets_manager: AppSecretsManager | None = None
151
151
  self._websocket_session_manager: Any | None = None # WebSocketSessionManager
152
+ self._websocket_ticket_store: Any | None = None # WebSocketTicketStore
152
153
 
153
154
  # Store app read_scopes mapping for validation
154
155
  self._app_read_scopes: dict[str, list[str]] = {}
@@ -210,6 +211,13 @@ class MongoDBEngine:
210
211
  encryption_service=self._encryption_service,
211
212
  )
212
213
 
214
+ # Initialize WebSocket ticket store (in-memory, no dependencies needed)
215
+ # Tickets are preferred for multi-app SSO setups (short-lived, single-use)
216
+ from ..auth.websocket_tickets import WebSocketTicketStore
217
+
218
+ self._websocket_ticket_store = WebSocketTicketStore()
219
+ logger.info("WebSocket ticket store initialized")
220
+
213
221
  # Set up component managers
214
222
  self._app_registration_manager = AppRegistrationManager(
215
223
  mongo_db=self._connection_manager.mongo_db,
@@ -265,6 +273,16 @@ class MongoDBEngine:
265
273
  """Check if Ray is enabled and initialized."""
266
274
  return self.enable_ray and self.ray_actor is not None
267
275
 
276
+ @property
277
+ def connection_manager(self):
278
+ """
279
+ Get the connection manager.
280
+
281
+ Returns:
282
+ ConnectionManager instance
283
+ """
284
+ return self._connection_manager
285
+
268
286
  @property
269
287
  def mongo_client(self) -> AsyncIOMotorClient:
270
288
  """
@@ -709,6 +727,25 @@ class MongoDBEngine:
709
727
  app: FastAPI application instance
710
728
  slug: App slug
711
729
  """
730
+ # CRITICAL: Ensure websocket_ticket_store is available
731
+ # Ticket authentication is required for WebSocket connections
732
+ if not self._websocket_ticket_store:
733
+ error_msg = (
734
+ f"WebSocket routes cannot be registered for app '{slug}': "
735
+ "websocket_ticket_store is not available. "
736
+ "WebSocket authentication requires ticket store to be initialized."
737
+ )
738
+ contextual_logger.error(error_msg)
739
+ raise RuntimeError(error_msg)
740
+
741
+ # Ensure ticket store is in app state (may have been set in create_app)
742
+ if (
743
+ not hasattr(app.state, "websocket_ticket_store")
744
+ or app.state.websocket_ticket_store is None
745
+ ):
746
+ app.state.websocket_ticket_store = self._websocket_ticket_store
747
+ contextual_logger.debug(f"WebSocket ticket store stored in app state for '{slug}'")
748
+
712
749
  # Check if WebSockets are configured for this app
713
750
  websockets_config = self.get_websocket_config(slug)
714
751
  if not websockets_config:
@@ -915,9 +952,9 @@ class MongoDBEngine:
915
952
  return get_embedding_service_for_app(slug, self)
916
953
 
917
954
  @property
918
- def _apps(self) -> dict[str, Any]:
955
+ def apps(self) -> dict[str, Any]:
919
956
  """
920
- Get the apps dictionary (for backward compatibility with tests).
957
+ Get all registered apps.
921
958
 
922
959
  Returns:
923
960
  Dictionary of registered apps
@@ -927,7 +964,27 @@ class MongoDBEngine:
927
964
  """
928
965
  if not self._app_registration_manager:
929
966
  raise RuntimeError("MongoDBEngine not initialized. Call initialize() first.")
930
- return self._app_registration_manager._apps
967
+ return self._app_registration_manager.apps
968
+
969
+ @property
970
+ def websocket_ticket_store(self):
971
+ """
972
+ Get the WebSocket ticket store.
973
+
974
+ Returns:
975
+ WebSocketTicketStore instance or None if not initialized
976
+ """
977
+ return self._websocket_ticket_store
978
+
979
+ @property
980
+ def websocket_session_manager(self):
981
+ """
982
+ Get the WebSocket session manager.
983
+
984
+ Returns:
985
+ WebSocketSessionManager instance or None if not initialized
986
+ """
987
+ return self._websocket_session_manager
931
988
 
932
989
  def list_apps(self) -> list[str]:
933
990
  """
@@ -1100,6 +1157,73 @@ class MongoDBEngine:
1100
1157
  # FastAPI Integration Methods
1101
1158
  # =========================================================================
1102
1159
 
1160
+ async def _register_websocket_endpoints(self, app: "FastAPI", engine: "MongoDBEngine") -> None:
1161
+ """Register WebSocket ticket and session endpoints."""
1162
+ # Register WebSocket ticket endpoint AFTER initialization
1163
+ # (ticket store is now available)
1164
+ if engine.websocket_ticket_store:
1165
+ app.state.websocket_ticket_store = engine.websocket_ticket_store
1166
+ logger.info("WebSocket ticket store stored in app state")
1167
+
1168
+ # Set global ticket store for WebSocket authentication (works with routers)
1169
+ from ..routing.websockets import set_global_websocket_ticket_store
1170
+
1171
+ set_global_websocket_ticket_store(engine.websocket_ticket_store)
1172
+ logger.info("Global WebSocket ticket store set for multi-app authentication")
1173
+
1174
+ # Register WebSocket ticket endpoint
1175
+ from ..auth.websocket_tickets import create_websocket_ticket_endpoint
1176
+
1177
+ ticket_endpoint = create_websocket_ticket_endpoint(engine.websocket_ticket_store)
1178
+ app.post("/auth/ticket")(ticket_endpoint)
1179
+ logger.info("WebSocket ticket endpoint registered at /auth/ticket")
1180
+
1181
+ # Register WebSocket session endpoint AFTER initialization
1182
+ # (session manager is now available)
1183
+ if engine.websocket_session_manager:
1184
+ app.state.websocket_session_manager = engine.websocket_session_manager
1185
+ logger.info("WebSocket session manager stored in app state")
1186
+
1187
+ # Register WebSocket session endpoint
1188
+ from ..auth.websocket_sessions import create_websocket_session_endpoint
1189
+
1190
+ session_endpoint = create_websocket_session_endpoint(engine.websocket_session_manager)
1191
+ app.get("/auth/websocket-session")(session_endpoint)
1192
+ logger.info("WebSocket session endpoint registered at /auth/websocket-session")
1193
+
1194
+ async def _configure_websocket_ticket_ttl(
1195
+ self, app: "FastAPI", app_manifest: dict[str, Any], slug: str
1196
+ ) -> None:
1197
+ """Configure WebSocket ticket TTL from manifest."""
1198
+ websockets_config = app_manifest.get("websockets", {})
1199
+ if not websockets_config:
1200
+ return
1201
+
1202
+ from ..auth.websocket_tickets import WebSocketTicketStore
1203
+
1204
+ ticket_ttl_values: list[int] = []
1205
+ for endpoint_config in websockets_config.values():
1206
+ if isinstance(endpoint_config, dict):
1207
+ ticket_ttl = endpoint_config.get("ticket_ttl_seconds")
1208
+ if ticket_ttl is not None:
1209
+ ticket_ttl_values.append(ticket_ttl)
1210
+
1211
+ if ticket_ttl_values:
1212
+ configured_ticket_ttl = min(ticket_ttl_values) # Use minimum for maximum security
1213
+ # Reinitialize ticket store if needed
1214
+ ticket_store = self._websocket_ticket_store
1215
+ if ticket_store is None or ticket_store.ticket_ttl != configured_ticket_ttl:
1216
+ self._websocket_ticket_store = WebSocketTicketStore(
1217
+ ticket_ttl_seconds=configured_ticket_ttl
1218
+ )
1219
+ logger.info(
1220
+ f"WebSocket ticket store initialized with TTL: "
1221
+ f"{configured_ticket_ttl}s (from app '{slug}' manifest)"
1222
+ )
1223
+ # Update app state if ticket store was already set
1224
+ if hasattr(app.state, "websocket_ticket_store"):
1225
+ app.state.websocket_ticket_store = self._websocket_ticket_store
1226
+
1103
1227
  def create_app(
1104
1228
  self,
1105
1229
  slug: str,
@@ -1199,8 +1323,15 @@ class MongoDBEngine:
1199
1323
  if not is_sub_app:
1200
1324
  await engine.initialize()
1201
1325
 
1326
+ # Register WebSocket endpoints
1327
+ await self._register_websocket_endpoints(app, engine)
1328
+
1202
1329
  # Load and register manifest
1203
1330
  app_manifest = await engine.load_manifest(manifest_path)
1331
+
1332
+ # Configure WebSocket ticket TTL from manifest
1333
+ await self._configure_websocket_ticket_ttl(app, app_manifest, slug)
1334
+
1204
1335
  await engine.register_app(app_manifest)
1205
1336
 
1206
1337
  # Auto-detect multi-site mode from manifest
@@ -1235,11 +1366,11 @@ class MongoDBEngine:
1235
1366
  f"Sub-app '{slug}' uses shared auth but user_pool not found. "
1236
1367
  "Initializing now (parent should have initialized it)."
1237
1368
  )
1238
- await engine._initialize_shared_user_pool(app, app_manifest)
1369
+ await self._initialize_shared_user_pool(app, app_manifest)
1239
1370
  else:
1240
1371
  logger.debug(f"Sub-app '{slug}' using shared user_pool from parent app")
1241
1372
  else:
1242
- await engine._initialize_shared_user_pool(app, app_manifest)
1373
+ await self._initialize_shared_user_pool(app, app_manifest)
1243
1374
  else:
1244
1375
  logger.info(f"Per-app auth mode for '{slug}'")
1245
1376
  # Auto-retrieve app token for "app" mode
@@ -1505,6 +1636,10 @@ class MongoDBEngine:
1505
1636
  f"(require_role={auth_config.get('require_role')})"
1506
1637
  )
1507
1638
 
1639
+ # NOTE: WebSocket ticket endpoint registration is moved to lifespan context manager
1640
+ # (after engine.initialize()) because ticket store is only available after initialization.
1641
+ # This ensures consistency with create_multi_app() behavior.
1642
+
1508
1643
  # Add CSRF middleware (after auth - auto-enabled for shared mode)
1509
1644
  # CSRF protection is enabled by default for shared auth mode
1510
1645
  # SKIP for sub-apps in multi-app setups - parent app handles CSRF
@@ -1512,8 +1647,17 @@ class MongoDBEngine:
1512
1647
  if csrf_config and not is_sub_app: # Don't add CSRF to child apps
1513
1648
  from ..auth.csrf import create_csrf_middleware
1514
1649
 
1650
+ # Add ticket endpoint to public routes (it handles its own auth)
1651
+ public_routes = auth_config.get("public_routes", [])
1652
+ public_routes_with_ticket = list(public_routes) + ["/auth/ticket"]
1653
+
1654
+ csrf_config_with_routes = {
1655
+ **auth_config,
1656
+ "public_routes": public_routes_with_ticket,
1657
+ }
1658
+
1515
1659
  csrf_middleware = create_csrf_middleware(
1516
- manifest_auth=auth_config,
1660
+ manifest_auth=csrf_config_with_routes,
1517
1661
  )
1518
1662
  app.add_middleware(csrf_middleware)
1519
1663
  logger.info(f"CSRFMiddleware added for '{slug}'")
@@ -2142,6 +2286,7 @@ class MongoDBEngine:
2142
2286
  )
2143
2287
 
2144
2288
  # Check if any app uses shared auth and collect public routes for CSRF exemption
2289
+ # Also collect ticket TTL values from websocket configs
2145
2290
  has_shared_auth = False
2146
2291
  all_public_routes = [
2147
2292
  "/health",
@@ -2149,6 +2294,7 @@ class MongoDBEngine:
2149
2294
  "/openapi.json",
2150
2295
  "/_mdb/routes",
2151
2296
  ] # Base exempt routes
2297
+ ticket_ttl_values: list[int] = [] # Collect ticket TTLs from all apps
2152
2298
  for app_config in apps:
2153
2299
  try:
2154
2300
  manifest_path = app_config["manifest"]
@@ -2168,9 +2314,45 @@ class MongoDBEngine:
2168
2314
  prefixed_route = f"{path_prefix.rstrip('/')}/{route}"
2169
2315
  if prefixed_route not in all_public_routes:
2170
2316
  all_public_routes.append(prefixed_route)
2317
+
2318
+ # Collect ticket TTL from websocket configs
2319
+ websockets_config = app_manifest_pre.get("websockets", {})
2320
+ if websockets_config:
2321
+ for endpoint_config in websockets_config.values():
2322
+ if isinstance(endpoint_config, dict):
2323
+ ticket_ttl = endpoint_config.get("ticket_ttl_seconds")
2324
+ if ticket_ttl is not None:
2325
+ ticket_ttl_values.append(ticket_ttl)
2171
2326
  except (FileNotFoundError, json.JSONDecodeError, KeyError) as e:
2172
2327
  logger.warning(f"Could not check auth mode for app '{app_config.get('slug')}': {e}")
2173
2328
 
2329
+ # Determine ticket TTL: use minimum from app configs (most secure), or default
2330
+ from ..auth.websocket_tickets import DEFAULT_TICKET_TTL_SECONDS
2331
+
2332
+ if ticket_ttl_values:
2333
+ configured_ticket_ttl = min(ticket_ttl_values) # Use minimum for maximum security
2334
+ logger.info(
2335
+ f"Ticket TTL configured from app manifests: {configured_ticket_ttl}s "
2336
+ f"(found values: {ticket_ttl_values}, using minimum)"
2337
+ )
2338
+ else:
2339
+ configured_ticket_ttl = DEFAULT_TICKET_TTL_SECONDS
2340
+ logger.debug(
2341
+ f"No ticket TTL specified in app manifests, using default: {configured_ticket_ttl}s"
2342
+ )
2343
+
2344
+ # Reinitialize ticket store with configured TTL if different from current
2345
+ if (
2346
+ self._websocket_ticket_store is None
2347
+ or self._websocket_ticket_store.ticket_ttl != configured_ticket_ttl
2348
+ ):
2349
+ from ..auth.websocket_tickets import WebSocketTicketStore
2350
+
2351
+ self._websocket_ticket_store = WebSocketTicketStore(
2352
+ ticket_ttl_seconds=configured_ticket_ttl
2353
+ )
2354
+ logger.info(f"WebSocket ticket store initialized with TTL: {configured_ticket_ttl}s")
2355
+
2174
2356
  # Validate hooks before creating lifespan (fail fast)
2175
2357
  for app_config in apps:
2176
2358
  slug = app_config.get("slug", "unknown")
@@ -2291,10 +2473,49 @@ class MongoDBEngine:
2291
2473
  logger.debug(f"No WebSocket configuration found for app '{slug}'")
2292
2474
  return
2293
2475
 
2476
+ # CRITICAL: Check if session manager is required and available
2477
+ # Some endpoints require session keys (csrf_required=True), which need
2478
+ # session manager
2479
+ requires_session_manager = False
2480
+ for _endpoint_name, endpoint_config in websockets_config.items():
2481
+ auth_config = endpoint_config.get("auth", {})
2482
+ if isinstance(auth_config, dict):
2483
+ csrf_required = auth_config.get("csrf_required", True) # Default to True
2484
+ if csrf_required:
2485
+ requires_session_manager = True
2486
+ break
2487
+
2488
+ if requires_session_manager and not engine.websocket_session_manager:
2489
+ error_msg = (
2490
+ f"WebSocket routes cannot be registered for app '{slug}': "
2491
+ "websocket_session_manager is not available. "
2492
+ "WebSocket endpoints with csrf_required=True require "
2493
+ "session manager to be initialized. "
2494
+ "Set MDB_ENGINE_MASTER_KEY environment variable to enable "
2495
+ "session manager."
2496
+ )
2497
+ logger.error(error_msg)
2498
+ raise RuntimeError(error_msg)
2499
+
2500
+ # CRITICAL: Ensure websocket_ticket_store is available
2501
+ # Ticket authentication is required for WebSocket connections
2502
+ if not engine.websocket_ticket_store:
2503
+ error_msg = (
2504
+ f"WebSocket routes cannot be registered for app '{slug}': "
2505
+ "websocket_ticket_store is not available. "
2506
+ "WebSocket authentication requires ticket store to be initialized."
2507
+ )
2508
+ logger.error(error_msg)
2509
+ raise RuntimeError(error_msg)
2510
+
2294
2511
  # Store WebSocket config in parent app state for CSRF middleware to access
2295
2512
  if not hasattr(parent_app.state, "websocket_configs"):
2296
2513
  parent_app.state.websocket_configs = {}
2297
2514
  parent_app.state.websocket_configs[slug] = websockets_config
2515
+ logger.info(
2516
+ f"✅ Stored WebSocket config for '{slug}' in parent app state "
2517
+ f"({len(websockets_config)} endpoint(s))"
2518
+ )
2298
2519
 
2299
2520
  try:
2300
2521
  from fastapi import APIRouter
@@ -2335,16 +2556,51 @@ class MongoDBEngine:
2335
2556
  ping_interval=ping_interval,
2336
2557
  )
2337
2558
 
2338
- # Register on parent app with full path
2339
- ws_router = APIRouter()
2340
- ws_router.websocket(full_ws_path)(handler)
2341
- parent_app.include_router(ws_router)
2559
+ # Register on parent app with full path using FastAPI's
2560
+ # proper WebSocket registration
2561
+ # We register BEFORE mounting apps to ensure WebSocket
2562
+ # routes are checked first
2563
+ try:
2564
+ # Use FastAPI's APIRouter approach (same as single-app mode)
2565
+ # This maintains FastAPI features
2566
+ # (dependency injection, OpenAPI docs, etc.)
2567
+ ws_router = APIRouter()
2568
+ ws_router.websocket(full_ws_path)(handler)
2569
+
2570
+ # Include router BEFORE mounting child app to ensure route priority
2571
+ parent_app.include_router(ws_router)
2572
+
2573
+ logger.info(
2574
+ f"✅ Registered WebSocket route '{full_ws_path}' "
2575
+ f"using FastAPI APIRouter "
2576
+ f"(registered before app mount to ensure priority)"
2577
+ )
2578
+ except (
2579
+ ValueError,
2580
+ RuntimeError,
2581
+ AttributeError,
2582
+ TypeError,
2583
+ ) as fastapi_error:
2584
+ logger.error(
2585
+ f"❌ Failed to register WebSocket route "
2586
+ f"'{full_ws_path}' with FastAPI: {fastapi_error}",
2587
+ exc_info=True,
2588
+ )
2589
+ raise
2342
2590
 
2343
2591
  logger.info(
2344
2592
  f"✅ Registered WebSocket route '{full_ws_path}' "
2345
2593
  f"for mounted app '{slug}' (mounted at '{path_prefix}', "
2346
2594
  f"auth: {require_auth}, ping: {ping_interval}s)"
2347
2595
  )
2596
+ import sys
2597
+
2598
+ print(
2599
+ f"✅ [ROUTE REGISTRATION] WebSocket route '{full_ws_path}' "
2600
+ f"registered for '{slug}' using FastAPI APIRouter",
2601
+ file=sys.stderr,
2602
+ flush=True,
2603
+ )
2348
2604
 
2349
2605
  # Verify route was actually registered
2350
2606
  registered_routes = [
@@ -2352,6 +2608,26 @@ class MongoDBEngine:
2352
2608
  for r in parent_app.routes
2353
2609
  if hasattr(r, "path") and full_ws_path in str(getattr(r, "path", ""))
2354
2610
  ]
2611
+
2612
+ # CRITICAL: Log all WebSocket routes to verify registration
2613
+ # FastAPI APIRouter creates routes of type 'APIWebSocketRoute'
2614
+ all_ws_routes = [
2615
+ (r.path, type(r).__name__)
2616
+ for r in parent_app.routes
2617
+ if hasattr(r, "path")
2618
+ and ("ws" in str(r.path).lower() or hasattr(r, "endpoint"))
2619
+ ]
2620
+ import sys
2621
+
2622
+ print(
2623
+ f"📋 [ROUTE VERIFICATION] All WebSocket-like routes: {all_ws_routes}",
2624
+ file=sys.stderr,
2625
+ flush=True,
2626
+ )
2627
+ logger.info(
2628
+ f"📋 [ROUTE VERIFICATION] All WebSocket-like routes: {all_ws_routes}"
2629
+ )
2630
+
2355
2631
  if registered_routes:
2356
2632
  registered_count += 1
2357
2633
  logger.debug(
@@ -2403,6 +2679,40 @@ class MongoDBEngine:
2403
2679
  # Initialize engine
2404
2680
  await engine.initialize()
2405
2681
 
2682
+ # Register WebSocket ticket endpoint AFTER initialization
2683
+ # (ticket store is now available)
2684
+ if engine.websocket_ticket_store:
2685
+ app.state.websocket_ticket_store = engine.websocket_ticket_store
2686
+ logger.info("WebSocket ticket store stored in parent app state")
2687
+
2688
+ # Set global ticket store for WebSocket authentication (works with routers)
2689
+ from ..routing.websockets import set_global_websocket_ticket_store
2690
+
2691
+ set_global_websocket_ticket_store(engine.websocket_ticket_store)
2692
+ logger.info("Global WebSocket ticket store set for multi-app authentication")
2693
+
2694
+ # Register WebSocket ticket endpoint on parent app
2695
+ from ..auth.websocket_tickets import create_websocket_ticket_endpoint
2696
+
2697
+ ticket_endpoint = create_websocket_ticket_endpoint(engine.websocket_ticket_store)
2698
+ app.post("/auth/ticket")(ticket_endpoint)
2699
+ logger.info("WebSocket ticket endpoint registered at /auth/ticket")
2700
+
2701
+ # Register WebSocket session endpoint AFTER initialization
2702
+ # (session manager is now available)
2703
+ if engine.websocket_session_manager:
2704
+ app.state.websocket_session_manager = engine.websocket_session_manager
2705
+ logger.info("WebSocket session manager stored in parent app state")
2706
+
2707
+ # Register WebSocket session endpoint on parent app
2708
+ from ..auth.websocket_sessions import create_websocket_session_endpoint
2709
+
2710
+ session_endpoint = create_websocket_session_endpoint(
2711
+ engine.websocket_session_manager
2712
+ )
2713
+ app.get("/auth/websocket-session")(session_endpoint)
2714
+ logger.info("WebSocket session endpoint registered at /auth/websocket-session")
2715
+
2406
2716
  # Initialize shared user pool once if any app uses shared auth
2407
2717
  if has_shared_auth:
2408
2718
  logger.info("Initializing shared user pool for multi-app deployment")
@@ -2414,7 +2724,7 @@ class MongoDBEngine:
2414
2724
  app_manifest_pre = json.load(f)
2415
2725
  auth_config = app_manifest_pre.get("auth", {})
2416
2726
  if auth_config.get("mode") == "shared":
2417
- await engine._initialize_shared_user_pool(app, app_manifest_pre)
2727
+ await self._initialize_shared_user_pool(app, app_manifest_pre)
2418
2728
  shared_user_pool_initialized = True
2419
2729
  logger.info("Shared user pool initialized for multi-app deployment")
2420
2730
  break
@@ -2509,6 +2819,11 @@ class MongoDBEngine:
2509
2819
  )
2510
2820
  logger.debug(f"Shared WebSocket session manager with child app '{slug}'")
2511
2821
 
2822
+ # Share WebSocket ticket store with child app
2823
+ if hasattr(app.state, "websocket_ticket_store"):
2824
+ child_app.state.websocket_ticket_store = app.state.websocket_ticket_store
2825
+ logger.debug(f"Shared WebSocket ticket store with child app '{slug}'")
2826
+
2512
2827
  # Add middleware for app context helpers
2513
2828
  from starlette.middleware.base import BaseHTTPMiddleware
2514
2829
  from starlette.requests import Request
@@ -2596,15 +2911,17 @@ class MongoDBEngine:
2596
2911
  child_app.add_middleware(middleware_class)
2597
2912
  logger.debug(f"Added AppContextMiddleware to child app '{slug}'")
2598
2913
 
2599
- # Mount child app at path prefix
2914
+ # CRITICAL FIX: Register WebSocket routes on parent app BEFORE mounting
2915
+ # This ensures WebSocket routes are checked before mounted app routes
2916
+ # Mounted apps create catch-all routes that intercept /app-slug/* paths
2917
+ await _register_websocket_routes(app, app_manifest_data, slug, path_prefix)
2918
+
2919
+ # Mount child app at path prefix (AFTER WebSocket routes are registered)
2600
2920
  app.mount(path_prefix, child_app)
2601
2921
 
2602
2922
  # CRITICAL FIX: Merge CORS config from child app to parent app
2603
2923
  await _merge_cors_config_to_parent(app, child_app, app_manifest_data, slug)
2604
2924
 
2605
- # CRITICAL FIX: Register WebSocket routes on parent app with full path
2606
- await _register_websocket_routes(app, app_manifest_data, slug, path_prefix)
2607
-
2608
2925
  # Update existing entry instead of appending
2609
2926
  entry = _find_mounted_app_entry(slug)
2610
2927
  if entry:
@@ -2733,6 +3050,11 @@ class MongoDBEngine:
2733
3050
  "manifest_path": str(manifest_path),
2734
3051
  }
2735
3052
  )
3053
+ # Always re-raise RuntimeError for critical failures
3054
+ # (like missing session manager)
3055
+ # These are configuration errors that should fail fast
3056
+ if isinstance(e, RuntimeError) and "websocket_session_manager" in str(e):
3057
+ raise RuntimeError(error_msg) from e
2736
3058
  if strict:
2737
3059
  raise RuntimeError(error_msg) from e
2738
3060
  continue
@@ -2826,11 +3148,78 @@ class MongoDBEngine:
2826
3148
  logger.debug("Set default CORS config on parent app for WebSocket origin validation")
2827
3149
 
2828
3150
  # Store app reference in engine for get_mounted_apps()
2829
- engine._multi_app_instance = parent_app
3151
+ self._multi_app_instance = parent_app
2830
3152
 
2831
- # Add request scope middleware
3153
+ # Add diagnostic ASGI middleware FIRST (outermost - runs before everything)
3154
+ # This will catch WebSocket upgrades before any other middleware
2832
3155
  from starlette.middleware.base import BaseHTTPMiddleware
2833
3156
 
3157
+ class DiagnosticMiddleware(BaseHTTPMiddleware):
3158
+ """Diagnostic middleware to log ALL requests, especially WebSocket upgrades."""
3159
+
3160
+ async def dispatch(self, request, call_next):
3161
+ path = request.url.path
3162
+ method = request.method
3163
+ upgrade_header = request.headers.get("upgrade", "").lower()
3164
+ connection_header = request.headers.get("connection", "").lower()
3165
+ origin_header = request.headers.get("origin")
3166
+
3167
+ # Log WebSocket upgrade attempts IMMEDIATELY
3168
+ if upgrade_header == "websocket" or "websocket" in path.lower():
3169
+ import sys
3170
+
3171
+ print(
3172
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket upgrade detected: "
3173
+ f"{method} {path}, upgrade={upgrade_header}, "
3174
+ f"connection={connection_header}, origin={origin_header}",
3175
+ file=sys.stderr,
3176
+ flush=True,
3177
+ )
3178
+ logger.info(
3179
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket upgrade: {method} {path}, "
3180
+ f"origin={origin_header}"
3181
+ )
3182
+
3183
+ try:
3184
+ response = await call_next(request)
3185
+
3186
+ # Log response for WebSocket upgrades
3187
+ if upgrade_header == "websocket" or "websocket" in path.lower():
3188
+ import sys
3189
+
3190
+ print(
3191
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket response: "
3192
+ f"{method} {path} -> {response.status_code}",
3193
+ file=sys.stderr,
3194
+ flush=True,
3195
+ )
3196
+ logger.info(
3197
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket response: "
3198
+ f"{method} {path} -> {response.status_code}"
3199
+ )
3200
+
3201
+ return response
3202
+ except (RuntimeError, ConnectionError, ValueError, AttributeError) as e:
3203
+ if upgrade_header == "websocket" or "websocket" in path.lower():
3204
+ import sys
3205
+
3206
+ print(
3207
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket exception: "
3208
+ f"{method} {path} -> {type(e).__name__}: {e}",
3209
+ file=sys.stderr,
3210
+ flush=True,
3211
+ )
3212
+ logger.error(
3213
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket exception: "
3214
+ f"{method} {path} -> {type(e).__name__}: {e}",
3215
+ exc_info=True,
3216
+ )
3217
+ raise
3218
+
3219
+ parent_app.add_middleware(DiagnosticMiddleware)
3220
+ logger.debug("DiagnosticMiddleware added for parent app (outermost layer)")
3221
+
3222
+ # Add request scope middleware
2834
3223
  from ..di import ScopeManager
2835
3224
 
2836
3225
  class RequestScopeMiddleware(BaseHTTPMiddleware):
@@ -2854,32 +3243,27 @@ class MongoDBEngine:
2854
3243
  from ..auth.csrf import create_csrf_middleware
2855
3244
 
2856
3245
  # Create CSRF middleware with default config (will use parent app's CORS config)
2857
- # Exempt routes that don't need CSRF (health checks, public routes from child apps)
2858
- # all_public_routes includes base routes + child app public routes with path prefixes
2859
- # Add WebSocket session endpoint to public routes (it handles its own auth)
2860
- public_routes_with_session_endpoint = list(all_public_routes) + [
2861
- "/auth/websocket-session"
3246
+ # Exempt routes that don't need CSRF (health checks, public routes
3247
+ # from child apps)
3248
+ # all_public_routes includes base routes + child app public routes
3249
+ # with path prefixes
3250
+ # Add WebSocket session and ticket endpoints to public routes
3251
+ # (they handle their own auth)
3252
+ public_routes_with_websocket_endpoints = list(all_public_routes) + [
3253
+ "/auth/websocket-session",
3254
+ "/auth/ticket",
2862
3255
  ]
2863
3256
  parent_csrf_config = {
2864
3257
  "csrf_protection": True,
2865
- "public_routes": public_routes_with_session_endpoint,
3258
+ "public_routes": public_routes_with_websocket_endpoints,
2866
3259
  }
2867
3260
  csrf_middleware = create_csrf_middleware(parent_csrf_config)
2868
3261
  parent_app.add_middleware(csrf_middleware)
2869
3262
 
2870
- # Store WebSocket session manager in app state for CSRF middleware and endpoints
2871
- if self._websocket_session_manager:
2872
- parent_app.state.websocket_session_manager = self._websocket_session_manager
2873
- logger.info("WebSocket session manager stored in parent app state")
2874
-
2875
- # Register WebSocket session endpoint on parent app
2876
- from ..auth.websocket_sessions import create_websocket_session_endpoint
2877
-
2878
- session_endpoint = create_websocket_session_endpoint(
2879
- self._websocket_session_manager
2880
- )
2881
- parent_app.get("/auth/websocket-session")(session_endpoint)
2882
- logger.info("WebSocket session endpoint registered at /auth/websocket-session")
3263
+ # NOTE: WebSocket ticket and session endpoint registrations are moved to lifespan
3264
+ # context manager (after engine.initialize()) because they're only available after
3265
+ # initialization. The CSRF middleware still needs to know about these routes, so
3266
+ # they're added to public_routes_with_websocket_endpoints above.
2883
3267
 
2884
3268
  logger.info("CSRFMiddleware added to parent app for WebSocket origin validation")
2885
3269
 
@@ -2904,6 +3288,21 @@ class MongoDBEngine:
2904
3288
  # Read CORS config from app.state (may have been merged from child apps)
2905
3289
  cors_config = getattr(request.app.state, "cors_config", {})
2906
3290
 
3291
+ # CRITICAL: Log WebSocket upgrade requests to see if CORS
3292
+ # is intercepting them
3293
+ upgrade_header = request.headers.get("upgrade", "").lower()
3294
+ if upgrade_header == "websocket" or "websocket" in request.url.path.lower():
3295
+ import sys
3296
+
3297
+ print(
3298
+ f"🌐 [CORS MIDDLEWARE] WebSocket upgrade: "
3299
+ f"{request.method} {request.url.path}, "
3300
+ f"origin={request.headers.get('origin')}, "
3301
+ f"cors_enabled={cors_config.get('enabled', False)}",
3302
+ file=sys.stderr,
3303
+ flush=True,
3304
+ )
3305
+
2907
3306
  if not cors_config.get("enabled", False):
2908
3307
  # CORS not enabled, pass through
2909
3308
  return await call_next(request)
@@ -2980,6 +3379,46 @@ class MongoDBEngine:
2980
3379
  except ImportError:
2981
3380
  logger.warning("CORS middleware not available")
2982
3381
 
3382
+ # Wrap parent app in ASGI wrapper to intercept WebSocket connections at ASGI level
3383
+ # This must be done AFTER all middleware and routes are registered
3384
+ from starlette.types import ASGIApp, Receive, Scope, Send
3385
+
3386
+ class WebSocketASGIWrapper:
3387
+ """ASGI wrapper to intercept WebSocket connections before FastAPI routing."""
3388
+
3389
+ def __init__(self, app: ASGIApp):
3390
+ self.app = app
3391
+ # Delegate attribute access to underlying app
3392
+ self.__dict__.update(app.__dict__)
3393
+
3394
+ def __getattr__(self, name):
3395
+ # Delegate any missing attributes to underlying app
3396
+ return getattr(self.app, name)
3397
+
3398
+ async def __call__(self, scope: Scope, receive: Receive, send: Send):
3399
+ # Intercept WebSocket connections at ASGI level
3400
+ if scope["type"] == "websocket":
3401
+ import sys
3402
+
3403
+ path = scope.get("path", "unknown")
3404
+ headers_dict = {k.decode(): v.decode() for k, v in scope.get("headers", [])}
3405
+ origin = headers_dict.get("origin", "missing")
3406
+ upgrade = headers_dict.get("upgrade", "missing")
3407
+ connection = headers_dict.get("connection", "missing")
3408
+ print(
3409
+ f"🌐 [ASGI WEBSOCKET] Intercepted at ASGI level: path={path}, "
3410
+ f"origin={origin}, upgrade={upgrade}, connection={connection}",
3411
+ file=sys.stderr,
3412
+ flush=True,
3413
+ )
3414
+ logger.info(f"🌐 [ASGI WEBSOCKET] Intercepted: path={path}, origin={origin}")
3415
+
3416
+ # Call the actual app
3417
+ await self.app(scope, receive, send)
3418
+
3419
+ # Wrap the app (but keep reference to original for internal use)
3420
+ WebSocketASGIWrapper(parent_app) # Wrapped for WebSocket support
3421
+
2983
3422
  # Add unified health check endpoint
2984
3423
  @parent_app.get("/health")
2985
3424
  async def health_check():
@@ -3266,7 +3705,163 @@ class MongoDBEngine:
3266
3705
 
3267
3706
  logger.info(f"Multi-app parent created with {len(apps)} app(s) configured")
3268
3707
 
3269
- return parent_app
3708
+ # CRITICAL: Wrap the FastAPI app in an ASGI wrapper to intercept WebSocket connections
3709
+ # BEFORE FastAPI's routing handles them. This will catch rejections at the framework level.
3710
+ from starlette.types import ASGIApp, Receive, Scope, Send
3711
+
3712
+ class WebSocketASGIInterceptor:
3713
+ """ASGI wrapper to intercept WebSocket connections at the ASGI level."""
3714
+
3715
+ def __init__(self, app: ASGIApp):
3716
+ self.app = app
3717
+ # Delegate attribute access to underlying FastAPI app
3718
+ # This allows app.routes, etc. to work
3719
+ # Note: We don't copy state here - we delegate it via property
3720
+ # to ensure changes to app.state are always visible
3721
+ for key, value in app.__dict__.items():
3722
+ if key != "state": # Don't copy state - delegate it
3723
+ setattr(self, key, value)
3724
+
3725
+ @property
3726
+ def state(self):
3727
+ # Always delegate state access to underlying app
3728
+ # This ensures changes to app.state are immediately visible
3729
+ return self.app.state
3730
+
3731
+ def __getattr__(self, name):
3732
+ # Delegate any missing attributes to underlying app
3733
+ return getattr(self.app, name)
3734
+
3735
+ @property
3736
+ def __class__(self):
3737
+ # Make isinstance() checks work by returning the underlying app's class
3738
+ # This allows isinstance(wrapper, FastAPI) to return True
3739
+ return type(self.app)
3740
+
3741
+ async def __call__(self, scope: Scope, receive: Receive, send: Send):
3742
+ # Intercept WebSocket connections at ASGI level (before FastAPI routing)
3743
+ if scope["type"] == "websocket":
3744
+ import sys
3745
+
3746
+ path = scope.get("path", "unknown")
3747
+ headers = {k.decode(): v.decode() for k, v in scope.get("headers", [])}
3748
+ origin = headers.get("origin", "missing")
3749
+ upgrade = headers.get("upgrade", "missing")
3750
+ connection = headers.get("connection", "missing")
3751
+
3752
+ print(
3753
+ f"🌐 [ASGI INTERCEPTOR] WebSocket connection at ASGI level: "
3754
+ f"path={path}, origin={origin}, upgrade={upgrade}, connection={connection}",
3755
+ file=sys.stderr,
3756
+ flush=True,
3757
+ )
3758
+ logger.info(f"🌐 [ASGI INTERCEPTOR] WebSocket: path={path}, origin={origin}")
3759
+
3760
+ # Wrap send to catch ALL messages to see what FastAPI is doing
3761
+ async def intercepted_send(message):
3762
+ import sys
3763
+
3764
+ msg_type = message.get("type", "unknown")
3765
+ print(
3766
+ f"🌐 [ASGI INTERCEPTOR] Message type: {msg_type}, "
3767
+ f"keys: {list(message.keys())}",
3768
+ file=sys.stderr,
3769
+ flush=True,
3770
+ )
3771
+
3772
+ if msg_type == "websocket.close":
3773
+ print(
3774
+ f"🌐 [ASGI INTERCEPTOR] WebSocket closed: "
3775
+ f"code={message.get('code', 'unknown')}, "
3776
+ f"reason={message.get('reason', 'unknown')}",
3777
+ file=sys.stderr,
3778
+ flush=True,
3779
+ )
3780
+ logger.warning(
3781
+ f"🌐 [ASGI INTERCEPTOR] WebSocket closed: "
3782
+ f"code={message.get('code')}, reason={message.get('reason')}"
3783
+ )
3784
+ elif msg_type == "websocket.accept":
3785
+ print(
3786
+ "🌐 [ASGI INTERCEPTOR] WebSocket ACCEPTED!",
3787
+ file=sys.stderr,
3788
+ flush=True,
3789
+ )
3790
+ logger.info("🌐 [ASGI INTERCEPTOR] WebSocket ACCEPTED!")
3791
+ elif msg_type == "websocket.http.response.start":
3792
+ status = message.get("status", "unknown")
3793
+ print(
3794
+ f"🌐 [ASGI INTERCEPTOR] HTTP response: status={status}",
3795
+ file=sys.stderr,
3796
+ flush=True,
3797
+ )
3798
+ logger.warning(f"🌐 [ASGI INTERCEPTOR] HTTP response: status={status}")
3799
+
3800
+ await send(message)
3801
+
3802
+ # Wrap receive to see what FastAPI is receiving
3803
+ async def intercepted_receive():
3804
+ msg = await receive()
3805
+ import sys
3806
+
3807
+ msg_type = msg.get("type", "unknown")
3808
+ print(
3809
+ f"🌐 [ASGI INTERCEPTOR] Received message: type={msg_type}, "
3810
+ f"keys: {list(msg.keys())}",
3811
+ file=sys.stderr,
3812
+ flush=True,
3813
+ )
3814
+ if msg_type == "websocket.connect":
3815
+ print(
3816
+ "🌐 [ASGI INTERCEPTOR] WebSocket CONNECT received!",
3817
+ file=sys.stderr,
3818
+ flush=True,
3819
+ )
3820
+ return msg
3821
+
3822
+ # Check if route exists before calling app
3823
+ import sys
3824
+
3825
+ if hasattr(self.app, "routes"):
3826
+ ws_routes = [
3827
+ r for r in self.app.routes if hasattr(r, "path") and "ws" in str(r.path)
3828
+ ]
3829
+ print(
3830
+ f"🌐 [ASGI INTERCEPTOR] Found {len(ws_routes)} WebSocket route(s): "
3831
+ f"{[r.path for r in ws_routes]}",
3832
+ file=sys.stderr,
3833
+ flush=True,
3834
+ )
3835
+
3836
+ try:
3837
+ await self.app(scope, intercepted_receive, intercepted_send)
3838
+ except (
3839
+ RuntimeError,
3840
+ ConnectionError,
3841
+ OSError,
3842
+ ValueError,
3843
+ AttributeError,
3844
+ ) as e:
3845
+ import sys
3846
+
3847
+ print(
3848
+ f"🌐 [ASGI INTERCEPTOR] Exception during WebSocket handling: "
3849
+ f"{type(e).__name__}: {e}",
3850
+ file=sys.stderr,
3851
+ flush=True,
3852
+ )
3853
+ logger.error(
3854
+ f"🌐 [ASGI INTERCEPTOR] Exception: {type(e).__name__}: {e}",
3855
+ exc_info=True,
3856
+ )
3857
+ raise
3858
+ else:
3859
+ # Non-WebSocket requests pass through normally
3860
+ await self.app(scope, receive, send)
3861
+
3862
+ # Re-enable ASGI interceptor to debug WebSocket connection issues
3863
+ # This will show us exactly what's happening at the ASGI level
3864
+ return WebSocketASGIInterceptor(parent_app)
3270
3865
 
3271
3866
  def get_mounted_apps(self, app: Optional["FastAPI"] = None) -> list[dict[str, Any]]:
3272
3867
  """