Tighten budget-window batch state contract

Author: anth-volk

projects/policyengine-api-simulation/src/modal/budget_window_state.py

@@ -44,12 +44,15 @@ def create_initial_batch_state(
         batch_job_id=batch_job_id,
         status="submitted",
         country=request.country,
+        region=request.region,
         version=resolved_version,
+        target=request.target,
         resolved_app_name=resolved_app_name,
         policyengine_bundle=bundle,
         start_year=request.start_year,
         window_size=request.window_size,
         max_parallel=request.max_parallel,
+        request_payload=request.model_dump(exclude={"telemetry"}, mode="json"),
         years=years,
         queued_years=list(years),
         running_years=[],

projects/policyengine-api-simulation/src/modal/gateway/models.py

@@ -2,7 +2,7 @@
 Pydantic models for the Gateway API.
 """
 
-from typing import ClassVar, Literal, Optional
+from typing import Any, ClassVar, Literal, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
@@ -80,11 +80,12 @@ class BudgetWindowBatchRequest(GatewayRequestBase):
 
     MAX_YEARS: ClassVar[int] = 75
     MAX_END_YEAR: ClassVar[int] = 2099
+    MAX_PARALLEL: ClassVar[int] = 3
 
     region: str
     start_year: str
     window_size: int = Field(ge=1, le=MAX_YEARS)
-    max_parallel: int = Field(default=3, ge=1)
+    max_parallel: int = Field(default=MAX_PARALLEL, ge=1, le=MAX_PARALLEL)
     target: Literal["general"] = "general"
 
     @field_validator("start_year")
@@ -182,12 +183,15 @@ class BudgetWindowBatchState(BaseModel):
     batch_job_id: str
     status: str
     country: str
+    region: str
     version: str
+    target: Literal["general"] = "general"
     resolved_app_name: str
     policyengine_bundle: PolicyEngineBundle
     start_year: str
     window_size: int
     max_parallel: int
+    request_payload: dict[str, Any] = Field(default_factory=dict)
     years: list[str] = Field(default_factory=list)
     queued_years: list[str] = Field(default_factory=list)
     running_years: list[str] = Field(default_factory=list)

projects/policyengine-api-simulation/tests/gateway/test_budget_window_state.py

@@ -3,6 +3,8 @@
 from src.modal.budget_window_state import (
     build_batch_status_response,
     create_initial_batch_state,
+    get_batch_job_state,
+    put_batch_job_state,
 )
 from src.modal.gateway.models import BudgetWindowBatchRequest, PolicyEngineBundle
 
@@ -14,6 +16,9 @@ def test_create_initial_batch_state_builds_queued_years_and_run_id():
         start_year="2026",
         window_size=3,
         max_parallel=2,
+        dataset="enhanced_cps_2024",
+        scope="macro",
+        reform={},
         _telemetry={
             "run_id": "batch-run-123",
             "process_id": "proc-123",
@@ -31,8 +36,13 @@ def test_create_initial_batch_state_builds_queued_years_and_run_id():
 
     assert state.batch_job_id == "fc-parent-123"
     assert state.status == "submitted"
+    assert state.region == "us"
+    assert state.target == "general"
     assert state.years == ["2026", "2027", "2028"]
     assert state.queued_years == ["2026", "2027", "2028"]
+    assert state.request_payload["dataset"] == "enhanced_cps_2024"
+    assert state.request_payload["scope"] == "macro"
+    assert state.request_payload["reform"] == {}
     assert state.run_id == "batch-run-123"
 
 
@@ -61,3 +71,32 @@ def test_build_batch_status_response_computes_progress_from_completed_years():
     assert response.completed_years == ["2026", "2027"]
     assert response.running_years == ["2028"]
     assert response.queued_years == ["2029"]
+
+
+def test_batch_state_round_trips_through_modal_dict(mock_modal):
+    request = BudgetWindowBatchRequest(
+        country="us",
+        region="state/ca",
+        start_year="2026",
+        window_size=2,
+        max_parallel=2,
+        scope="macro",
+        reform={"foo": True},
+    )
+
+    state = create_initial_batch_state(
+        batch_job_id="fc-parent-123",
+        request=request,
+        resolved_version="1.500.0",
+        resolved_app_name="policyengine-simulation-us1-500-0-uk2-66-0",
+        bundle=PolicyEngineBundle(model_version="1.500.0"),
+    )
+    put_batch_job_state(state)
+
+    restored = get_batch_job_state("fc-parent-123")
+
+    assert restored is not None
+    assert restored.region == "state/ca"
+    assert restored.target == "general"
+    assert restored.request_payload["scope"] == "macro"
+    assert restored.request_payload["reform"] == {"foo": True}

projects/policyengine-api-simulation/tests/gateway/test_endpoints.py

@@ -527,3 +527,22 @@ def test__given_non_general_target__then_budget_window_submit_returns_422(
 
         assert response.status_code == 422
         assert response.json()["detail"][0]["loc"] == ["body", "target"]
+
+    def test__given_max_parallel_above_active_limit__then_budget_window_submit_returns_422(
+        self, mock_modal, client: TestClient
+    ):
+        response = client.post(
+            "/simulate/economy/budget-window",
+            json={
+                "country": "us",
+                "region": "us",
+                "scope": "macro",
+                "reform": {},
+                "start_year": "2026",
+                "window_size": 3,
+                "max_parallel": 4,
+            },
+        )
+
+        assert response.status_code == 422
+        assert response.json()["detail"][0]["loc"] == ["body", "max_parallel"]

projects/policyengine-api-simulation/tests/gateway/test_models.py

@@ -379,6 +379,16 @@ def test_budget_window_batch_request_rejects_non_general_target(self):
                 target="cliff",
             )
 
+    def test_budget_window_batch_request_rejects_max_parallel_above_active_limit(self):
+        with pytest.raises(ValidationError):
+            BudgetWindowBatchRequest(
+                country="us",
+                region="us",
+                start_year="2026",
+                window_size=3,
+                max_parallel=4,
+            )
+
     def test_budget_window_batch_request_accepts_extra_simulation_fields(self):
         request = BudgetWindowBatchRequest(
             country="us",