diff --git a/changelog.d/request-id-propagation.fixed.md b/changelog.d/request-id-propagation.fixed.md new file mode 100644 index 000000000..380ffbd45 --- /dev/null +++ b/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. diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py index f062363f7..8d3a2b0c3 100644 --- a/policyengine_api/asgi_factory.py +++ b/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: diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 42ffdb74b..cea087cb9 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/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 = { diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py index aa3c575aa..93c5bbb5a 100644 --- a/policyengine_api/migration_logging.py +++ b/policyengine_api/migration_logging.py @@ -3,10 +3,8 @@ 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, @@ -14,6 +12,10 @@ 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, diff --git a/policyengine_api/request_context.py b/policyengine_api/request_context.py new file mode 100644 index 000000000..feff6194f --- /dev/null +++ b/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() diff --git a/tests/conftest.py b/tests/conftest.py index cf5a9aa70..06d171e37 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/to_refactor/api/test_rest_client_context.py b/tests/to_refactor/api/test_rest_client_context.py new file mode 100644 index 000000000..8b0096bdd --- /dev/null +++ b/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() diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 681c996d1..b1deb55ad 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -12,6 +12,7 @@ import httpx import pytest +from flask import Flask, g sys.modules.setdefault( "policyengine_api.gcp_logging", @@ -28,10 +29,14 @@ from policyengine_api.libs.simulation_entrypoint import ( # noqa: E402 ModalBudgetWindowBatchExecution, ModalSimulationExecution, - SimulationEntrypointClient, SimulationAPIModal, + SimulationEntrypointClient, resolve_simulation_entrypoint_url, ) +from policyengine_api.request_context import ( # noqa: E402 + REQUEST_ID_HEADER, + _asgi_request_id, +) from tests.fixtures.libs.simulation_entrypoint import ( # noqa: E402 MOCK_BATCH_JOB_ID, @@ -73,6 +78,54 @@ ) +class RequestRecordingHTTPXClient: + """Minimal HTTPX-compatible client that executes request event hooks.""" + + instances = [] + + def __init__(self, *, event_hooks=None, **kwargs): + self.event_hooks = event_hooks or {} + self.requests = [] + type(self).instances.append(self) + + def _response(self, method, url, json=None): + request = httpx.Request(method, url, json=json) + for hook in self.event_hooks.get("request", []): + hook(request) + self.requests.append(request) + + path = request.url.path + if method == "POST" and path.endswith("/comparison"): + payload = MOCK_SUBMIT_RESPONSE_SUCCESS + status_code = 202 + elif method == "POST" and path.endswith("/budget-window"): + payload = MOCK_BATCH_SUBMIT_RESPONSE_SUCCESS + status_code = 202 + elif "/budget-window-jobs/" in path: + payload = MOCK_BATCH_POLL_RESPONSE_RUNNING + status_code = 202 + elif "/jobs/" in path: + payload = MOCK_POLL_RESPONSE_RUNNING + status_code = 202 + elif "/versions/" in path: + payload = { + "latest": "1.459.0", + "1.459.0": MOCK_RESOLVED_APP_NAME, + } + status_code = 200 + else: + payload = MOCK_HEALTH_RESPONSE + status_code = 200 + + return httpx.Response(status_code, request=request, json=payload) + + def post(self, url, json=None): + return self._response("POST", url, json=json) + + def get(self, url): + return self._response("GET", url) + + @pytest.fixture(autouse=True) def clear_gateway_auth_env(monkeypatch): """Isolate unit tests from gateway-auth env injected during Docker builds.""" @@ -88,8 +141,11 @@ def test_generic_client_name_retains_old_class_alias(): def test_legacy_simulation_api_modules_export_entrypoint_aliases(): - from policyengine_api.libs import simulation_api, simulation_api_modal - from policyengine_api.libs import simulation_entrypoint + from policyengine_api.libs import ( + simulation_api, + simulation_api_modal, + simulation_entrypoint, + ) assert simulation_api.SimulationAPIClient is SimulationEntrypointClient assert simulation_api_modal.SimulationAPIClient is SimulationEntrypointClient @@ -310,6 +366,123 @@ def test__given_partial_gateway_auth_env_vars__then_raises( with pytest.raises(GatewayAuthError): SimulationAPIModal() + def test__given_client_initialized__then_installs_one_request_id_hook( + self, mock_httpx_client + ): + from policyengine_api.libs.simulation_entrypoint import httpx as modal_httpx + + SimulationAPIModal() + + _, kwargs = modal_httpx.Client.call_args + assert list(kwargs["event_hooks"]) == ["request"] + assert len(kwargs["event_hooks"]["request"]) == 1 + + def test__given_flask_request__then_hook_uses_current_request_id( + self, mock_httpx_client + ): + from policyengine_api.libs.simulation_entrypoint import httpx as modal_httpx + + SimulationAPIModal() + hook = modal_httpx.Client.call_args.kwargs["event_hooks"]["request"][0] + request = httpx.Request("GET", MOCK_MODAL_BASE_URL) + app = Flask("request-id-test") + + token = _asgi_request_id.set("asgi-request-id") + try: + with app.test_request_context(): + g.request_id = "flask-request-id" + hook(request) + finally: + _asgi_request_id.reset(token) + + assert request.headers[REQUEST_ID_HEADER] == "flask-request-id" + + def test__given_asgi_request__then_hook_uses_current_request_id( + self, mock_httpx_client + ): + from policyengine_api.libs.simulation_entrypoint import httpx as modal_httpx + + SimulationAPIModal() + hook = modal_httpx.Client.call_args.kwargs["event_hooks"]["request"][0] + request = httpx.Request("GET", MOCK_MODAL_BASE_URL) + token = _asgi_request_id.set("asgi-request-id") + try: + hook(request) + finally: + _asgi_request_id.reset(token) + + assert request.headers[REQUEST_ID_HEADER] == "asgi-request-id" + + def test__given_no_request_context__then_hook_omits_request_id( + self, monkeypatch, mock_modal_logger + ): + from policyengine_api.libs import simulation_entrypoint as module + + RequestRecordingHTTPXClient.instances.clear() + monkeypatch.setattr( + module.httpx, + "Client", + RequestRecordingHTTPXClient, + ) + api = SimulationAPIModal() + + assert api.health_check() is True + request = RequestRecordingHTTPXClient.instances[-1].requests[-1] + assert REQUEST_ID_HEADER not in request.headers + + @pytest.mark.parametrize( + ("entrypoint", "url_env_name"), + [ + ("old_gateway_direct", "OLD_SIMULATION_GATEWAY_URL"), + ( + "cloud_run_simulation_entrypoint", + "SIMULATION_ENTRYPOINT_URL", + ), + ], + ) + def test__given_request_context__then_all_calls_forward_request_id( + self, + monkeypatch, + mock_modal_logger, + entrypoint, + url_env_name, + ): + from policyengine_api.libs import simulation_entrypoint as module + + RequestRecordingHTTPXClient.instances.clear() + monkeypatch.setenv("SIM_ENTRYPOINT", entrypoint) + monkeypatch.setenv(url_env_name, MOCK_MODAL_BASE_URL) + monkeypatch.setattr( + module.httpx, + "Client", + RequestRecordingHTTPXClient, + ) + api = SimulationAPIModal() + app = Flask("request-id-all-calls") + + with app.test_request_context(): + g.request_id = "flask-request-id" + api.run(MOCK_SIMULATION_PAYLOAD) + api.run_budget_window_batch(MOCK_SIMULATION_PAYLOAD) + api.get_execution_by_id(MOCK_MODAL_JOB_ID) + api.get_budget_window_batch_by_id(MOCK_BATCH_JOB_ID) + api.resolve_app_name("us", "1.459.0") + assert api.health_check() is True + + requests = RequestRecordingHTTPXClient.instances[-1].requests + assert {request.url.path for request in requests} == { + "/simulate/economy/comparison", + "/simulate/economy/budget-window", + f"/jobs/{MOCK_MODAL_JOB_ID}", + f"/budget-window-jobs/{MOCK_BATCH_JOB_ID}", + "/versions/us", + "/health", + } + assert all( + request.headers[REQUEST_ID_HEADER] == "flask-request-id" + for request in requests + ) + class TestRun: def test__given_valid_payload__then_returns_execution_with_job_id( self, diff --git a/tests/unit/routes/test_migration_context_logging.py b/tests/unit/routes/test_migration_context_logging.py index 9e3b9939b..34c2597d2 100644 --- a/tests/unit/routes/test_migration_context_logging.py +++ b/tests/unit/routes/test_migration_context_logging.py @@ -1,11 +1,15 @@ +from types import SimpleNamespace from unittest.mock import patch from fastapi.testclient import TestClient from flask import Flask, Response - from policyengine_api.asgi_factory import create_asgi_app from policyengine_api.migration_flags import BACKEND_RESPONSE_HEADER from policyengine_api.migration_logging import register_migration_request_logging +from policyengine_api.request_context import ( + REQUEST_ID_HEADER, + current_request_id, +) def _app(): @@ -16,6 +20,10 @@ def _app(): def readiness_check(): return Response("OK", status=200, mimetype="text/plain") + @app.route("/request-id") + def request_id(): + return Response(current_request_id(), status=200, mimetype="text/plain") + register_migration_request_logging(app) return app @@ -45,6 +53,84 @@ def test_request_logging_includes_migration_context(): assert log_payload["migration"]["route_impl"] == "flask_fallback" +def test_flask_preserves_policyengine_request_id_in_context_log_and_response(): + with patch("policyengine_api.migration_logging.logger") as mock_logger: + response = ( + _app() + .test_client() + .get( + "/request-id", + headers={REQUEST_ID_HEADER: "request-123"}, + ) + ) + + assert response.status_code == 200 + assert response.text == "request-123" + assert response.headers[REQUEST_ID_HEADER] == "request-123" + assert mock_logger.log_struct.call_args.args[0]["request_id"] == "request-123" + + +def test_flask_generates_one_request_id_for_context_log_and_response(): + with ( + patch( + "policyengine_api.request_context.uuid.uuid4", + return_value=SimpleNamespace(hex="generated-request-id"), + ) as generate_request_id, + patch("policyengine_api.migration_logging.logger") as mock_logger, + ): + response = _app().test_client().get("/request-id") + + assert response.status_code == 200 + assert response.text == "generated-request-id" + assert response.headers[REQUEST_ID_HEADER] == "generated-request-id" + assert ( + mock_logger.log_struct.call_args.args[0]["request_id"] == "generated-request-id" + ) + generate_request_id.assert_called_once_with() + + +def test_flask_does_not_accept_x_request_id_as_an_alias(): + with ( + patch( + "policyengine_api.request_context.uuid.uuid4", + return_value=SimpleNamespace(hex="generated-request-id"), + ), + patch("policyengine_api.migration_logging.logger") as mock_logger, + ): + response = ( + _app() + .test_client() + .get( + "/request-id", + headers={"X-Request-ID": "legacy-request-id"}, + ) + ) + + assert response.text == "generated-request-id" + assert response.headers[REQUEST_ID_HEADER] == "generated-request-id" + assert ( + mock_logger.log_struct.call_args.args[0]["request_id"] == "generated-request-id" + ) + + +def test_asgi_generated_request_id_reaches_mounted_flask_unchanged(): + with ( + patch( + "policyengine_api.request_context.uuid.uuid4", + return_value=SimpleNamespace(hex="generated-request-id"), + ) as generate_request_id, + patch("policyengine_api.migration_logging.logger") as mock_logger, + ): + response = TestClient(create_asgi_app(_app())).get("/request-id") + + assert response.text == "generated-request-id" + assert response.headers[REQUEST_ID_HEADER] == "generated-request-id" + assert ( + mock_logger.log_struct.call_args.args[0]["request_id"] == "generated-request-id" + ) + generate_request_id.assert_called_once_with() + + def test_request_logging_failure_does_not_change_response(): with patch( "policyengine_api.migration_logging.logger.log_struct", @@ -73,7 +159,7 @@ def test_request_logging_runs_for_fastapi_native_health_routes(monkeypatch): with patch("policyengine_api.migration_logging.logger") as mock_logger: response = TestClient(create_asgi_app(_app())).get( "/health", - headers={"X-Request-ID": "request-123"}, + headers={REQUEST_ID_HEADER: "request-123"}, ) assert response.status_code == 200 @@ -89,6 +175,16 @@ def test_request_logging_runs_for_fastapi_native_health_routes(monkeypatch): assert log_payload["migration"]["route_impl"] == "flask_fallback" +def test_fastapi_native_response_includes_policyengine_request_id(): + with patch("policyengine_api.migration_logging.logger"): + response = TestClient(create_asgi_app(_app())).get( + "/health", + headers={REQUEST_ID_HEADER: "request-123"}, + ) + + assert response.headers[REQUEST_ID_HEADER] == "request-123" + + def test_fastapi_native_logging_failure_does_not_change_response(): with patch( "policyengine_api.migration_logging.logger.log_struct", diff --git a/tests/unit/test_asgi_factory.py b/tests/unit/test_asgi_factory.py index ebb9bf866..ab388e5b4 100644 --- a/tests/unit/test_asgi_factory.py +++ b/tests/unit/test_asgi_factory.py @@ -1,6 +1,9 @@ import importlib import json import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -8,6 +11,10 @@ from flask import Flask, Response, jsonify, make_response, request from flask_cors import CORS from policyengine_api.asgi_factory import _add_vary_origin, create_asgi_app +from policyengine_api.request_context import ( + REQUEST_ID_HEADER, + current_request_id, +) from starlette.responses import Response as ASGIResponse @@ -62,6 +69,114 @@ def test_native_health_route_is_fastapi_json(): assert response.headers["content-type"].startswith("application/json") +def test_native_route_preserves_request_id_in_context_log_and_response(): + observed_context_ids = [] + + def capture_request(**kwargs): + observed_context_ids.append(current_request_id()) + + with patch( + "policyengine_api.asgi_factory.log_migration_request", + side_effect=capture_request, + ) as log_request: + response = TestClient(create_asgi_app(create_test_wsgi_app())).get( + "/health", + headers={REQUEST_ID_HEADER: "request-123"}, + ) + + assert response.headers[REQUEST_ID_HEADER] == "request-123" + assert log_request.call_args.kwargs["request_id"] == "request-123" + assert observed_context_ids == ["request-123"] + + +def test_native_route_generates_one_request_id_for_context_log_and_response(): + observed_context_ids = [] + + def capture_request(**kwargs): + observed_context_ids.append(current_request_id()) + + with ( + patch( + "policyengine_api.request_context.uuid.uuid4", + return_value=SimpleNamespace(hex="generated-request-id"), + ) as generate_request_id, + patch( + "policyengine_api.asgi_factory.log_migration_request", + side_effect=capture_request, + ) as log_request, + ): + response = TestClient(create_asgi_app(create_test_wsgi_app())).get("/health") + + assert response.headers[REQUEST_ID_HEADER] == "generated-request-id" + assert log_request.call_args.kwargs["request_id"] == "generated-request-id" + assert observed_context_ids == ["generated-request-id"] + generate_request_id.assert_called_once_with() + + +def test_native_route_does_not_accept_x_request_id_as_an_alias(): + with ( + patch( + "policyengine_api.request_context.uuid.uuid4", + return_value=SimpleNamespace(hex="generated-request-id"), + ), + patch("policyengine_api.asgi_factory.log_migration_request") as log_request, + ): + response = TestClient(create_asgi_app(create_test_wsgi_app())).get( + "/health", + headers={"X-Request-ID": "legacy-request-id"}, + ) + + assert response.headers[REQUEST_ID_HEADER] == "generated-request-id" + assert log_request.call_args.kwargs["request_id"] == "generated-request-id" + + +def test_asgi_request_context_is_reset_after_response(): + assert current_request_id() is None + + response = TestClient(create_asgi_app(create_test_wsgi_app())).get( + "/health", + headers={REQUEST_ID_HEADER: "request-123"}, + ) + + assert response.status_code == 200 + assert current_request_id() is None + + +def test_concurrent_asgi_requests_do_not_leak_request_ids(): + barrier = threading.Barrier(2) + observed_context_ids = {} + observed_lock = threading.Lock() + + def capture_request(**kwargs): + barrier.wait(timeout=5) + with observed_lock: + observed_context_ids[kwargs["request_id"]] = current_request_id() + + app = create_asgi_app(create_test_wsgi_app()) + + def make_request(request_id): + with TestClient(app) as client: + return client.get( + "/health", + headers={REQUEST_ID_HEADER: request_id}, + ) + + with ( + patch( + "policyengine_api.asgi_factory.log_migration_request", + side_effect=capture_request, + ), + ThreadPoolExecutor(max_workers=2) as executor, + ): + responses = list(executor.map(make_request, ("request-one", "request-two"))) + + assert [response.status_code for response in responses] == [200, 200] + assert observed_context_ids == { + "request-one": "request-one", + "request-two": "request-two", + } + + @pytest.mark.parametrize( ("existing_vary", "expected_vary"), [