Add improvements

gateway/app/middleware/auth.py

@@ -82,13 +82,16 @@ class AuthMiddleware(BaseHTTPMiddleware):
         # For SSE endpoint with demo_session_id in query params, validate it here
         if request.url.path == "/api/events" and demo_session_query and not hasattr(request.state, "is_demo_session"):
             logger.info(f"SSE endpoint with demo_session_id query param: {demo_session_query}")
-            # Validate demo session via demo-session service
+            # Validate demo session via demo-session service using JWT service token
             import httpx
             try:
+                # Create service token for gateway-to-demo-session communication
+                service_token = jwt_handler.create_service_token(service_name="gateway")
+
                 async with httpx.AsyncClient() as client:
                     response = await client.get(
                         f"http://demo-session-service:8000/api/v1/demo/sessions/{demo_session_query}",
-                        headers={"X-Internal-API-Key": "dev-internal-key-change-in-production"}
+                        headers={"Authorization": f"Bearer {service_token}"}
                     )
                     if response.status_code == 200:
                         session_data = response.json()
@@ -161,22 +164,27 @@ class AuthMiddleware(BaseHTTPMiddleware):
 
         # ✅ STEP 4: Verify tenant access if this is a tenant-scoped route
         if tenant_id and is_tenant_scoped_path(request.url.path):
-            # Use TenantAccessManager for gateway-level verification with caching
-            if self.redis_client and tenant_access_manager.redis_client is None:
-                tenant_access_manager.redis_client = self.redis_client
+            # Skip tenant access verification for service tokens (services have admin access)
+            if user_context.get("type") != "service":
+                # Use TenantAccessManager for gateway-level verification with caching
+                if self.redis_client and tenant_access_manager.redis_client is None:
+                    tenant_access_manager.redis_client = self.redis_client
 
-            has_access = await tenant_access_manager.verify_basic_tenant_access(
-                user_context["user_id"],
-                tenant_id
-            )
-
-            if not has_access:
-                logger.warning(f"User {user_context['email']} denied access to tenant {tenant_id}")
-                return JSONResponse(
-                    status_code=403,
-                    content={"detail": f"Access denied to tenant {tenant_id}"}
+                has_access = await tenant_access_manager.verify_basic_tenant_access(
+                    user_context["user_id"],
+                    tenant_id
                 )
 
+                if not has_access:
+                    logger.warning(f"User {user_context['email']} denied access to tenant {tenant_id}")
+                    return JSONResponse(
+                        status_code=403,
+                        content={"detail": f"Access denied to tenant {tenant_id}"}
+                    )
+            else:
+                logger.debug(f"Service token granted access to tenant {tenant_id}", 
+                           service=user_context.get("service"))
+
             # Get tenant subscription tier and inject into user context
             # NEW: Use JWT data if available, skip HTTP call
             if user_context.get("subscription_from_jwt"):
@@ -365,6 +373,12 @@ class AuthMiddleware(BaseHTTPMiddleware):
             except Exception as e:
                 logger.warning("Token freshness check setup failed", error=str(e))
 
+        # FIX: Validate service tokens with tenant context for tenant-scoped routes
+        if token_type == "service" and payload.get("tenant_id"):
+            # Service tokens with tenant context are valid for tenant-scoped operations
+            logger.debug("Service token with tenant context validated",
+                       service=payload.get("service"), tenant_id=payload.get("tenant_id"))
+
         return True
 
     def _validate_jwt_integrity(self, payload: Dict[str, Any]) -> bool:
@@ -469,7 +483,13 @@ class AuthMiddleware(BaseHTTPMiddleware):
             base_context["role"] = "admin"
             base_context["user_id"] = f"{service_name}-service"
             base_context["email"] = f"{service_name}-service@internal"
-            logger.debug(f"Service authentication: {payload['service']}")
+            
+            # FIX: Service tokens with tenant context should use that tenant_id
+            if payload.get("tenant_id"):
+                base_context["tenant_id"] = payload["tenant_id"]
+                logger.debug(f"Service authentication with tenant context: {service_name}, tenant_id: {payload['tenant_id']}")
+            else:
+                logger.debug(f"Service authentication: {service_name}")
     
         return base_context
 
@@ -556,18 +576,30 @@ class AuthMiddleware(BaseHTTPMiddleware):
         Inject user and tenant context headers for downstream services
         ENHANCED: Added logging to verify header injection
         """
-        # Log what we're injecting for debugging
-        logger.debug(
-            "Injecting context headers",
+        # 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', []))}")
+        
+        # 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()
         ))
@@ -607,10 +639,17 @@ class AuthMiddleware(BaseHTTPMiddleware):
 
         # 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", "")

gateway/app/middleware/demo_middleware.py

@@ -304,14 +304,27 @@ class DemoMiddleware(BaseHTTPMiddleware):
         return response
 
     async def _get_session_info(self, session_id: str) -> Optional[dict]:
-        """Get session information from demo service"""
+        """Get session information from demo service using JWT service token"""
         try:
+            # Create JWT service token for gateway-to-demo-session communication
+            from shared.auth.jwt_handler import JWTHandler
+            from app.core.config import settings
+
+            jwt_handler = JWTHandler(settings.JWT_SECRET_KEY, settings.JWT_ALGORITHM)
+            service_token = jwt_handler.create_service_token(service_name="gateway")
+
             async with httpx.AsyncClient(timeout=5.0) as client:
                 response = await client.get(
-                    f"{self.demo_session_url}/api/v1/demo/sessions/{session_id}"
+                    f"{self.demo_session_url}/api/v1/demo/sessions/{session_id}",
+                    headers={"Authorization": f"Bearer {service_token}"}
                 )
                 if response.status_code == 200:
                     return response.json()
+                else:
+                    logger.warning("Demo session fetch failed",
+                                 session_id=session_id,
+                                 status_code=response.status_code,
+                                 response_text=response.text[:200] if hasattr(response, 'text') else '')
                 return None
         except Exception as e:
             logger.error("Failed to get session info", session_id=session_id, error=str(e))