forward bearer token from authorization header to OpenAI (#19)

* forward bearer token from authorization header to OpenAI

* format

Author: kelsey-wong

services/chat-backend/src/app.py

@@ -3,16 +3,17 @@
 # setup logging, before importing anything else (and before creating other loggers)
 setup_logging()
 
-import json
 import logging
 
-from fastapi import FastAPI, Request, Response  # noqa: E402
+from fastapi import FastAPI, Request  # noqa: E402
 from fastapi.middleware.cors import CORSMiddleware  # noqa: E402
 from fastapi.responses import JSONResponse  # noqa: E402
 from fastapi.routing import APIRoute  # noqa: E402
+from fastapi.openapi.utils import get_openapi
 
 from src.core.config import get_settings  # noqa: E402
 from src.middleware import AuditMiddleware  # noqa: E402
+from src.middleware.request_context import RequestContextMiddleware  # noqa: E402
 from src.routers.main import api_router  # noqa: E402
 from src.utils.openapi_converter import convert_openapi_spec  # noqa: E402
 
@@ -35,6 +36,38 @@ def custom_generate_unique_id(route: APIRoute) -> str:
     generate_unique_id_function=custom_generate_unique_id,
 )
 
+def custom_openapi():
+    if app.openapi_schema:
+        return app.openapi_schema
+    
+    openapi_schema = get_openapi(
+        title=app.title,
+        version="1.0.0",
+        description=app.description,
+        routes=app.routes,
+    )
+    
+    # Add security scheme for Bearer token
+    openapi_schema["components"]["securitySchemes"] = {
+        "BearerAuth": {
+            "type": "http",
+            "scheme": "bearer",
+            "bearerFormat": "JWT",
+            "description": "Enter your bearer token for Azure AD authentication"
+        }
+    }
+    
+    # Apply security globally to all endpoints
+    for path in openapi_schema["paths"]:
+        for method in openapi_schema["paths"][path]:
+            if method != "parameters":
+                openapi_schema["paths"][path][method]["security"] = [{"BearerAuth": []}]
+    
+    app.openapi_schema = openapi_schema
+    return app.openapi_schema
+
+app.openapi = custom_openapi
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -43,6 +76,7 @@ def custom_generate_unique_id(route: APIRoute) -> str:
     allow_headers=["*"],
 )
 
+app.add_middleware(RequestContextMiddleware)
 app.add_middleware(AuditMiddleware)
 app.include_router(api_router, prefix=config.API_PREFIX)
 

services/chat-backend/src/middleware/request_context.py

@@ -0,0 +1,17 @@
+from fastapi import Request
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import Response
+
+from src.utils.request_context import set_authorization_header
+
+
+class RequestContextMiddleware(BaseHTTPMiddleware):
+    """Middleware to extract request headers and store them in context variables."""
+    
+    async def dispatch(self, request: Request, call_next):
+        # Extract Authorization header
+        auth_header = request.headers.get("Authorization")
+        set_authorization_header(auth_header)
+        
+        response = await call_next(request)
+        return response

services/chat-backend/src/services/azure.py

@@ -4,6 +4,7 @@
 from azure.core.exceptions import ClientAuthenticationError
 from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
 from tlm.config.types import ModelProvider
+from src.utils.request_context import get_authorization_header
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +28,12 @@ def _is_azure_model(model_provider: ModelProvider) -> bool:
 
 
 def get_azure_ad_token() -> str | None:
+    auth_header = get_authorization_header()
+    if auth_header:
+        # Bearer token
+        logger.info(f"using bearer token: {auth_header}")
+        return auth_header.split(" ")[1]
+
     try:
         managed_identity_client_id = os.getenv("AZURE_CLIENT_ID")
 

services/chat-backend/src/utils/request_context.py

@@ -0,0 +1,14 @@
+from contextvars import ContextVar
+from typing import Optional
+
+authorization_header: ContextVar[Optional[str]] = ContextVar("authorization_header", default=None)
+
+
+def get_authorization_header() -> Optional[str]:
+    """Get the Authorization header value from the current request context."""
+    return authorization_header.get()
+
+
+def set_authorization_header(value: Optional[str]) -> None:
+    """Set the Authorization header value in the current request context."""
+    authorization_header.set(value)