Propagate PolicyEngine request IDs to the simulation service

changelog.d/request-id-propagation.fixed.md

@@ -0,0 +1,2 @@
+Propagate `X-PolicyEngine-Request-Id` from API requests to every synchronous
+simulation gateway request for end-to-end log correlation.

policyengine_api/asgi_factory.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import time
-import uuid
 from typing import Literal
 
 from a2wsgi import WSGIMiddleware
@@ -14,7 +13,13 @@
     get_api_host_backend,
 )
 from policyengine_api.migration_logging import log_migration_request
+from policyengine_api.request_context import (
+    REQUEST_ID_HEADER,
+    _asgi_request_id,
+    generate_request_id,
+)
 from pydantic import BaseModel
+from starlette.datastructures import MutableHeaders
 from starlette.middleware.gzip import GZipMiddleware
 
 FASTAPI_NATIVE_LOGGED_PATHS = frozenset(
@@ -63,26 +68,32 @@ def create_asgi_app(wsgi_app) -> FastAPI:
     @app.middleware("http")
     async def add_cors_for_native_routes(request, call_next):
         started_at = time.time()
-        request_id = request.headers.get("X-Request-ID") or uuid.uuid4().hex
-        response = await call_next(request)
-        if BACKEND_RESPONSE_HEADER not in response.headers:
-            response.headers[BACKEND_RESPONSE_HEADER] = get_api_host_backend()
-        origin = request.headers.get("origin")
-        if origin and "access-control-allow-origin" not in response.headers:
-            response.headers["Access-Control-Allow-Origin"] = origin
-            _add_vary_origin(response)
-        if request.url.path in FASTAPI_NATIVE_LOGGED_PATHS:
-            try:
-                log_migration_request(
-                    request_id=request_id,
-                    method=request.method,
-                    path=request.url.path,
-                    status_code=response.status_code,
-                    started_at=started_at,
-                )
-            except Exception:
-                pass
-        return response
+        request_id = request.headers.get(REQUEST_ID_HEADER) or generate_request_id()
+        MutableHeaders(scope=request.scope)[REQUEST_ID_HEADER] = request_id
+        context_token = _asgi_request_id.set(request_id)
+        try:
+            response = await call_next(request)
+            response.headers[REQUEST_ID_HEADER] = request_id
+            if BACKEND_RESPONSE_HEADER not in response.headers:
+                response.headers[BACKEND_RESPONSE_HEADER] = get_api_host_backend()
+            origin = request.headers.get("origin")
+            if origin and "access-control-allow-origin" not in response.headers:
+                response.headers["Access-Control-Allow-Origin"] = origin
+                _add_vary_origin(response)
+            if request.url.path in FASTAPI_NATIVE_LOGGED_PATHS:
+                try:
+                    log_migration_request(
+                        request_id=request_id,
+                        method=request.method,
+                        path=request.url.path,
+                        status_code=response.status_code,
+                        started_at=started_at,
+                    )
+                except Exception:
+                    pass
+            return response
+        finally:
+            _asgi_request_id.reset(context_token)
 
     @app.get("/health", response_model=HealthResponse)
     def health() -> HealthResponse:

policyengine_api/libs/simulation_entrypoint.py

@@ -16,6 +16,10 @@
     gateway_auth_required,
 )
 from policyengine_api.migration_flags import get_sim_entrypoint
+from policyengine_api.request_context import (
+    REQUEST_ID_HEADER,
+    current_request_id,
+)
 
 
 def _required_base_url(env_name: str) -> str:
@@ -49,6 +53,14 @@ def resolve_simulation_entrypoint_url(entrypoint: str | None = None) -> str:
     raise ValueError(f"Unsupported simulation entrypoint: {selected_entrypoint!r}")
 
 
+def _attach_current_request_id(request: httpx.Request) -> None:
+    """Attach the active API request ID to an outbound simulation request."""
+
+    request_id = current_request_id()
+    if request_id is not None:
+        request.headers[REQUEST_ID_HEADER] = request_id
+
+
 @dataclass
 class ModalSimulationExecution:
     """
@@ -124,7 +136,11 @@ def __init__(self, entrypoint: str | None = None):
                 file=sys.stderr,
                 flush=True,
             )
-        self.client = httpx.Client(timeout=30.0, auth=auth)
+        self.client = httpx.Client(
+            timeout=30.0,
+            auth=auth,
+            event_hooks={"request": [_attach_current_request_id]},
+        )
 
     def _normalize_submission_payload(self, payload: dict) -> dict:
         modal_payload = {

policyengine_api/migration_logging.py

@@ -3,17 +3,19 @@
 from __future__ import annotations
 
 import time
-import uuid
 
 import flask
-
 from policyengine_api.gcp_logging import logger
 from policyengine_api.migration_flags import (
     BACKEND_RESPONSE_HEADER,
     get_api_host_backend,
     get_migration_log_context,
     infer_route_group,
 )
+from policyengine_api.request_context import (
+    REQUEST_ID_HEADER,
+    generate_request_id,
+)
 
 
 def register_migration_request_logging(app: flask.Flask) -> None:
@@ -23,15 +25,18 @@ def register_migration_request_logging(app: flask.Flask) -> None:
     def set_request_migration_context():
         flask.g.request_started_at = time.time()
         flask.g.request_id = (
-            flask.request.headers.get("X-Request-ID") or uuid.uuid4().hex
+            flask.request.headers.get(REQUEST_ID_HEADER) or generate_request_id()
         )
 
     @app.after_request
     def log_request_migration_context(response):
         response.headers[BACKEND_RESPONSE_HEADER] = get_api_host_backend()
+        request_id = getattr(flask.g, "request_id", None)
+        if request_id is not None:
+            response.headers[REQUEST_ID_HEADER] = request_id
         try:
             log_migration_request(
-                request_id=getattr(flask.g, "request_id", None),
+                request_id=request_id,
                 method=flask.request.method,
                 path=flask.request.path,
                 status_code=response.status_code,

policyengine_api/request_context.py

@@ -0,0 +1,29 @@
+"""Request-scoped correlation identifiers shared by Flask and ASGI."""
+
+from __future__ import annotations
+
+import uuid
+from contextvars import ContextVar
+
+import flask
+
+REQUEST_ID_HEADER = "X-PolicyEngine-Request-Id"
+
+_asgi_request_id: ContextVar[str | None] = ContextVar(
+    "policyengine_api_request_id",
+    default=None,
+)
+
+
+def generate_request_id() -> str:
+    """Generate an API request correlation identifier."""
+
+    return uuid.uuid4().hex
+
+
+def current_request_id() -> str | None:
+    """Return the correlation identifier for the current request, if any."""
+
+    if flask.has_request_context():
+        return getattr(flask.g, "request_id", None)
+    return _asgi_request_id.get()

tests/conftest.py

@@ -44,5 +44,4 @@ def client():
     with running(["redis-server"], 3):
         redis_client = redis.Redis()
         redis_client.ping()
-        with app.test_client() as test_client:
-            yield test_client
+        yield app.test_client()

tests/to_refactor/api/test_rest_client_context.py

@@ -0,0 +1,10 @@
+import flask
+
+
+def test_rest_client_does_not_preserve_request_context(rest_client):
+    assert not flask.has_request_context()
+
+    response = rest_client.get("/readiness-check")
+
+    assert response.status_code == 200
+    assert not flask.has_request_context()