MCP Tool Governance

┌──────────────────────────────────────────────────────────────────┐
│  Orchestrator / Agent                                             │
│  (LangChain, AutoGen, custom, etc.)                               │
│                                                                   │
│   1. Agent decides to call a tool (LLM completion)                │
│   2. MCP client invokes tool via transport                        │
└────────────────────┬─────────────────────────────────────────────┘
                     │  MCP tool/call request
                     ▼
┌──────────────────────────────────────────────────────────────────┐
│  Cognisafe MCP Intercept Layer                                    │
│  - Logs tool name + arguments → llm_requests                     │
│  - Tags: agent_name, session_id, tool_name                        │
│  - Runs synchronous policy checks (keyword_list scorer)           │
│  - Fires threat_detected webhook if policy violated               │
│  - Passes through (or blocks) to MCP server                       │
└────────────────────┬─────────────────────────────────────────────┘
                     │
                     ▼
┌──────────────────────────────────────────────────────────────────┐
│  MCP Server                                                       │
│  - filesystem, database, API, code executor, etc.                 │
│  - Executes tool; returns result                                  │
└────────────────────┬─────────────────────────────────────────────┘
                     │  Tool result
                     ▼
┌──────────────────────────────────────────────────────────────────┐
│  Cognisafe intercept (response path)                              │
│  - Logs tool result → response_body                               │
│  - Scores result for PII / sensitive data patterns                │
└──────────────────────────────────────────────────────────────────┘

import time
import json
import httpx
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

COGNISAFE_API_URL = "http://localhost:8000"
COGNISAFE_API_KEY = "csk_..."
PROJECT_ID = "mcp-agent"
AGENT_NAME = "file-assistant"


async def log_tool_call(
    tool_name: str,
    arguments: dict,
    result: dict,
    duration_ms: float,
    session_id: str,
):
    """Fire-and-forget log to Cognisafe API."""
    payload = {
        "project_id": PROJECT_ID,
        "agent_name": AGENT_NAME,
        "session_id": session_id,
        "model": "mcp-tool",
        "request_body": {
            "type": "mcp_tool_call",
            "tool_name": tool_name,
            "arguments": arguments,
        },
        "response_body": result,
        "latency_ms": duration_ms,
    }
    async with httpx.AsyncClient() as client:
        await client.post(
            f"{COGNISAFE_API_URL}/internal/log",
            json=payload,
            headers={"Authorization": f"Bearer {COGNISAFE_API_KEY}"},
            timeout=2.0,
        )


class GovernedMCPSession:
    """
    Wraps an MCP ClientSession to log every tool invocation to Cognisafe
    and apply synchronous policy checks before forwarding to the MCP server.
    """

    BLOCKED_TOOLS = {"execute_command", "run_script", "drop_table"}

    def __init__(self, session: ClientSession, session_id: str):
        self._session = session
        self._session_id = session_id

    async def call_tool(self, tool_name: str, arguments: dict):
        # Synchronous policy gate — block before the tool runs
        if tool_name in self.BLOCKED_TOOLS:
            raise PermissionError(
                f"Tool '{tool_name}' is blocked by governance policy. "
                f"Session: {self._session_id}"
            )

        start = time.monotonic()
        result = await self._session.call_tool(tool_name, arguments)
        duration_ms = (time.monotonic() - start) * 1000

        # Non-blocking log — do not await in the hot path
        import asyncio
        asyncio.create_task(
            log_tool_call(
                tool_name=tool_name,
                arguments=arguments,
                result=result.model_dump() if hasattr(result, "model_dump") else result,
                duration_ms=duration_ms,
                session_id=self._session_id,
            )
        )

        return result


# Usage
async def main():
    server_params = StdioServerParameters(
        command="python",
        args=["mcp_server.py"],
    )

    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()

            governed = GovernedMCPSession(session, session_id="sess_abc123")

            # This call is logged to Cognisafe and checked against policy
            result = await governed.call_tool(
                "read_file",
                {"path": "/data/reports/q1.csv"},
            )
            print(result)

- name: mcp_pii_in_args
  type: regex
  severity: high
  description: "Detects PII patterns in MCP tool call arguments"
  patterns:
    - '\b\d{3}-\d{2}-\d{4}\b'          # SSN
    - '\b[A-Z]{2}\d{6}[A-Z]\b'         # UK passport
    - '\b\d{13,16}\b'                    # Payment card
    - '[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}'  # Email
  match_on: request_body

HITL_REQUIRED_TOOLS = {"create_payment", "issue_refund", "wire_transfer"}

async def governed_call(governed_session, tool_name, arguments):
    if tool_name in HITL_REQUIRED_TOOLS:
        approval = await request_human_approval(
            tool_name=tool_name,
            arguments=arguments,
            approver_channel="slack://ops-approvals",
        )
        if not approval.approved:
            raise PermissionError(f"Human declined {tool_name} invocation")

    return await governed_session.call_tool(tool_name, arguments)