Fix simulation request correlation and service metadata

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

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."""

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"

libs/policyengine-simulation-observability/tests/test_observability.py

@@ -1,3 +1,4 @@
+import inspect
 import json
 import os
 
@@ -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,11 +194,89 @@ 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"
     assert record["route"] == "/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

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

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
                 },