From 4489fd3040d14ba93a2cdac821699a9b210a79ed Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 30 Jul 2026 21:10:08 +0300 Subject: [PATCH 1/2] Test simulation request correlation metadata --- .../tests/test_observability.py | 102 +++++++++++++++++- .../tests/test_app.py | 72 +++++++++++-- .../tests/test_backend.py | 7 +- .../test_standalone_simulation_contract.py | 7 ++ .../tests/test_observability.py | 41 +++++++ 5 files changed, 212 insertions(+), 17 deletions(-) create mode 100644 projects/policyengine-simulation-gateway/tests/test_observability.py diff --git a/libs/policyengine-simulation-observability/tests/test_observability.py b/libs/policyengine-simulation-observability/tests/test_observability.py index e0347dcb1..60057dbb8 100644 --- a/libs/policyengine-simulation-observability/tests/test_observability.py +++ b/libs/policyengine-simulation-observability/tests/test_observability.py @@ -1,4 +1,5 @@ import json +import inspect import os import pytest @@ -14,7 +15,7 @@ from policyengine_simulation_observability.observability import ( LOG_DESTINATIONS, - SERVICE_NAME, + _environment, configure_process_observability, init_process_observability, init_simulation_observability, @@ -23,6 +24,7 @@ OBSERVABILITY_ENV_KEYS = ( + "OBSERVABILITY_ENVIRONMENT", "OBSERVABILITY_PLATFORM", "OBSERVABILITY_SERVICE_ROLE", "OBSERVABILITY_RUNTIME_ROLE", @@ -32,6 +34,10 @@ "OTEL_ENABLED", "GOOGLE_CLOUD_PROJECT", "MODAL_ENVIRONMENT", + "DEPLOYMENT_ENVIRONMENT", + "APP_ENVIRONMENT", + "APP_ENV", + "ENVIRONMENT", ) @@ -143,10 +149,14 @@ def test_init_simulation_observability_forces_stdout_and_disables_exports(): os.environ["MODAL_ENVIRONMENT"] = "staging" app = FastAPI() - runtime = init_simulation_observability(app, service_role="modal_gateway") + runtime = init_simulation_observability( + app, + service_name="policyengine-simulation-gateway", + service_role="modal_gateway", + ) assert app.state.policyengine_observability is runtime - assert runtime.config.service_name == SERVICE_NAME + assert runtime.config.service_name == "policyengine-simulation-gateway" assert runtime.config.service_role == "modal_gateway" assert runtime.config.environment == "staging" assert runtime.config.log_destinations == LOG_DESTINATIONS @@ -161,7 +171,11 @@ def test_fastapi_observability_emits_structured_request_log(monkeypatch): def health(): return {"status": "healthy"} - init_simulation_observability(app, service_role="api") + init_simulation_observability( + app, + service_name="policyengine-simulation-entry", + service_role="api", + ) records = [] monkeypatch.setattr( @@ -180,7 +194,7 @@ def health(): record = records[0] assert record["schema_version"] == "policyengine.observability.request.v1" assert record["event"] == "http_request_completed" - assert record["service_name"] == SERVICE_NAME + assert record["service_name"] == "policyengine-simulation-entry" assert record["service_role"] == "api" assert record["request_id"] == "request-123" assert record["method"] == "GET" @@ -188,3 +202,81 @@ def health(): assert record["status_code"] == 200 assert record["logfire_status"] == "legacy_candidate_for_replacement" assert record["logfire_replacement_candidate"] == "policyengine-observability" + + +def test_fastapi_service_name_is_a_required_keyword_only_argument(): + parameter = inspect.signature(init_simulation_observability).parameters[ + "service_name" + ] + + assert parameter.kind is inspect.Parameter.KEYWORD_ONLY + assert parameter.default is inspect.Parameter.empty + + +@pytest.mark.parametrize("app_environment", ["beta", "prod"]) +def test_app_environment_configures_cloud_run_runtime(app_environment): + os.environ["APP_ENVIRONMENT"] = app_environment + + app = FastAPI() + runtime = init_simulation_observability( + app, + service_name="policyengine-simulation-entry", + service_role="simulation_entry", + ) + + assert runtime.config.environment == app_environment + + +@pytest.mark.parametrize( + ("environment", "expected"), + [ + ( + { + "OBSERVABILITY_ENVIRONMENT": "observability", + "MODAL_ENVIRONMENT": "modal", + "DEPLOYMENT_ENVIRONMENT": "deployment", + "APP_ENVIRONMENT": "app-environment", + "APP_ENV": "app-env", + "ENVIRONMENT": "generic", + }, + "observability", + ), + ( + { + "MODAL_ENVIRONMENT": "modal", + "DEPLOYMENT_ENVIRONMENT": "deployment", + "APP_ENVIRONMENT": "app-environment", + "APP_ENV": "app-env", + "ENVIRONMENT": "generic", + }, + "modal", + ), + ( + { + "DEPLOYMENT_ENVIRONMENT": "deployment", + "APP_ENVIRONMENT": "app-environment", + "APP_ENV": "app-env", + "ENVIRONMENT": "generic", + }, + "deployment", + ), + ( + { + "APP_ENVIRONMENT": "app-environment", + "APP_ENV": "app-env", + "ENVIRONMENT": "generic", + }, + "app-environment", + ), + ( + {"APP_ENV": "app-env", "ENVIRONMENT": "generic"}, + "app-env", + ), + ({"ENVIRONMENT": "generic"}, "generic"), + ({}, "local"), + ], +) +def test_environment_precedence(environment, expected): + os.environ.update(environment) + + assert _environment() == expected diff --git a/projects/policyengine-simulation-entry/tests/test_app.py b/projects/policyengine-simulation-entry/tests/test_app.py index 7c6190902..bbfa7bdea 100644 --- a/projects/policyengine-simulation-entry/tests/test_app.py +++ b/projects/policyengine-simulation-entry/tests/test_app.py @@ -184,11 +184,26 @@ def test_versions_and_ping_are_public_proxy_routes(client, backend): assert client.post("/ping", json={"value": 1}).json() == {"incremented": 2} -def test_request_id_is_propagated_and_returned(client, backend): - result = client.get("/versions", headers={"X-Request-ID": "request-123"}) +def test_request_id_is_propagated_logged_and_returned( + client, + backend, + caplog, +): + with caplog.at_level(logging.INFO): + result = client.get( + "/versions", + headers={REQUEST_ID_HEADER: "request-123"}, + ) - assert result.headers["x-request-id"] == "request-123" + assert result.headers[REQUEST_ID_HEADER] == "request-123" + assert "x-request-id" not in result.headers assert backend.requests[-1].request_id == "request-123" + request_record = next( + record + for record in caplog.records + if record.getMessage() == "simulation_entry_request" + ) + assert request_record.request_id == "request-123" def test_generated_request_id_is_shared_with_observability(): @@ -214,13 +229,39 @@ async def request(self, *args, **kwargs): with TestClient(app) as client: result = client.get("/versions") - request_id = result.headers["x-request-id"] + request_id = result.headers[REQUEST_ID_HEADER] assert request_id - assert result.headers[REQUEST_ID_HEADER] == request_id + assert "x-request-id" not in result.headers assert backend.requests[-1].request_id == request_id assert backend.observability_request_id == request_id +def test_x_request_id_is_not_an_alias(client, backend): + result = client.get( + "/versions", + headers={"X-Request-ID": "must-not-be-used"}, + ) + + request_id = result.headers[REQUEST_ID_HEADER] + assert request_id + assert request_id != "must-not-be-used" + assert backend.requests[-1].request_id == request_id + assert "x-request-id" not in result.headers + + +def test_entrypoint_uses_its_own_service_name(backend): + app = create_app( + settings=make_settings(), + backend=backend, + auth_dependency=lambda: None, + ) + + assert ( + app.state.policyengine_observability.config.service_name + == "policyengine-simulation-entry" + ) + + def test_request_validation_matches_shared_contract(client): result = client.post( "/simulate/economy/comparison", @@ -251,12 +292,17 @@ async def request(self, *args, **kwargs): from fastapi.testclient import TestClient with TestClient(app) as client: - result = client.get("/jobs/job-1") + result = client.get( + "/jobs/job-1", + headers={REQUEST_ID_HEADER: "request-backend-failure"}, + ) assert result.status_code == expected_status assert result.json() == {"detail": expected_detail} assert result.headers["retry-after"] == "10" assert result.headers["x-policyengine-simulation-backend"] == "old_gateway" + assert result.headers[REQUEST_ID_HEADER] == "request-backend-failure" + assert "x-request-id" not in result.headers def test_unexpected_failure_preserves_correlation_headers_and_request_log(caplog): @@ -275,12 +321,13 @@ async def request(self, *args, **kwargs): with caplog.at_level(logging.INFO), TestClient(app) as client: result = client.get( "/jobs/job-1", - headers={"X-Request-ID": "request-500"}, + headers={REQUEST_ID_HEADER: "request-500"}, ) assert result.status_code == 500 assert result.text == "Internal Server Error" - assert result.headers["x-request-id"] == "request-500" + assert result.headers[REQUEST_ID_HEADER] == "request-500" + assert "x-request-id" not in result.headers assert ( result.headers["x-policyengine-simulation-revision"] == "simulation-entry-test-revision" @@ -295,14 +342,19 @@ async def request(self, *args, **kwargs): assert request_record.job_id == "job-1" -@pytest.mark.parametrize("status_code", [400, 404, 409, 500]) +@pytest.mark.parametrize("status_code", [400, 404, 409, 500, 502]) def test_upstream_error_status_and_body_are_preserved(client, backend, status_code): backend.responses[("GET", "/jobs/job-1")] = response( status_code, {"detail": "upstream response"}, ) - result = client.get("/jobs/job-1") + result = client.get( + "/jobs/job-1", + headers={REQUEST_ID_HEADER: "request-upstream-error"}, + ) assert result.status_code == status_code assert result.json() == {"detail": "upstream response"} + assert result.headers[REQUEST_ID_HEADER] == "request-upstream-error" + assert "x-request-id" not in result.headers diff --git a/projects/policyengine-simulation-entry/tests/test_backend.py b/projects/policyengine-simulation-entry/tests/test_backend.py index 3b66da246..6c0ceee99 100644 --- a/projects/policyengine-simulation-entry/tests/test_backend.py +++ b/projects/policyengine-simulation-entry/tests/test_backend.py @@ -4,6 +4,7 @@ import httpx import pytest +from policyengine_observability import REQUEST_ID_HEADER from policyengine_simulation_entry.backend import ( BackendAuthenticationError, @@ -58,7 +59,7 @@ def handler(request: httpx.Request) -> httpx.Response: headers={ "Retry-After": "2", "Set-Cookie": "must-not-pass", - "X-Request-ID": "upstream-request", + REQUEST_ID_HEADER: "upstream-request", }, ) @@ -79,12 +80,14 @@ def handler(request: httpx.Request) -> httpx.Response: upstream = requests[-1] assert upstream.headers["authorization"] == "Bearer service-token" - assert upstream.headers["x-request-id"] == "caller-request" + assert upstream.headers[REQUEST_ID_HEADER] == "caller-request" + assert "x-request-id" not in upstream.headers assert json.loads(result.content) == { "job_id": "fc-123", "status": "submitted", } assert result.headers["retry-after"] == "2" + assert result.headers[REQUEST_ID_HEADER] == "upstream-request" assert "set-cookie" not in result.headers diff --git a/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py b/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py index 06f2fdfc2..58b3545ad 100644 --- a/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py +++ b/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py @@ -29,6 +29,13 @@ ) +def test_standalone_executor_uses_its_own_service_name(): + assert ( + app.state.policyengine_observability.config.service_name + == "policyengine-simulation-executor" + ) + + def test_standalone_package_runtime_does_not_import_unpackaged_modal_source(): for module_name in PACKAGED_RUNTIME_MODULES: module = import_module(module_name) diff --git a/projects/policyengine-simulation-gateway/tests/test_observability.py b/projects/policyengine-simulation-gateway/tests/test_observability.py new file mode 100644 index 000000000..4f6cd0efb --- /dev/null +++ b/projects/policyengine-simulation-gateway/tests/test_observability.py @@ -0,0 +1,41 @@ +"""Gateway observability identity and request-correlation tests.""" + +import json + +from policyengine_observability import REQUEST_ID_HEADER +from policyengine_observability.runtime import REQUEST_LOGGER +from policyengine_simulation_gateway.testing import create_gateway_app + + +def test_gateway_reads_policyengine_request_id_into_observability_context( + monkeypatch, +): + records = [] + monkeypatch.setattr( + REQUEST_LOGGER, + "info", + lambda message: records.append(json.loads(message)), + ) + app = create_gateway_app() + + from fastapi.testclient import TestClient + + response = TestClient(app).get( + "/health", + headers={REQUEST_ID_HEADER: "request-through-entrypoint"}, + ) + + assert response.status_code == 200 + assert len(records) == 1 + assert records[0]["request_id"] == "request-through-entrypoint" + assert records[0]["service_name"] == "policyengine-simulation-gateway" + assert records[0]["service_role"] == "modal_gateway" + + +def test_gateway_uses_its_own_service_name(): + app = create_gateway_app() + + assert ( + app.state.policyengine_observability.config.service_name + == "policyengine-simulation-gateway" + ) From c03d5a9b9ea38db6d8e2911795d44c12da4dcddc Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 30 Jul 2026 21:29:32 +0300 Subject: [PATCH 2/2] Fix simulation request correlation metadata --- .../README.md | 9 ++++++--- .../__init__.py | 3 ++- .../observability.py | 19 +++++++++++++++---- .../tests/test_observability.py | 2 +- .../src/policyengine_simulation_entry/app.py | 16 +++++++--------- .../policyengine_simulation_entry/backend.py | 11 ++++++++--- .../policyengine_simulation_executor/main.py | 6 +++++- .../policyengine_simulation_gateway/app.py | 6 +++++- .../testing.py | 6 +++++- 9 files changed, 54 insertions(+), 24 deletions(-) diff --git a/libs/policyengine-simulation-observability/README.md b/libs/policyengine-simulation-observability/README.md index f3be9bf47..7635e90f5 100644 --- a/libs/policyengine-simulation-observability/README.md +++ b/libs/policyengine-simulation-observability/README.md @@ -1,8 +1,11 @@ # policyengine-simulation-observability -Shared observability plumbing for the simulation service (gateway and -executor): the `policyengine-observability` configuration wrapper, legacy -Logfire helpers, error redaction/reporting, and the telemetry envelope. +Shared observability plumbing for the simulation entrypoint, gateway, and +executor: the `policyengine-observability` configuration wrapper, legacy +Logfire helpers, error redaction/reporting, and the telemetry envelope. Each +FastAPI service supplies its own service name, while +`X-PolicyEngine-Request-Id` is the single request-correlation header across +service boundaries. This lives in its own lib because it has a different reason to change than the gateway↔executor contract: Logfire is retained only while a replacement diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/__init__.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/__init__.py index a4f3cd139..f65a80295 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/__init__.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/__init__.py @@ -1 +1,2 @@ -"""Shared observability plumbing for the simulation gateway and executor.""" +"""Shared observability plumbing for the simulation entrypoint, gateway, +and executor.""" diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py index fc31be389..aac5399f4 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py @@ -132,16 +132,21 @@ def configure_process_observability( def init_simulation_observability( app: FastAPI, *, + service_name: str, service_role: str = "api", ) -> ObservabilityRuntime: service_role = _service_role(service_role) platform = _platform() - config = _config(service_role=service_role, platform=platform) + config = _config( + service_name=service_name, + service_role=service_role, + platform=platform, + ) return init_fastapi_observability( app, config=config, runtime=ObservabilityRuntime(config, segment_registry=SegmentName), - service_name=SERVICE_NAME, + service_name=service_name, service_role=service_role, span_prefix=SPAN_PREFIX, segment_registry=SegmentName, @@ -155,7 +160,11 @@ def init_process_observability( ) -> ObservabilityRuntime: service_role = _service_role(service_role) platform = _platform() - config = _config(service_role=service_role, platform=platform) + config = _config( + service_name=SERVICE_NAME, + service_role=service_role, + platform=platform, + ) runtime = ObservabilityRuntime(config, segment_registry=SegmentName) runtime.configure() set_observability_runtime(runtime) @@ -182,11 +191,12 @@ def logfire_replacement_attributes() -> dict[str, str]: def _config( *, + service_name: str, service_role: str, platform: str, ) -> ObservabilityConfig: config = ObservabilityConfig.from_env( - service_name=SERVICE_NAME, + service_name=service_name, service_role=service_role, span_prefix=SPAN_PREFIX, extra_metric_attribute_keys=SIMULATION_METRIC_ATTRIBUTE_KEYS, @@ -206,6 +216,7 @@ def _environment() -> str: os.getenv("OBSERVABILITY_ENVIRONMENT") or os.getenv("MODAL_ENVIRONMENT") or os.getenv("DEPLOYMENT_ENVIRONMENT") + or os.getenv("APP_ENVIRONMENT") or os.getenv("APP_ENV") or os.getenv("ENVIRONMENT") or "local" diff --git a/libs/policyengine-simulation-observability/tests/test_observability.py b/libs/policyengine-simulation-observability/tests/test_observability.py index 60057dbb8..bc3c9a70d 100644 --- a/libs/policyengine-simulation-observability/tests/test_observability.py +++ b/libs/policyengine-simulation-observability/tests/test_observability.py @@ -1,5 +1,5 @@ -import json import inspect +import json import os import pytest diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index b4c1b4d92..f3bbe511f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -125,18 +125,16 @@ async def lifespan(_: FastAPI): platform="cloud_run", service_role="simulation_entry", ) - init_simulation_observability(app, service_role="simulation_entry") + init_simulation_observability( + app, + service_name="policyengine-simulation-entry", + service_role="simulation_entry", + ) @app.middleware("http") async def request_context(request: Request, call_next): headers = MutableHeaders(scope=request.scope) - request_id = ( - headers.get("x-request-id") - or headers.get(REQUEST_ID_HEADER) - or str(uuid.uuid4()) - ) - headers["x-request-id"] = request_id - # policyengine-observability currently reads its legacy header name. + request_id = headers.get(REQUEST_ID_HEADER) or str(uuid.uuid4()) headers[REQUEST_ID_HEADER] = request_id request.state.request_id = request_id started = time.monotonic() @@ -156,9 +154,9 @@ async def request_context(request: Request, call_next): content="Internal Server Error", status_code=500, media_type="text/plain", + headers={REQUEST_ID_HEADER: request_id}, ) elapsed_ms = round((time.monotonic() - started) * 1000, 2) - response.headers["X-Request-ID"] = request_id if runtime_settings.revision: response.headers["X-PolicyEngine-Simulation-Revision"] = ( runtime_settings.revision diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py index 67343e470..065547535 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py @@ -9,6 +9,7 @@ import httpx from pydantic import ValidationError +from policyengine_observability import REQUEST_ID_HEADER from policyengine_simulation_contract.json_types import JsonObject from policyengine_simulation_entry.config import Settings @@ -21,7 +22,7 @@ "content-type", "etag", "retry-after", - "x-request-id", + REQUEST_ID_HEADER.lower(), } ) @@ -181,7 +182,7 @@ async def request( token = await token_provider.get_token() headers = {"Authorization": f"Bearer {token}"} if request_id: - headers["X-Request-ID"] = request_id + headers[REQUEST_ID_HEADER] = request_id try: response = await client.request( @@ -216,7 +217,11 @@ async def request( status_code=response.status_code, content=response.content, headers={ - key: value + ( + REQUEST_ID_HEADER + if key.lower() == REQUEST_ID_HEADER.lower() + else key + ): value for key, value in response.headers.items() if key.lower() in SAFE_RESPONSE_HEADERS }, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py index de2fd8733..4c915b7ab 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py @@ -33,7 +33,11 @@ async def lifespan(app: FastAPI): title="policyengine-simulation-executor", summary="Policyengine simulation api", ) -init_simulation_observability(app, service_role="api") +init_simulation_observability( + app, + service_name="policyengine-simulation-executor", + service_role="api", +) # attach the api defined in the app package initialize(app=app) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py index ade6c8558..8546ac0db 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py @@ -93,7 +93,11 @@ def web_app(): modal_app_name="policyengine-simulation-gateway", modal_function_name="web_app", ) - init_simulation_observability(api, service_role="modal_gateway") + init_simulation_observability( + api, + service_name="policyengine-simulation-gateway", + service_role="modal_gateway", + ) configure_logfire("policyengine-simulation-gateway") # Startup guard: crash the container if GATEWAY_AUTH_DISABLED is set in diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/testing.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/testing.py index 703723ce3..71b520028 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/testing.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/testing.py @@ -27,7 +27,11 @@ def create_gateway_app(*, authenticate: bool = True) -> FastAPI: description="Test instance for unit tests", version="0.0.1", ) - init_simulation_observability(app, service_role="modal_gateway") + init_simulation_observability( + app, + service_name="policyengine-simulation-gateway", + service_role="modal_gateway", + ) app.include_router(router) if authenticate: app.dependency_overrides[require_auth] = lambda: None