Bound public model work requests

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)

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)
 

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:

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)
 # ---------------------------------------------------------------------------