scale_forecasting.mcp

Built-in, zero-extra JSON-RPC 2.0 Model Context Protocol (MCP) server over stdio (python -m scale_forecasting.mcp) exposing 7 resources and 9 tools across offline model benchmarking, RunConfig validation, multi-family DAG planning, Airflow DAG emission, and live BigQuery registry inspection/review/repair.

See the AI Agents, Portable Skills & Built-in MCP Server guide for client configuration and tool examples.

mcp

Built-in Model Context Protocol (MCP) server for scale-forecasting over stdio.

Exposes the platform's configuration schema, canonical catalogs (37 models, 21 metrics, 7 runtimes, 4 GPU types, 13 extras, 5 views, 62 shipped configs), environment readiness probe, offline model playground, RunConfig validator and DAG planner, Airflow DAG emitter, and live BigQuery registry inspection/review/repair tools to any MCP-compatible AI coding agent (Google Antigravity agy CLI & IDE, Claude Code, Cursor, VS Code GitHub Copilot, Windsurf, OpenAI Codex CLI).

Zero extra dependencies are required: the server runs on a bare pip install scale-forecasting core installation using the Python standard library and Pydantic. Tools that query or mutate Google Cloud (inspect_registry, review_run, launch_or_repair_run, and cloud modes of plan_execution) check for [gcp] and SF_* environment variables first and return structured remediation steps if either is absent. Cloud job submission and mutating repair verbs are gated behind the explicit --allow-launch CLI flag.

CLI usage::

# Start stdio MCP server (read/plan/playground enabled; cloud launches safety-locked)
python -m scale_forecasting.mcp

# Start stdio MCP server with cloud job launch & repair unlocked
python -m scale_forecasting.mcp --allow-launch

# Print environment readiness JSON and exit
python -m scale_forecasting.mcp --probe-env

McpServer

McpServer(*, allow_launch=False, repo_root=None)

JSON-RPC 2.0 Model Context Protocol server for scale-forecasting.

Source code in src/scale_forecasting/mcp.py

def __init__(self, *, allow_launch: bool = False, repo_root: Path | None = None) -> None:
    self.allow_launch = allow_launch
    self.repo_root = repo_root or Path(__file__).resolve().parents[2]

call_tool

call_tool(name, arguments=None)

Dispatch a tool call and wrap the result in an MCP tool response.

Source code in src/scale_forecasting/mcp.py

def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> dict[str, Any]:
    """Dispatch a tool call and wrap the result in an MCP tool response."""
    args = arguments or {}
    try:
        payload = self._dispatch_tool(name, args)
        return {
            "content": [{"type": "text", "text": _dumps(payload)}],
            "isError": False,
        }
    except (ScaleForecastError, ValueError, KeyError, FileNotFoundError) as exc:
        err_payload = {
            "error": type(exc).__name__,
            "message": str(exc),
            "environment": probe_environment(self.repo_root),
        }
        return {
            "content": [{"type": "text", "text": _dumps(err_payload)}],
            "isError": True,
        }
    except Exception as exc:  # noqa: BLE001 - MCP tool boundary must return structured JSON-RPC result
        err_payload = {
            "error": type(exc).__name__,
            "message": str(exc),
        }
        return {
            "content": [{"type": "text", "text": _dumps(err_payload)}],
            "isError": True,
        }

handle_message

handle_message(message)

Process one JSON-RPC 2.0 request or notification and return a response dict (or None).

Source code in src/scale_forecasting/mcp.py

def handle_message(self, message: dict[str, Any]) -> dict[str, Any] | None:
    """Process one JSON-RPC 2.0 request or notification and return a response dict (or None)."""
    method = message.get("method", "")
    msg_id = message.get("id")
    params = message.get("params") or {}

    # Notifications have no `id` and expect no response.
    if msg_id is None:
        return None

    try:
        if method == "initialize":
            client_version = params.get("protocolVersion") or MCP_PROTOCOL_VERSION
            result: dict[str, Any] = {
                "protocolVersion": client_version,
                "capabilities": {
                    "resources": {"subscribe": False, "listChanged": False},
                    "tools": {"listChanged": False},
                },
                "serverInfo": {
                    "name": "scale-forecasting",
                    "version": _pkg_version("scale-forecasting"),
                },
            }
        elif method == "ping":
            result = {}
        elif method == "resources/list":
            result = {"resources": self.list_resources()}
        elif method == "resources/read":
            uri = str(params.get("uri", ""))
            result = self.read_resource(uri)
        elif method == "tools/list":
            result = {"tools": self.list_tools()}
        elif method == "tools/call":
            tool_name = str(params.get("name", ""))
            tool_args = params.get("arguments") or {}
            result = self.call_tool(tool_name, tool_args)
        elif method == "prompts/list":
            result = {"prompts": []}
        else:
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "error": {
                    "code": -32601,
                    "message": f"Method not found: {method}",
                },
            }
        return {"jsonrpc": "2.0", "id": msg_id, "result": result}
    except Exception as exc:  # noqa: BLE001 - JSON-RPC protocol handler must return structured error object
        return {
            "jsonrpc": "2.0",
            "id": msg_id,
            "error": {
                "code": -32602,
                "message": str(exc),
            },
        }

list_resources

list_resources()

Return the 7 MCP resources exposed by this server.

Source code in src/scale_forecasting/mcp.py

def list_resources(self) -> list[dict[str, str]]:
    """Return the 7 MCP resources exposed by this server."""
    return [
        {
            "uri": "forecast://environment",
            "name": "Active Environment & Readiness Report",
            "description": (
                "Installed extras (13), available vs. missing models (37), SF_* environment "
                "variables, ADC status, and readiness across all 5 execution paths."
            ),
            "mimeType": "application/json",
        },
        {
            "uri": "forecast://schema/run-config",
            "name": "RunConfig JSON Schema (Draft 2020-12)",
            "description": "Complete JSON Schema for RunConfig and all nested blocks.",
            "mimeType": "application/schema+json",
        },
        {
            "uri": "forecast://catalog/models",
            "name": "Forecasting Models Catalog (37 Models)",
            "description": (
                "All 37 registered models across 5 families with runtime, package, extra, "
                "covariate flags, explainability, and local availability."
            ),
            "mimeType": "application/json",
        },
        {
            "uri": "forecast://catalog/metrics",
            "name": "Evaluation Metrics Catalog (21 Metrics)",
            "description": (
                "All 16 point and 5 interval evaluation metrics with direction and "
                "requirements."
            ),
            "mimeType": "application/json",
        },
        {
            "uri": "forecast://catalog/runtimes",
            "name": "Compute Runtimes & GPU Catalog (7 Runtimes, 4 GPUs)",
            "description": (
                "All 7 Google Cloud runtimes, sub-modes, scaling models, and 4 GPU types."
            ),
            "mimeType": "application/json",
        },
        {
            "uri": "forecast://catalog/configs",
            "name": "Shipped Demo & Smoke Configurations (62 Configs)",
            "description": (
                "All 20 root demo configs and 42 smoke configs with run_ids and models."
            ),
            "mimeType": "application/json",
        },
        {
            "uri": "forecast://catalog/views",
            "name": "BigQuery Source Tables, Registry Tables & Views",
            "description": (
                "The 4 source tables, 5 registry tables, and 5 analytical SQL views."
            ),
            "mimeType": "application/json",
        },
    ]

list_tools

list_tools()

Return the 9 MCP tools exposed by this server.

Source code in src/scale_forecasting/mcp.py

def list_tools(self) -> list[dict[str, Any]]:
    """Return the 9 MCP tools exposed by this server."""
    return [
        {
            "name": "probe_environment",
            "description": (
                "Inspect the active Python environment: installed optional extras (out of 13), "
                "locally available vs. missing-extra models (out of 37), SF_* environment "
                "variables, .env.infra presence, ADC credentials, and readiness across all 5 "
                "execution paths. Call this first to adapt recommendations to the machine."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {},
                "additionalProperties": False,
            },
        },
        {
            "name": "list_catalog",
            "description": (
                "Query the canonical catalogs of models (37), metrics (21), runtimes (7), "
                "GPUs (4), dependency extras (13), BigQuery views (5), or shipped configs (62)."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "enum": [
                            "all",
                            "models",
                            "metrics",
                            "runtimes",
                            "gpus",
                            "extras",
                            "views",
                            "configs",
                        ],
                        "default": "all",
                    },
                    "family": {
                        "type": "string",
                        "description": (
                            "Optional model family filter "
                            "(statistical, ml, deep_learning, automl, native)."
                        ),
                    },
                    "available_only": {
                        "type": "boolean",
                        "default": False,
                        "description": (
                            "When true, filter models to those installed in the active "
                            "environment."
                        ),
                    },
                },
                "additionalProperties": False,
            },
        },
        {
            "name": "inspect_config_schema",
            "description": (
                "Inspect the JSON Schema and validation rules for RunConfig or a specific "
                "nested config block (DataConfig, FeaturesConfig, BacktestConfig, "
                "OutputConfig, HpoConfig, EnsembleConfig, HierarchyConfig, ComputeConfig, "
                "FamilyCompute, EnsembleCompute, ProfileConfig, CapacityConfig, "
                "CapacityServicePolicy, RetryResources)."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "section": {
                        "type": "string",
                        "description": (
                            "Optional section or Pydantic class name to inspect (e.g., "
                            "'compute', 'backtest', 'FamilyCompute'). Omit for overview."
                        ),
                    }
                },
                "additionalProperties": False,
            },
        },
        {
            "name": "validate_and_dry_run",
            "description": (
                "Validate a RunConfig (dict, JSON string, or file path) offline, compute its "
                "deterministic content-addressed run_id, estimate cell/fit workload, resolve "
                "per-family compute placement and DAG nodes, and check whether any requested "
                "models require an uninstalled optional extra."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "config": {
                        "description": (
                            "RunConfig dict, JSON string, or relative/absolute file path."
                        ),
                    },
                    "n_series": {
                        "type": "integer",
                        "description": "Optional series_limit override.",
                    },
                    "ignore_unavailable_models": {
                        "type": "boolean",
                        "default": False,
                        "description": (
                            "Filter models to those installed locally before planning."
                        ),
                    },
                },
                "required": ["config"],
                "additionalProperties": False,
            },
        },
        {
            "name": "run_playground",
            "description": (
                "Run a real forecasting model offline on deterministic synthetic time-series "
                "panel data (`worker.run_cell`), optionally with expanding-window backtesting, "
                "and return per-series evaluation metrics and sample forecast rows."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "model": {
                        "type": "string",
                        "default": "holtwinters",
                        "description": (
                            "Registered Python model name "
                            "(e.g., 'holtwinters', 'theta', 'naive_seasonal')."
                        ),
                    },
                    "horizon": {
                        "type": "integer",
                        "default": 14,
                    },
                    "n_series": {
                        "type": "integer",
                        "default": 3,
                    },
                    "backtest": {
                        "type": "boolean",
                        "default": True,
                    },
                    "n_folds": {
                        "type": "integer",
                        "default": 3,
                    },
                    "scenarios": {
                        "type": "object",
                        "description": (
                            "Optional counterfactual covariate shock map "
                            "(e.g., {'promo_week': {'price_index': '*0.85'}})."
                        ),
                    },
                    "price_col": {
                        "type": "string",
                        "description": (
                            "Optional price column name for revenue impact simulation."
                        ),
                    },
                    "include_inventory_policy": {
                        "type": "boolean",
                        "default": False,
                        "description": (
                            "When true, compute operational (s, S) / (R, S) inventory & "
                            "replenishment policy (ROP, safety_stock, order_up_to_level, "
                            "recommended_order_qty, Newsvendor cost) from the forward "
                            "forecast and backtest residuals."
                        ),
                    },
                    "inventory_params": {
                        "type": "object",
                        "description": (
                            "Optional keyword overrides for compute_inventory_policy "
                            "(lead_time, lead_time_std, review_period, service_level, "
                            "fill_rate, on_hand, on_order, backorders, holding_cost, "
                            "stockout_cost, min_order_qty, order_multiple)."
                        ),
                    },
                    "export_report_path": {
                        "type": "string",
                        "description": (
                            "Optional file path (.html) to export a self-contained "
                            "interactive HTML report."
                        ),
                    },
                },
                "additionalProperties": False,
            },
        },
        {
            "name": "plan_execution",
            "description": (
                "Plan or emit execution artifacts across the platform's execution paths: "
                "'dry_run' (offline plan + CLI/gcloud launch commands + idempotency check when "
                "SF_* is configured), 'emit_airflow' (render Cloud Composer 3 / Airflow DAG "
                "Python source offline), 'feasibility' (query live BigQuery series lengths vs. "
                "backtest geometry), 'quota' (query regional CPU/GPU quota meters), or "
                "'stage_only' (upload config + code archive to GCS; requires --allow-launch)."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "config": {
                        "description": "RunConfig dict, JSON string, or file path.",
                    },
                    "mode": {
                        "type": "string",
                        "enum": [
                            "dry_run",
                            "emit_airflow",
                            "feasibility",
                            "quota",
                            "stage_only",
                        ],
                        "default": "dry_run",
                    },
                    "force": {
                        "type": "boolean",
                        "default": False,
                    },
                    "with_retry": {
                        "type": "boolean",
                        "default": False,
                        "description": (
                            "For mode='emit_airflow': include the surgical retry node before "
                            "barrier ensembling."
                        ),
                    },
                    "emit_out": {
                        "type": "string",
                        "description": (
                            "Optional file path to write the rendered Airflow DAG when "
                            "mode='emit_airflow'."
                        ),
                    },
                },
                "required": ["config"],
                "additionalProperties": False,
            },
        },
        {
            "name": "inspect_registry",
            "description": (
                "Read-only inspection of the live BigQuery run registry (`[gcp]` + `SF_*`): "
                "'doctor' (table row counts, stuck RUNNING runs, orphaned GCS artifacts), "
                "'recent_runs' (latest runs from v_run_summary with time ledger and "
                "overhead_seconds), 'probe_run' (reconcile live runtime status for a run_id), "
                "or 'retry_preview' (preview surgical cell-repair decision table)."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["doctor", "recent_runs", "probe_run", "retry_preview"],
                        "default": "doctor",
                    },
                    "run_id": {
                        "type": "string",
                        "description": (
                            "Target run_id (required for 'probe_run' and 'retry_preview')."
                        ),
                    },
                    "limit": {
                        "type": "integer",
                        "default": 20,
                        "description": "Maximum rows for 'recent_runs'.",
                    },
                    "job": {
                        "type": "string",
                        "description": "Optional family filter for 'probe_run'.",
                    },
                },
                "additionalProperties": False,
            },
        },
        {
            "name": "review_run",
            "description": (
                "Read-only data-science evaluation of a completed or partial run in BigQuery "
                "(`review.review_run` + `review.calibration_report`): best model overall and "
                "per family, ensemble lift over best base model, leaderboard metrics, and "
                "point-forecast arm / prediction interval coverage calibration."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "run_id": {
                        "type": "string",
                        "description": "Content-addressed run_id (<slug>-<12hex>) to review.",
                    },
                    "include_calibration": {
                        "type": "boolean",
                        "default": True,
                    },
                    "include_leaderboard": {
                        "type": "boolean",
                        "default": True,
                    },
                    "include_fva": {
                        "type": "boolean",
                        "default": False,
                        "description": (
                            "When true, include Forecast Value Added (FVA) & Diebold-Mariano "
                            "statistical significance table."
                        ),
                    },
                    "include_horizon_decay": {
                        "type": "boolean",
                        "default": False,
                        "description": "When true, include lead-time horizon accuracy decay.",
                    },
                    "include_bias_report": {
                        "type": "boolean",
                        "default": False,
                        "description": (
                            "When true, include directional bias & Tracking Signal audit."
                        ),
                    },
                    "include_exceptions": {
                        "type": "boolean",
                        "default": False,
                        "description": (
                            "When true, include Top-N operational exception triage queue."
                        ),
                    },
                    "include_inventory_policy": {
                        "type": "boolean",
                        "default": False,
                        "description": (
                            "When true, include operational (s, S) / (R, S) inventory & "
                            "replenishment policy table (ROP, safety_stock, "
                            "order_up_to_level, recommended_order_qty, Newsvendor cost)."
                        ),
                    },
                    "inventory_params": {
                        "type": "object",
                        "description": (
                            "Optional keyword overrides for compute_inventory_policy "
                            "(lead_time, lead_time_std, review_period, service_level, "
                            "fill_rate, on_hand, on_order, backorders, holding_cost, "
                            "stockout_cost, min_order_qty, order_multiple, model_type)."
                        ),
                    },
                    "export_report_path": {
                        "type": "string",
                        "description": (
                            "Optional file path (.html) to export a self-contained "
                            "interactive HTML report."
                        ),
                    },
                },
                "required": ["run_id"],
                "additionalProperties": False,
            },
        },
        {
            "name": "launch_or_repair_run",
            "description": (
                "Safety-gated cloud execution and repair tool (requires the MCP server to be "
                "started with `--allow-launch`). Actions: 'run' (launch multi-family DAG), "
                "'retry' (submit surgical repair of failed/missing cells), 'settle' (repair "
                "stale job status rows from runtime probes), 'cancel' (stop in-flight cloud "
                "jobs), 'close_runs' (finalize abandoned RUNNING headers with terminal jobs)."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["run", "retry", "settle", "cancel", "close_runs"],
                    },
                    "config": {
                        "description": (
                            "RunConfig dict, JSON string, or file path (for 'run', or for "
                            "'retry'/'settle'/'cancel' when run_id is not passed)."
                        ),
                    },
                    "run_id": {
                        "type": "string",
                        "description": "Target run_id (for 'retry', 'settle', or 'cancel').",
                    },
                    "force": {
                        "type": "boolean",
                        "default": False,
                    },
                    "reason": {
                        "type": "string",
                        "default": "",
                    },
                    "n_series": {
                        "type": "integer",
                    },
                    "job": {
                        "type": "string",
                    },
                },
                "required": ["action"],
                "additionalProperties": False,
            },
        },
    ]

read_resource

read_resource(uri)

Read one MCP resource by URI.

Source code in src/scale_forecasting/mcp.py

def read_resource(self, uri: str) -> dict[str, Any]:
    """Read one MCP resource by URI."""
    if uri == "forecast://environment":
        text = _dumps(probe_environment(self.repo_root))
        mime = "application/json"
    elif uri == "forecast://schema/run-config":
        text = render_run_config_json_schema()
        mime = "application/schema+json"
    elif uri == "forecast://catalog/models":
        text = _dumps(_build_models_catalog())
        mime = "application/json"
    elif uri == "forecast://catalog/metrics":
        text = _dumps(_build_metrics_catalog())
        mime = "application/json"
    elif uri == "forecast://catalog/runtimes":
        text = _dumps({"runtimes": list(RUNTIME_CATALOG), "gpus": list(GPU_CATALOG)})
        mime = "application/json"
    elif uri == "forecast://catalog/configs":
        text = _dumps(_build_configs_catalog(self.repo_root))
        mime = "application/json"
    elif uri == "forecast://catalog/views":
        text = _dumps(_build_views_catalog())
        mime = "application/json"
    else:
        raise ValueError(f"Unknown resource URI: {uri!r}")

    return {
        "contents": [
            {
                "uri": uri,
                "mimeType": mime,
                "text": text,
            }
        ]
    }

serve_stdio

serve_stdio(stdin=None, stdout=None)

Serve MCP JSON-RPC 2.0 requests over stdio (newline-delimited or Content-Length).

Source code in src/scale_forecasting/mcp.py

def serve_stdio(
    self,
    stdin: BinaryIO | None = None,
    stdout: BinaryIO | None = None,
) -> None:
    """Serve MCP JSON-RPC 2.0 requests over `stdio` (newline-delimited or Content-Length)."""
    in_stream: BinaryIO = stdin if stdin is not None else sys.stdin.buffer
    out_stream: BinaryIO = stdout if stdout is not None else sys.stdout.buffer

    while True:
        line_bytes = in_stream.readline()
        if not line_bytes:
            break
        stripped = line_bytes.strip()
        if not stripped:
            continue

        use_content_length = False
        if stripped.lower().startswith(b"content-length:"):
            use_content_length = True
            length = int(stripped.split(b":", 1)[1].strip())
            while True:
                hdr_bytes = in_stream.readline()
                if not hdr_bytes or not hdr_bytes.strip():
                    break
            body_bytes = in_stream.read(length)
        else:
            body_bytes = stripped

        try:
            request = json.loads(body_bytes.decode("utf-8"))
        except json.JSONDecodeError as exc:
            err_resp = {
                "jsonrpc": "2.0",
                "id": None,
                "error": {"code": -32700, "message": f"Parse error: {exc}"},
            }
            _write_frame(out_stream, err_resp, use_content_length=use_content_length)
            continue

        response = self.handle_message(request)
        if response is not None:
            _write_frame(out_stream, response, use_content_length=use_content_length)

main

main(argv=None)

CLI entry point for python -m scale_forecasting.mcp.

Source code in src/scale_forecasting/mcp.py

def main(argv: list[str] | None = None) -> int:
    """CLI entry point for ``python -m scale_forecasting.mcp``."""
    _configure_stderr_logging()
    parser = argparse.ArgumentParser(
        prog="scale_forecasting.mcp",
        description="Run the scale-forecasting Model Context Protocol (MCP) server over stdio.",
    )
    parser.add_argument(
        "--allow-launch",
        action="store_true",
        help=(
            "Unlock cloud job submission and mutating registry repair tools "
            "(`launch_or_repair_run` and `stage_only`)."
        ),
    )
    parser.add_argument(
        "--probe-env",
        action="store_true",
        help="Print the environment readiness JSON report to stdout and exit.",
    )
    args = parser.parse_args(argv)

    if args.probe_env:
        print(_dumps(probe_environment()))
        return 0

    server = McpServer(allow_launch=args.allow_launch)
    server.serve_stdio()
    return 0