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

"""Closed, minimized contract for MCP calls made by governed Agents."""

from __future__ import annotations

import copy
import re
import unicodedata
from typing import Any


class InvocationContractError(ValueError):
    """Raised when an MCP invocation exceeds its declared contract."""


_ALLOWED = frozenset(
    {"interface_type", "tool_name", "action", "arguments_digest", "evidence_refs"}
)
_FORBIDDEN_KEYS = frozenset(
    {
        "url", "uri", "endpoint", "path", "file_path", "filename", "command",
        "shell", "script", "secret", "token", "password", "api_key", "raw_rows",
        "rows", "records", "arguments", "payload",
    }
)
_TOOL = re.compile(r"^[A-Za-z][A-Za-z0-9_.:-]{1,159}$")
_DIGEST = re.compile(r"^[a-f0-9]{64}$")


def _closed(value: Any, allowed: frozenset[str], label: str) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise InvocationContractError(f"{label}_invalid")
    for key in value:
        normalized = unicodedata.normalize("NFKC", str(key))
        if normalized != key or normalized.casefold() in _FORBIDDEN_KEYS:
            raise InvocationContractError("forbidden_contract_key")
    unknown = set(value) - allowed
    if unknown:
        raise InvocationContractError("unknown_contract_field")
    return copy.deepcopy(value)


def _digest(value: Any, label: str) -> str:
    if not isinstance(value, str) or not _DIGEST.fullmatch(value):
        raise InvocationContractError(f"{label}_invalid")
    return value


def normalize_mcp_invocation(value: Any) -> dict[str, Any]:
    """Accept hash-only, bounded MCP metadata; never raw arguments or locations."""
    body = _closed(value, _ALLOWED, "invocation")
    if body.get("interface_type") != "mcp":
        raise InvocationContractError("interface_type_invalid")
    tool_name = body.get("tool_name")
    if not isinstance(tool_name, str) or not _TOOL.fullmatch(tool_name):
        raise InvocationContractError("tool_name_invalid")
    if body.get("action") not in {"read", "suggest", "execute"}:
        raise InvocationContractError("action_invalid")
    refs = body.get("evidence_refs")
    if not isinstance(refs, list) or len(refs) > 20:
        raise InvocationContractError("evidence_refs_invalid")
    normalized_refs = []
    for item in refs:
        ref = _closed(item, frozenset({"evidence_id", "digest"}), "evidence_ref")
        evidence_id = ref.get("evidence_id")
        if not isinstance(evidence_id, str) or not re.fullmatch(r"[A-Za-z0-9_.:-]{1,120}", evidence_id):
            raise InvocationContractError("evidence_id_invalid")
        normalized_refs.append({"evidence_id": evidence_id, "digest": _digest(ref.get("digest"), "evidence_digest")})
    return {
        "interface_type": "mcp",
        "tool_name": tool_name,
        "action": body["action"],
        "arguments_digest": _digest(body.get("arguments_digest"), "arguments_digest"),
        "evidence_refs": normalized_refs,
    }