Add improvements 2

gateway/app/middleware/subscription.py

@@ -15,6 +15,7 @@ import asyncio
 from datetime import datetime, timezone
 
 from app.core.config import settings
+from app.core.header_manager import header_manager
 from app.utils.subscription_error_responses import create_upgrade_required_response
 
 logger = structlog.get_logger()
@@ -178,7 +179,10 @@ class SubscriptionMiddleware(BaseHTTPMiddleware):
             r'/api/v1/subscriptions/.*',  # Subscription management itself
             r'/api/v1/tenants/[^/]+/members.*',  # Basic tenant info
             r'/docs.*',
-            r'/openapi\.json'
+            r'/openapi\.json',
+            # Training monitoring endpoints (WebSocket and status checks)
+            r'/api/v1/tenants/[^/]+/training/jobs/.*/live.*',  # WebSocket endpoint
+            r'/api/v1/tenants/[^/]+/training/jobs/.*/status.*',  # Status polling endpoint
         ]
 
         # Skip OPTIONS requests (CORS preflight)
@@ -275,21 +279,11 @@ class SubscriptionMiddleware(BaseHTTPMiddleware):
                         'current_tier': current_tier
                     }
 
-            # Use the same authentication pattern as gateway routes for fallback
-            headers = dict(request.headers)
-            headers.pop("host", None)
-
+            # Use unified HeaderManager for consistent header handling
+            headers = header_manager.get_all_headers_for_proxy(request)
+            
             # Extract user_id for logging (fallback path)
-            user_id = 'unknown'
-            # Add user context headers if available
-            if hasattr(request.state, 'user') and request.state.user:
-                user = request.state.user
-                user_id = str(user.get('user_id', 'unknown'))
-                headers["x-user-id"] = user_id
-                headers["x-user-email"] = str(user.get('email', ''))
-                headers["x-user-role"] = str(user.get('role', 'user'))
-                headers["x-user-full-name"] = str(user.get('full_name', ''))
-                headers["x-tenant-id"] = str(user.get('tenant_id', ''))
+            user_id = header_manager.get_header_value(request, 'x-user-id', 'unknown')
 
             # Call tenant service fast tier endpoint with caching (fallback for old tokens)
             timeout_config = httpx.Timeout(