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