mdb-engine 0.6.0__py3-none-any.whl → 0.7.1__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
@@ -36,6 +36,11 @@ from typing import TYPE_CHECKING, Any, Optional
36
36
  from motor.motor_asyncio import AsyncIOMotorClient
37
37
  from pymongo.errors import PyMongoError
38
38
 
39
+ try:
40
+ from openai import OpenAIError
41
+ except ImportError:
42
+ OpenAIError = RuntimeError
43
+
39
44
  if TYPE_CHECKING:
40
45
  from fastapi import FastAPI
41
46
 
@@ -149,6 +154,7 @@ class MongoDBEngine:
149
154
  self._encryption_service: EnvelopeEncryptionService | None = None
150
155
  self._app_secrets_manager: AppSecretsManager | None = None
151
156
  self._websocket_session_manager: Any | None = None # WebSocketSessionManager
157
+ self._websocket_ticket_store: Any | None = None # WebSocketTicketStore
152
158
 
153
159
  # Store app read_scopes mapping for validation
154
160
  self._app_read_scopes: dict[str, list[str]] = {}
@@ -210,6 +216,13 @@ class MongoDBEngine:
210
216
  encryption_service=self._encryption_service,
211
217
  )
212
218
 
219
+ # Initialize WebSocket ticket store (in-memory, no dependencies needed)
220
+ # Tickets are preferred for multi-app SSO setups (short-lived, single-use)
221
+ from ..auth.websocket_tickets import WebSocketTicketStore
222
+
223
+ self._websocket_ticket_store = WebSocketTicketStore()
224
+ logger.info("WebSocket ticket store initialized")
225
+
213
226
  # Set up component managers
214
227
  self._app_registration_manager = AppRegistrationManager(
215
228
  mongo_db=self._connection_manager.mongo_db,
@@ -265,6 +278,16 @@ class MongoDBEngine:
265
278
  """Check if Ray is enabled and initialized."""
266
279
  return self.enable_ray and self.ray_actor is not None
267
280
 
281
+ @property
282
+ def connection_manager(self):
283
+ """
284
+ Get the connection manager.
285
+
286
+ Returns:
287
+ ConnectionManager instance
288
+ """
289
+ return self._connection_manager
290
+
268
291
  @property
269
292
  def mongo_client(self) -> AsyncIOMotorClient:
270
293
  """
@@ -709,6 +732,25 @@ class MongoDBEngine:
709
732
  app: FastAPI application instance
710
733
  slug: App slug
711
734
  """
735
+ # CRITICAL: Ensure websocket_ticket_store is available
736
+ # Ticket authentication is required for WebSocket connections
737
+ if not self._websocket_ticket_store:
738
+ error_msg = (
739
+ f"WebSocket routes cannot be registered for app '{slug}': "
740
+ "websocket_ticket_store is not available. "
741
+ "WebSocket authentication requires ticket store to be initialized."
742
+ )
743
+ contextual_logger.error(error_msg)
744
+ raise RuntimeError(error_msg)
745
+
746
+ # Ensure ticket store is in app state (may have been set in create_app)
747
+ if (
748
+ not hasattr(app.state, "websocket_ticket_store")
749
+ or app.state.websocket_ticket_store is None
750
+ ):
751
+ app.state.websocket_ticket_store = self._websocket_ticket_store
752
+ contextual_logger.debug(f"WebSocket ticket store stored in app state for '{slug}'")
753
+
712
754
  # Check if WebSockets are configured for this app
713
755
  websockets_config = self.get_websocket_config(slug)
714
756
  if not websockets_config:
@@ -915,9 +957,9 @@ class MongoDBEngine:
915
957
  return get_embedding_service_for_app(slug, self)
916
958
 
917
959
  @property
918
- def _apps(self) -> dict[str, Any]:
960
+ def apps(self) -> dict[str, Any]:
919
961
  """
920
- Get the apps dictionary (for backward compatibility with tests).
962
+ Get all registered apps.
921
963
 
922
964
  Returns:
923
965
  Dictionary of registered apps
@@ -927,7 +969,27 @@ class MongoDBEngine:
927
969
  """
928
970
  if not self._app_registration_manager:
929
971
  raise RuntimeError("MongoDBEngine not initialized. Call initialize() first.")
930
- return self._app_registration_manager._apps
972
+ return self._app_registration_manager.apps
973
+
974
+ @property
975
+ def websocket_ticket_store(self):
976
+ """
977
+ Get the WebSocket ticket store.
978
+
979
+ Returns:
980
+ WebSocketTicketStore instance or None if not initialized
981
+ """
982
+ return self._websocket_ticket_store
983
+
984
+ @property
985
+ def websocket_session_manager(self):
986
+ """
987
+ Get the WebSocket session manager.
988
+
989
+ Returns:
990
+ WebSocketSessionManager instance or None if not initialized
991
+ """
992
+ return self._websocket_session_manager
931
993
 
932
994
  def list_apps(self) -> list[str]:
933
995
  """
@@ -1100,6 +1162,73 @@ class MongoDBEngine:
1100
1162
  # FastAPI Integration Methods
1101
1163
  # =========================================================================
1102
1164
 
1165
+ async def _register_websocket_endpoints(self, app: "FastAPI", engine: "MongoDBEngine") -> None:
1166
+ """Register WebSocket ticket and session endpoints."""
1167
+ # Register WebSocket ticket endpoint AFTER initialization
1168
+ # (ticket store is now available)
1169
+ if engine.websocket_ticket_store:
1170
+ app.state.websocket_ticket_store = engine.websocket_ticket_store
1171
+ logger.info("WebSocket ticket store stored in app state")
1172
+
1173
+ # Set global ticket store for WebSocket authentication (works with routers)
1174
+ from ..routing.websockets import set_global_websocket_ticket_store
1175
+
1176
+ set_global_websocket_ticket_store(engine.websocket_ticket_store)
1177
+ logger.info("Global WebSocket ticket store set for multi-app authentication")
1178
+
1179
+ # Register WebSocket ticket endpoint
1180
+ from ..auth.websocket_tickets import create_websocket_ticket_endpoint
1181
+
1182
+ ticket_endpoint = create_websocket_ticket_endpoint(engine.websocket_ticket_store)
1183
+ app.post("/auth/ticket")(ticket_endpoint)
1184
+ logger.info("WebSocket ticket endpoint registered at /auth/ticket")
1185
+
1186
+ # Register WebSocket session endpoint AFTER initialization
1187
+ # (session manager is now available)
1188
+ if engine.websocket_session_manager:
1189
+ app.state.websocket_session_manager = engine.websocket_session_manager
1190
+ logger.info("WebSocket session manager stored in app state")
1191
+
1192
+ # Register WebSocket session endpoint
1193
+ from ..auth.websocket_sessions import create_websocket_session_endpoint
1194
+
1195
+ session_endpoint = create_websocket_session_endpoint(engine.websocket_session_manager)
1196
+ app.get("/auth/websocket-session")(session_endpoint)
1197
+ logger.info("WebSocket session endpoint registered at /auth/websocket-session")
1198
+
1199
+ async def _configure_websocket_ticket_ttl(
1200
+ self, app: "FastAPI", app_manifest: dict[str, Any], slug: str
1201
+ ) -> None:
1202
+ """Configure WebSocket ticket TTL from manifest."""
1203
+ websockets_config = app_manifest.get("websockets", {})
1204
+ if not websockets_config:
1205
+ return
1206
+
1207
+ from ..auth.websocket_tickets import WebSocketTicketStore
1208
+
1209
+ ticket_ttl_values: list[int] = []
1210
+ for endpoint_config in websockets_config.values():
1211
+ if isinstance(endpoint_config, dict):
1212
+ ticket_ttl = endpoint_config.get("ticket_ttl_seconds")
1213
+ if ticket_ttl is not None:
1214
+ ticket_ttl_values.append(ticket_ttl)
1215
+
1216
+ if ticket_ttl_values:
1217
+ configured_ticket_ttl = min(ticket_ttl_values) # Use minimum for maximum security
1218
+ # Reinitialize ticket store if needed
1219
+ ticket_store = self._websocket_ticket_store
1220
+ if ticket_store is None or ticket_store.ticket_ttl != configured_ticket_ttl:
1221
+ self._websocket_ticket_store = WebSocketTicketStore(
1222
+ ticket_ttl_seconds=configured_ticket_ttl
1223
+ )
1224
+ logger.info(
1225
+ f"WebSocket ticket store initialized with TTL: "
1226
+ f"{configured_ticket_ttl}s (from app '{slug}' manifest)"
1227
+ )
1228
+ # Update app state if ticket store was already set
1229
+ if hasattr(app.state, "websocket_ticket_store"):
1230
+ app.state.websocket_ticket_store = self._websocket_ticket_store
1231
+
1103
1232
  def create_app(
1104
1233
  self,
1105
1234
  slug: str,
@@ -1199,8 +1328,15 @@ class MongoDBEngine:
1199
1328
  if not is_sub_app:
1200
1329
  await engine.initialize()
1201
1330
 
1331
+ # Register WebSocket endpoints
1332
+ await self._register_websocket_endpoints(app, engine)
1333
+
1202
1334
  # Load and register manifest
1203
1335
  app_manifest = await engine.load_manifest(manifest_path)
1336
+
1337
+ # Configure WebSocket ticket TTL from manifest
1338
+ await self._configure_websocket_ticket_ttl(app, app_manifest, slug)
1339
+
1204
1340
  await engine.register_app(app_manifest)
1205
1341
 
1206
1342
  # Auto-detect multi-site mode from manifest
@@ -1235,11 +1371,11 @@ class MongoDBEngine:
1235
1371
  f"Sub-app '{slug}' uses shared auth but user_pool not found. "
1236
1372
  "Initializing now (parent should have initialized it)."
1237
1373
  )
1238
- await engine._initialize_shared_user_pool(app, app_manifest)
1374
+ await self._initialize_shared_user_pool(app, app_manifest)
1239
1375
  else:
1240
1376
  logger.debug(f"Sub-app '{slug}' using shared user_pool from parent app")
1241
1377
  else:
1242
- await engine._initialize_shared_user_pool(app, app_manifest)
1378
+ await self._initialize_shared_user_pool(app, app_manifest)
1243
1379
  else:
1244
1380
  logger.info(f"Per-app auth mode for '{slug}'")
1245
1381
  # Auto-retrieve app token for "app" mode
@@ -1505,6 +1641,10 @@ class MongoDBEngine:
1505
1641
  f"(require_role={auth_config.get('require_role')})"
1506
1642
  )
1507
1643
 
1644
+ # NOTE: WebSocket ticket endpoint registration is moved to lifespan context manager
1645
+ # (after engine.initialize()) because ticket store is only available after initialization.
1646
+ # This ensures consistency with create_multi_app() behavior.
1647
+
1508
1648
  # Add CSRF middleware (after auth - auto-enabled for shared mode)
1509
1649
  # CSRF protection is enabled by default for shared auth mode
1510
1650
  # SKIP for sub-apps in multi-app setups - parent app handles CSRF
@@ -1512,8 +1652,17 @@ class MongoDBEngine:
1512
1652
  if csrf_config and not is_sub_app: # Don't add CSRF to child apps
1513
1653
  from ..auth.csrf import create_csrf_middleware
1514
1654
 
1655
+ # Add ticket endpoint to public routes (it handles its own auth)
1656
+ public_routes = auth_config.get("public_routes", [])
1657
+ public_routes_with_ticket = list(public_routes) + ["/auth/ticket"]
1658
+
1659
+ csrf_config_with_routes = {
1660
+ **auth_config,
1661
+ "public_routes": public_routes_with_ticket,
1662
+ }
1663
+
1515
1664
  csrf_middleware = create_csrf_middleware(
1516
- manifest_auth=auth_config,
1665
+ manifest_auth=csrf_config_with_routes,
1517
1666
  )
1518
1667
  app.add_middleware(csrf_middleware)
1519
1668
  logger.info(f"CSRFMiddleware added for '{slug}'")
@@ -2142,6 +2291,7 @@ class MongoDBEngine:
2142
2291
  )
2143
2292
 
2144
2293
  # Check if any app uses shared auth and collect public routes for CSRF exemption
2294
+ # Also collect ticket TTL values from websocket configs
2145
2295
  has_shared_auth = False
2146
2296
  all_public_routes = [
2147
2297
  "/health",
@@ -2149,6 +2299,7 @@ class MongoDBEngine:
2149
2299
  "/openapi.json",
2150
2300
  "/_mdb/routes",
2151
2301
  ] # Base exempt routes
2302
+ ticket_ttl_values: list[int] = [] # Collect ticket TTLs from all apps
2152
2303
  for app_config in apps:
2153
2304
  try:
2154
2305
  manifest_path = app_config["manifest"]
@@ -2168,9 +2319,45 @@ class MongoDBEngine:
2168
2319
  prefixed_route = f"{path_prefix.rstrip('/')}/{route}"
2169
2320
  if prefixed_route not in all_public_routes:
2170
2321
  all_public_routes.append(prefixed_route)
2322
+
2323
+ # Collect ticket TTL from websocket configs
2324
+ websockets_config = app_manifest_pre.get("websockets", {})
2325
+ if websockets_config:
2326
+ for endpoint_config in websockets_config.values():
2327
+ if isinstance(endpoint_config, dict):
2328
+ ticket_ttl = endpoint_config.get("ticket_ttl_seconds")
2329
+ if ticket_ttl is not None:
2330
+ ticket_ttl_values.append(ticket_ttl)
2171
2331
  except (FileNotFoundError, json.JSONDecodeError, KeyError) as e:
2172
2332
  logger.warning(f"Could not check auth mode for app '{app_config.get('slug')}': {e}")
2173
2333
 
2334
+ # Determine ticket TTL: use minimum from app configs (most secure), or default
2335
+ from ..auth.websocket_tickets import DEFAULT_TICKET_TTL_SECONDS
2336
+
2337
+ if ticket_ttl_values:
2338
+ configured_ticket_ttl = min(ticket_ttl_values) # Use minimum for maximum security
2339
+ logger.info(
2340
+ f"Ticket TTL configured from app manifests: {configured_ticket_ttl}s "
2341
+ f"(found values: {ticket_ttl_values}, using minimum)"
2342
+ )
2343
+ else:
2344
+ configured_ticket_ttl = DEFAULT_TICKET_TTL_SECONDS
2345
+ logger.debug(
2346
+ f"No ticket TTL specified in app manifests, using default: {configured_ticket_ttl}s"
2347
+ )
2348
+
2349
+ # Reinitialize ticket store with configured TTL if different from current
2350
+ if (
2351
+ self._websocket_ticket_store is None
2352
+ or self._websocket_ticket_store.ticket_ttl != configured_ticket_ttl
2353
+ ):
2354
+ from ..auth.websocket_tickets import WebSocketTicketStore
2355
+
2356
+ self._websocket_ticket_store = WebSocketTicketStore(
2357
+ ticket_ttl_seconds=configured_ticket_ttl
2358
+ )
2359
+ logger.info(f"WebSocket ticket store initialized with TTL: {configured_ticket_ttl}s")
2360
+
2174
2361
  # Validate hooks before creating lifespan (fail fast)
2175
2362
  for app_config in apps:
2176
2363
  slug = app_config.get("slug", "unknown")
@@ -2291,10 +2478,49 @@ class MongoDBEngine:
2291
2478
  logger.debug(f"No WebSocket configuration found for app '{slug}'")
2292
2479
  return
2293
2480
 
2481
+ # CRITICAL: Check if session manager is required and available
2482
+ # Some endpoints require session keys (csrf_required=True), which need
2483
+ # session manager
2484
+ requires_session_manager = False
2485
+ for _endpoint_name, endpoint_config in websockets_config.items():
2486
+ auth_config = endpoint_config.get("auth", {})
2487
+ if isinstance(auth_config, dict):
2488
+ csrf_required = auth_config.get("csrf_required", True) # Default to True
2489
+ if csrf_required:
2490
+ requires_session_manager = True
2491
+ break
2492
+
2493
+ if requires_session_manager and not engine.websocket_session_manager:
2494
+ error_msg = (
2495
+ f"WebSocket routes cannot be registered for app '{slug}': "
2496
+ "websocket_session_manager is not available. "
2497
+ "WebSocket endpoints with csrf_required=True require "
2498
+ "session manager to be initialized. "
2499
+ "Set MDB_ENGINE_MASTER_KEY environment variable to enable "
2500
+ "session manager."
2501
+ )
2502
+ logger.error(error_msg)
2503
+ raise RuntimeError(error_msg)
2504
+
2505
+ # CRITICAL: Ensure websocket_ticket_store is available
2506
+ # Ticket authentication is required for WebSocket connections
2507
+ if not engine.websocket_ticket_store:
2508
+ error_msg = (
2509
+ f"WebSocket routes cannot be registered for app '{slug}': "
2510
+ "websocket_ticket_store is not available. "
2511
+ "WebSocket authentication requires ticket store to be initialized."
2512
+ )
2513
+ logger.error(error_msg)
2514
+ raise RuntimeError(error_msg)
2515
+
2294
2516
  # Store WebSocket config in parent app state for CSRF middleware to access
2295
2517
  if not hasattr(parent_app.state, "websocket_configs"):
2296
2518
  parent_app.state.websocket_configs = {}
2297
2519
  parent_app.state.websocket_configs[slug] = websockets_config
2520
+ logger.info(
2521
+ f"✅ Stored WebSocket config for '{slug}' in parent app state "
2522
+ f"({len(websockets_config)} endpoint(s))"
2523
+ )
2298
2524
 
2299
2525
  try:
2300
2526
  from fastapi import APIRouter
@@ -2335,16 +2561,51 @@ class MongoDBEngine:
2335
2561
  ping_interval=ping_interval,
2336
2562
  )
2337
2563
 
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)
2564
+ # Register on parent app with full path using FastAPI's
2565
+ # proper WebSocket registration
2566
+ # We register BEFORE mounting apps to ensure WebSocket
2567
+ # routes are checked first
2568
+ try:
2569
+ # Use FastAPI's APIRouter approach (same as single-app mode)
2570
+ # This maintains FastAPI features
2571
+ # (dependency injection, OpenAPI docs, etc.)
2572
+ ws_router = APIRouter()
2573
+ ws_router.websocket(full_ws_path)(handler)
2574
+
2575
+ # Include router BEFORE mounting child app to ensure route priority
2576
+ parent_app.include_router(ws_router)
2577
+
2578
+ logger.info(
2579
+ f"✅ Registered WebSocket route '{full_ws_path}' "
2580
+ f"using FastAPI APIRouter "
2581
+ f"(registered before app mount to ensure priority)"
2582
+ )
2583
+ except (
2584
+ ValueError,
2585
+ RuntimeError,
2586
+ AttributeError,
2587
+ TypeError,
2588
+ ) as fastapi_error:
2589
+ logger.error(
2590
+ f"❌ Failed to register WebSocket route "
2591
+ f"'{full_ws_path}' with FastAPI: {fastapi_error}",
2592
+ exc_info=True,
2593
+ )
2594
+ raise
2342
2595
 
2343
2596
  logger.info(
2344
2597
  f"✅ Registered WebSocket route '{full_ws_path}' "
2345
2598
  f"for mounted app '{slug}' (mounted at '{path_prefix}', "
2346
2599
  f"auth: {require_auth}, ping: {ping_interval}s)"
2347
2600
  )
2601
+ import sys
2602
+
2603
+ print(
2604
+ f"✅ [ROUTE REGISTRATION] WebSocket route '{full_ws_path}' "
2605
+ f"registered for '{slug}' using FastAPI APIRouter",
2606
+ file=sys.stderr,
2607
+ flush=True,
2608
+ )
2348
2609
 
2349
2610
  # Verify route was actually registered
2350
2611
  registered_routes = [
@@ -2352,6 +2613,26 @@ class MongoDBEngine:
2352
2613
  for r in parent_app.routes
2353
2614
  if hasattr(r, "path") and full_ws_path in str(getattr(r, "path", ""))
2354
2615
  ]
2616
+
2617
+ # CRITICAL: Log all WebSocket routes to verify registration
2618
+ # FastAPI APIRouter creates routes of type 'APIWebSocketRoute'
2619
+ all_ws_routes = [
2620
+ (r.path, type(r).__name__)
2621
+ for r in parent_app.routes
2622
+ if hasattr(r, "path")
2623
+ and ("ws" in str(r.path).lower() or hasattr(r, "endpoint"))
2624
+ ]
2625
+ import sys
2626
+
2627
+ print(
2628
+ f"📋 [ROUTE VERIFICATION] All WebSocket-like routes: {all_ws_routes}",
2629
+ file=sys.stderr,
2630
+ flush=True,
2631
+ )
2632
+ logger.info(
2633
+ f"📋 [ROUTE VERIFICATION] All WebSocket-like routes: {all_ws_routes}"
2634
+ )
2635
+
2355
2636
  if registered_routes:
2356
2637
  registered_count += 1
2357
2638
  logger.debug(
@@ -2396,13 +2677,47 @@ class MongoDBEngine:
2396
2677
  )
2397
2678
 
2398
2679
  @asynccontextmanager
2399
- async def lifespan(app: FastAPI):
2680
+ async def lifespan(app: FastAPI): # noqa: C901
2400
2681
  """Lifespan context manager for parent app."""
2401
2682
  nonlocal mounted_apps, shared_user_pool_initialized
2402
2683
 
2403
2684
  # Initialize engine
2404
2685
  await engine.initialize()
2405
2686
 
2687
+ # Register WebSocket ticket endpoint AFTER initialization
2688
+ # (ticket store is now available)
2689
+ if engine.websocket_ticket_store:
2690
+ app.state.websocket_ticket_store = engine.websocket_ticket_store
2691
+ logger.info("WebSocket ticket store stored in parent app state")
2692
+
2693
+ # Set global ticket store for WebSocket authentication (works with routers)
2694
+ from ..routing.websockets import set_global_websocket_ticket_store
2695
+
2696
+ set_global_websocket_ticket_store(engine.websocket_ticket_store)
2697
+ logger.info("Global WebSocket ticket store set for multi-app authentication")
2698
+
2699
+ # Register WebSocket ticket endpoint on parent app
2700
+ from ..auth.websocket_tickets import create_websocket_ticket_endpoint
2701
+
2702
+ ticket_endpoint = create_websocket_ticket_endpoint(engine.websocket_ticket_store)
2703
+ app.post("/auth/ticket")(ticket_endpoint)
2704
+ logger.info("WebSocket ticket endpoint registered at /auth/ticket")
2705
+
2706
+ # Register WebSocket session endpoint AFTER initialization
2707
+ # (session manager is now available)
2708
+ if engine.websocket_session_manager:
2709
+ app.state.websocket_session_manager = engine.websocket_session_manager
2710
+ logger.info("WebSocket session manager stored in parent app state")
2711
+
2712
+ # Register WebSocket session endpoint on parent app
2713
+ from ..auth.websocket_sessions import create_websocket_session_endpoint
2714
+
2715
+ session_endpoint = create_websocket_session_endpoint(
2716
+ engine.websocket_session_manager
2717
+ )
2718
+ app.get("/auth/websocket-session")(session_endpoint)
2719
+ logger.info("WebSocket session endpoint registered at /auth/websocket-session")
2720
+
2406
2721
  # Initialize shared user pool once if any app uses shared auth
2407
2722
  if has_shared_auth:
2408
2723
  logger.info("Initializing shared user pool for multi-app deployment")
@@ -2414,7 +2729,7 @@ class MongoDBEngine:
2414
2729
  app_manifest_pre = json.load(f)
2415
2730
  auth_config = app_manifest_pre.get("auth", {})
2416
2731
  if auth_config.get("mode") == "shared":
2417
- await engine._initialize_shared_user_pool(app, app_manifest_pre)
2732
+ await self._initialize_shared_user_pool(app, app_manifest_pre)
2418
2733
  shared_user_pool_initialized = True
2419
2734
  logger.info("Shared user pool initialized for multi-app deployment")
2420
2735
  break
@@ -2509,6 +2824,11 @@ class MongoDBEngine:
2509
2824
  )
2510
2825
  logger.debug(f"Shared WebSocket session manager with child app '{slug}'")
2511
2826
 
2827
+ # Share WebSocket ticket store with child app
2828
+ if hasattr(app.state, "websocket_ticket_store"):
2829
+ child_app.state.websocket_ticket_store = app.state.websocket_ticket_store
2830
+ logger.debug(f"Shared WebSocket ticket store with child app '{slug}'")
2831
+
2512
2832
  # Add middleware for app context helpers
2513
2833
  from starlette.middleware.base import BaseHTTPMiddleware
2514
2834
  from starlette.requests import Request
@@ -2596,15 +2916,66 @@ class MongoDBEngine:
2596
2916
  child_app.add_middleware(middleware_class)
2597
2917
  logger.debug(f"Added AppContextMiddleware to child app '{slug}'")
2598
2918
 
2599
- # Mount child app at path prefix
2919
+ await _register_websocket_routes(app, app_manifest_data, slug, path_prefix)
2920
+
2921
+ memory_config = app_manifest_data.get("memory_config")
2922
+ if memory_config and memory_config.get("enabled", False):
2923
+ if engine._service_initializer: # noqa: SLF001
2924
+ try:
2925
+ await engine._service_initializer.initialize_memory_service( # noqa: SLF001
2926
+ slug, memory_config
2927
+ )
2928
+ logger.info(
2929
+ f"Memory service initialized for mounted app '{slug}' "
2930
+ f"in multi-app context"
2931
+ )
2932
+ except OpenAIError as e:
2933
+ logger.warning(
2934
+ f"Memory service initialization skipped for mounted app "
2935
+ f"'{slug}': OpenAI API error. {e}",
2936
+ extra={"app_slug": slug, "error": str(e)},
2937
+ )
2938
+ except (
2939
+ ImportError,
2940
+ AttributeError,
2941
+ TypeError,
2942
+ ValueError,
2943
+ RuntimeError,
2944
+ ConnectionError,
2945
+ OSError,
2946
+ ) as e:
2947
+ error_msg = str(e).lower()
2948
+ error_type = type(e).__name__
2949
+ is_api_key_error = (
2950
+ "api_key" in error_msg
2951
+ or "api key" in error_msg
2952
+ or "openai" in error_type.lower()
2953
+ )
2954
+ if is_api_key_error:
2955
+ logger.warning(
2956
+ f"Memory service initialization skipped for mounted app "
2957
+ f"'{slug}': Missing API key or configuration. {e}",
2958
+ extra={"app_slug": slug, "error": str(e)},
2959
+ )
2960
+ else:
2961
+ logger.error(
2962
+ f"Failed to initialize memory service for mounted app "
2963
+ f"'{slug}': {e}",
2964
+ exc_info=True,
2965
+ extra={"app_slug": slug, "error": str(e)},
2966
+ )
2967
+ else:
2968
+ logger.warning(
2969
+ f"Memory service requested for '{slug}' but "
2970
+ f"service_initializer is not available"
2971
+ )
2972
+
2973
+ # Mount child app at path prefix (AFTER WebSocket routes are registered)
2600
2974
  app.mount(path_prefix, child_app)
2601
2975
 
2602
2976
  # CRITICAL FIX: Merge CORS config from child app to parent app
2603
2977
  await _merge_cors_config_to_parent(app, child_app, app_manifest_data, slug)
2604
2978
 
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
2979
  # Update existing entry instead of appending
2609
2980
  entry = _find_mounted_app_entry(slug)
2610
2981
  if entry:
@@ -2733,6 +3104,11 @@ class MongoDBEngine:
2733
3104
  "manifest_path": str(manifest_path),
2734
3105
  }
2735
3106
  )
3107
+ # Always re-raise RuntimeError for critical failures
3108
+ # (like missing session manager)
3109
+ # These are configuration errors that should fail fast
3110
+ if isinstance(e, RuntimeError) and "websocket_session_manager" in str(e):
3111
+ raise RuntimeError(error_msg) from e
2736
3112
  if strict:
2737
3113
  raise RuntimeError(error_msg) from e
2738
3114
  continue
@@ -2826,11 +3202,78 @@ class MongoDBEngine:
2826
3202
  logger.debug("Set default CORS config on parent app for WebSocket origin validation")
2827
3203
 
2828
3204
  # Store app reference in engine for get_mounted_apps()
2829
- engine._multi_app_instance = parent_app
3205
+ self._multi_app_instance = parent_app
2830
3206
 
2831
- # Add request scope middleware
3207
+ # Add diagnostic ASGI middleware FIRST (outermost - runs before everything)
3208
+ # This will catch WebSocket upgrades before any other middleware
2832
3209
  from starlette.middleware.base import BaseHTTPMiddleware
2833
3210
 
3211
+ class DiagnosticMiddleware(BaseHTTPMiddleware):
3212
+ """Diagnostic middleware to log ALL requests, especially WebSocket upgrades."""
3213
+
3214
+ async def dispatch(self, request, call_next):
3215
+ path = request.url.path
3216
+ method = request.method
3217
+ upgrade_header = request.headers.get("upgrade", "").lower()
3218
+ connection_header = request.headers.get("connection", "").lower()
3219
+ origin_header = request.headers.get("origin")
3220
+
3221
+ # Log WebSocket upgrade attempts IMMEDIATELY
3222
+ if upgrade_header == "websocket" or "websocket" in path.lower():
3223
+ import sys
3224
+
3225
+ print(
3226
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket upgrade detected: "
3227
+ f"{method} {path}, upgrade={upgrade_header}, "
3228
+ f"connection={connection_header}, origin={origin_header}",
3229
+ file=sys.stderr,
3230
+ flush=True,
3231
+ )
3232
+ logger.info(
3233
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket upgrade: {method} {path}, "
3234
+ f"origin={origin_header}"
3235
+ )
3236
+
3237
+ try:
3238
+ response = await call_next(request)
3239
+
3240
+ # Log response for WebSocket upgrades
3241
+ if upgrade_header == "websocket" or "websocket" in path.lower():
3242
+ import sys
3243
+
3244
+ print(
3245
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket response: "
3246
+ f"{method} {path} -> {response.status_code}",
3247
+ file=sys.stderr,
3248
+ flush=True,
3249
+ )
3250
+ logger.info(
3251
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket response: "
3252
+ f"{method} {path} -> {response.status_code}"
3253
+ )
3254
+
3255
+ return response
3256
+ except (RuntimeError, ConnectionError, ValueError, AttributeError) as e:
3257
+ if upgrade_header == "websocket" or "websocket" in path.lower():
3258
+ import sys
3259
+
3260
+ print(
3261
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket exception: "
3262
+ f"{method} {path} -> {type(e).__name__}: {e}",
3263
+ file=sys.stderr,
3264
+ flush=True,
3265
+ )
3266
+ logger.error(
3267
+ f"🔬 [DIAGNOSTIC MIDDLEWARE] WebSocket exception: "
3268
+ f"{method} {path} -> {type(e).__name__}: {e}",
3269
+ exc_info=True,
3270
+ )
3271
+ raise
3272
+
3273
+ parent_app.add_middleware(DiagnosticMiddleware)
3274
+ logger.debug("DiagnosticMiddleware added for parent app (outermost layer)")
3275
+
3276
+ # Add request scope middleware
2834
3277
  from ..di import ScopeManager
2835
3278
 
2836
3279
  class RequestScopeMiddleware(BaseHTTPMiddleware):
@@ -2854,32 +3297,27 @@ class MongoDBEngine:
2854
3297
  from ..auth.csrf import create_csrf_middleware
2855
3298
 
2856
3299
  # 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"
3300
+ # Exempt routes that don't need CSRF (health checks, public routes
3301
+ # from child apps)
3302
+ # all_public_routes includes base routes + child app public routes
3303
+ # with path prefixes
3304
+ # Add WebSocket session and ticket endpoints to public routes
3305
+ # (they handle their own auth)
3306
+ public_routes_with_websocket_endpoints = list(all_public_routes) + [
3307
+ "/auth/websocket-session",
3308
+ "/auth/ticket",
2862
3309
  ]
2863
3310
  parent_csrf_config = {
2864
3311
  "csrf_protection": True,
2865
- "public_routes": public_routes_with_session_endpoint,
3312
+ "public_routes": public_routes_with_websocket_endpoints,
2866
3313
  }
2867
3314
  csrf_middleware = create_csrf_middleware(parent_csrf_config)
2868
3315
  parent_app.add_middleware(csrf_middleware)
2869
3316
 
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")
3317
+ # NOTE: WebSocket ticket and session endpoint registrations are moved to lifespan
3318
+ # context manager (after engine.initialize()) because they're only available after
3319
+ # initialization. The CSRF middleware still needs to know about these routes, so
3320
+ # they're added to public_routes_with_websocket_endpoints above.
2883
3321
 
2884
3322
  logger.info("CSRFMiddleware added to parent app for WebSocket origin validation")
2885
3323
 
@@ -2904,6 +3342,21 @@ class MongoDBEngine:
2904
3342
  # Read CORS config from app.state (may have been merged from child apps)
2905
3343
  cors_config = getattr(request.app.state, "cors_config", {})
2906
3344
 
3345
+ # CRITICAL: Log WebSocket upgrade requests to see if CORS
3346
+ # is intercepting them
3347
+ upgrade_header = request.headers.get("upgrade", "").lower()
3348
+ if upgrade_header == "websocket" or "websocket" in request.url.path.lower():
3349
+ import sys
3350
+
3351
+ print(
3352
+ f"🌐 [CORS MIDDLEWARE] WebSocket upgrade: "
3353
+ f"{request.method} {request.url.path}, "
3354
+ f"origin={request.headers.get('origin')}, "
3355
+ f"cors_enabled={cors_config.get('enabled', False)}",
3356
+ file=sys.stderr,
3357
+ flush=True,
3358
+ )
3359
+
2907
3360
  if not cors_config.get("enabled", False):
2908
3361
  # CORS not enabled, pass through
2909
3362
  return await call_next(request)
@@ -2980,6 +3433,46 @@ class MongoDBEngine:
2980
3433
  except ImportError:
2981
3434
  logger.warning("CORS middleware not available")
2982
3435
 
3436
+ # Wrap parent app in ASGI wrapper to intercept WebSocket connections at ASGI level
3437
+ # This must be done AFTER all middleware and routes are registered
3438
+ from starlette.types import ASGIApp, Receive, Scope, Send
3439
+
3440
+ class WebSocketASGIWrapper:
3441
+ """ASGI wrapper to intercept WebSocket connections before FastAPI routing."""
3442
+
3443
+ def __init__(self, app: ASGIApp):
3444
+ self.app = app
3445
+ # Delegate attribute access to underlying app
3446
+ self.__dict__.update(app.__dict__)
3447
+
3448
+ def __getattr__(self, name):
3449
+ # Delegate any missing attributes to underlying app
3450
+ return getattr(self.app, name)
3451
+
3452
+ async def __call__(self, scope: Scope, receive: Receive, send: Send):
3453
+ # Intercept WebSocket connections at ASGI level
3454
+ if scope["type"] == "websocket":
3455
+ import sys
3456
+
3457
+ path = scope.get("path", "unknown")
3458
+ headers_dict = {k.decode(): v.decode() for k, v in scope.get("headers", [])}
3459
+ origin = headers_dict.get("origin", "missing")
3460
+ upgrade = headers_dict.get("upgrade", "missing")
3461
+ connection = headers_dict.get("connection", "missing")
3462
+ print(
3463
+ f"🌐 [ASGI WEBSOCKET] Intercepted at ASGI level: path={path}, "
3464
+ f"origin={origin}, upgrade={upgrade}, connection={connection}",
3465
+ file=sys.stderr,
3466
+ flush=True,
3467
+ )
3468
+ logger.info(f"🌐 [ASGI WEBSOCKET] Intercepted: path={path}, origin={origin}")
3469
+
3470
+ # Call the actual app
3471
+ await self.app(scope, receive, send)
3472
+
3473
+ # Wrap the app (but keep reference to original for internal use)
3474
+ WebSocketASGIWrapper(parent_app) # Wrapped for WebSocket support
3475
+
2983
3476
  # Add unified health check endpoint
2984
3477
  @parent_app.get("/health")
2985
3478
  async def health_check():
@@ -3266,7 +3759,163 @@ class MongoDBEngine:
3266
3759
 
3267
3760
  logger.info(f"Multi-app parent created with {len(apps)} app(s) configured")
3268
3761
 
3269
- return parent_app
3762
+ # CRITICAL: Wrap the FastAPI app in an ASGI wrapper to intercept WebSocket connections
3763
+ # BEFORE FastAPI's routing handles them. This will catch rejections at the framework level.
3764
+ from starlette.types import ASGIApp, Receive, Scope, Send
3765
+
3766
+ class WebSocketASGIInterceptor:
3767
+ """ASGI wrapper to intercept WebSocket connections at the ASGI level."""
3768
+
3769
+ def __init__(self, app: ASGIApp):
3770
+ self.app = app
3771
+ # Delegate attribute access to underlying FastAPI app
3772
+ # This allows app.routes, etc. to work
3773
+ # Note: We don't copy state here - we delegate it via property
3774
+ # to ensure changes to app.state are always visible
3775
+ for key, value in app.__dict__.items():
3776
+ if key != "state": # Don't copy state - delegate it
3777
+ setattr(self, key, value)
3778
+
3779
+ @property
3780
+ def state(self):
3781
+ # Always delegate state access to underlying app
3782
+ # This ensures changes to app.state are immediately visible
3783
+ return self.app.state
3784
+
3785
+ def __getattr__(self, name):
3786
+ # Delegate any missing attributes to underlying app
3787
+ return getattr(self.app, name)
3788
+
3789
+ @property
3790
+ def __class__(self):
3791
+ # Make isinstance() checks work by returning the underlying app's class
3792
+ # This allows isinstance(wrapper, FastAPI) to return True
3793
+ return type(self.app)
3794
+
3795
+ async def __call__(self, scope: Scope, receive: Receive, send: Send):
3796
+ # Intercept WebSocket connections at ASGI level (before FastAPI routing)
3797
+ if scope["type"] == "websocket":
3798
+ import sys
3799
+
3800
+ path = scope.get("path", "unknown")
3801
+ headers = {k.decode(): v.decode() for k, v in scope.get("headers", [])}
3802
+ origin = headers.get("origin", "missing")
3803
+ upgrade = headers.get("upgrade", "missing")
3804
+ connection = headers.get("connection", "missing")
3805
+
3806
+ print(
3807
+ f"🌐 [ASGI INTERCEPTOR] WebSocket connection at ASGI level: "
3808
+ f"path={path}, origin={origin}, upgrade={upgrade}, connection={connection}",
3809
+ file=sys.stderr,
3810
+ flush=True,
3811
+ )
3812
+ logger.info(f"🌐 [ASGI INTERCEPTOR] WebSocket: path={path}, origin={origin}")
3813
+
3814
+ # Wrap send to catch ALL messages to see what FastAPI is doing
3815
+ async def intercepted_send(message):
3816
+ import sys
3817
+
3818
+ msg_type = message.get("type", "unknown")
3819
+ print(
3820
+ f"🌐 [ASGI INTERCEPTOR] Message type: {msg_type}, "
3821
+ f"keys: {list(message.keys())}",
3822
+ file=sys.stderr,
3823
+ flush=True,
3824
+ )
3825
+
3826
+ if msg_type == "websocket.close":
3827
+ print(
3828
+ f"🌐 [ASGI INTERCEPTOR] WebSocket closed: "
3829
+ f"code={message.get('code', 'unknown')}, "
3830
+ f"reason={message.get('reason', 'unknown')}",
3831
+ file=sys.stderr,
3832
+ flush=True,
3833
+ )
3834
+ logger.warning(
3835
+ f"🌐 [ASGI INTERCEPTOR] WebSocket closed: "
3836
+ f"code={message.get('code')}, reason={message.get('reason')}"
3837
+ )
3838
+ elif msg_type == "websocket.accept":
3839
+ print(
3840
+ "🌐 [ASGI INTERCEPTOR] WebSocket ACCEPTED!",
3841
+ file=sys.stderr,
3842
+ flush=True,
3843
+ )
3844
+ logger.info("🌐 [ASGI INTERCEPTOR] WebSocket ACCEPTED!")
3845
+ elif msg_type == "websocket.http.response.start":
3846
+ status = message.get("status", "unknown")
3847
+ print(
3848
+ f"🌐 [ASGI INTERCEPTOR] HTTP response: status={status}",
3849
+ file=sys.stderr,
3850
+ flush=True,
3851
+ )
3852
+ logger.warning(f"🌐 [ASGI INTERCEPTOR] HTTP response: status={status}")
3853
+
3854
+ await send(message)
3855
+
3856
+ # Wrap receive to see what FastAPI is receiving
3857
+ async def intercepted_receive():
3858
+ msg = await receive()
3859
+ import sys
3860
+
3861
+ msg_type = msg.get("type", "unknown")
3862
+ print(
3863
+ f"🌐 [ASGI INTERCEPTOR] Received message: type={msg_type}, "
3864
+ f"keys: {list(msg.keys())}",
3865
+ file=sys.stderr,
3866
+ flush=True,
3867
+ )
3868
+ if msg_type == "websocket.connect":
3869
+ print(
3870
+ "🌐 [ASGI INTERCEPTOR] WebSocket CONNECT received!",
3871
+ file=sys.stderr,
3872
+ flush=True,
3873
+ )
3874
+ return msg
3875
+
3876
+ # Check if route exists before calling app
3877
+ import sys
3878
+
3879
+ if hasattr(self.app, "routes"):
3880
+ ws_routes = [
3881
+ r for r in self.app.routes if hasattr(r, "path") and "ws" in str(r.path)
3882
+ ]
3883
+ print(
3884
+ f"🌐 [ASGI INTERCEPTOR] Found {len(ws_routes)} WebSocket route(s): "
3885
+ f"{[r.path for r in ws_routes]}",
3886
+ file=sys.stderr,
3887
+ flush=True,
3888
+ )
3889
+
3890
+ try:
3891
+ await self.app(scope, intercepted_receive, intercepted_send)
3892
+ except (
3893
+ RuntimeError,
3894
+ ConnectionError,
3895
+ OSError,
3896
+ ValueError,
3897
+ AttributeError,
3898
+ ) as e:
3899
+ import sys
3900
+
3901
+ print(
3902
+ f"🌐 [ASGI INTERCEPTOR] Exception during WebSocket handling: "
3903
+ f"{type(e).__name__}: {e}",
3904
+ file=sys.stderr,
3905
+ flush=True,
3906
+ )
3907
+ logger.error(
3908
+ f"🌐 [ASGI INTERCEPTOR] Exception: {type(e).__name__}: {e}",
3909
+ exc_info=True,
3910
+ )
3911
+ raise
3912
+ else:
3913
+ # Non-WebSocket requests pass through normally
3914
+ await self.app(scope, receive, send)
3915
+
3916
+ # Re-enable ASGI interceptor to debug WebSocket connection issues
3917
+ # This will show us exactly what's happening at the ASGI level
3918
+ return WebSocketASGIInterceptor(parent_app)
3270
3919
 
3271
3920
  def get_mounted_apps(self, app: Optional["FastAPI"] = None) -> list[dict[str, Any]]:
3272
3921
  """