Remote smoke tests: run obr_shock and microsim score_reform through jobs

integration/tests/remote_jobs.py

@@ -0,0 +1,92 @@
+"""Run a slow hosted tool through start_job / get_job_result.
+
+Modal abandons any web request after 150 seconds and answers HTTP 303, which
+the MCP client refuses ("Redirect to .../mcp not followed"). obr_shock and a
+microsim score_reform used to fit under that on a warm server; on a cold one
+they do not, so test_obr_shock_cgg has failed on most deploys and daily
+validations since October, and test_score_reform_microsim_scoreresult on most
+since the 30 Oct deploy. Both tools are on jobs.JOB_TOOLS, so the remote tests reach
+them the way a real caller should: as a job, polled in short calls that each
+stay well inside the ceiling.
+
+Kept out of test_remote_mcp.py so it can be unit-tested without a network
+(test_remote_jobs.py); test_remote_mcp.py is skipped unless
+POLICYENGINE_MACRO_REMOTE_TESTS=1.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+
+import json
+
+POLL_WAIT_SECONDS = 30  # jobs.DEFAULT_WAIT_SECONDS: short polls are the reliable ones
+POLL_TIMEOUT_SECONDS = 180  # one poll: 30s server-side wait plus a cold start
+MAX_POLL_ERRORS = 5
+
+
+async def run_via_job(call, tool: str, arguments: dict, deadline: float,
+                      clock=time.monotonic, sleep=asyncio.sleep) -> dict:
+    """Start ``tool`` as a hosted job and return its result.
+
+    ``call(tool, arguments)`` is an async MCP call that returns the decoded
+    JSON and raises AssertionError when the server reports a tool error.
+    A transport error on a poll is retried (the job runs on regardless of the
+    connection watching it); a tool error, or a job still running at
+    ``deadline`` seconds, fails with the job id in the message.
+    """
+    started = await call("start_job", {"tool": tool, "arguments": arguments})
+    assert started["status"] == "running", started
+    job_id = started["job_id"]
+    assert job_id, started
+
+    stop = clock() + deadline
+    errors = 0
+    out = None
+    while clock() < stop:
+        try:
+            out = await asyncio.wait_for(
+                call("get_job_result",
+                     {"job_id": job_id, "wait_seconds": POLL_WAIT_SECONDS}),
+                timeout=POLL_TIMEOUT_SECONDS,
+            )
+        except AssertionError:
+            raise
+        except Exception as exc:  # noqa: BLE001 - a failed poll is not a failed job
+            errors += 1
+            assert errors <= MAX_POLL_ERRORS, (
+                f"{tool} job {job_id}: {errors} polls in a row failed: {exc!r}"
+            )
+            await sleep(5)
+            continue
+        errors = 0
+        if out["status"] == "done":
+            return out["result"]
+        assert out["status"] == "running", out
+    raise AssertionError(f"{tool} job {job_id} did not finish within {deadline:.0f}s: {out}")
+
+
+def mcp_caller(url: str, timeout: float = 300):
+    """An async ``call(tool, arguments)`` for run_via_job, one session per call.
+
+    The reply is judged only after the MCP contexts have closed. Both
+    streamablehttp_client and ClientSession run anyio task groups, which wrap
+    anything raised inside them in an ExceptionGroup; an AssertionError for a
+    tool error raised in there would reach run_via_job looking like a
+    transport error and be retried instead of failing the test.
+    """
+
+    async def call(tool: str, arguments: dict) -> dict:
+        from mcp import ClientSession
+        from mcp.client.streamable_http import streamablehttp_client
+
+        async with streamablehttp_client(url, timeout=timeout) as (read, write, _):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                result = await session.call_tool(tool, arguments)
+        text = "\n".join(c.text for c in result.content if getattr(c, "text", None))
+        assert not result.isError, f"{tool} error: {text[:1000]}"
+        return json.loads(text)
+
+    return call

integration/tests/test_remote_jobs.py

@@ -0,0 +1,159 @@
+"""Unit tests for remote_jobs.run_via_job, with a scripted server.
+
+No network: the remote suite that uses this helper only runs against the live
+deployment, so the polling contract is pinned here instead.
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from remote_jobs import MAX_POLL_ERRORS, mcp_caller, run_via_job
+
+RUNNING = {"job_id": "fc-x", "status": "running", "waited_seconds": 30}
+RESULT = {"budgetary_impact_bn": 7.1}
+DONE = {"job_id": "fc-x", "status": "done", "result": RESULT}
+
+
+class Server:
+    def __init__(self, polls, clock=None):
+        self.polls = list(polls)
+        self.calls = []
+        self.clock = clock
+
+    async def __call__(self, tool, arguments):
+        self.calls.append((tool, arguments))
+        if tool == "start_job":
+            return {"job_id": "fc-x", "status": "running", "next_step": "get_job_result"}
+        if self.clock is not None:
+            self.clock[0] += 30
+        reply = self.polls.pop(0)
+        if isinstance(reply, BaseException):
+            raise reply
+        return reply
+
+
+async def _no_sleep(_):
+    return None
+
+
+def run(server, deadline=600, clock=None):
+    kw = {"sleep": _no_sleep}
+    if clock is not None:
+        kw["clock"] = lambda: clock[0]
+    return asyncio.run(run_via_job(server, "obr_shock", {"var": "CGG"}, deadline, **kw))
+
+
+def test_starts_a_job_and_polls_until_done():
+    server = Server([RUNNING, DONE])
+    assert run(server) == RESULT
+    assert server.calls[0] == ("start_job", {"tool": "obr_shock", "arguments": {"var": "CGG"}})
+    assert [tool for tool, _ in server.calls[1:]] == ["get_job_result"] * 2
+    assert all(args["wait_seconds"] <= 60 for _, args in server.calls[1:])
+
+
+def test_the_tool_is_never_called_directly():
+    server = Server([DONE])
+    run(server)
+    assert "obr_shock" not in [tool for tool, _ in server.calls]
+
+
+def test_a_transport_error_on_a_poll_is_retried():
+    server = Server([RUNNING, OSError("303 not followed"), DONE])
+    assert run(server) == RESULT
+
+
+def test_repeated_transport_errors_fail_with_the_job_id():
+    server = Server([OSError("down")] * (MAX_POLL_ERRORS + 1))
+    with pytest.raises(AssertionError, match="fc-x"):
+        run(server)
+
+
+def test_a_job_error_fails_without_polling_again():
+    server = Server([AssertionError("job raised"), DONE])
+    with pytest.raises(AssertionError, match="job raised"):
+        run(server)
+    assert server.polls == [DONE]
+
+
+def test_a_job_still_running_at_the_deadline_fails_with_its_id():
+    clock = [0.0]
+    server = Server([RUNNING] * 50, clock=clock)
+    with pytest.raises(AssertionError, match="fc-x did not finish within 90s"):
+        run(server, deadline=90, clock=clock)
+
+
+def _fake_mcp(monkeypatch, replies):
+    """Replace the MCP client with fakes that hold a live anyio task group,
+    as the real streamablehttp_client and ClientSession do."""
+    import contextlib
+    import sys
+    import types
+
+    import anyio
+
+    @contextlib.asynccontextmanager
+    async def background_group():
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(anyio.sleep_forever)
+            try:
+                yield
+            finally:
+                tg.cancel_scope.cancel()
+
+    @contextlib.asynccontextmanager
+    async def streamablehttp_client(url, timeout):
+        async with background_group():
+            yield (None, None, None)
+
+    class ClientSession:
+        def __init__(self, read, write):
+            self._group = background_group()
+
+        async def __aenter__(self):
+            await self._group.__aenter__()
+            return self
+
+        async def __aexit__(self, *exc):
+            return await self._group.__aexit__(*exc)
+
+        async def initialize(self):
+            return None
+
+        async def call_tool(self, tool, arguments):
+            is_error, text = replies[tool]
+            return types.SimpleNamespace(
+                isError=is_error, content=[types.SimpleNamespace(text=text)]
+            )
+
+    mcp = types.ModuleType("mcp")
+    mcp.ClientSession = ClientSession
+    transport = types.ModuleType("mcp.client.streamable_http")
+    transport.streamablehttp_client = streamablehttp_client
+    monkeypatch.setitem(sys.modules, "mcp", mcp)
+    monkeypatch.setitem(sys.modules, "mcp.client", types.ModuleType("mcp.client"))
+    monkeypatch.setitem(sys.modules, "mcp.client.streamable_http", transport)
+
+
+def test_a_failed_job_fails_at_once_through_real_task_groups(monkeypatch):
+    """A tool error must surface as the job's failure, not be retried as if
+    the connection had dropped."""
+    _fake_mcp(monkeypatch, {
+        "start_job": (False, '{"job_id": "fc-x", "status": "running"}'),
+        "get_job_result": (True, "ValueError: reform rejected"),
+    })
+    with pytest.raises(AssertionError, match="reform rejected"):
+        asyncio.run(run_via_job(mcp_caller("http://fake"), "obr_shock", {},
+                                deadline=600, sleep=_no_sleep))
+
+
+def test_a_finished_job_returns_its_result_through_real_task_groups(monkeypatch):
+    _fake_mcp(monkeypatch, {
+        "start_job": (False, '{"job_id": "fc-x", "status": "running"}'),
+        "get_job_result": (False, '{"status": "done", "result": {"ok": 1}}'),
+    })
+    out = asyncio.run(run_via_job(mcp_caller("http://fake"), "obr_shock", {},
+                                  deadline=600, sleep=_no_sleep))
+    assert out == {"ok": 1}

integration/tests/test_remote_mcp.py

@@ -24,6 +24,7 @@
     "https://policyengine--policyengine-macro-mcp-serve.modal.run/mcp",
 )
 
+from remote_jobs import mcp_caller, run_via_job
 from tool_surface import GOLDEN_TOOL_COUNT, GOLDEN_TOOLS, assert_surface
 
 EXPECTED_TOOLS = set(GOLDEN_TOOLS)
@@ -220,8 +221,11 @@ async def test_model_summary_returns_fevd():
 
 @pytest.mark.anyio
 async def test_obr_shock_cgg():
-    out = await _call(
-        "obr_shock", {"var": "CGG", "shock": 1250, "periods": 4}
+    # Through a job: direct, a cold obr_shock passes Modal's 150s ceiling and
+    # the client gets a 303 it cannot follow (failing on most runs since October).
+    out = await run_via_job(
+        mcp_caller(URL), "obr_shock", {"var": "CGG", "shock": 1250, "periods": 4},
+        deadline=300,
     )
     assert len(out["results"]) >= 4
     # £1.25bn/quarter of extra spending should raise GDP by ~£1.25bn/quarter
@@ -294,17 +298,16 @@ async def test_score_reform_obr_refuses_corp_tax_actionably():
 async def test_score_reform_microsim_scoreresult():
     """score_reform(model='microsim') on the hosted server: a basic-rate
     rise must raise revenue, and the common ScoreResult block (#10) must be
-    present and coherent (units, basis, distributional)."""
-    import asyncio
-
-    out = await asyncio.wait_for(
-        _call(
-            "score_reform",
-            {"country": "uk",
-             "reform": {"gov.hmrc.income_tax.rates.uk[0].rate": 0.21},
-             "model": "microsim"},
-        ),
-        timeout=300,
+    present and coherent (units, basis, distributional).
+
+    Through a job: since the 30 Oct 2026 deploy a direct call usually passes
+    Modal's 150s ceiling and the client gets a 303 it cannot follow."""
+    out = await run_via_job(
+        mcp_caller(URL), "score_reform",
+        {"country": "uk",
+         "reform": {"gov.hmrc.income_tax.rates.uk[0].rate": 0.21},
+         "model": "microsim"},
+        deadline=300,
     )
     assert out["budgetary_impact_bn"] > 0, out["budgetary_impact_bn"]
     score = out["score"]