Fix missing simulation job polling status

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

@@ -30,6 +30,7 @@ class MockFunctionCall:
     """Mock for Modal FunctionCall returned by spawn."""
 
     registry = {}
+    from_id_errors = {}
 
     def __init__(self, object_id: str = "mock-job-id-123"):
         self.object_id = object_id
@@ -47,6 +48,8 @@ def get(self, timeout: int = 0):
 
     @classmethod
     def from_id(cls, object_id: str):
+        if object_id in cls.from_id_errors:
+            raise cls.from_id_errors[object_id]
         if object_id not in cls.registry:
             raise KeyError(object_id)
         return cls.registry[object_id]
@@ -64,6 +67,11 @@ def __init__(self):
     def bind(self, app_name: str, func_name: str) -> "BoundMockFunction":
         return BoundMockFunction(self, app_name, func_name)
 
+    def call_for(self, object_id: str) -> MockFunctionCall:
+        call = MockFunctionCall(object_id=object_id)
+        self.last_call = call
+        return call
+
 
 class BoundMockFunction:
     """Function handle returned by Modal.Function.from_name."""
@@ -84,6 +92,14 @@ def spawn(self, payload: dict) -> MockFunctionCall:
         return self.recorder.last_call
 
 
+class MockModalException:
+    class NotFoundError(Exception):
+        pass
+
+    class OutputExpiredError(Exception):
+        pass
+
+
 @pytest.fixture
 def mock_modal(monkeypatch):
     """Patch Modal calls in the gateway endpoints module."""
@@ -93,6 +109,7 @@ def mock_modal(monkeypatch):
     mock_func = MockFunction()
     mock_dicts = {}
     MockFunctionCall.registry = {}
+    MockFunctionCall.from_id_errors = {}
 
     class MockModalDict:
         @staticmethod
@@ -113,11 +130,14 @@ class MockModal:
         Dict = MockModalDict
         Function = MockModalFunction
         FunctionCall = MockFunctionCall
+        exception = MockModalException
 
     monkeypatch.setattr(endpoints, "modal", MockModal)
     monkeypatch.setattr(budget_window_state, "modal", MockModal)
 
     return {
         "func": mock_func,
         "dicts": mock_dicts,
+        "function_call": MockFunctionCall,
+        "exception": MockModalException,
     }

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

@@ -61,6 +61,24 @@ def _job_metadata_store():
     return modal.Dict.from_name(JOB_METADATA_DICT_NAME, create_if_missing=True)
 
 
+def _modal_exception_class(name: str):
+    exception_module = getattr(modal, "exception", None)
+    if exception_module is None:
+        return None
+    return getattr(exception_module, name, None)
+
+
+def _is_modal_exception(exc: BaseException, name: str) -> bool:
+    exception_class = _modal_exception_class(name)
+    return exception_class is not None and isinstance(exc, exception_class)
+
+
+def _is_modal_job_not_found(exc: BaseException) -> bool:
+    return _is_modal_exception(exc, "NotFoundError") or _is_modal_exception(
+        exc, "OutputExpiredError"
+    )
+
+
 def _build_policyengine_bundle(
     country: str, resolved_version: str, payload: dict
 ) -> PolicyEngineBundle:
@@ -267,12 +285,16 @@ async def get_job_status(job_id: str):
         - 500 with status="failed" and error on failure
         - 404 if job_id not found
     """
-    try:
-        call = modal.FunctionCall.from_id(job_id)
-    except Exception:
+    job_metadata = _job_metadata_store().get(job_id)
+    if job_metadata is None:
         raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
 
-    job_metadata = _job_metadata_store().get(job_id)
+    try:
+        call = modal.FunctionCall.from_id(job_id)
+    except Exception as exc:
+        if _is_modal_job_not_found(exc):
+            raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
+        raise
 
     try:
         result = call.get(timeout=0)
@@ -282,6 +304,8 @@ async def get_job_status(job_id: str):
     except TimeoutError:
         return running_job_response(job_metadata)
     except Exception as exc:
+        if _is_modal_job_not_found(exc):
+            raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
         redacted = log_and_redact_exception(
             exc,
             scope="simulation_job_status",

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

@@ -433,6 +433,143 @@ def test__given_submitted_job_with_telemetry__then_polling_echoes_run_id(
         assert response.status_code == 200
         assert response.json()["run_id"] == "run-123"
 
+    def test__given_unknown_job_id__then_polling_returns_404(
+        self, mock_modal, client: TestClient
+    ):
+        """
+        Given a job id that the gateway never issued
+        When polling job status
+        Then the gateway returns 404 before asking Modal for a call result.
+        """
+        response = client.get("/jobs/unknown-job-id")
+
+        assert response.status_code == 404
+        assert response.json()["detail"] == "Job not found: unknown-job-id"
+
+    def test__given_lazy_modal_call_without_metadata__then_polling_returns_404(
+        self, mock_modal, client: TestClient
+    ):
+        """
+        Given Modal can construct a FunctionCall handle for an arbitrary id
+        When the gateway has no metadata for that id
+        Then the gateway still treats it as not found.
+        """
+        mock_modal["func"].call_for("auth-smoke-probe-does-not-exist")
+
+        response = client.get("/jobs/auth-smoke-probe-does-not-exist")
+
+        assert response.status_code == 404
+        assert (
+            response.json()["detail"]
+            == "Job not found: auth-smoke-probe-does-not-exist"
+        )
+
+    def test__given_running_job__then_polling_returns_202(
+        self, mock_modal, client: TestClient
+    ):
+        mock_modal["dicts"]["simulation-api-us-versions"] = {
+            "latest": "1.500.0",
+            "1.500.0": "policyengine-simulation-us1-500-0-uk2-66-0",
+        }
+
+        submit_response = client.post(
+            "/simulate/economy/comparison",
+            json={
+                "country": "us",
+                "scope": "macro",
+                "reform": {},
+            },
+        )
+        job_id = submit_response.json()["job_id"]
+        mock_modal["func"].last_call.running = True
+
+        response = client.get(f"/jobs/{job_id}")
+
+        assert response.status_code == 202
+        assert response.json()["status"] == "running"
+
+    def test__given_expired_modal_output__then_polling_returns_404(
+        self, mock_modal, client: TestClient
+    ):
+        mock_modal["dicts"]["simulation-api-us-versions"] = {
+            "latest": "1.500.0",
+            "1.500.0": "policyengine-simulation-us1-500-0-uk2-66-0",
+        }
+
+        submit_response = client.post(
+            "/simulate/economy/comparison",
+            json={
+                "country": "us",
+                "scope": "macro",
+                "reform": {},
+            },
+        )
+        job_id = submit_response.json()["job_id"]
+        mock_modal["func"].last_call.error = mock_modal[
+            "exception"
+        ].OutputExpiredError()
+
+        # Modal's FastAPI job queue example maps OutputExpiredError to 404:
+        # https://modal.com/docs/guide/job-queue#integration-with-web-frameworks
+        response = client.get(f"/jobs/{job_id}")
+
+        assert response.status_code == 404
+        assert response.json()["detail"] == f"Job not found: {job_id}"
+
+    def test__given_modal_call_not_found__then_polling_returns_404(
+        self, mock_modal, client: TestClient
+    ):
+        mock_modal["dicts"]["simulation-api-us-versions"] = {
+            "latest": "1.500.0",
+            "1.500.0": "policyengine-simulation-us1-500-0-uk2-66-0",
+        }
+
+        submit_response = client.post(
+            "/simulate/economy/comparison",
+            json={
+                "country": "us",
+                "scope": "macro",
+                "reform": {},
+            },
+        )
+        job_id = submit_response.json()["job_id"]
+        mock_modal["function_call"].from_id_errors[job_id] = mock_modal[
+            "exception"
+        ].NotFoundError()
+
+        response = client.get(f"/jobs/{job_id}")
+
+        assert response.status_code == 404
+        assert response.json()["detail"] == f"Job not found: {job_id}"
+
+    def test__given_worker_error__then_polling_returns_redacted_500(
+        self, mock_modal, client: TestClient
+    ):
+        mock_modal["dicts"]["simulation-api-us-versions"] = {
+            "latest": "1.500.0",
+            "1.500.0": "policyengine-simulation-us1-500-0-uk2-66-0",
+        }
+
+        submit_response = client.post(
+            "/simulate/economy/comparison",
+            json={
+                "country": "us",
+                "scope": "macro",
+                "reform": {},
+            },
+        )
+        job_id = submit_response.json()["job_id"]
+        mock_modal["func"].last_call.error = RuntimeError("worker crashed")
+
+        response = client.get(f"/jobs/{job_id}")
+
+        assert response.status_code == 500
+        body = response.json()
+        assert body["status"] == "failed"
+        assert body["error"].startswith("Simulation failed")
+        assert "correlation_id=" in body["error"]
+        assert "worker crashed" not in body["error"]
+
 
 class TestBudgetWindowBatchEndpoints:
     """Tests for budget-window batch gateway endpoints."""