DataOps-platform/app/core/mcp/runtime.py

"""Fail-closed stdio bootstrap for DataOps MCP servers."""

from __future__ import annotations

import argparse
import importlib
import os
import re
import uuid
from collections.abc import Mapping

from app.core.mcp.identity import AgentIdentity
from app.core.mcp.servers import create_context_mcp, create_scheduling_mcp

ALLOWED_ROLES = {"viewer", "editor", "scheduler", "admin"}
TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9_.:*@-]{1,200}$")


def _required(environ, name):
    value = str(environ.get(name) or "").strip()
    if not value:
        raise ValueError(f"{name} is required")
    return value


def _bounded_set(environ, name, *, allowed=None):
    values = {
        item.strip() for item in _required(environ, name).split(",") if item.strip()
    }
    if not values or len(values) > 50:
        raise ValueError(f"{name} must contain between 1 and 50 values")
    if any(not TOKEN_PATTERN.fullmatch(value) for value in values):
        raise ValueError(f"{name} contains an invalid value")
    if allowed is not None and not values <= allowed:
        raise ValueError(f"{name} contains an unsupported role")
    return frozenset(values)


def identity_from_env(environ: Mapping[str, str] | None = None):
    source = environ if environ is not None else os.environ
    subject = _required(source, "DATAOPS_MCP_SUBJECT")
    if not TOKEN_PATTERN.fullmatch(subject):
        raise ValueError("DATAOPS_MCP_SUBJECT contains an invalid value")
    correlation_id = _required(source, "DATAOPS_MCP_CORRELATION_ID")
    try:
        correlation_id = str(uuid.UUID(correlation_id))
    except ValueError as exc:
        raise ValueError("DATAOPS_MCP_CORRELATION_ID must be a UUID") from exc
    return AgentIdentity(
        subject=subject,
        roles=_bounded_set(source, "DATAOPS_MCP_ROLES", allowed=ALLOWED_ROLES),
        business_domains=_bounded_set(source, "DATAOPS_MCP_BUSINESS_DOMAINS"),
        environments=_bounded_set(source, "DATAOPS_MCP_ENVIRONMENTS"),
        correlation_id=correlation_id,
    )


def load_factory(path, *, allow_test_factory=False):
    if not isinstance(path, str) or ":" not in path:
        raise ValueError("MCP factory must use module:function syntax")
    module_name, function_name = path.rsplit(":", 1)
    allowed_prefixes = ("app.",)
    if allow_test_factory:
        allowed_prefixes += ("tests.mcp.",)
    if not module_name.startswith(allowed_prefixes):
        raise ValueError("MCP factory module is outside the allowlist")
    if not function_name.isidentifier():
        raise ValueError("MCP factory function is invalid")
    module = importlib.import_module(module_name)
    factory = getattr(module, function_name, None)
    if not callable(factory):
        raise ValueError("MCP factory is not callable")
    return factory


def build_server(
    *,
    kind,
    factory_path,
    environ=None,
    allow_test_factory=False,
):
    identity = identity_from_env(environ)
    service = load_factory(factory_path, allow_test_factory=allow_test_factory)()
    if kind == "context":
        return create_context_mcp(service, identity=identity)
    if kind == "scheduling":
        return create_scheduling_mcp(service, identity=identity)
    raise ValueError("MCP kind must be context or scheduling")


def main(argv=None):
    parser = argparse.ArgumentParser(description="Run a DataOps MCP server")
    parser.add_argument("--kind", choices=("context", "scheduling"), required=True)
    parser.add_argument(
        "--factory",
        default=os.getenv("DATAOPS_MCP_FACTORY"),
        required=os.getenv("DATAOPS_MCP_FACTORY") is None,
    )
    parser.add_argument(
        "--allow-test-factory",
        action="store_true",
        help=argparse.SUPPRESS,
    )
    arguments = parser.parse_args(argv)
    server = build_server(
        kind=arguments.kind,
        factory_path=arguments.factory,
        allow_test_factory=arguments.allow_test_factory,
    )
    server.run(transport="stdio")