Add improvements 2

gateway/app/middleware/auth.py

@@ -14,6 +14,7 @@ import httpx
 import json
 
 from app.core.config import settings
+from app.core.header_manager import header_manager
 from shared.auth.jwt_handler import JWTHandler
 from shared.auth.tenant_access import tenant_access_manager, extract_tenant_id_from_path, is_tenant_scoped_path
 
@@ -60,15 +61,8 @@ class AuthMiddleware(BaseHTTPMiddleware):
         if request.method == "OPTIONS":
             return await call_next(request)
 
-        # SECURITY: Remove any incoming x-subscription-* headers
-        # These will be re-injected from verified JWT only
-        sanitized_headers = [
-            (k, v) for k, v in request.headers.raw
-            if not k.decode().lower().startswith('x-subscription-')
-            and not k.decode().lower().startswith('x-user-')
-            and not k.decode().lower().startswith('x-tenant-')
-        ]
-        request.headers.__dict__["_list"] = sanitized_headers
+        # SECURITY: Remove any incoming sensitive headers using HeaderManager
+        header_manager.sanitize_incoming_headers(request)
 
         # Skip authentication for public routes
         if self._is_public_route(request.url.path):
@@ -573,109 +567,13 @@ class AuthMiddleware(BaseHTTPMiddleware):
 
     async def _inject_context_headers(self, request: Request, user_context: Dict[str, Any], tenant_id: Optional[str] = None):
         """
-        Inject user and tenant context headers for downstream services
-        ENHANCED: Added logging to verify header injection
+        Inject user and tenant context headers for downstream services using unified HeaderManager
         """
-        # Enhanced logging for debugging
-        logger.info(
-            "🔧 Injecting context headers",
-            user_id=user_context.get("user_id"),
-            user_type=user_context.get("type", ""),
-            service_name=user_context.get("service", ""),
-            role=user_context.get("role", ""),
-            tenant_id=tenant_id,
-            is_demo=user_context.get("is_demo", False),
-            demo_session_id=user_context.get("demo_session_id", ""),
-            path=request.url.path
-        )
-
-        # Add user context headers
-        logger.debug(f"DEBUG: Injecting headers for user: {user_context.get('user_id')}, is_demo: {user_context.get('is_demo', False)}")
-        logger.debug(f"DEBUG: request.headers object id: {id(request.headers)}, _list id: {id(request.headers.__dict__.get('_list', []))}")
+        # Use unified HeaderManager for consistent header injection
+        injected_headers = header_manager.inject_context_headers(request, user_context, tenant_id)
         
-        # Store headers in request.state for cross-middleware access
-        request.state.injected_headers = {
-            "x-user-id": user_context["user_id"],
-            "x-user-email": user_context["email"],
-            "x-user-role": user_context.get("role", "user")
-        }
-        
-        request.headers.__dict__["_list"].append((
-            b"x-user-id", user_context["user_id"].encode()
-        ))
-        request.headers.__dict__["_list"].append((
-            b"x-user-email", user_context["email"].encode()
-        ))
-
-        user_role = user_context.get("role", "user")
-        request.headers.__dict__["_list"].append((
-            b"x-user-role", user_role.encode()
-        ))
-
-        user_type = user_context.get("type", "")
-        if user_type:
-            request.headers.__dict__["_list"].append((
-                b"x-user-type", user_type.encode()
-        ))
-
-        service_name = user_context.get("service", "")
-        if service_name:
-            request.headers.__dict__["_list"].append((
-                b"x-service-name", service_name.encode()
-        ))
-
-        # Add tenant context if available
-        if tenant_id:
-            request.headers.__dict__["_list"].append((
-                b"x-tenant-id", tenant_id.encode()
-        ))
-
-        # Add subscription tier if available
-        subscription_tier = user_context.get("subscription_tier", "")
-        if subscription_tier:
-            request.headers.__dict__["_list"].append((
-                b"x-subscription-tier", subscription_tier.encode()
-            ))
-
-        # Add is_demo flag for demo sessions
-        is_demo = user_context.get("is_demo", False)
-        logger.debug(f"DEBUG: is_demo value: {is_demo}, type: {type(is_demo)}")
-        if is_demo:
-            logger.info(f"🎭 Adding demo session headers",
-                       demo_session_id=user_context.get("demo_session_id", ""),
-                       demo_account_type=user_context.get("demo_account_type", ""),
-                       path=request.url.path)
-            request.headers.__dict__["_list"].append((
-                b"x-is-demo", b"true"
-            ))
-        else:
-            logger.debug(f"DEBUG: Not adding demo headers because is_demo is: {is_demo}")
-
-            # Add demo session context headers for backend services
-            demo_session_id = user_context.get("demo_session_id", "")
-            if demo_session_id:
-                request.headers.__dict__["_list"].append((
-                    b"x-demo-session-id", demo_session_id.encode()
-                ))
-
-            demo_account_type = user_context.get("demo_account_type", "")
-            if demo_account_type:
-                request.headers.__dict__["_list"].append((
-                    b"x-demo-account-type", demo_account_type.encode()
-                ))
-
         # Add hierarchical access headers if tenant context exists
         if tenant_id:
-            tenant_access_type = getattr(request.state, 'tenant_access_type', 'direct')
-            can_view_children = getattr(request.state, 'can_view_children', False)
-
-            request.headers.__dict__["_list"].append((
-                b"x-tenant-access-type", tenant_access_type.encode()
-            ))
-            request.headers.__dict__["_list"].append((
-                b"x-can-view-children", str(can_view_children).encode()
-            ))
-
             # If this is hierarchical access, include parent tenant ID
             # Get parent tenant ID from the auth service if available
             try:
@@ -689,17 +587,16 @@ class AuthMiddleware(BaseHTTPMiddleware):
                         hierarchy_data = response.json()
                         parent_tenant_id = hierarchy_data.get("parent_tenant_id")
                         if parent_tenant_id:
-                            request.headers.__dict__["_list"].append((
-                                b"x-parent-tenant-id", parent_tenant_id.encode()
-                            ))
+                            # Add parent tenant ID using HeaderManager for consistency
+                            header_name = header_manager.STANDARD_HEADERS['parent_tenant_id']
+                            header_value = str(parent_tenant_id)
+                            header_manager.add_header_for_middleware(request, header_name, header_value)
+                            logger.info(f"Added parent tenant ID header: {parent_tenant_id}")
             except Exception as e:
                 logger.warning(f"Failed to get parent tenant ID: {e}")
                 pass
-
-        # Add gateway identification
-        request.headers.__dict__["_list"].append((
-            b"x-forwarded-by", b"bakery-gateway"
-        ))
+        
+        return injected_headers
 
     async def _get_tenant_subscription_tier(self, tenant_id: str, request: Request) -> Optional[str]:
         """

gateway/app/middleware/rate_limiting.py

@@ -45,8 +45,17 @@ class APIRateLimitMiddleware(BaseHTTPMiddleware):
             return await call_next(request)
 
         try:
-            # Get subscription tier
-            subscription_tier = await self._get_subscription_tier(tenant_id, request)
+            # Get subscription tier from headers (added by AuthMiddleware)
+            subscription_tier = request.headers.get("x-subscription-tier")
+            
+            if not subscription_tier:
+                # Fallback: get from request state if headers not available
+                subscription_tier = getattr(request.state, "subscription_tier", None)
+            
+            if not subscription_tier:
+                # Final fallback: get from tenant service (should rarely happen)
+                subscription_tier = await self._get_subscription_tier(tenant_id, request)
+                logger.warning(f"Subscription tier not found in headers or state, fetched from tenant service: {subscription_tier}")
 
             # Get quota limit for tier
             quota_limit = self._get_quota_limit(subscription_tier)

gateway/app/middleware/request_id.py

@@ -9,6 +9,8 @@ from fastapi import Request
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import Response
 
+from app.core.header_manager import header_manager
+
 logger = structlog.get_logger()
 
 
@@ -40,11 +42,9 @@ class RequestIDMiddleware(BaseHTTPMiddleware):
         # Bind request ID to structured logger context
         logger_ctx = logger.bind(request_id=request_id)
 
-        # Inject request ID header for downstream services
-        # This is done by modifying the headers that will be forwarded
-        request.headers.__dict__["_list"].append((
-            b"x-request-id", request_id.encode()
-        ))
+        # Inject request ID header for downstream services using HeaderManager
+        # Note: This runs early in middleware chain, so we use add_header_for_middleware
+        header_manager.add_header_for_middleware(request, "x-request-id", request_id)
 
         # Log request start
         logger_ctx.info(

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(