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

@@ -8,6 +8,7 @@
 
 import httpx
 from policyengine_observability import instrument_httpx
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 from policyengine_api.gcp_logging import logger
 from policyengine_api.libs.gateway_auth import (
     GatewayAuthError,
@@ -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: Optional[str] = None
+    progress: Optional[int] = Field(default=None, ge=0, le=100)
+    progress_max: Optional[int] = Field(default=None, ge=0, le=100)
+    submitted_at: Optional[str] = None
+    started_at: Optional[str] = None
+    stage_started_at: Optional[str] = None
+    segments: Optional[SimulationRunSegments] = None
+
+    @classmethod
+    def from_status(cls, data: dict) -> Optional["SimulationRunProgress"]:
+        """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: Optional[str] = None
     policyengine_bundle: Optional[dict] = None
     resolved_app_name: Optional[str] = None
+    run_progress: Optional[SimulationRunProgress] = 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

@@ -101,13 +101,17 @@ 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 worker reports them, so other bodies keep their shape.
+    run_progress = getattr(economic_impact_result, "run_progress", None)
+    if result_dict["status"] == "computing" and isinstance(run_progress, dict):
+        body.update(run_progress)
+    return _json_response(body)
 
 
 @economy_bp.route(

policyengine_api/services/economy_service.py

@@ -23,7 +23,10 @@
 from policyengine_api.data.v1_models import ReformImpact
 from policyengine_api.data.places import validate_place_code
 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,
     set_runtime_context,
@@ -128,6 +131,9 @@ class EconomicImpactResult(BaseModel):
     status: ImpactStatus
     data: Optional[dict] = None
     message: Optional[str] = None
+    # Stage fields relayed from the Simulation Entrypoint while computing.
+    # Kept out of to_dict(); the route adds them only when present.
+    run_progress: Optional[dict] = None
 
     model_config = {"frozen": True}  # Make model immutable
 
@@ -149,11 +155,11 @@ def completed(cls, data: dict) -> "EconomicImpactResult":
         return cls(status=ImpactStatus.OK, data=data)
 
     @classmethod
-    def computing(cls) -> "EconomicImpactResult":
+    def computing(cls, run_progress: Optional[dict] = 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":
@@ -1278,7 +1284,14 @@ def _handle_execution_state(
                 {"message": "Sim API execution is still running"},
                 severity="INFO",
             )
-            return EconomicImpactResult.computing()
+            run_progress = getattr(execution, "run_progress", None)
+            return EconomicImpactResult.computing(
+                run_progress=(
+                    run_progress.to_dict()
+                    if isinstance(run_progress, SimulationRunProgress)
+                    else None
+                )
+            )
 
         else:
             raise ValueError(f"Unexpected sim API execution state: {execution_state}")

tests/unit/libs/test_simulation_entrypoint.py

@@ -26,6 +26,7 @@
     ModalSimulationExecution,
     SimulationAPIModal,
     SimulationEntrypointClient,
+    SimulationRunProgress,
     resolve_simulation_entrypoint_url,
 )
 from policyengine_api.request_context import (  # noqa: E402
@@ -961,6 +962,84 @@ 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_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,
@@ -1008,6 +1087,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,

tests/unit/routes/test_economy_dataset_validation.py

@@ -98,6 +98,56 @@ def test_economy_route_keeps_computing_response_as_http_200(mock_get_economic_im
     response = client.get("/us/economy/123/over/456?region=us&time_period=2026")
 
     assert response.status_code == 200
+    assert json.loads(response.data) == {
+        "status": "computing",
+        "message": None,
+        "result": None,
+    }
+
+
+@patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact")
+def test_economy_route_relays_stage_fields_while_computing(mock_get_economic_impact):
+    mock_result = Mock()
+    mock_result.to_dict.return_value = {
+        "status": "computing",
+        "data": None,
+        "message": None,
+    }
+    mock_result.run_progress = {
+        "stage": "reform",
+        "progress": 67,
+        "progress_max": 96,
+        "stage_started_at": "2026-10-13T12:01:00Z",
+    }
+    mock_get_economic_impact.return_value = mock_result
+    client = _client_with_economy_blueprint()
+
+    response = client.get("/us/economy/123/over/456?region=us&time_period=2026")
+
+    assert response.status_code == 200
+    assert json.loads(response.data) == {
+        "status": "computing",
+        "message": None,
+        "result": None,
+        "stage": "reform",
+        "progress": 67,
+        "progress_max": 96,
+        "stage_started_at": "2026-10-13T12:01:00Z",
+    }
+
+
+@patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact")
+def test_economy_route_never_adds_stage_fields_to_completed_results(
+    mock_get_economic_impact,
+):
+    mock_result = _mock_economic_result()
+    mock_result.run_progress = {"stage": "outputs", "progress": 97}
+    mock_get_economic_impact.return_value = mock_result
+    client = _client_with_economy_blueprint()
+
+    response = client.get("/us/economy/123/over/456?region=us&time_period=2026")
+
+    assert set(json.loads(response.data)) == {"status", "message", "result"}
 
 
 @patch(