From f67f1f9a0120201bbdd22924dfcb15145734d902 Mon Sep 17 00:00:00 2001 From: Vahid Ahmadi Date: Sun, 26 Jul 2026 15:43:27 +0200 Subject: [PATCH] Bound public model work requests --- integration/src/policyengine_macro/core.py | 46 +++++++++++++++---- .../src/policyengine_macro/mcp_server.py | 28 +++++++---- integration/tests/test_core.py | 26 +++++++++++ integration/tests/test_wiring.py | 18 ++++++++ 4 files changed, 100 insertions(+), 18 deletions(-) diff --git a/integration/src/policyengine_macro/core.py b/integration/src/policyengine_macro/core.py index d494e2a..728c28d 100644 --- a/integration/src/policyengine_macro/core.py +++ b/integration/src/policyengine_macro/core.py @@ -167,6 +167,9 @@ def obr_shock( closure solves to all-zero deltas, which would read as a misleading "no effect" result rather than a mis-specified run. """ + periods = _bounded_int( + "periods", periods, _OBR_MIN_PERIODS, _OBR_MAX_PERIODS + ) run_reform = _import_obr() if investment_closure is None: known = {v["var"]: v["investment_closure"] for v in OBR_VARIABLES} @@ -177,7 +180,7 @@ def obr_shock( name=name, var=var, shock=float(shock), - periods=int(periods), + periods=periods, investment_closure=bool(investment_closure), ) rows = _obr_result_rows(df) @@ -1081,6 +1084,27 @@ def _covid_dummies(index) -> np.ndarray: # residual ESS shortfall honestly instead of hiding it. First-call runtime is # a couple of minutes end to end and results are cached in-process. _SVAR_DEFAULT_DRAWS = 2000 +_SVAR_MIN_DRAWS = 50 +_SVAR_MAX_DRAWS = 10_000 +_SVAR_MIN_HORIZONS = 1 +_SVAR_MAX_HORIZONS = 40 +_OBR_MIN_PERIODS = 1 +_OBR_MAX_PERIODS = 40 + + +def _bounded_int(name: str, value: int, minimum: int, maximum: int) -> int: + """Normalize a public integer input and reject unsafe work requests.""" + try: + normalized = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be an integer") from exc + if normalized != value: + raise ValueError(f"{name} must be an integer") + if not minimum <= normalized <= maximum: + raise ValueError( + f"{name} must be between {minimum} and {maximum}; got {value}" + ) + return normalized def _estimate(draws: int = _SVAR_DEFAULT_DRAWS, seed: int = 0) -> dict: @@ -1143,10 +1167,14 @@ def svar_forecast(horizons: int = 12, Returns median and 68/90 percent bands per future quarter. Bands combine parameter and shock uncertainty. Cached in-process by (horizons, draws). """ - key = (int(horizons), int(draws)) + horizons = _bounded_int( + "horizons", horizons, _SVAR_MIN_HORIZONS, _SVAR_MAX_HORIZONS + ) + draws = _bounded_int("draws", draws, _SVAR_MIN_DRAWS, _SVAR_MAX_DRAWS) + key = (horizons, draws) if key in _FORECAST_CACHE: return _FORECAST_CACHE[key] - est = _estimate(int(draws)) + est = _estimate(draws) analysis, forecast = est["modules"] y_full = est["y_full"] rng = np.random.default_rng(1) @@ -1156,13 +1184,13 @@ def svar_forecast(horizons: int = 12, n_paths = 5 for i, (d, _B) in enumerate(est["pairs"]): for _ in range(n_paths): - path = forecast.sample_forecast(d, y_full, horizons=int(horizons), rng=rng) + path = forecast.sample_forecast(d, y_full, horizons=horizons, rng=rng) yoy_paths.append(forecast.yoy(np.vstack([tail, path]))) pw.append(est["weights"][i]) bands = analysis.aggregate(yoy_paths, weights=np.asarray(pw)) last_q = est["df_full"].index[-1] - quarters = [str(last_q + h) for h in range(1, int(horizons) + 1)] + quarters = [str(last_q + h) for h in range(1, horizons + 1)] def _series(idx: int) -> list[dict]: return [ @@ -1174,7 +1202,7 @@ def _series(idx: int) -> list[dict]: "lo90": round(float(bands["lo90"][h, idx]), 3), "hi90": round(float(bands["hi90"][h, idx]), 3), } - for h in range(int(horizons)) + for h in range(horizons) ] out = { @@ -1187,8 +1215,8 @@ def _series(idx: int) -> list[dict]: estimation_sample=est["estimation_sample"], ), "forecast_origin": str(last_q), - "horizons": int(horizons), - "draws": int(draws), + "horizons": horizons, + "draws": draws, "accepted_draws": est["n_accepted"], "ess": round(est["ess"], 1), "warnings": list(est["warnings"]), @@ -1202,7 +1230,7 @@ def _series(idx: int) -> list[dict]: def svar_latest_shocks(draws: int = _SVAR_DEFAULT_DRAWS) -> dict: """P(sign) of the 6 identified structural shocks in the latest data quarter.""" - key = int(draws) + key = _bounded_int("draws", draws, _SVAR_MIN_DRAWS, _SVAR_MAX_DRAWS) if key in _SHOCKS_CACHE: return _SHOCKS_CACHE[key] est = _estimate(key) diff --git a/integration/src/policyengine_macro/mcp_server.py b/integration/src/policyengine_macro/mcp_server.py index c051085..38c39e5 100644 --- a/integration/src/policyengine_macro/mcp_server.py +++ b/integration/src/policyengine_macro/mcp_server.py @@ -7,7 +7,10 @@ from __future__ import annotations +from typing import Annotated + from mcp.server.fastmcp import FastMCP +from pydantic import Field from policyengine_macro import core from policyengine_macro import capabilities @@ -147,7 +150,7 @@ def score_reform( def obr_shock( var: str, shock: float, - periods: int = 12, + periods: Annotated[int, Field(ge=1, le=40)] = 12, name: str | None = None, investment_closure: bool | None = None, ) -> dict: @@ -164,7 +167,8 @@ def obr_shock( shock: Shock size, applied each quarter. UNITS DEPEND ON THE VARIABLE: CGG and CGIPS are in £ million per quarter (e.g. 1250 = £5bn/year); TCPRO is a rate change in decimal (e.g. -0.05 = 5 percentage point cut). - periods: Number of quarters the shock is applied (default 12 = 3 years). + periods: Number of quarters the shock is applied (1-40; default 12 = + 3 years). name: Optional label for the reform. investment_closure: Omit to use the safe per-variable default (True for TCPRO, False otherwise). It activates the cost-of-capital @@ -365,7 +369,10 @@ def hank_summary() -> dict: @mcp.tool() -def forecast_uk(horizons: int = 12, draws: int = 2000) -> dict: +def forecast_uk( + horizons: Annotated[int, Field(ge=1, le=40)] = 12, + draws: Annotated[int, Field(ge=50, le=10_000)] = 2000, +) -> dict: """Forecast UK YoY GDP growth and CPI inflation with the UK SVAR model. Estimates a Bayesian VAR (1992Q1-2025Q1 sample, sign-identified structural @@ -376,18 +383,21 @@ def forecast_uk(horizons: int = 12, draws: int = 2000) -> dict: importance-weight ESS below 100) with a recommended draw count. Args: - horizons: Forecast horizon in quarters (default 12 = 3 years). + horizons: Forecast horizon in quarters (1-40; default 12 = 3 years). draws: Posterior draws (default 2000: ~135 accepted draws, ESS ~65, a couple of minutes on first call; ~3500 draws reaches ESS >= 100; 500 is faster but yields ~35 accepted, ESS ~14, and a warning). - Results are cached in-process, so repeated calls with the same + Must be between 50 and 10,000. Results are cached in-process, so + repeated calls with the same (horizons, draws) are instant. """ return core.svar_forecast(horizons=horizons, draws=draws) @mcp.tool() -def latest_shocks(draws: int = 2000) -> dict: +def latest_shocks( + draws: Annotated[int, Field(ge=50, le=10_000)] = 2000, +) -> dict: """Structural-shock reading for the latest data quarter from the UK SVAR. For each of the 6 identified shocks (world demand/energy/supply, UK @@ -398,9 +408,9 @@ def latest_shocks(draws: int = 2000) -> dict: than 100 accepted draws or importance-weight ESS below 100). Args: - draws: Posterior draws (default 2000, ~2 minutes on first call; can - be raised for precision). Cached in-process, so repeat calls are - instant. + draws: Posterior draws (50-10,000; default 2000, ~2 minutes on first + call; can be raised for precision). Cached in-process, so repeat + calls are instant. """ return core.svar_latest_shocks(draws=draws) diff --git a/integration/tests/test_core.py b/integration/tests/test_core.py index bdc0948..1afd045 100644 --- a/integration/tests/test_core.py +++ b/integration/tests/test_core.py @@ -16,6 +16,32 @@ def test_list_variables(): json.dumps(vars_) +@pytest.mark.parametrize("periods", [0, 41, 1.5, "many"]) +def test_obr_shock_rejects_unsafe_periods_before_import(periods): + with pytest.raises(ValueError, match="periods"): + core.obr_shock(var="CGG", shock=1250, periods=periods) + + +@pytest.mark.parametrize( + ("kwargs", "field"), + [ + ({"horizons": 0, "draws": 100}, "horizons"), + ({"horizons": 41, "draws": 100}, "horizons"), + ({"horizons": 4, "draws": 49}, "draws"), + ({"horizons": 4, "draws": 10_001}, "draws"), + ], +) +def test_svar_forecast_rejects_unsafe_work_before_import(kwargs, field): + with pytest.raises(ValueError, match=field): + core.svar_forecast(**kwargs) + + +@pytest.mark.parametrize("draws", [0, 49, 10_001, "many"]) +def test_latest_shocks_rejects_unsafe_draws_before_import(draws): + with pytest.raises(ValueError, match="draws"): + core.svar_latest_shocks(draws=draws) + + def test_summary_parses(): s = core.svar_summary() if "error" in s: diff --git a/integration/tests/test_wiring.py b/integration/tests/test_wiring.py index 0ac77e4..da1d0b9 100644 --- a/integration/tests/test_wiring.py +++ b/integration/tests/test_wiring.py @@ -74,6 +74,24 @@ def test_mcp_tool_schema_exposes_expected_params(tool, params): assert params <= props, f"{tool} missing {params - props}" +def test_expensive_mcp_inputs_publish_safe_schema_bounds(): + tools = _registered_tools() + obr_periods = tools["obr_shock"].inputSchema["properties"]["periods"] + forecast = tools["forecast_uk"].inputSchema["properties"] + shocks_draws = tools["latest_shocks"].inputSchema["properties"]["draws"] + + assert (obr_periods["minimum"], obr_periods["maximum"]) == (1, 40) + assert (forecast["horizons"]["minimum"], forecast["horizons"]["maximum"]) == ( + 1, + 40, + ) + assert (forecast["draws"]["minimum"], forecast["draws"]["maximum"]) == ( + 50, + 10_000, + ) + assert (shocks_draws["minimum"], shocks_draws["maximum"]) == (50, 10_000) + + # --------------------------------------------------------------------------- # MCP thin wrappers dispatch to core (instant tools, no heavy solve) # ---------------------------------------------------------------------------