Relay simulation run stage and progress on computing economy responses

changelog.d/relay-simulation-run-progress.added.md

@@ -0,0 +1 @@
+Annual economy responses relay the simulation entrypoint's optional stage fields (`stage`, `progress`, `progress_max`, timestamps and `segments`) while a run is computing.

policyengine_api/libs/simulation_entrypoint.py

@@ -7,6 +7,7 @@
 
 import httpx
 from policyengine_observability import instrument_httpx
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError
 
 from policyengine_api.gcp_logging import logger
 from policyengine_api.libs.gateway_auth import (
@@ -69,6 +70,45 @@ def _attach_current_request_id(request: httpx.Request) -> None:
         request.headers[OBSERVABILITY_ID_HEADER] = observability_id
 
 
+class SimulationRunSegments(BaseModel):
+    model_config = ConfigDict(extra="ignore", frozen=True)
+
+    completed: int = Field(ge=0)
+    total: int = Field(ge=1)
+
+
+class SimulationRunProgress(BaseModel):
+    """Stage fields of a running Simulation Entrypoint job.
+
+    The entrypoint adds these optional fields to ``GET /jobs/{job_id}`` for
+    jobs on workers that report stages (policyengine-sim-api#781). The API
+    relays them while an economy run is computing. ``stage`` stays a string:
+    later entrypoint versions may add stages, and clients must accept them.
+    """
+
+    model_config = ConfigDict(extra="ignore", frozen=True)
+
+    stage: str | None = None
+    progress: int | None = Field(default=None, ge=0, le=100)
+    progress_max: int | None = Field(default=None, ge=0, le=100)
+    submitted_at: AwareDatetime | None = None
+    started_at: AwareDatetime | None = None
+    stage_started_at: AwareDatetime | None = None
+    segments: SimulationRunSegments | None = None
+
+    @classmethod
+    def from_status(cls, data: dict) -> "SimulationRunProgress | None":
+        """Read the stage fields; never fail a poll because of them."""
+        try:
+            run_progress = cls.model_validate(data)
+        except ValidationError:
+            return None
+        return run_progress if run_progress.to_dict() else None
+
+    def to_dict(self) -> dict:
+        return self.model_dump(mode="json", exclude_none=True)
+
+
 @dataclass
 class ModalSimulationExecution:
     """
@@ -81,6 +121,7 @@ class ModalSimulationExecution:
     error: str | None = None
     policyengine_bundle: dict | None = None
     resolved_app_name: str | None = None
+    run_progress: SimulationRunProgress | None = None
 
     @property
     def name(self) -> str:
@@ -441,6 +482,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution:
                 error=data.get("error"),
                 policyengine_bundle=data.get("policyengine_bundle"),
                 resolved_app_name=data.get("resolved_app_name"),
+                run_progress=SimulationRunProgress.from_status(data),
             )
 
         except httpx.HTTPStatusError as e:

policyengine_api/openapi_spec.yaml

@@ -879,7 +879,7 @@ paths:
         # Query parameters are generated from policyengine_api.query_parameters.
       responses:
         200:
-          description: Annual economic impact, progress, or persisted error state. Canonical completed results include actual baseline/reform receipts containing the selected SPM configuration.
+          description: Annual economic impact, progress, or persisted error state. Canonical completed results include actual baseline/reform receipts containing the selected SPM configuration. While computing, the optional stage fields come from the simulation entrypoint and appear only when its worker reports stages.
           content:
             application/json:
               schema:
@@ -897,6 +897,41 @@ paths:
                     properties:
                       spm_provenance:
                         $ref: '#/components/schemas/SPMComparisonProvenance'
+                  stage:
+                    type: string
+                    description: Current stage of a computing run, such as queued, loading_data, baseline, reform, simulating or outputs. Later versions may add stages; clients must accept a value they do not know.
+                  progress:
+                    type: integer
+                    minimum: 0
+                    maximum: 100
+                    description: Coarse completion percent at the start of the current stage. It changes only at stage boundaries.
+                  progress_max:
+                    type: integer
+                    minimum: 0
+                    maximum: 100
+                    description: Highest percent the run reaches before its next stage starts. A client can move a time-based estimate between progress and progress_max.
+                  submitted_at:
+                    type: string
+                    format: date-time
+                    description: When the simulation entrypoint accepted the run.
+                  started_at:
+                    type: string
+                    format: date-time
+                    description: When a worker started the run.
+                  stage_started_at:
+                    type: string
+                    format: date-time
+                    description: When the current stage started.
+                  segments:
+                    type: object
+                    description: Completed region segments of a segmented US national run.
+                    properties:
+                      completed:
+                        type: integer
+                        minimum: 0
+                      total:
+                        type: integer
+                        minimum: 1
         400:
           $ref: '#/components/responses/SPMValidationError'
         404:

policyengine_api/routes/economy_routes.py

@@ -103,13 +103,19 @@ def get_economic_impact(country_id: str, policy_id: int, baseline_policy_id: int
             result=None,
         )
 
-    return _json_response(
-        {
-            "status": result_dict["status"],
-            "message": None,
-            "result": result_dict["data"],
-        }
-    )
+    body = {
+        "status": result_dict["status"],
+        "message": None,
+        "result": result_dict["data"],
+    }
+    # Stage fields from the Simulation Entrypoint, only while computing and
+    # only when the entrypoint reports them, so other bodies keep their shape.
+    if (
+        result_dict["status"] == "computing"
+        and economic_impact_result.run_progress is not None
+    ):
+        body.update(economic_impact_result.run_progress.to_dict())
+    return _json_response(body)
 
 
 @economy_bp.route(

policyengine_api/services/economy_service.py

@@ -25,7 +25,10 @@
 from policyengine_api.data.places import validate_place_code
 from policyengine_api.data.v1_models import ReformImpact
 from policyengine_api.gcp_logging import logger
-from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint
+from policyengine_api.libs.simulation_entrypoint import (
+    SimulationRunProgress,
+    simulation_entrypoint,
+)
 from policyengine_api.observability import (
     runtime as observability_runtime,
 )
@@ -131,6 +134,9 @@ class EconomicImpactResult(BaseModel):
     status: ImpactStatus
     data: dict | None = None
     message: str | None = None
+    # Stage fields relayed from the Simulation Entrypoint while computing.
+    # Kept out of to_dict(); the route adds them only when present.
+    run_progress: SimulationRunProgress | None = None
 
     model_config = {"frozen": True}  # Make model immutable
 
@@ -152,11 +158,13 @@ def completed(cls, data: dict) -> "EconomicImpactResult":
         return cls(status=ImpactStatus.OK, data=data)
 
     @classmethod
-    def computing(cls) -> "EconomicImpactResult":
+    def computing(
+        cls, run_progress: SimulationRunProgress | None = None
+    ) -> "EconomicImpactResult":
         """
         Create an EconomicImpactResult for a computing impact calculation.
         """
-        return cls(status=ImpactStatus.COMPUTING, data=None)
+        return cls(status=ImpactStatus.COMPUTING, data=None, run_progress=run_progress)
 
     @classmethod
     def error(cls, message: str) -> "EconomicImpactResult":
@@ -1276,7 +1284,9 @@ def _handle_execution_state(
                 {"message": "Sim API execution is still running"},
                 severity="INFO",
             )
-            return EconomicImpactResult.computing()
+            return EconomicImpactResult.computing(
+                run_progress=execution.run_progress if execution is not None else None
+            )
 
         raise ValueError(f"Unexpected sim API execution state: {execution_state}")
 

tests/fixtures/services/economy_service.py

@@ -264,6 +264,7 @@ def create_mock_modal_execution(
     result=None,
     error=None,
     policyengine_bundle=None,
+    run_progress=None,
 ):
     """
     Helper function to create mock Modal execution objects.
@@ -292,6 +293,7 @@ def create_mock_modal_execution(
     mock_execution.error = error
     mock_execution.policyengine_bundle = policyengine_bundle or MOCK_RUNTIME_BUNDLE
     mock_execution.resolved_app_name = MOCK_RESOLVED_APP_NAME
+    mock_execution.run_progress = run_progress
     return mock_execution
 
 

tests/unit/libs/test_simulation_entrypoint.py

@@ -26,6 +26,7 @@
     ModalSimulationExecution,
     SimulationAPIModal,
     SimulationEntrypointClient,
+    SimulationRunProgress,
     resolve_simulation_entrypoint_url,
 )
 from policyengine_api.observability.identifiers import (  # noqa: E402
@@ -960,6 +961,131 @@ def test__given_failed_job__then_returns_error(
             assert execution.status == MODAL_EXECUTION_STATUS_FAILED
             assert execution.error == "Simulation timed out"
 
+        def test__given_running_job_with_stage_fields__then_reads_run_progress(
+            self,
+            mock_httpx_client,
+            mock_modal_logger,
+        ):
+            # Given
+            mock_httpx_client.get.return_value = create_mock_httpx_response(
+                status_code=202,
+                json_data={
+                    **MOCK_POLL_RESPONSE_RUNNING,
+                    "stage": "simulating",
+                    "progress": 5,
+                    "progress_max": 96,
+                    "submitted_at": "2026-10-13T12:00:00Z",
+                    "started_at": "2026-10-13T12:00:10Z",
+                    "stage_started_at": "2026-10-13T12:00:17Z",
+                    "segments": {"completed": 19, "total": 20},
+                    "stage_durations": {"loading_data": 7.0},
+                },
+            )
+            api = SimulationAPIModal()
+
+            # When
+            execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID)
+
+            # Then
+            assert execution.status == MODAL_EXECUTION_STATUS_RUNNING
+            assert execution.run_progress.to_dict() == {
+                "stage": "simulating",
+                "progress": 5,
+                "progress_max": 96,
+                "submitted_at": "2026-10-13T12:00:00Z",
+                "started_at": "2026-10-13T12:00:10Z",
+                "stage_started_at": "2026-10-13T12:00:17Z",
+                "segments": {"completed": 19, "total": 20},
+            }
+
+        def test__given_running_job_without_stage_fields__then_has_no_run_progress(
+            self,
+            mock_httpx_client,
+            mock_modal_logger,
+        ):
+            # Given
+            mock_httpx_client.get.return_value = create_mock_httpx_response(
+                status_code=202,
+                json_data=MOCK_POLL_RESPONSE_RUNNING,
+            )
+            api = SimulationAPIModal()
+
+            # When
+            execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID)
+
+            # Then
+            assert execution.run_progress is None
+
+        def test__given_only_a_submission_time__then_relays_it(
+            self,
+            mock_httpx_client,
+            mock_modal_logger,
+        ):
+            # Given - a job on a worker app without stages, after
+            # policyengine-sim-api#781 deploys its gateway
+            mock_httpx_client.get.return_value = create_mock_httpx_response(
+                status_code=202,
+                json_data={
+                    **MOCK_POLL_RESPONSE_RUNNING,
+                    "submitted_at": "2026-10-13T12:00:00Z",
+                },
+            )
+            api = SimulationAPIModal()
+
+            # When
+            execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID)
+
+            # Then
+            assert execution.run_progress.to_dict() == {
+                "submitted_at": "2026-10-13T12:00:00Z"
+            }
+
+        def test__given_a_malformed_timestamp__then_drops_the_stage_fields(
+            self,
+            mock_httpx_client,
+            mock_modal_logger,
+        ):
+            # Given
+            mock_httpx_client.get.return_value = create_mock_httpx_response(
+                status_code=202,
+                json_data={
+                    **MOCK_POLL_RESPONSE_RUNNING,
+                    "stage": "reform",
+                    "stage_started_at": "yesterday",
+                },
+            )
+            api = SimulationAPIModal()
+
+            # When
+            execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID)
+
+            # Then
+            assert execution.status == MODAL_EXECUTION_STATUS_RUNNING
+            assert execution.run_progress is None
+
+        def test__given_malformed_stage_fields__then_still_returns_status(
+            self,
+            mock_httpx_client,
+            mock_modal_logger,
+        ):
+            # Given
+            mock_httpx_client.get.return_value = create_mock_httpx_response(
+                status_code=202,
+                json_data={
+                    **MOCK_POLL_RESPONSE_RUNNING,
+                    "stage": "reform",
+                    "progress": 150,
+                },
+            )
+            api = SimulationAPIModal()
+
+            # When
+            execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID)
+
+            # Then
+            assert execution.status == MODAL_EXECUTION_STATUS_RUNNING
+            assert execution.run_progress is None
+
         def test__given_job_id__then_polls_correct_endpoint(
             self,
             mock_httpx_client,
@@ -1007,6 +1133,17 @@ def test__given_network_error__then_raises_exception(
             log_payload = mock_modal_logger.log_struct.call_args.args[0]
             assert MOCK_MODAL_JOB_ID in log_payload["message"]
 
+    class TestSimulationRunProgress:
+        def test__given_unknown_stage__then_keeps_it(self):
+            run_progress = SimulationRunProgress.from_status(
+                {"status": "running", "stage": "a_future_stage", "progress": 40}
+            )
+
+            assert run_progress.to_dict() == {
+                "stage": "a_future_stage",
+                "progress": 40,
+            }
+
     class TestGetBudgetWindowBatchById:
         def test__given_running_batch__then_returns_running_status(
             self,