diff --git a/integration/tests/remote_jobs.py b/integration/tests/remote_jobs.py new file mode 100644 index 00000000..89eaeb88 --- /dev/null +++ b/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 diff --git a/integration/tests/test_remote_jobs.py b/integration/tests/test_remote_jobs.py new file mode 100644 index 00000000..f4fe9dbc --- /dev/null +++ b/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} diff --git a/integration/tests/test_remote_mcp.py b/integration/tests/test_remote_mcp.py index c7edf519..9ff14110 100644 --- a/integration/tests/test_remote_mcp.py +++ b/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"]