From f534a2b078699c1573b3429304167e15c572afce Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Fri, 24 Jul 2026 15:43:59 +0530 Subject: [PATCH 01/16] feat(mcp): host an MCP server for Unstract API deployments Exposes an Unstract API deployment to coding agents over the Model Context Protocol, so an agent can run document extraction as a tool call instead of hand-rolling HTTP requests. Follows the hosted-MCP pattern used in the mfbt backend: a hand-rolled JSON-RPC 2.0 endpoint mounted in the existing app (not a separate service), with a declarative tool registry behind it. Endpoint mirrors the deployment's own REST URL and reuses its API key, so there is no second credential to mint or revoke: POST /deployment/api/// # REST POST /mcp/// # MCP Tools: readMeFirst, getApiInfo, extractDocument, getExecutionStatus. Auth goes through the same DeploymentHelper validation the REST endpoint uses, and execution through ExecutionRequestSerializer, so the MCP surface cannot drift from the REST one on who is allowed in or what input is valid. All auth failures answer identically (401, no detail) so the endpoint cannot be used to enumerate deployment names. OAuth 2.1 with dynamic client registration is deliberately not implemented; bearer auth covers Claude Code and API clients, and OAuth would be additive. Tests: 20 passing (protocol/dispatch in the unit tier, auth boundary in the integration tier). Registers the mcp-server-auth critical path. Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/backend/base_urls.py | 2 + backend/backend/settings/base.py | 8 + backend/mcp_v2/README.md | 109 ++++++++ backend/mcp_v2/__init__.py | 0 backend/mcp_v2/apps.py | 5 + backend/mcp_v2/constants.py | 41 +++ backend/mcp_v2/context.py | 27 ++ backend/mcp_v2/exceptions.py | 15 ++ backend/mcp_v2/registry.py | 154 +++++++++++ backend/mcp_v2/tests/__init__.py | 0 backend/mcp_v2/tests/test_mcp_auth.py | 159 +++++++++++ backend/mcp_v2/tests/test_mcp_protocol.py | 216 +++++++++++++++ backend/mcp_v2/tests/test_tool_errors.py | 44 ++++ backend/mcp_v2/tools/__init__.py | 0 backend/mcp_v2/tools/execution.py | 305 ++++++++++++++++++++++ backend/mcp_v2/tools/info.py | 83 ++++++ backend/mcp_v2/urls.py | 31 +++ backend/mcp_v2/views.py | 292 +++++++++++++++++++++ tests/critical_paths.yaml | 5 + 19 files changed, 1496 insertions(+) create mode 100644 backend/mcp_v2/README.md create mode 100644 backend/mcp_v2/__init__.py create mode 100644 backend/mcp_v2/apps.py create mode 100644 backend/mcp_v2/constants.py create mode 100644 backend/mcp_v2/context.py create mode 100644 backend/mcp_v2/exceptions.py create mode 100644 backend/mcp_v2/registry.py create mode 100644 backend/mcp_v2/tests/__init__.py create mode 100644 backend/mcp_v2/tests/test_mcp_auth.py create mode 100644 backend/mcp_v2/tests/test_mcp_protocol.py create mode 100644 backend/mcp_v2/tests/test_tool_errors.py create mode 100644 backend/mcp_v2/tools/__init__.py create mode 100644 backend/mcp_v2/tools/execution.py create mode 100644 backend/mcp_v2/tools/info.py create mode 100644 backend/mcp_v2/urls.py create mode 100644 backend/mcp_v2/views.py diff --git a/backend/backend/base_urls.py b/backend/backend/base_urls.py index 8450add04d..353a253cec 100644 --- a/backend/backend/base_urls.py +++ b/backend/backend/base_urls.py @@ -22,6 +22,8 @@ f"{settings.API_DEPLOYMENT_PATH_PREFIX}/pipeline/", include("pipeline_v2.public_api_urls"), ), + # Hosted MCP server, authenticated by the API deployment's own key + path(f"{settings.MCP_PATH_PREFIX}/", include("mcp_v2.urls")), path("", include("health.urls")), # Internal API for worker communication path("internal/", include("backend.internal_base_urls")), diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 94071ac3b8..55c069998a 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -118,6 +118,9 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | API_DEPLOYMENT_PATH_PREFIX = os.environ.get( "API_DEPLOYMENT_PATH_PREFIX", "deployment" ).strip("/") +# Mount point for the hosted MCP server. Changing this invalidates the MCP +# endpoint URLs already configured in users' MCP clients. +MCP_PATH_PREFIX = os.environ.get("MCP_PATH_PREFIX", "mcp").strip("/") # Maximum file size for presigned URLs in API deployments (in MB) API_DEPL_PRESIGNED_URL_MAX_FILE_SIZE_MB = int( @@ -379,6 +382,7 @@ def filter(self, record): "pipeline_v2", "platform_settings_v2", "api_v2", + "mcp_v2", "usage_v2", "notification_v2", "prompt_studio.prompt_profile_manager_v2", @@ -650,6 +654,10 @@ def filter(self, record): # White lists workflow-api-deployment path WHITELISTED_PATHS.append(f"/{API_DEPLOYMENT_PATH_PREFIX}") +# White lists the hosted MCP server, which authenticates with the API +# deployment's own key rather than a session. +WHITELISTED_PATHS.append(f"/{MCP_PATH_PREFIX}") + # Whitelisting health check API WHITELISTED_PATHS.append("/health") diff --git a/backend/mcp_v2/README.md b/backend/mcp_v2/README.md new file mode 100644 index 0000000000..7e47f45ac0 --- /dev/null +++ b/backend/mcp_v2/README.md @@ -0,0 +1,109 @@ +# Hosted MCP server + +Exposes an Unstract API deployment to coding agents over the +[Model Context Protocol](https://modelcontextprotocol.io), so an agent can run +document extraction as a tool call instead of hand-rolling HTTP requests. + +The server is hosted inside the existing backend — it is a Django app served by +the same gunicorn process, not a separate service to deploy or scale. + +## Endpoint + +An MCP session is scoped to exactly one API deployment, and mirrors the URL +shape of that deployment's REST endpoint: + +``` +POST /deployment/api/// # REST +POST /mcp/// # MCP +``` + +Authentication is the deployment's **existing API key** — the same key used for +REST execution, managed from the same place in the UI. There is no separate MCP +credential to mint or revoke. + +``` +Authorization: Bearer +``` + +For MCP clients that cannot attach custom headers, the key may instead be given +as a path segment: + +``` +POST /mcp//// +``` + +The path key takes precedence over the header when both are present. + +`GET` on either URL returns server identity for clients that probe before +connecting. It is unauthenticated and reveals nothing about the deployment. + +## Connecting + +With Claude Code: + +```bash +claude mcp add --transport http unstract \ + https:///mcp/// \ + --header "Authorization: Bearer " +``` + +## Tools + +| Tool | Purpose | +| --- | --- | +| `readMeFirst` | Orientation guide, built from the live deployment. Call first. | +| `getApiInfo` | Name, description, workflow and active state of the deployment. | +| `extractDocument` | Run extraction over document URLs. **Consumes quota.** | +| `getExecutionStatus` | Poll for the result of a pending extraction. | + +Documents are passed as URLs and fetched server-side, so they must be reachable +by the backend — typically a pre-signed URL. Extraction is asynchronous: +`extractDocument` returns an `execution_id` when it does not finish within the +timeout, and the agent polls `getExecutionStatus` with it. + +## Adding a tool + +Write a handler taking `MCPContext` as its first argument, then register it in +`registry.py` with a JSON schema: + +```python +registry.register( + MCPTool( + name="myTool", + description="What it does, written for an LLM to read.", + input_schema={"type": "object", "properties": {...}, "required": [...]}, + handler=my_tool, + writes=False, + ) +) +``` + +Tool descriptions are prompts, not documentation — they are the only guidance +the calling agent gets. Say what the tool does, when to use it, and what it +costs. + +Raise `MCPToolError` for failures the agent can act on (bad arguments, inactive +deployment, rate limit); the message reaches the agent verbatim, so write it as +an instruction. Any other exception is logged and reported generically so +internal detail does not leak to the client. + +## Design notes + +- **Auth reuses the deployment key path.** `_resolve_context` calls the same + `DeploymentHelper` validation the REST endpoint uses, so the two surfaces + cannot drift apart on who is allowed in. +- **Execution reuses `ExecutionRequestSerializer`.** URL validation and the + file-count cap live there; reimplementing them for MCP would let the MCP + surface quietly diverge from the REST one. +- **All auth failures answer identically** (401, no detail), so the endpoint + cannot be used to enumerate deployment names. +- **Tool errors are JSON-RPC results, not protocol errors.** Clients treat + protocol errors as unrecoverable transport faults; an agent-fixable problem + comes back as `isError: true` content it can read and retry. + +## Not implemented + +OAuth 2.1 with dynamic client registration, which MCP defines for browser-based +one-click connectors. Header/path bearer auth covers Claude Code and +API clients. Adding OAuth is additive — it would mount discovery endpoints +alongside this router without changing the transport. diff --git a/backend/mcp_v2/__init__.py b/backend/mcp_v2/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/mcp_v2/apps.py b/backend/mcp_v2/apps.py new file mode 100644 index 0000000000..fce3156468 --- /dev/null +++ b/backend/mcp_v2/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class MCPConfig(AppConfig): + name = "mcp_v2" diff --git a/backend/mcp_v2/constants.py b/backend/mcp_v2/constants.py new file mode 100644 index 0000000000..baa838651f --- /dev/null +++ b/backend/mcp_v2/constants.py @@ -0,0 +1,41 @@ +"""Constants for the hosted MCP server.""" + + +class MCPServer: + """Identity advertised to MCP clients during `initialize`.""" + + NAME: str = "unstract" + VERSION: str = "0.1.0" + # Protocol revision this transport implements. Clients that request a + # different revision are answered with this one, per the MCP spec's + # version-negotiation rules. + PROTOCOL_VERSION: str = "2024-11-05" + + +class JSONRPC: + """JSON-RPC 2.0 wire constants.""" + + VERSION: str = "2.0" + + # Standard error codes (JSON-RPC 2.0 spec). + PARSE_ERROR: int = -32700 + INVALID_REQUEST: int = -32600 + METHOD_NOT_FOUND: int = -32601 + INVALID_PARAMS: int = -32602 + INTERNAL_ERROR: int = -32603 + + # Implementation-defined codes (-32000 to -32099 reserved for servers). + UNAUTHORIZED: int = -32001 + TOOL_EXECUTION_ERROR: int = -32002 + + +class MCPMethod: + """MCP methods handled by this transport.""" + + INITIALIZE: str = "initialize" + PING: str = "ping" + TOOLS_LIST: str = "tools/list" + TOOLS_CALL: str = "tools/call" + + # Notifications carry no `id` and expect no response body. + NOTIFICATION_PREFIX: str = "notifications/" diff --git a/backend/mcp_v2/context.py b/backend/mcp_v2/context.py new file mode 100644 index 0000000000..13e0d33a26 --- /dev/null +++ b/backend/mcp_v2/context.py @@ -0,0 +1,27 @@ +"""Resolved request context handed to MCP tool handlers. + +Auth, organization resolution and deployment lookup all happen once in the +transport view. Tools receive the result and never repeat that work. +""" + +from dataclasses import dataclass + +from api_v2.models import APIDeployment + + +@dataclass(frozen=True) +class MCPContext: + """Everything a tool needs about the caller's deployment. + + Attributes: + api: The API deployment this MCP server session is scoped to. + api_key: The validated API key used to authenticate the request. + Carried through because downstream execution helpers record it + against the execution. + org_name: Organization identifier taken from the URL, matching the + value already placed in the state store by the auth layer. + """ + + api: APIDeployment + api_key: str + org_name: str diff --git a/backend/mcp_v2/exceptions.py b/backend/mcp_v2/exceptions.py new file mode 100644 index 0000000000..3200cba49d --- /dev/null +++ b/backend/mcp_v2/exceptions.py @@ -0,0 +1,15 @@ +"""Exceptions raised by MCP tool handlers.""" + + +class MCPToolError(Exception): + """A tool failed in a way the calling agent can understand and act on. + + Raised for conditions the agent could plausibly recover from — bad + arguments, an inactive deployment, a rate limit. The transport turns these + into a JSON-RPC error carrying the message verbatim, so the message is + prompt-facing: state what went wrong and what the agent should do next. + + Unexpected exceptions are deliberately *not* funnelled through this type; + they are logged and reported generically so internal detail does not leak + to the client. + """ diff --git a/backend/mcp_v2/registry.py b/backend/mcp_v2/registry.py new file mode 100644 index 0000000000..cb03f29f1d --- /dev/null +++ b/backend/mcp_v2/registry.py @@ -0,0 +1,154 @@ +"""Tool registry for the hosted MCP server. + +Tools are plain callables registered with the JSON schema MCP clients need in +order to call them. Each tool receives the resolved deployment context as its +first argument, so tool implementations never re-do auth or org resolution. +""" + +import logging +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class MCPTool: + """A single tool exposed over the MCP transport. + + Attributes: + name: Tool name as seen by the MCP client. + description: Prompt-facing description. Written for an LLM audience — + it is the only guidance the calling agent gets. + input_schema: JSON schema for the tool's arguments. + handler: Callable invoked as ``handler(context, **arguments)``. + writes: True when the tool has side effects (consumes quota, starts an + execution). Read-only tools are safe to retry; write tools are not. + """ + + name: str + description: str + input_schema: dict[str, Any] + handler: Callable[..., Any] + writes: bool = False + + def to_mcp_schema(self) -> dict[str, Any]: + """Serialize to the shape returned by `tools/list`.""" + return { + "name": self.name, + "description": self.description, + "inputSchema": self.input_schema, + } + + +@dataclass +class MCPToolRegistry: + """Ordered name -> tool mapping. + + Ordering is preserved so `tools/list` presents tools in the order they were + registered; agents weight earlier tools more heavily, and `readMeFirst` + needs to come first. + """ + + _tools: dict[str, MCPTool] = field(default_factory=dict) + + def register(self, tool: MCPTool) -> None: + if tool.name in self._tools: + raise ValueError(f"Duplicate MCP tool registration: '{tool.name}'") + self._tools[tool.name] = tool + + def get(self, name: str) -> MCPTool | None: + return self._tools.get(name) + + def names(self) -> list[str]: + return list(self._tools) + + def list_schemas(self) -> list[dict[str, Any]]: + return [tool.to_mcp_schema() for tool in self._tools.values()] + + +def build_registry() -> MCPToolRegistry: + """Build the registry of tools exposed by the Unstract MCP server. + + Imported lazily inside the function so that registering a tool cannot + trigger Django model imports at module-import time. + """ + from mcp_v2.tools.execution import ( + extract_document, + extract_document_schema, + get_execution_status, + get_execution_status_schema, + ) + from mcp_v2.tools.info import get_api_info, get_api_info_schema, read_me_first + + registry = MCPToolRegistry() + + registry.register( + MCPTool( + name="readMeFirst", + description=( + "START HERE. Returns a guide to this MCP server: what the " + "connected Unstract API deployment does, the available tools, " + "and the recommended call sequence. Takes no arguments." + ), + input_schema={"type": "object", "properties": {}, "required": []}, + handler=read_me_first, + ) + ) + registry.register( + MCPTool( + name="getApiInfo", + description=( + "Get details of the Unstract API deployment this MCP server is " + "connected to: its display name, description, the workflow it " + "runs, and whether it is active. Call this to learn what kind of " + "document the deployment expects before extracting. " + "Takes no arguments." + ), + input_schema=get_api_info_schema(), + handler=get_api_info, + ) + ) + registry.register( + MCPTool( + name="extractDocument", + description=( + "Run the connected Unstract API deployment over one or more " + "documents and return the structured extraction result.\n\n" + "Documents are supplied as URLs, which Unstract fetches " + "server-side. Extraction is asynchronous: when it does not " + "finish within `timeout` seconds this returns " + "`execution_status: PENDING` along with an `execution_id`. Poll " + "`getExecutionStatus` with that id to collect the result.\n\n" + "This consumes the organization's extraction quota — do not call " + "it speculatively or retry a call that already returned an " + "execution_id." + ), + input_schema=extract_document_schema(), + handler=extract_document, + writes=True, + ) + ) + registry.register( + MCPTool( + name="getExecutionStatus", + description=( + "Fetch the status and, once available, the result of an " + "extraction previously started by `extractDocument`. Pass the " + "`execution_id` that call returned.\n\n" + "Returns an `execution_status` of PENDING, EXECUTING, COMPLETED " + "or ERROR. Poll while it is PENDING or EXECUTING, leaving a few " + "seconds between calls." + ), + input_schema=get_execution_status_schema(), + handler=get_execution_status, + ) + ) + + return registry + + +# Single shared registry. The tool set is static, so building it once at import +# time is safe and keeps per-request work down. +TOOL_REGISTRY = build_registry() diff --git a/backend/mcp_v2/tests/__init__.py b/backend/mcp_v2/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/mcp_v2/tests/test_mcp_auth.py b/backend/mcp_v2/tests/test_mcp_auth.py new file mode 100644 index 0000000000..c92fffda07 --- /dev/null +++ b/backend/mcp_v2/tests/test_mcp_auth.py @@ -0,0 +1,159 @@ +"""Critical path ``mcp-server-auth``: the hosted MCP endpoint rejects +unauthenticated callers before any tool runs. + +``POST /mcp///`` is unauthenticated at the DRF layer — like the +deployment execution endpoint it mirrors, it is guarded entirely by the API key +check inside the view. A regression there exposes every registered tool, +including the one that spends the organization's extraction quota. These tests +assert rejection lands before the tool handler is reached. Needs a live DB +(integration tier). +""" + +from __future__ import annotations + +import json +import uuid +from unittest.mock import patch + +import pytest +from account_v2.models import Organization +from api_v2.models import APIDeployment, APIKey +from django.test import TestCase +from rest_framework.test import APIRequestFactory +from utils.user_context import UserContext +from workflow_manager.workflow_v2.models.workflow import Workflow + +from mcp_v2.views import MCPServerView + +ORG_ID = "org-mcp" + + +def rpc_body(response): + """Decode the JSON-RPC envelope from a rendered response.""" + return json.loads(response.content) + + +class MCPServerAuthTest(TestCase): + def setUp(self) -> None: + self.org = Organization.objects.create( + name=ORG_ID, display_name="Org MCP", organization_id=ORG_ID + ) + UserContext.set_organization_identifier(ORG_ID) + workflow = Workflow.objects.create(workflow_name="wf-mcp", is_active=True) + + self.api = APIDeployment.objects.create(api_name="live-api", workflow=workflow) + self.key = APIKey.objects.create(api=self.api) + self.inactive_key = APIKey.objects.create(api=self.api, is_active=False) + + self.inactive_api = APIDeployment.objects.create( + api_name="dead-api", workflow=workflow, is_active=False + ) + self.other_api = APIDeployment.objects.create( + api_name="other-api", workflow=workflow + ) + self.other_key = APIKey.objects.create(api=self.other_api) + + self.view = MCPServerView.as_view() + self.factory = APIRequestFactory() + + def _post( + self, + api_name: str, + auth: str | None, + org: str = ORG_ID, + payload: dict | None = None, + path_key: str | None = None, + ): + headers = {"HTTP_AUTHORIZATION": auth} if auth is not None else {} + body = payload if payload is not None else { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/list", + } + request = self.factory.post( + f"/mcp/{org}/{api_name}/", body, format="json", **headers + ) + kwargs = {"org_name": org, "api_name": api_name} + if path_key is not None: + kwargs["api_key"] = path_key + return self.view(request, **kwargs) + + @pytest.mark.critical_path("mcp-server-auth") + @patch("mcp_v2.views.TOOL_REGISTRY.get") + def test_bad_credentials_rejected_before_dispatch(self, registry_get) -> None: + cases = [ + ("missing header", "live-api", None, ORG_ID), + ("no bearer prefix", "live-api", f"Token {self.key.api_key}", ORG_ID), + ("empty bearer", "live-api", "Bearer ", ORG_ID), + ("unknown key", "live-api", f"Bearer {uuid.uuid4()}", ORG_ID), + ("not a uuid", "live-api", "Bearer not-a-uuid", ORG_ID), + ("inactive key", "live-api", f"Bearer {self.inactive_key.api_key}", ORG_ID), + ( + "key of another api", + "live-api", + f"Bearer {self.other_key.api_key}", + ORG_ID, + ), + ("unknown api", "ghost-api", f"Bearer {self.key.api_key}", ORG_ID), + ("inactive api", "dead-api", f"Bearer {self.key.api_key}", ORG_ID), + ("wrong org", "live-api", f"Bearer {self.key.api_key}", "no-such-org"), + ] + for label, api_name, auth, org in cases: + with self.subTest(label): + response = self._post(api_name, auth, org) + assert response.status_code == 401, response.content + assert rpc_body(response)["error"]["code"] == -32001 + + registry_get.assert_not_called() + + @pytest.mark.critical_path("mcp-server-auth") + def test_valid_key_reaches_tool_listing(self) -> None: + """Guard the inverse: a check that rejected everything would pass all + the rejection cases above. + """ + response = self._post("live-api", f"Bearer {self.key.api_key}") + + assert response.status_code == 200, response.content + tools = rpc_body(response)["result"]["tools"] + assert [tool["name"] for tool in tools] == [ + "readMeFirst", + "getApiInfo", + "extractDocument", + "getExecutionStatus", + ] + + @pytest.mark.critical_path("mcp-server-auth") + def test_api_key_in_url_path_authenticates(self) -> None: + """The path-key route exists for MCP clients that cannot set headers; + it must enforce exactly the same check as the header route. + """ + ok = self._post("live-api", auth=None, path_key=str(self.key.api_key)) + assert ok.status_code == 200, ok.content + + rejected = self._post("live-api", auth=None, path_key=str(uuid.uuid4())) + assert rejected.status_code == 401, rejected.content + + @pytest.mark.critical_path("mcp-server-auth") + def test_path_key_wins_over_valid_header(self) -> None: + """The path key takes precedence, so a bad path key must not be + rescued by a good header — otherwise the precedence rule would be a + way to smuggle an unchecked credential past the check. + """ + response = self._post( + "live-api", + auth=f"Bearer {self.key.api_key}", + path_key=str(uuid.uuid4()), + ) + assert response.status_code == 401, response.content + + @pytest.mark.critical_path("mcp-server-auth") + def test_get_probe_does_not_leak_deployment_details(self) -> None: + """The unauthenticated GET probe advertises the server, not the + deployment behind it. + """ + request = self.factory.get(f"/mcp/{ORG_ID}/live-api/") + response = self.view(request, org_name=ORG_ID, api_name="live-api") + + assert response.status_code == 200 + assert response.data["name"] == "unstract" + assert "live-api" not in json.dumps(response.data) diff --git a/backend/mcp_v2/tests/test_mcp_protocol.py b/backend/mcp_v2/tests/test_mcp_protocol.py new file mode 100644 index 0000000000..e29b273f52 --- /dev/null +++ b/backend/mcp_v2/tests/test_mcp_protocol.py @@ -0,0 +1,216 @@ +"""JSON-RPC protocol behaviour of the hosted MCP server. + +Auth is stubbed here so these stay in the unit tier; the auth boundary itself +is covered by ``test_mcp_auth``. What is checked is the wire contract an MCP +client depends on: envelope shape, method routing, and the error mapping that +decides whether a client retries, re-prompts, or gives up. +""" + +from __future__ import annotations + +import json +from unittest.mock import patch + +from django.test import SimpleTestCase +from rest_framework.test import APIRequestFactory + +from mcp_v2.constants import JSONRPC, MCPServer +from mcp_v2.context import MCPContext +from mcp_v2.exceptions import MCPToolError +from mcp_v2.registry import TOOL_REGISTRY +from mcp_v2.views import MCPServerView + +ORG_ID = "org-mcp" +API_NAME = "live-api" + + +class FakeWorkflow: + id = "11111111-1111-1111-1111-111111111111" + workflow_name = "wf-mcp" + + +class FakeAPI: + """Stands in for an APIDeployment so these tests need no database.""" + + id = "22222222-2222-2222-2222-222222222222" + display_name = "Invoice Extractor" + api_name = API_NAME + description = "Extracts invoice fields" + is_active = True + workflow = FakeWorkflow() + workflow_id = FakeWorkflow.id + + +class MCPProtocolTest(SimpleTestCase): + def setUp(self) -> None: + self.view = MCPServerView.as_view() + self.factory = APIRequestFactory() + self.context = MCPContext( + api=FakeAPI(), api_key="test-key", org_name=ORG_ID + ) + + def _failing_tool(self, error: Exception, name: str = "getApiInfo"): + """Swap a registered tool for one that raises. + + The registry stores handler references captured at import, so patching + the tool's module attribute would not affect an already-built registry — + the substitution has to happen at the registry lookup. + """ + from dataclasses import replace + + def boom(*args, **kwargs): + raise error + + original = TOOL_REGISTRY.get(name) + return patch.object( + TOOL_REGISTRY, "get", return_value=replace(original, handler=boom) + ) + + def _call(self, body): + """POST a JSON-RPC body with auth stubbed out.""" + request = self.factory.post(f"/mcp/{ORG_ID}/{API_NAME}/", body, format="json") + with patch.object( + MCPServerView, "_resolve_context", return_value=self.context + ): + response = self.view( + request, org_name=ORG_ID, api_name=API_NAME, api_key="test-key" + ) + return response, json.loads(response.content) + + def test_initialize_returns_protocol_and_server_info(self) -> None: + response, body = self._call( + {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}} + ) + + assert response.status_code == 200 + assert body["id"] == 1 + result = body["result"] + assert result["protocolVersion"] == MCPServer.PROTOCOL_VERSION + assert result["serverInfo"]["name"] == MCPServer.NAME + # The agent is told which deployment it is attached to, so it does not + # have to guess what the tools will operate on. + assert "Invoice Extractor" in result["instructions"] + + def test_notification_gets_no_result_body(self) -> None: + """Notifications carry no id; answering one with a result would be a + protocol violation that some clients treat as a fatal error. + """ + response, body = self._call( + {"jsonrpc": "2.0", "method": "notifications/initialized"} + ) + + assert response.status_code == 202 + assert body == {} + + def test_ping_returns_empty_result(self) -> None: + _, body = self._call({"jsonrpc": "2.0", "id": 7, "method": "ping"}) + assert body["result"] == {} + + def test_wrong_jsonrpc_version_rejected(self) -> None: + _, body = self._call({"jsonrpc": "1.0", "id": 1, "method": "ping"}) + assert body["error"]["code"] == JSONRPC.INVALID_REQUEST + + def test_missing_method_rejected(self) -> None: + _, body = self._call({"jsonrpc": "2.0", "id": 1}) + assert body["error"]["code"] == JSONRPC.INVALID_REQUEST + + def test_unknown_method_rejected(self) -> None: + _, body = self._call( + {"jsonrpc": "2.0", "id": 1, "method": "resources/list"} + ) + assert body["error"]["code"] == JSONRPC.METHOD_NOT_FOUND + + def test_unknown_tool_lists_available_tools(self) -> None: + """The error names the real tools, so a mistaken agent can correct + itself without another round trip. + """ + _, body = self._call( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "deleteEverything"}, + } + ) + + assert body["error"]["code"] == JSONRPC.METHOD_NOT_FOUND + assert "readMeFirst" in body["error"]["data"] + + def test_tool_result_is_wrapped_in_mcp_content(self) -> None: + _, body = self._call( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "getApiInfo", "arguments": {}}, + } + ) + + result = body["result"] + assert result["isError"] is False + payload = json.loads(result["content"][0]["text"]) + assert payload["display_name"] == "Invoice Extractor" + assert payload["workflow"]["name"] == "wf-mcp" + + def test_tool_error_returns_actionable_result_not_protocol_error(self) -> None: + """An MCPToolError is the agent's problem to fix, so it comes back as + an isError result the agent can read — not a protocol error, which + many clients surface as an unrecoverable transport fault. + """ + with self._failing_tool(MCPToolError("Deployment is not active")): + _, body = self._call( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "getApiInfo", "arguments": {}}, + } + ) + + assert "error" not in body + assert body["result"]["isError"] is True + assert "not active" in body["result"]["content"][0]["text"] + + def test_unexpected_tool_failure_does_not_leak_internals(self) -> None: + with self._failing_tool( + RuntimeError("psycopg2 connection string user=admin") + ): + _, body = self._call( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "getApiInfo", "arguments": {}}, + } + ) + + assert body["error"]["code"] == JSONRPC.TOOL_EXECUTION_ERROR + assert "psycopg2" not in json.dumps(body) + + def test_unexpected_argument_reported_as_invalid_params(self) -> None: + _, body = self._call( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "getApiInfo", "arguments": {"nope": 1}}, + } + ) + + assert body["error"]["code"] == JSONRPC.INVALID_PARAMS + + def test_batch_request_rejected(self) -> None: + request = self.factory.post( + f"/mcp/{ORG_ID}/{API_NAME}/", + [{"jsonrpc": "2.0", "id": 1, "method": "ping"}], + format="json", + ) + with patch.object( + MCPServerView, "_resolve_context", return_value=self.context + ): + response = self.view( + request, org_name=ORG_ID, api_name=API_NAME, api_key="test-key" + ) + + body = json.loads(response.content) + assert body["error"]["code"] == JSONRPC.INVALID_REQUEST diff --git a/backend/mcp_v2/tests/test_tool_errors.py b/backend/mcp_v2/tests/test_tool_errors.py new file mode 100644 index 0000000000..5f4a9c8d28 --- /dev/null +++ b/backend/mcp_v2/tests/test_tool_errors.py @@ -0,0 +1,44 @@ +"""Tool-layer error messages, which are read by an agent rather than a human. + +A tool error is the agent's next prompt: if it does not name the offending +argument in plain text, the agent's most likely next move is to retry the same +call — and for `extractDocument` a retry spends the organization's quota. +""" + +from __future__ import annotations + +from django.test import SimpleTestCase +from rest_framework.exceptions import ValidationError + +from mcp_v2.tools.execution import _format_validation_error + + +class FormatValidationErrorTest(SimpleTestCase): + def test_flattens_nested_detail_without_errordetail_repr(self) -> None: + """DRF nests per-item errors under the field name and wraps each in an + ``ErrorDetail``; both would otherwise reach the agent verbatim. + """ + error = ValidationError({"presigned_urls": {0: ["Enter a valid URL."]}}) + + message = _format_validation_error(error) + + assert message == "presigned_urls.0: Enter a valid URL." + assert "ErrorDetail" not in message + + def test_reports_every_failing_field(self) -> None: + """Surfacing only the first failure would send the agent round the loop + once per bad argument. + """ + error = ValidationError( + {"timeout": ["Ensure this value is at most 300."], "tags": ["Invalid tag."]} + ) + + message = _format_validation_error(error) + + assert "timeout: Ensure this value is at most 300." in message + assert "tags: Invalid tag." in message + + def test_plain_message_passes_through_unchanged(self) -> None: + error = ValidationError("Something was wrong.") + + assert _format_validation_error(error) == "Something was wrong." diff --git a/backend/mcp_v2/tools/__init__.py b/backend/mcp_v2/tools/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/mcp_v2/tools/execution.py b/backend/mcp_v2/tools/execution.py new file mode 100644 index 0000000000..0b22c49dd7 --- /dev/null +++ b/backend/mcp_v2/tools/execution.py @@ -0,0 +1,305 @@ +"""Tools that run the connected API deployment and read back its results.""" + +import logging +import uuid +from typing import Any + +from api_v2.constants import ApiExecution +from api_v2.deployment_helper import DeploymentHelper +from api_v2.dto import DeploymentExecutionDTO +from api_v2.exceptions import RateLimitExceeded +from api_v2.rate_limiter import APIDeploymentRateLimiter +from api_v2.serializers import ExecutionQuerySerializer, ExecutionRequestSerializer +from rest_framework.exceptions import ValidationError +from utils.enums import CeleryTaskState +from workflow_manager.workflow_v2.dto import ExecutionResponse + +from mcp_v2.context import MCPContext +from mcp_v2.exceptions import MCPToolError + +logger = logging.getLogger(__name__) + +# Mirrors ExecutionRequestSerializer.MAX_FILES_ALLOWED. Declared in the JSON +# schema too, so a well-behaved client is stopped before the request is made. +MAX_DOCUMENTS = ExecutionRequestSerializer.MAX_FILES_ALLOWED + +# An agent polling with getExecutionStatus does not benefit from holding the +# HTTP request open for the full 300s the REST API permits, and a long-held +# MCP call reads as a hang. Default to a short wait and let the agent poll. +DEFAULT_TIMEOUT_SEC = 30 + + +def _format_validation_error(error: ValidationError) -> str: + """Flatten DRF validation detail into a sentence an agent can act on. + + ``error.detail`` stringifies with ``ErrorDetail(string=..., code=...)`` + wrappers and nested dict/list structure. That is noise the agent has to + parse around before it can see which argument was wrong, so flatten it to + ``field: message`` pairs. + """ + + def walk(node: Any, path: str = "") -> list[str]: + if isinstance(node, dict): + return [ + message + for key, value in node.items() + for message in walk(value, f"{path}.{key}" if path else str(key)) + ] + if isinstance(node, list): + return [message for item in node for message in walk(item, path)] + return [f"{path}: {node}" if path else str(node)] + + return "; ".join(walk(error.detail)) + + +def extract_document_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "document_urls": { + "type": "array", + "items": {"type": "string"}, + "minItems": 1, + "maxItems": MAX_DOCUMENTS, + "description": ( + "URLs of the documents to extract. Unstract fetches these " + "server-side, so each must be reachable by the Unstract " + "backend — typically a pre-signed URL." + ), + }, + "timeout": { + "type": "integer", + "minimum": -1, + "maximum": ApiExecution.MAXIMUM_TIMEOUT_IN_SEC, + "default": DEFAULT_TIMEOUT_SEC, + "description": ( + "Seconds to wait for the extraction to finish before " + "returning a pending result. Use -1 to return immediately " + "and poll with getExecutionStatus." + ), + }, + "include_metadata": { + "type": "boolean", + "default": False, + "description": "Include extraction metadata in the result.", + }, + "include_metrics": { + "type": "boolean", + "default": False, + "description": "Include token and cost metrics in the result.", + }, + "include_extracted_text": { + "type": "boolean", + "default": False, + "description": ( + "Include the full raw text extracted from each document, " + "alongside the structured output. This can be very large." + ), + }, + "tags": { + "type": "array", + "items": {"type": "string"}, + "description": "Tags to associate with this execution.", + }, + "llm_profile_id": { + "type": "string", + "description": ( + "UUID of an LLM profile to override the deployment's " + "default. Omit to use the configured default." + ), + }, + }, + "required": ["document_urls"], + } + + +def extract_document( + context: MCPContext, + document_urls: list[str], + timeout: int = DEFAULT_TIMEOUT_SEC, + include_metadata: bool = False, + include_metrics: bool = False, + include_extracted_text: bool = False, + tags: list[str] | None = None, + llm_profile_id: str | None = None, +) -> dict[str, Any]: + """Run the deployment's workflow over the supplied document URLs. + + Deliberately routed through ``ExecutionRequestSerializer`` rather than + calling the execution helper directly: that serializer owns URL validation + (HTTPS/endpoint restrictions) and the file-count cap, and reimplementing + those here would let the MCP surface drift away from the REST surface it + is meant to mirror. + """ + if not context.api.is_active: + raise MCPToolError( + f"API deployment '{context.api.display_name}' is not active. " + "Activate it in Unstract before extracting." + ) + + # `tags` is a comma-separated string on the REST surface; MCP clients + # produce JSON arrays, so convert rather than exposing the wire format. + payload: dict[str, Any] = { + ApiExecution.PRESIGNED_URLS: document_urls, + ApiExecution.TIMEOUT_FORM_DATA: timeout, + ApiExecution.INCLUDE_METADATA: include_metadata, + ApiExecution.INCLUDE_METRICS: include_metrics, + ApiExecution.INCLUDE_EXTRACTED_TEXT: include_extracted_text, + } + if tags: + payload[ApiExecution.TAGS] = ",".join(tags) + if llm_profile_id: + payload[ApiExecution.LLM_PROFILE_ID] = llm_profile_id + + serializer = ExecutionRequestSerializer( + data=payload, context={"api": context.api, "api_key": context.api_key} + ) + try: + serializer.is_valid(raise_exception=True) + except ValidationError as error: + raise MCPToolError( + f"Invalid arguments: {_format_validation_error(error)}" + ) from error + + validated = serializer.validated_data + presigned_urls = validated.get(ApiExecution.PRESIGNED_URLS, []) + + file_objs: list[Any] = [] + DeploymentHelper.load_presigned_files(presigned_urls, file_objs) + + organization = context.api.organization + execution_id = str(uuid.uuid4()) + + can_proceed, limit_info = APIDeploymentRateLimiter.check_and_acquire( + organization, execution_id + ) + if not can_proceed: + raise MCPToolError( + "Rate limit exceeded for this organization " + f"({limit_info['current_usage']}/{limit_info['limit']} " + f"{limit_info['limit_type']}). Retry once running extractions finish." + ) + + try: + response = DeploymentHelper.execute_workflow( + organization_name=context.org_name, + api=context.api, + file_objs=file_objs, + timeout=validated.get(ApiExecution.TIMEOUT_FORM_DATA), + include_metadata=validated.get(ApiExecution.INCLUDE_METADATA), + include_metrics=validated.get(ApiExecution.INCLUDE_METRICS), + include_extracted_text=validated.get(ApiExecution.INCLUDE_EXTRACTED_TEXT), + use_file_history=validated.get(ApiExecution.USE_FILE_HISTORY, False), + tag_names=validated.get(ApiExecution.TAGS, []), + llm_profile_id=validated.get(ApiExecution.LLM_PROFILE_ID), + execution_id=execution_id, + ) + except RateLimitExceeded: + # Slot was acquired above, so a limit raised from deeper in the stack + # is a different limit; still release ours before surfacing it. + APIDeploymentRateLimiter.release_slot(organization, execution_id) + raise + except Exception as error: + APIDeploymentRateLimiter.release_slot(organization, execution_id) + logger.exception( + f"MCP extractDocument failed for api '{context.api.api_name}': {error}" + ) + raise MCPToolError(f"Extraction failed: {error}") from error + + result = dict(response) + # Surface the execution id unconditionally. On the pending path it is the + # agent's only handle for polling, and execute_workflow does not guarantee + # it is present in the response body. + result.setdefault("execution_id", execution_id) + return result + + +def get_execution_status_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "execution_id": { + "type": "string", + "description": ( + "The execution_id returned by a previous extractDocument call." + ), + }, + "include_metadata": { + "type": "boolean", + "default": False, + "description": "Include extraction metadata in the result.", + }, + "include_metrics": { + "type": "boolean", + "default": False, + "description": "Include token and cost metrics in the result.", + }, + "include_extracted_text": { + "type": "boolean", + "default": False, + "description": ( + "Include the full raw text extracted from each document. " + "This can be very large." + ), + }, + }, + "required": ["execution_id"], + } + + +def get_execution_status( + context: MCPContext, + execution_id: str, + include_metadata: bool = False, + include_metrics: bool = False, + include_extracted_text: bool = False, +) -> dict[str, Any]: + """Return the status, and once ready the result, of an extraction.""" + serializer = ExecutionQuerySerializer( + data={ + ApiExecution.EXECUTION_ID: execution_id, + ApiExecution.INCLUDE_METADATA: include_metadata, + ApiExecution.INCLUDE_METRICS: include_metrics, + ApiExecution.INCLUDE_EXTRACTED_TEXT: include_extracted_text, + } + ) + try: + serializer.is_valid(raise_exception=True) + except ValidationError as error: + raise MCPToolError( + f"Invalid arguments: {_format_validation_error(error)}" + ) from error + + validated = serializer.validated_data + response: ExecutionResponse = DeploymentHelper.get_execution_status( + validated.get(ApiExecution.EXECUTION_ID) + ) + + if response.result_acknowledged: + # The REST surface answers 406 here. An agent cannot act on a status + # code, so say plainly that the result is gone and not retrievable. + return { + "execution_status": response.execution_status, + "message": ( + "This result was already acknowledged and is no longer " + "retrievable. Start a new extraction if you need it again." + ), + } + + if response.execution_status == CeleryTaskState.COMPLETED.value: + deployment_execution_dto = DeploymentExecutionDTO( + api=context.api, api_key=context.api_key + ) + DeploymentHelper.process_completed_execution( + response=response, + deployment_execution_dto=deployment_execution_dto, + include_metadata=validated.get(ApiExecution.INCLUDE_METADATA), + include_metrics=validated.get(ApiExecution.INCLUDE_METRICS), + include_extracted_text=validated.get(ApiExecution.INCLUDE_EXTRACTED_TEXT), + ) + + return { + "execution_id": execution_id, + "execution_status": response.execution_status, + "result": response.result, + } diff --git a/backend/mcp_v2/tools/info.py b/backend/mcp_v2/tools/info.py new file mode 100644 index 0000000000..06cc218e83 --- /dev/null +++ b/backend/mcp_v2/tools/info.py @@ -0,0 +1,83 @@ +"""Read-only tools describing the connected API deployment.""" + +import logging +from typing import Any + +from mcp_v2.context import MCPContext + +logger = logging.getLogger(__name__) + + +def read_me_first(context: MCPContext) -> dict[str, Any]: + """Return the orientation guide for a coding agent. + + Deliberately built from the live deployment rather than a static string — + the agent should learn the name of the deployment it is actually talking + to, not a generic description of Unstract. + """ + return { + "server": "Unstract MCP Server", + "connected_deployment": { + "name": context.api.display_name, + "description": context.api.description or None, + "is_active": context.api.is_active, + }, + "what_this_does": ( + "Unstract runs LLM-driven extraction over unstructured documents " + "(PDFs, scans, images) and returns structured JSON. This MCP " + "server is scoped to a single API deployment, which encapsulates " + "one extraction workflow — the prompts and output schema are fixed " + "by that deployment, so you supply documents, not instructions." + ), + "tools": [ + { + "name": "getApiInfo", + "purpose": "Learn what the connected deployment extracts.", + }, + { + "name": "extractDocument", + "purpose": "Run extraction over document URLs. Consumes quota.", + }, + { + "name": "getExecutionStatus", + "purpose": "Poll for the result of a pending extraction.", + }, + ], + "recommended_workflow": [ + "1. Call getApiInfo to confirm the deployment is active and see " + "what it is meant to process.", + "2. Call extractDocument with the URL(s) of the document(s).", + "3. If the response has execution_status COMPLETED, the result is " + "already in the response — you are done.", + "4. Otherwise poll getExecutionStatus with the returned " + "execution_id until the status is COMPLETED or ERROR, pausing a " + "few seconds between polls.", + ], + "notes": [ + "Documents are fetched server-side from the URLs you provide, so " + "they must be reachable by Unstract (for example a pre-signed URL).", + "extractDocument consumes the organization's extraction quota. " + "Never call it speculatively, and never retry a call that already " + "returned an execution_id — poll instead.", + ], + } + + +def get_api_info_schema() -> dict[str, Any]: + return {"type": "object", "properties": {}, "required": []} + + +def get_api_info(context: MCPContext) -> dict[str, Any]: + """Return metadata about the connected API deployment.""" + api = context.api + return { + "id": str(api.id), + "display_name": api.display_name, + "api_name": api.api_name, + "description": api.description or None, + "is_active": api.is_active, + "workflow": { + "id": str(api.workflow_id), + "name": api.workflow.workflow_name, + }, + } diff --git a/backend/mcp_v2/urls.py b/backend/mcp_v2/urls.py new file mode 100644 index 0000000000..6a6ce1cdb8 --- /dev/null +++ b/backend/mcp_v2/urls.py @@ -0,0 +1,31 @@ +"""URLs for the hosted MCP server. + +Mounted alongside the API deployment execution endpoint so an MCP endpoint is +the same shape as the REST endpoint for the same deployment: + + POST /deployment/api/// (REST) + POST /mcp/// (MCP) +""" + +from django.urls import re_path + +from mcp_v2.views import MCPServerView + +mcp_server = MCPServerView.as_view() + +urlpatterns = [ + re_path( + r"^(?P[\w-]+)/(?P[\w-]+)/?$", + mcp_server, + name="mcp_server", + ), + # API key as a path segment, for MCP clients that cannot attach an + # Authorization header to the request. The key is a UUID, so the pattern + # cannot collide with the header-authenticated route above. + re_path( + r"^(?P[\w-]+)/(?P[\w-]+)/" + r"(?P[0-9a-fA-F-]{36})/?$", + mcp_server, + name="mcp_server_with_key", + ), +] diff --git a/backend/mcp_v2/views.py b/backend/mcp_v2/views.py new file mode 100644 index 0000000000..125bc32733 --- /dev/null +++ b/backend/mcp_v2/views.py @@ -0,0 +1,292 @@ +"""JSON-RPC 2.0 transport for the hosted Unstract MCP server. + +Mirrors the shape of the API deployment execution endpoint: the URL carries the +organization and the deployment name, and the deployment's own API key +authenticates the caller. An MCP session is therefore scoped to exactly one API +deployment, and reuses that deployment's existing key management. +""" + +import json +import logging +from typing import Any + +from api_v2.deployment_helper import DeploymentHelper +from django.http import JsonResponse +from rest_framework import status, views +from rest_framework.request import Request +from rest_framework.response import Response + +from mcp_v2.constants import JSONRPC, MCPMethod, MCPServer +from mcp_v2.context import MCPContext +from mcp_v2.exceptions import MCPToolError +from mcp_v2.registry import TOOL_REGISTRY + +logger = logging.getLogger(__name__) + + +def _rpc_result(request_id: Any, result: Any) -> JsonResponse: + """Build a JSON-RPC success response. + + JsonResponse rather than DRF's Response so the body is emitted verbatim: + content negotiation must not turn a JSON-RPC envelope into the browsable + API renderer for a client that sent a permissive Accept header. + """ + return JsonResponse({"jsonrpc": JSONRPC.VERSION, "id": request_id, "result": result}) + + +def _rpc_error( + request_id: Any, code: int, message: str, data: Any = None +) -> JsonResponse: + """Build a JSON-RPC error response. + + Always HTTP 200: transport-level success with an application-level error is + exactly what JSON-RPC models, and clients read the envelope, not the status + code. The one exception is authentication, which must be a real 401 so MCP + clients can react to it (see `_auth_error`). + """ + error: dict[str, Any] = {"code": code, "message": message} + if data is not None: + error["data"] = data + return JsonResponse({"jsonrpc": JSONRPC.VERSION, "id": request_id, "error": error}) + + +def _auth_error(message: str) -> JsonResponse: + """Build a 401 for a failed or missing credential. + + Unlike other failures this carries a real HTTP status, because a client + that cannot authenticate needs to distinguish "wrong key" from "tool + failed" before it has a usable session. + """ + response = JsonResponse( + { + "jsonrpc": JSONRPC.VERSION, + "id": None, + "error": {"code": JSONRPC.UNAUTHORIZED, "message": message}, + }, + status=status.HTTP_401_UNAUTHORIZED, + ) + response["WWW-Authenticate"] = 'Bearer realm="unstract-mcp"' + return response + + +class MCPServerView(views.APIView): + """MCP JSON-RPC endpoint for a single API deployment. + + Authentication is the deployment's API key, accepted either as a bearer + token or — for MCP clients that cannot attach custom headers — as a path + segment. This is a public, key-authenticated endpoint like the deployment + execution endpoint it sits beside, so session auth does not apply. + """ + + authentication_classes: list = [] + permission_classes: list = [] + + def initialize_request(self, request: Request, *args: Any, **kwargs: Any) -> Request: + """Skip CSRF, matching the public API deployment endpoint.""" + request.csrf_processing_done = True + return super().initialize_request(request, *args, **kwargs) + + def get( + self, request: Request, org_name: str, api_name: str, api_key: str | None = None + ) -> Response: + """Advertise server identity. + + Clients probe with GET to check connectivity before opening a session. + Deliberately unauthenticated and free of deployment detail — it reveals + only that an MCP server is mounted here. + """ + return Response( + { + "name": MCPServer.NAME, + "version": MCPServer.VERSION, + "protocolVersion": MCPServer.PROTOCOL_VERSION, + "transport": "http", + "authMethods": ["bearer"], + } + ) + + def post( + self, request: Request, org_name: str, api_name: str, api_key: str | None = None + ) -> JsonResponse: + """Handle a single JSON-RPC request.""" + # `api_key` in the path takes precedence when present; otherwise fall + # back to the Authorization header. + if not api_key: + header = request.headers.get("Authorization", "") + if header.startswith("Bearer "): + api_key = header.split(" ", 1)[1].strip() + + if not api_key: + return _auth_error("Missing API key") + + context = self._resolve_context( + org_name=org_name, api_name=api_name, api_key=api_key + ) + if context is None: + return _auth_error("Invalid API key or unknown API deployment") + + body = request.data + if not isinstance(body, dict): + # Batch requests are a valid part of JSON-RPC 2.0 but are not used + # by MCP clients; rejecting them explicitly beats a confusing + # downstream AttributeError. + return _rpc_error( + None, + JSONRPC.INVALID_REQUEST, + "Invalid Request", + "Expected a single JSON-RPC object", + ) + + request_id = body.get("id") + method = body.get("method") + + if body.get("jsonrpc") != JSONRPC.VERSION: + return _rpc_error( + request_id, + JSONRPC.INVALID_REQUEST, + "Invalid Request", + "Only JSON-RPC 2.0 is supported", + ) + if not method: + return _rpc_error( + request_id, JSONRPC.INVALID_REQUEST, "Invalid Request", "Missing method" + ) + + return self._dispatch( + method=method, + request_id=request_id, + params=body.get("params") or {}, + context=context, + ) + + def _resolve_context( + self, org_name: str, api_name: str, api_key: str + ) -> MCPContext | None: + """Authenticate the key against the named deployment. + + Returns None for every failure mode — unknown deployment, wrong key, + malformed key — so the caller answers all of them identically and the + endpoint cannot be used to probe which deployment names exist. + """ + # Pin the organization before touching org-scoped managers; the + # deployment lookup below filters on it. + DeploymentHelper.validate_parameters( + self.request, api_name=api_name, org_name=org_name + ) + + api_deployment = DeploymentHelper.get_deployment_by_api_name(api_name=api_name) + try: + DeploymentHelper.validate_api(api_deployment=api_deployment, api_key=api_key) + except Exception as error: + logger.warning( + f"MCP auth rejected for org '{org_name}', api '{api_name}': {error}" + ) + return None + + return MCPContext(api=api_deployment, api_key=api_key, org_name=org_name) + + def _dispatch( + self, method: str, request_id: Any, params: dict[str, Any], context: MCPContext + ) -> JsonResponse: + """Route a JSON-RPC method to its handler.""" + if method == MCPMethod.INITIALIZE: + return _rpc_result( + request_id, + { + "protocolVersion": MCPServer.PROTOCOL_VERSION, + "capabilities": {"tools": {"listChanged": False}}, + "serverInfo": { + "name": MCPServer.NAME, + "version": MCPServer.VERSION, + }, + "instructions": ( + "Unstract runs LLM-driven extraction over unstructured " + "documents and returns structured JSON. This server is " + "scoped to the API deployment " + f"'{context.api.display_name}'. Call readMeFirst before " + "any other tool." + ), + }, + ) + + if method.startswith(MCPMethod.NOTIFICATION_PREFIX): + # Notifications (e.g. notifications/initialized) carry no id and + # must not receive a result body. + return JsonResponse({}, status=status.HTTP_202_ACCEPTED) + + if method == MCPMethod.PING: + return _rpc_result(request_id, {}) + + if method == MCPMethod.TOOLS_LIST: + return _rpc_result(request_id, {"tools": TOOL_REGISTRY.list_schemas()}) + + if method == MCPMethod.TOOLS_CALL: + return self._call_tool(request_id=request_id, params=params, context=context) + + return _rpc_error( + request_id, + JSONRPC.METHOD_NOT_FOUND, + "Method not found", + f"Unsupported method '{method}'", + ) + + def _call_tool( + self, request_id: Any, params: dict[str, Any], context: MCPContext + ) -> JsonResponse: + """Invoke a registered tool and wrap its result in MCP content format.""" + tool_name = params.get("name") + if not tool_name: + return _rpc_error( + request_id, JSONRPC.INVALID_PARAMS, "Invalid params", "Missing tool name" + ) + + tool = TOOL_REGISTRY.get(tool_name) + if tool is None: + return _rpc_error( + request_id, + JSONRPC.METHOD_NOT_FOUND, + "Method not found", + f"Tool '{tool_name}' not found. Available tools: {TOOL_REGISTRY.names()}", + ) + + arguments = params.get("arguments") or {} + if not isinstance(arguments, dict): + return _rpc_error( + request_id, + JSONRPC.INVALID_PARAMS, + "Invalid params", + "'arguments' must be an object", + ) + + try: + result = tool.handler(context, **arguments) + except MCPToolError as error: + # Expected, agent-actionable failure: the message is written for + # the agent, so pass it through as an error result rather than a + # protocol error. isError=True lets the agent see and retry it. + return _rpc_result(request_id, _tool_content(str(error), is_error=True)) + except TypeError as error: + # Almost always the agent passing arguments the tool does not + # accept; report it as bad params rather than a server fault. + logger.warning(f"MCP tool '{tool_name}' called with bad arguments: {error}") + return _rpc_error( + request_id, JSONRPC.INVALID_PARAMS, "Invalid params", str(error) + ) + except Exception as error: + logger.exception(f"MCP tool '{tool_name}' failed: {error}") + return _rpc_error( + request_id, + JSONRPC.TOOL_EXECUTION_ERROR, + "Tool execution failed", + f"Tool '{tool_name}' failed unexpectedly. " + "Contact your Unstract administrator if this persists.", + ) + + return _rpc_result( + request_id, _tool_content(json.dumps(result, indent=2, default=str)) + ) + + +def _tool_content(text: str, is_error: bool = False) -> dict[str, Any]: + """Wrap tool output in the MCP `tools/call` content envelope.""" + return {"content": [{"type": "text", "text": text}], "isError": is_error} diff --git a/tests/critical_paths.yaml b/tests/critical_paths.yaml index 00429280e6..7e8a7a97d1 100644 --- a/tests/critical_paths.yaml +++ b/tests/critical_paths.yaml @@ -66,6 +66,11 @@ paths: covered_by: [e2e-api-deployment] proof: marker + - id: mcp-server-auth + description: "Unauthenticated or mis-scoped hosted-MCP calls are rejected before any tool runs." + covered_by: [integration-backend] + proof: marker + - id: prompt-studio-author description: "Create a Prompt Studio project and add a prompt to it." covered_by: [integration-backend] From 763239a51d058ee8a7740bf9196579fa8c577eed Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Fri, 24 Jul 2026 16:02:10 +0530 Subject: [PATCH 02/16] feat(mcp): cover the extraction happy path, fix what it caught MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The transport and auth tests never reached the extraction helpers on a successful call, so the mapping from tool kwargs into execute_workflow was executed by nothing — a renamed kwarg would have passed every test and failed on the first real extraction. Adding that coverage surfaced three real problems: - Tool descriptions promised documents could be passed as any reachable URL. The execution serializer accepts *only* S3 pre-signed URLs, so an agent following the description would have failed every call. Descriptions, JSON schema and README now state the S3 requirement. - The tags schema advertised an unbounded array; the serializer caps it at one. Both limits are now sourced from the serializers rather than restated, so the advertised schema tracks what is actually enforced. - Documents were downloaded before the rate-limit slot was taken, so a call about to be rejected still pulled every document over the network. The slot is now acquired first, with the fetch inside the try block so a failed fetch still releases it. Also converts a RateLimitExceeded raised from deeper in the stack into an agent-readable message instead of letting it reach the generic "failed unexpectedly" branch. Tests: 34 passing (was 20). Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/mcp_v2/README.md | 5 +- backend/mcp_v2/registry.py | 6 +- backend/mcp_v2/tests/test_tool_execution.py | 341 ++++++++++++++++++++ backend/mcp_v2/tools/execution.py | 42 ++- backend/mcp_v2/tools/info.py | 5 +- 5 files changed, 380 insertions(+), 19 deletions(-) create mode 100644 backend/mcp_v2/tests/test_tool_execution.py diff --git a/backend/mcp_v2/README.md b/backend/mcp_v2/README.md index 7e47f45ac0..146e44cb11 100644 --- a/backend/mcp_v2/README.md +++ b/backend/mcp_v2/README.md @@ -56,8 +56,9 @@ claude mcp add --transport http unstract \ | `extractDocument` | Run extraction over document URLs. **Consumes quota.** | | `getExecutionStatus` | Poll for the result of a pending extraction. | -Documents are passed as URLs and fetched server-side, so they must be reachable -by the backend — typically a pre-signed URL. Extraction is asynchronous: +Documents are passed as **S3 pre-signed URLs** and fetched server-side; the +execution serializer rejects anything else, so a plain public link will not +work. Extraction is asynchronous: `extractDocument` returns an `execution_id` when it does not finish within the timeout, and the agent polls `getExecutionStatus` with it. diff --git a/backend/mcp_v2/registry.py b/backend/mcp_v2/registry.py index cb03f29f1d..238adbbc8d 100644 --- a/backend/mcp_v2/registry.py +++ b/backend/mcp_v2/registry.py @@ -116,8 +116,10 @@ def build_registry() -> MCPToolRegistry: description=( "Run the connected Unstract API deployment over one or more " "documents and return the structured extraction result.\n\n" - "Documents are supplied as URLs, which Unstract fetches " - "server-side. Extraction is asynchronous: when it does not " + "Documents are supplied as S3 pre-signed URLs, which Unstract " + "fetches server-side — ordinary public links are rejected, so " + "upload to S3 and pre-sign first if the document is not " + "already there. Extraction is asynchronous: when it does not " "finish within `timeout` seconds this returns " "`execution_status: PENDING` along with an `execution_id`. Poll " "`getExecutionStatus` with that id to collect the result.\n\n" diff --git a/backend/mcp_v2/tests/test_tool_execution.py b/backend/mcp_v2/tests/test_tool_execution.py new file mode 100644 index 0000000000..46ab0ad3d4 --- /dev/null +++ b/backend/mcp_v2/tests/test_tool_execution.py @@ -0,0 +1,341 @@ +"""The success path of the extraction tools. + +The transport and auth tests never reach these helpers on a successful call, so +without this file the argument mapping from tool kwargs into +``execute_workflow`` is executed by nothing — a renamed or mis-mapped kwarg +would pass every other test and fail on the first real extraction. The +``assert_called_once_with`` below is the point of these tests. + +Mocked at the ``DeploymentHelper`` boundary so no database or Celery is needed. +""" + +from __future__ import annotations + +from unittest.mock import patch + +from django.test import SimpleTestCase +from rest_framework.exceptions import ValidationError +from workflow_manager.workflow_v2.dto import ExecutionResponse + +from mcp_v2.context import MCPContext +from mcp_v2.exceptions import MCPToolError +from mcp_v2.tools.execution import extract_document, get_execution_status + +# The execution serializer accepts only HTTPS S3 pre-signed URLs; using a +# realistic one here keeps these tests honest about what the tool accepts. +DOC_URL = "https://my-bucket.s3.us-east-1.amazonaws.com/invoice.pdf?X-Amz-Signature=abc" +EXECUTION_ID = "33333333-3333-3333-3333-333333333333" + + +class FakeOrganization: + organization_id = "org-mcp" + + +class FakeAPI: + display_name = "Invoice Extractor" + api_name = "live-api" + is_active = True + organization = FakeOrganization() + + +def make_context(active: bool = True) -> MCPContext: + api = FakeAPI() + api.is_active = active + return MCPContext(api=api, api_key="test-key", org_name="org-mcp") + + +class ExtractDocumentTest(SimpleTestCase): + def setUp(self) -> None: + self.context = make_context() + + def _run(self, execute_return=None, **kwargs): + """Drive extract_document with the whole execution stack stubbed.""" + with ( + patch( + "mcp_v2.tools.execution.DeploymentHelper.load_presigned_files" + ) as load_files, + patch( + "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + return_value=(True, {}), + ), + patch( + "mcp_v2.tools.execution.APIDeploymentRateLimiter.release_slot" + ) as release, + patch( + "mcp_v2.tools.execution.DeploymentHelper.execute_workflow", + return_value=execute_return + if execute_return is not None + else {"execution_status": "COMPLETED", "message": "ok"}, + ) as execute, + ): + result = extract_document(self.context, **kwargs) + return result, execute, load_files, release + + def test_arguments_map_onto_execute_workflow(self) -> None: + """Guards the tool-kwarg -> execute_workflow mapping. A rename on + either side breaks here rather than in production. + """ + result, execute, load_files, _ = self._run( + document_urls=[DOC_URL], + timeout=45, + include_metadata=True, + include_metrics=True, + include_extracted_text=True, + tags=["alpha"], + ) + + load_files.assert_called_once() + assert load_files.call_args.args[0] == [DOC_URL] + + execute.assert_called_once() + kwargs = execute.call_args.kwargs + assert kwargs["organization_name"] == "org-mcp" + assert kwargs["api"] is self.context.api + assert kwargs["timeout"] == 45 + assert kwargs["include_metadata"] is True + assert kwargs["include_metrics"] is True + assert kwargs["include_extracted_text"] is True + # Agents send a JSON array; the execution layer expects parsed names. + assert kwargs["tag_names"] == ["alpha"] + assert kwargs["execution_id"] == result["execution_id"] + + def test_result_always_carries_execution_id(self) -> None: + """On the pending path the id is the agent's only handle for polling, + and execute_workflow does not guarantee it in the body. + """ + result, execute, _, _ = self._run( + execute_return={"execution_status": "PENDING"}, document_urls=[DOC_URL] + ) + + assert result["execution_status"] == "PENDING" + assert result["execution_id"] == execute.call_args.kwargs["execution_id"] + + def test_existing_execution_id_in_response_is_preserved(self) -> None: + """Never overwrite an id the execution layer reported itself.""" + result, _, _, _ = self._run( + execute_return={"execution_status": "COMPLETED", "execution_id": "from-core"}, + document_urls=[DOC_URL], + ) + + assert result["execution_id"] == "from-core" + + def test_defaults_are_conservative(self) -> None: + """Defaults must not silently opt an agent into large or costly output.""" + _, execute, _, _ = self._run(document_urls=[DOC_URL]) + + kwargs = execute.call_args.kwargs + assert kwargs["include_metadata"] is False + assert kwargs["include_metrics"] is False + assert kwargs["include_extracted_text"] is False + assert kwargs["tag_names"] == [] + + def test_inactive_deployment_rejected_before_quota_is_spent(self) -> None: + with patch( + "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + ) as execute: + with self.assertRaises(MCPToolError): + extract_document(make_context(active=False), document_urls=[DOC_URL]) + + execute.assert_not_called() + + def test_non_s3_url_rejected_with_an_explanation(self) -> None: + """Only S3 pre-signed URLs are accepted. The tool description says so, + and this pins the behaviour the description promises — an agent handed + an ordinary link must be told why, not just that it failed. + """ + with patch( + "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + ) as execute: + with self.assertRaises(MCPToolError) as caught: + extract_document( + self.context, document_urls=["https://example.com/invoice.pdf"] + ) + + execute.assert_not_called() + assert "S3" in str(caught.exception) + + def test_invalid_url_rejected_before_quota_is_spent(self) -> None: + with ( + patch( + "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire" + ) as acquire, + patch( + "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + ) as execute, + ): + with self.assertRaises(MCPToolError): + extract_document(self.context, document_urls=["not-a-url"]) + + acquire.assert_not_called() + execute.assert_not_called() + + def test_rate_limited_call_downloads_nothing(self) -> None: + """The slot is taken before the documents are fetched, so a rejected + call must not pull every document over the network first. + """ + with ( + patch( + "mcp_v2.tools.execution.DeploymentHelper.load_presigned_files" + ) as load_files, + patch( + "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + return_value=( + False, + {"current_usage": 5, "limit": 5, "limit_type": "organization"}, + ), + ), + patch( + "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + ) as execute, + ): + with self.assertRaises(MCPToolError) as caught: + extract_document(self.context, document_urls=[DOC_URL]) + + load_files.assert_not_called() + execute.assert_not_called() + assert "5/5" in str(caught.exception) + + def test_failed_download_releases_the_rate_limit_slot(self) -> None: + """The fetch happens after the slot is taken, so a fetch failure has to + give the slot back or the org silently loses throughput. + """ + with ( + patch( + "mcp_v2.tools.execution.DeploymentHelper.load_presigned_files", + side_effect=RuntimeError("403 Forbidden"), + ), + patch( + "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + return_value=(True, {}), + ), + patch( + "mcp_v2.tools.execution.APIDeploymentRateLimiter.release_slot" + ) as release, + patch( + "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + ) as execute, + ): + with self.assertRaises(MCPToolError): + extract_document(self.context, document_urls=[DOC_URL]) + + release.assert_called_once() + execute.assert_not_called() + + def test_failed_execution_releases_the_rate_limit_slot(self) -> None: + """A leaked slot degrades the org's throughput permanently, so the + release must survive the failure path. + """ + with ( + patch("mcp_v2.tools.execution.DeploymentHelper.load_presigned_files"), + patch( + "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + return_value=(True, {}), + ), + patch( + "mcp_v2.tools.execution.APIDeploymentRateLimiter.release_slot" + ) as release, + patch( + "mcp_v2.tools.execution.DeploymentHelper.execute_workflow", + side_effect=RuntimeError("boom"), + ), + ): + with self.assertRaises(MCPToolError): + extract_document(self.context, document_urls=[DOC_URL]) + + release.assert_called_once() + + +class GetExecutionStatusTest(SimpleTestCase): + def setUp(self) -> None: + self.context = make_context() + + def _run(self, response: ExecutionResponse, **kwargs): + with ( + patch( + "mcp_v2.tools.execution.ExecutionQuerySerializer.is_valid", + return_value=True, + ), + patch( + "mcp_v2.tools.execution.ExecutionQuerySerializer.validated_data", + { + "execution_id": EXECUTION_ID, + "include_metadata": kwargs.get("include_metadata", False), + "include_metrics": kwargs.get("include_metrics", False), + "include_extracted_text": kwargs.get( + "include_extracted_text", False + ), + }, + ), + patch( + "mcp_v2.tools.execution.DeploymentHelper.get_execution_status", + return_value=response, + ), + patch( + "mcp_v2.tools.execution.DeploymentHelper.process_completed_execution" + ) as process, + ): + result = get_execution_status( + self.context, execution_id=EXECUTION_ID, **kwargs + ) + return result, process + + def test_completed_execution_is_enriched_and_returned(self) -> None: + response = ExecutionResponse( + workflow_id="wf", + execution_id=EXECUTION_ID, + execution_status="COMPLETED", + result=[{"file": "invoice.pdf", "result": {"output": {"total": 42}}}], + ) + + result, process = self._run(response, include_metadata=True) + + # Enrichment is what turns the raw execution into the API-shaped + # result; skipping it would silently return a thinner payload. + process.assert_called_once() + assert process.call_args.kwargs["include_metadata"] is True + assert result["execution_status"] == "COMPLETED" + assert result["result"][0]["result"]["output"]["total"] == 42 + assert result["execution_id"] == EXECUTION_ID + + def test_pending_execution_is_not_enriched(self) -> None: + response = ExecutionResponse( + workflow_id="wf", + execution_id=EXECUTION_ID, + execution_status="EXECUTING", + result=None, + ) + + result, process = self._run(response) + + process.assert_not_called() + assert result["execution_status"] == "EXECUTING" + + def test_acknowledged_result_explains_itself(self) -> None: + """The REST surface answers 406 here; an agent cannot act on a status + code, so the tool must say the result is gone and not retryable. + """ + response = ExecutionResponse( + workflow_id="wf", + execution_id=EXECUTION_ID, + execution_status="COMPLETED", + result=None, + ) + response.result_acknowledged = True + + result, process = self._run(response) + + process.assert_not_called() + assert "already acknowledged" in result["message"] + + def test_unknown_execution_id_is_an_agent_error(self) -> None: + """A bad id comes back as an actionable tool error, not a server + fault — the agent may simply have mistyped it. + """ + with patch( + "mcp_v2.tools.execution.ExecutionQuerySerializer.is_valid", + side_effect=ValidationError({"execution_id": ["Invalid execution_id."]}), + ): + with self.assertRaises(MCPToolError) as caught: + get_execution_status(self.context, execution_id="nope") + + assert "execution_id" in str(caught.exception) diff --git a/backend/mcp_v2/tools/execution.py b/backend/mcp_v2/tools/execution.py index 0b22c49dd7..b2af589010 100644 --- a/backend/mcp_v2/tools/execution.py +++ b/backend/mcp_v2/tools/execution.py @@ -11,6 +11,7 @@ from api_v2.rate_limiter import APIDeploymentRateLimiter from api_v2.serializers import ExecutionQuerySerializer, ExecutionRequestSerializer from rest_framework.exceptions import ValidationError +from tags.serializers import TagParamsSerializer from utils.enums import CeleryTaskState from workflow_manager.workflow_v2.dto import ExecutionResponse @@ -19,9 +20,12 @@ logger = logging.getLogger(__name__) -# Mirrors ExecutionRequestSerializer.MAX_FILES_ALLOWED. Declared in the JSON -# schema too, so a well-behaved client is stopped before the request is made. +# Mirrors the serializer limits into the JSON schema, so a well-behaved client +# is stopped before it spends a round trip on a request that cannot validate. +# Sourced from the serializers rather than copied, so the advertised schema +# tracks the limits actually enforced. MAX_DOCUMENTS = ExecutionRequestSerializer.MAX_FILES_ALLOWED +MAX_TAGS = TagParamsSerializer.MAX_TAGS_ALLOWED # An agent polling with getExecutionStatus does not benefit from holding the # HTTP request open for the full 300s the REST API permits, and a long-held @@ -62,9 +66,9 @@ def extract_document_schema() -> dict[str, Any]: "minItems": 1, "maxItems": MAX_DOCUMENTS, "description": ( - "URLs of the documents to extract. Unstract fetches these " - "server-side, so each must be reachable by the Unstract " - "backend — typically a pre-signed URL." + "S3 pre-signed URLs of the documents to extract. Unstract " + "fetches these server-side. Only S3 pre-signed URLs are " + "accepted — an ordinary public http(s) link is rejected." ), }, "timeout": { @@ -99,7 +103,12 @@ def extract_document_schema() -> dict[str, Any]: "tags": { "type": "array", "items": {"type": "string"}, - "description": "Tags to associate with this execution.", + "maxItems": MAX_TAGS, + "description": ( + f"Tags to associate with this execution (at most " + f"{MAX_TAGS}). Each tag must start with a letter and " + "contain only letters, numbers, underscores and hyphens." + ), }, "llm_profile_id": { "type": "string", @@ -164,12 +173,13 @@ def extract_document( validated = serializer.validated_data presigned_urls = validated.get(ApiExecution.PRESIGNED_URLS, []) - file_objs: list[Any] = [] - DeploymentHelper.load_presigned_files(presigned_urls, file_objs) - organization = context.api.organization execution_id = str(uuid.uuid4()) + # Take the rate-limit slot before downloading anything. The REST view + # fetches first because it already holds the uploaded files in the request + # body; here the fetch is ours to make, and doing it for a call we are + # about to reject would pull every document over the network for nothing. can_proceed, limit_info = APIDeploymentRateLimiter.check_and_acquire( organization, execution_id ) @@ -181,6 +191,8 @@ def extract_document( ) try: + file_objs: list[Any] = [] + DeploymentHelper.load_presigned_files(presigned_urls, file_objs) response = DeploymentHelper.execute_workflow( organization_name=context.org_name, api=context.api, @@ -194,11 +206,15 @@ def extract_document( llm_profile_id=validated.get(ApiExecution.LLM_PROFILE_ID), execution_id=execution_id, ) - except RateLimitExceeded: - # Slot was acquired above, so a limit raised from deeper in the stack - # is a different limit; still release ours before surfacing it. + except RateLimitExceeded as error: + # A limit raised from deeper in the stack is a different limit than the + # one checked above, but it is still the agent's to wait out — so + # convert it rather than letting it fall through to the generic + # "failed unexpectedly" branch below. APIDeploymentRateLimiter.release_slot(organization, execution_id) - raise + raise MCPToolError( + f"Rate limit exceeded: {error.detail}. Retry once running extractions finish." + ) from error except Exception as error: APIDeploymentRateLimiter.release_slot(organization, execution_id) logger.exception( diff --git a/backend/mcp_v2/tools/info.py b/backend/mcp_v2/tools/info.py index 06cc218e83..a5ad1da9a9 100644 --- a/backend/mcp_v2/tools/info.py +++ b/backend/mcp_v2/tools/info.py @@ -54,8 +54,9 @@ def read_me_first(context: MCPContext) -> dict[str, Any]: "few seconds between polls.", ], "notes": [ - "Documents are fetched server-side from the URLs you provide, so " - "they must be reachable by Unstract (for example a pre-signed URL).", + "Documents are fetched server-side from the URLs you provide, and " + "only S3 pre-signed URLs are accepted — an ordinary public link " + "is rejected. Upload to S3 and pre-sign it first if needed.", "extractDocument consumes the organization's extraction quota. " "Never call it speculatively, and never retry a call that already " "returned an execution_id — poll instead.", From 89f651d7d4a5be58cbf741296ae52d0adc3b8eb5 Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sat, 25 Jul 2026 17:25:02 +0530 Subject: [PATCH 03/16] refactor(mcp): rename mcp_v2 to mcp_server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `_v2` suffix in this codebase marks apps carried through the v0.93.0 v1->v2 multitenancy data migration (see backend/migrating/v2/README.md), not current naming convention. Every app added since — pg_queue, dashboard_metrics, platform_api, configuration — is unsuffixed. A brand-new app with no v1 predecessor should not claim the marker. Named `mcp_server` rather than plain `mcp` because `mcp` is also the PyPI package name for the MCP SDK (which mfbt's server uses). A bare `mcp` app shadows it on the import path: `import mcp` resolved to the Django app, so a later `pip install mcp` would break silently. Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/backend/base_urls.py | 2 +- backend/backend/settings/base.py | 2 +- backend/{mcp_v2 => mcp_server}/README.md | 0 backend/{mcp_v2 => mcp_server}/__init__.py | 0 backend/{mcp_v2 => mcp_server}/apps.py | 2 +- backend/{mcp_v2 => mcp_server}/constants.py | 0 backend/{mcp_v2 => mcp_server}/context.py | 0 backend/{mcp_v2 => mcp_server}/exceptions.py | 0 backend/{mcp_v2 => mcp_server}/registry.py | 4 +- .../{mcp_v2 => mcp_server}/tests/__init__.py | 0 .../tests/test_mcp_auth.py | 4 +- .../tests/test_mcp_protocol.py | 10 ++-- .../tests/test_tool_errors.py | 2 +- .../tests/test_tool_execution.py | 54 +++++++++---------- .../{mcp_v2 => mcp_server}/tools/__init__.py | 0 .../{mcp_v2 => mcp_server}/tools/execution.py | 4 +- backend/{mcp_v2 => mcp_server}/tools/info.py | 2 +- backend/{mcp_v2 => mcp_server}/urls.py | 2 +- backend/{mcp_v2 => mcp_server}/views.py | 8 +-- 19 files changed, 48 insertions(+), 48 deletions(-) rename backend/{mcp_v2 => mcp_server}/README.md (100%) rename backend/{mcp_v2 => mcp_server}/__init__.py (100%) rename backend/{mcp_v2 => mcp_server}/apps.py (72%) rename backend/{mcp_v2 => mcp_server}/constants.py (100%) rename backend/{mcp_v2 => mcp_server}/context.py (100%) rename backend/{mcp_v2 => mcp_server}/exceptions.py (100%) rename backend/{mcp_v2 => mcp_server}/registry.py (97%) rename backend/{mcp_v2 => mcp_server}/tests/__init__.py (100%) rename backend/{mcp_v2 => mcp_server}/tests/test_mcp_auth.py (98%) rename backend/{mcp_v2 => mcp_server}/tests/test_mcp_protocol.py (97%) rename backend/{mcp_v2 => mcp_server}/tests/test_tool_errors.py (96%) rename backend/{mcp_v2 => mcp_server}/tests/test_tool_execution.py (84%) rename backend/{mcp_v2 => mcp_server}/tools/__init__.py (100%) rename backend/{mcp_v2 => mcp_server}/tools/execution.py (99%) rename backend/{mcp_v2 => mcp_server}/tools/info.py (98%) rename backend/{mcp_v2 => mcp_server}/urls.py (95%) rename backend/{mcp_v2 => mcp_server}/views.py (98%) diff --git a/backend/backend/base_urls.py b/backend/backend/base_urls.py index 353a253cec..ee539e39de 100644 --- a/backend/backend/base_urls.py +++ b/backend/backend/base_urls.py @@ -23,7 +23,7 @@ include("pipeline_v2.public_api_urls"), ), # Hosted MCP server, authenticated by the API deployment's own key - path(f"{settings.MCP_PATH_PREFIX}/", include("mcp_v2.urls")), + path(f"{settings.MCP_PATH_PREFIX}/", include("mcp_server.urls")), path("", include("health.urls")), # Internal API for worker communication path("internal/", include("backend.internal_base_urls")), diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 55c069998a..85947ad807 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -382,7 +382,7 @@ def filter(self, record): "pipeline_v2", "platform_settings_v2", "api_v2", - "mcp_v2", + "mcp_server", "usage_v2", "notification_v2", "prompt_studio.prompt_profile_manager_v2", diff --git a/backend/mcp_v2/README.md b/backend/mcp_server/README.md similarity index 100% rename from backend/mcp_v2/README.md rename to backend/mcp_server/README.md diff --git a/backend/mcp_v2/__init__.py b/backend/mcp_server/__init__.py similarity index 100% rename from backend/mcp_v2/__init__.py rename to backend/mcp_server/__init__.py diff --git a/backend/mcp_v2/apps.py b/backend/mcp_server/apps.py similarity index 72% rename from backend/mcp_v2/apps.py rename to backend/mcp_server/apps.py index fce3156468..3447b53558 100644 --- a/backend/mcp_v2/apps.py +++ b/backend/mcp_server/apps.py @@ -2,4 +2,4 @@ class MCPConfig(AppConfig): - name = "mcp_v2" + name = "mcp_server" diff --git a/backend/mcp_v2/constants.py b/backend/mcp_server/constants.py similarity index 100% rename from backend/mcp_v2/constants.py rename to backend/mcp_server/constants.py diff --git a/backend/mcp_v2/context.py b/backend/mcp_server/context.py similarity index 100% rename from backend/mcp_v2/context.py rename to backend/mcp_server/context.py diff --git a/backend/mcp_v2/exceptions.py b/backend/mcp_server/exceptions.py similarity index 100% rename from backend/mcp_v2/exceptions.py rename to backend/mcp_server/exceptions.py diff --git a/backend/mcp_v2/registry.py b/backend/mcp_server/registry.py similarity index 97% rename from backend/mcp_v2/registry.py rename to backend/mcp_server/registry.py index 238adbbc8d..d211c00f96 100644 --- a/backend/mcp_v2/registry.py +++ b/backend/mcp_server/registry.py @@ -74,13 +74,13 @@ def build_registry() -> MCPToolRegistry: Imported lazily inside the function so that registering a tool cannot trigger Django model imports at module-import time. """ - from mcp_v2.tools.execution import ( + from mcp_server.tools.execution import ( extract_document, extract_document_schema, get_execution_status, get_execution_status_schema, ) - from mcp_v2.tools.info import get_api_info, get_api_info_schema, read_me_first + from mcp_server.tools.info import get_api_info, get_api_info_schema, read_me_first registry = MCPToolRegistry() diff --git a/backend/mcp_v2/tests/__init__.py b/backend/mcp_server/tests/__init__.py similarity index 100% rename from backend/mcp_v2/tests/__init__.py rename to backend/mcp_server/tests/__init__.py diff --git a/backend/mcp_v2/tests/test_mcp_auth.py b/backend/mcp_server/tests/test_mcp_auth.py similarity index 98% rename from backend/mcp_v2/tests/test_mcp_auth.py rename to backend/mcp_server/tests/test_mcp_auth.py index c92fffda07..5118fb3cf9 100644 --- a/backend/mcp_v2/tests/test_mcp_auth.py +++ b/backend/mcp_server/tests/test_mcp_auth.py @@ -23,7 +23,7 @@ from utils.user_context import UserContext from workflow_manager.workflow_v2.models.workflow import Workflow -from mcp_v2.views import MCPServerView +from mcp_server.views import MCPServerView ORG_ID = "org-mcp" @@ -79,7 +79,7 @@ def _post( return self.view(request, **kwargs) @pytest.mark.critical_path("mcp-server-auth") - @patch("mcp_v2.views.TOOL_REGISTRY.get") + @patch("mcp_server.views.TOOL_REGISTRY.get") def test_bad_credentials_rejected_before_dispatch(self, registry_get) -> None: cases = [ ("missing header", "live-api", None, ORG_ID), diff --git a/backend/mcp_v2/tests/test_mcp_protocol.py b/backend/mcp_server/tests/test_mcp_protocol.py similarity index 97% rename from backend/mcp_v2/tests/test_mcp_protocol.py rename to backend/mcp_server/tests/test_mcp_protocol.py index e29b273f52..e6b43df72c 100644 --- a/backend/mcp_v2/tests/test_mcp_protocol.py +++ b/backend/mcp_server/tests/test_mcp_protocol.py @@ -14,11 +14,11 @@ from django.test import SimpleTestCase from rest_framework.test import APIRequestFactory -from mcp_v2.constants import JSONRPC, MCPServer -from mcp_v2.context import MCPContext -from mcp_v2.exceptions import MCPToolError -from mcp_v2.registry import TOOL_REGISTRY -from mcp_v2.views import MCPServerView +from mcp_server.constants import JSONRPC, MCPServer +from mcp_server.context import MCPContext +from mcp_server.exceptions import MCPToolError +from mcp_server.registry import TOOL_REGISTRY +from mcp_server.views import MCPServerView ORG_ID = "org-mcp" API_NAME = "live-api" diff --git a/backend/mcp_v2/tests/test_tool_errors.py b/backend/mcp_server/tests/test_tool_errors.py similarity index 96% rename from backend/mcp_v2/tests/test_tool_errors.py rename to backend/mcp_server/tests/test_tool_errors.py index 5f4a9c8d28..f8a70823e2 100644 --- a/backend/mcp_v2/tests/test_tool_errors.py +++ b/backend/mcp_server/tests/test_tool_errors.py @@ -10,7 +10,7 @@ from django.test import SimpleTestCase from rest_framework.exceptions import ValidationError -from mcp_v2.tools.execution import _format_validation_error +from mcp_server.tools.execution import _format_validation_error class FormatValidationErrorTest(SimpleTestCase): diff --git a/backend/mcp_v2/tests/test_tool_execution.py b/backend/mcp_server/tests/test_tool_execution.py similarity index 84% rename from backend/mcp_v2/tests/test_tool_execution.py rename to backend/mcp_server/tests/test_tool_execution.py index 46ab0ad3d4..869fc5f2b3 100644 --- a/backend/mcp_v2/tests/test_tool_execution.py +++ b/backend/mcp_server/tests/test_tool_execution.py @@ -17,9 +17,9 @@ from rest_framework.exceptions import ValidationError from workflow_manager.workflow_v2.dto import ExecutionResponse -from mcp_v2.context import MCPContext -from mcp_v2.exceptions import MCPToolError -from mcp_v2.tools.execution import extract_document, get_execution_status +from mcp_server.context import MCPContext +from mcp_server.exceptions import MCPToolError +from mcp_server.tools.execution import extract_document, get_execution_status # The execution serializer accepts only HTTPS S3 pre-signed URLs; using a # realistic one here keeps these tests honest about what the tool accepts. @@ -52,17 +52,17 @@ def _run(self, execute_return=None, **kwargs): """Drive extract_document with the whole execution stack stubbed.""" with ( patch( - "mcp_v2.tools.execution.DeploymentHelper.load_presigned_files" + "mcp_server.tools.execution.DeploymentHelper.load_presigned_files" ) as load_files, patch( - "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire", return_value=(True, {}), ), patch( - "mcp_v2.tools.execution.APIDeploymentRateLimiter.release_slot" + "mcp_server.tools.execution.APIDeploymentRateLimiter.release_slot" ) as release, patch( - "mcp_v2.tools.execution.DeploymentHelper.execute_workflow", + "mcp_server.tools.execution.DeploymentHelper.execute_workflow", return_value=execute_return if execute_return is not None else {"execution_status": "COMPLETED", "message": "ok"}, @@ -131,7 +131,7 @@ def test_defaults_are_conservative(self) -> None: def test_inactive_deployment_rejected_before_quota_is_spent(self) -> None: with patch( - "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + "mcp_server.tools.execution.DeploymentHelper.execute_workflow" ) as execute: with self.assertRaises(MCPToolError): extract_document(make_context(active=False), document_urls=[DOC_URL]) @@ -144,7 +144,7 @@ def test_non_s3_url_rejected_with_an_explanation(self) -> None: an ordinary link must be told why, not just that it failed. """ with patch( - "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + "mcp_server.tools.execution.DeploymentHelper.execute_workflow" ) as execute: with self.assertRaises(MCPToolError) as caught: extract_document( @@ -157,10 +157,10 @@ def test_non_s3_url_rejected_with_an_explanation(self) -> None: def test_invalid_url_rejected_before_quota_is_spent(self) -> None: with ( patch( - "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire" + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire" ) as acquire, patch( - "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + "mcp_server.tools.execution.DeploymentHelper.execute_workflow" ) as execute, ): with self.assertRaises(MCPToolError): @@ -175,17 +175,17 @@ def test_rate_limited_call_downloads_nothing(self) -> None: """ with ( patch( - "mcp_v2.tools.execution.DeploymentHelper.load_presigned_files" + "mcp_server.tools.execution.DeploymentHelper.load_presigned_files" ) as load_files, patch( - "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire", return_value=( False, {"current_usage": 5, "limit": 5, "limit_type": "organization"}, ), ), patch( - "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + "mcp_server.tools.execution.DeploymentHelper.execute_workflow" ) as execute, ): with self.assertRaises(MCPToolError) as caught: @@ -201,18 +201,18 @@ def test_failed_download_releases_the_rate_limit_slot(self) -> None: """ with ( patch( - "mcp_v2.tools.execution.DeploymentHelper.load_presigned_files", + "mcp_server.tools.execution.DeploymentHelper.load_presigned_files", side_effect=RuntimeError("403 Forbidden"), ), patch( - "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire", return_value=(True, {}), ), patch( - "mcp_v2.tools.execution.APIDeploymentRateLimiter.release_slot" + "mcp_server.tools.execution.APIDeploymentRateLimiter.release_slot" ) as release, patch( - "mcp_v2.tools.execution.DeploymentHelper.execute_workflow" + "mcp_server.tools.execution.DeploymentHelper.execute_workflow" ) as execute, ): with self.assertRaises(MCPToolError): @@ -226,16 +226,16 @@ def test_failed_execution_releases_the_rate_limit_slot(self) -> None: release must survive the failure path. """ with ( - patch("mcp_v2.tools.execution.DeploymentHelper.load_presigned_files"), + patch("mcp_server.tools.execution.DeploymentHelper.load_presigned_files"), patch( - "mcp_v2.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire", return_value=(True, {}), ), patch( - "mcp_v2.tools.execution.APIDeploymentRateLimiter.release_slot" + "mcp_server.tools.execution.APIDeploymentRateLimiter.release_slot" ) as release, patch( - "mcp_v2.tools.execution.DeploymentHelper.execute_workflow", + "mcp_server.tools.execution.DeploymentHelper.execute_workflow", side_effect=RuntimeError("boom"), ), ): @@ -252,11 +252,11 @@ def setUp(self) -> None: def _run(self, response: ExecutionResponse, **kwargs): with ( patch( - "mcp_v2.tools.execution.ExecutionQuerySerializer.is_valid", + "mcp_server.tools.execution.ExecutionQuerySerializer.is_valid", return_value=True, ), patch( - "mcp_v2.tools.execution.ExecutionQuerySerializer.validated_data", + "mcp_server.tools.execution.ExecutionQuerySerializer.validated_data", { "execution_id": EXECUTION_ID, "include_metadata": kwargs.get("include_metadata", False), @@ -267,11 +267,11 @@ def _run(self, response: ExecutionResponse, **kwargs): }, ), patch( - "mcp_v2.tools.execution.DeploymentHelper.get_execution_status", + "mcp_server.tools.execution.DeploymentHelper.get_execution_status", return_value=response, ), patch( - "mcp_v2.tools.execution.DeploymentHelper.process_completed_execution" + "mcp_server.tools.execution.DeploymentHelper.process_completed_execution" ) as process, ): result = get_execution_status( @@ -332,7 +332,7 @@ def test_unknown_execution_id_is_an_agent_error(self) -> None: fault — the agent may simply have mistyped it. """ with patch( - "mcp_v2.tools.execution.ExecutionQuerySerializer.is_valid", + "mcp_server.tools.execution.ExecutionQuerySerializer.is_valid", side_effect=ValidationError({"execution_id": ["Invalid execution_id."]}), ): with self.assertRaises(MCPToolError) as caught: diff --git a/backend/mcp_v2/tools/__init__.py b/backend/mcp_server/tools/__init__.py similarity index 100% rename from backend/mcp_v2/tools/__init__.py rename to backend/mcp_server/tools/__init__.py diff --git a/backend/mcp_v2/tools/execution.py b/backend/mcp_server/tools/execution.py similarity index 99% rename from backend/mcp_v2/tools/execution.py rename to backend/mcp_server/tools/execution.py index b2af589010..40f00bcadc 100644 --- a/backend/mcp_v2/tools/execution.py +++ b/backend/mcp_server/tools/execution.py @@ -15,8 +15,8 @@ from utils.enums import CeleryTaskState from workflow_manager.workflow_v2.dto import ExecutionResponse -from mcp_v2.context import MCPContext -from mcp_v2.exceptions import MCPToolError +from mcp_server.context import MCPContext +from mcp_server.exceptions import MCPToolError logger = logging.getLogger(__name__) diff --git a/backend/mcp_v2/tools/info.py b/backend/mcp_server/tools/info.py similarity index 98% rename from backend/mcp_v2/tools/info.py rename to backend/mcp_server/tools/info.py index a5ad1da9a9..41fc08928b 100644 --- a/backend/mcp_v2/tools/info.py +++ b/backend/mcp_server/tools/info.py @@ -3,7 +3,7 @@ import logging from typing import Any -from mcp_v2.context import MCPContext +from mcp_server.context import MCPContext logger = logging.getLogger(__name__) diff --git a/backend/mcp_v2/urls.py b/backend/mcp_server/urls.py similarity index 95% rename from backend/mcp_v2/urls.py rename to backend/mcp_server/urls.py index 6a6ce1cdb8..89d42b3892 100644 --- a/backend/mcp_v2/urls.py +++ b/backend/mcp_server/urls.py @@ -9,7 +9,7 @@ from django.urls import re_path -from mcp_v2.views import MCPServerView +from mcp_server.views import MCPServerView mcp_server = MCPServerView.as_view() diff --git a/backend/mcp_v2/views.py b/backend/mcp_server/views.py similarity index 98% rename from backend/mcp_v2/views.py rename to backend/mcp_server/views.py index 125bc32733..b4599ad9ad 100644 --- a/backend/mcp_v2/views.py +++ b/backend/mcp_server/views.py @@ -16,10 +16,10 @@ from rest_framework.request import Request from rest_framework.response import Response -from mcp_v2.constants import JSONRPC, MCPMethod, MCPServer -from mcp_v2.context import MCPContext -from mcp_v2.exceptions import MCPToolError -from mcp_v2.registry import TOOL_REGISTRY +from mcp_server.constants import JSONRPC, MCPMethod, MCPServer +from mcp_server.context import MCPContext +from mcp_server.exceptions import MCPToolError +from mcp_server.registry import TOOL_REGISTRY logger = logging.getLogger(__name__) From 2ab12a70d8d00585c35d3a894b1d3e6600dad9f3 Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sat, 25 Jul 2026 18:11:21 +0530 Subject: [PATCH 04/16] feat(mcp): add an organization-scoped MCP server alongside the deployment one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The deployment MCP server can only ever expose one workflow: its credential is an API deployment key, which resolves to no user and grants nothing beyond that deployment. Reaching the rest of the platform needs a different credential — and Unstract already has one. PlatformApiKey (backend/platform_api/) is an org-scoped bearer key that CustomAuthMiddleware resolves to a service-account User, validating the key's organization against the URL and rejecting unknown, inactive and wrong-org keys before any view runs. So the new server does not authenticate at all; it is mounted where that middleware can do it. That placement is the security-critical detail. WHITELISTED_PATHS is matched with startswith, so the deployment server's /mcp/ prefix is exempt from the middleware while the tenant path is not: POST /mcp/// deployment key, view authenticates POST /api/v1/unstract//mcp/ platform key, middleware authenticates Moving the latter under the whitelisted prefix would silently remove all authentication, so a test asserts it is not whitelisted. Because that auth lives in middleware, the platform auth tests go through django.test.Client and the real URL. APIRequestFactory — correct for the deployment server, which owns its auth — bypasses middleware entirely and would have passed against a completely open endpoint. Shared JSON-RPC transport extracted to transport.py so the two servers cannot drift on protocol behaviour; the deployment server's URL, behaviour and 34 tests are unchanged. The platform tools are read-only by design (readMeFirst, whoami, listApiDeployments, listWorkflows, listPromptStudioProjects). Listings go through each model's for_user manager so platform sharing rules apply, and are capped with an explicit truncation note rather than silently cut. Two constraints are documented rather than worked around: - A read-tier key cannot use the server at all: the middleware gates tiers on HTTP method and every MCP call is a POST. Special-casing MCP paths in shared middleware is a maintainer decision, not this app's to make. - Service accounts bypass per-user sharing (for_user returns self.all()), so listings cover the whole org. whoami and readMeFirst both say so, since an agent assuming otherwise would misread a listing. check_tool_allowed refuses write tools below read_write, and a test pins the registry as read-only so adding the first write tool is a deliberate act. Tests: 50 in this app (was 34); full backend unit tier 310 passing. Registers the mcp-platform-auth critical path. Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/backend/urls_v2.py | 4 + backend/mcp_server/README.md | 124 +++++++- backend/mcp_server/context.py | 25 ++ backend/mcp_server/platform_urls.py | 21 ++ backend/mcp_server/platform_views.py | 127 ++++++++ backend/mcp_server/registry.py | 96 +++++- backend/mcp_server/tests/test_mcp_auth.py | 2 +- backend/mcp_server/tests/test_mcp_protocol.py | 10 +- .../mcp_server/tests/test_platform_auth.py | 164 ++++++++++ .../tests/test_platform_tier_guard.py | 76 +++++ .../mcp_server/tests/test_platform_tools.py | 148 +++++++++ backend/mcp_server/tools/platform.py | 169 +++++++++++ backend/mcp_server/transport.py | 280 ++++++++++++++++++ backend/mcp_server/views.py | 265 +++-------------- tests/critical_paths.yaml | 5 + 15 files changed, 1264 insertions(+), 252 deletions(-) create mode 100644 backend/mcp_server/platform_urls.py create mode 100644 backend/mcp_server/platform_views.py create mode 100644 backend/mcp_server/tests/test_platform_auth.py create mode 100644 backend/mcp_server/tests/test_platform_tier_guard.py create mode 100644 backend/mcp_server/tests/test_platform_tools.py create mode 100644 backend/mcp_server/tools/platform.py create mode 100644 backend/mcp_server/transport.py diff --git a/backend/backend/urls_v2.py b/backend/backend/urls_v2.py index fd91d3cc2f..be12a14d1a 100644 --- a/backend/backend/urls_v2.py +++ b/backend/backend/urls_v2.py @@ -65,4 +65,8 @@ path("execution/", include("workflow_manager.file_execution.urls")), path("metrics/", include("dashboard_metrics.urls")), path("platform-api/", include("platform_api.urls")), + # Organization-scoped MCP server. Mounted here rather than under + # MCP_PATH_PREFIX so CustomAuthMiddleware authenticates the platform API + # key — that prefix is whitelisted and would bypass auth entirely. + path("mcp/", include("mcp_server.platform_urls")), ] diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index 146e44cb11..4d8e862d4a 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -1,11 +1,31 @@ -# Hosted MCP server +# Hosted MCP servers -Exposes an Unstract API deployment to coding agents over the -[Model Context Protocol](https://modelcontextprotocol.io), so an agent can run -document extraction as a tool call instead of hand-rolling HTTP requests. +Exposes Unstract to coding agents over the +[Model Context Protocol](https://modelcontextprotocol.io), so an agent can +discover and run document extraction as tool calls instead of hand-rolling HTTP +requests. -The server is hosted inside the existing backend — it is a Django app served by -the same gunicorn process, not a separate service to deploy or scale. +Both servers are hosted inside the existing backend — Django views served by the +same gunicorn process, not separate services to deploy or scale. They share the +JSON-RPC transport in `transport.py` and differ only in how a caller is +authenticated and which tools they expose. + +| | Deployment server | Platform server | +| --- | --- | --- | +| Scope | one API deployment | one organization | +| Credential | that deployment's API key | a platform API key | +| Authenticated by | the view itself | `CustomAuthMiddleware` | +| URL | `/mcp///` | `/api/v1/unstract//mcp/` | +| Tools | extract, poll status | read-only discovery | + +The split is not cosmetic. A deployment key grants exactly one workflow and +resolves to no user, so it cannot authorize anything organization-wide; a +platform key resolves to a service-account user and is checked by the shared +auth middleware. Neither key works on the other server. + +--- + +# Deployment server ## Endpoint @@ -90,7 +110,7 @@ internal detail does not leak to the client. ## Design notes -- **Auth reuses the deployment key path.** `_resolve_context` calls the same +- **Auth reuses the deployment key path.** `resolve_context` calls the same `DeploymentHelper` validation the REST endpoint uses, so the two surfaces cannot drift apart on who is allowed in. - **Execution reuses `ExecutionRequestSerializer`.** URL validation and the @@ -102,9 +122,93 @@ internal detail does not leak to the client. protocol errors as unrecoverable transport faults; an agent-fixable problem comes back as `isError: true` content it can read and retry. +--- + +# Platform server + +Organization-scoped and **read-only**: it lets an agent discover what exists in +an Unstract organization — deployments, workflows, Prompt Studio projects — but +cannot run or change anything. + +## Endpoint + +``` +POST /api/v1/unstract//mcp/ +Authorization: Bearer +``` + +```bash +claude mcp add --transport http unstract-platform \ + https:///api/v1/unstract//mcp/ \ + --header "Authorization: Bearer " +``` + +Platform API keys are managed at `/api/v1/unstract//platform-api/keys/`. + +## Tools + +| Tool | Purpose | +| --- | --- | +| `readMeFirst` | Orientation, including what this credential can see. | +| `whoami` | The key's organization, permission tier and visibility scope. | +| `listApiDeployments` | Deployed extraction endpoints, with their `api_name`. | +| `listWorkflows` | The pipelines behind them. | +| `listPromptStudioProjects` | Where extraction prompts are authored. | + +To actually extract, an agent calls `listApiDeployments` to find an `api_name`, +then opens a **separate** session against the deployment server above. + +## Why the URL matters + +This server is mounted under the tenant prefix, next to `platform-api/`, and +**not** under `MCP_PATH_PREFIX`. `WHITELISTED_PATHS` is matched with +`startswith`, so `/mcp/...` is exempt from `CustomAuthMiddleware` while +`/api/v1/unstract//mcp/` is not — which is exactly how this server +inherits platform-key authentication. + +Moving these URLs under the whitelisted prefix would silently remove all +authentication. `test_platform_auth.py::test_the_endpoint_is_not_whitelisted` +guards against that. + +Because auth lives in middleware, its tests go through `django.test.Client` and +the real URL. A test using `APIRequestFactory` or calling the view directly +bypasses the middleware and would pass against a completely open endpoint. + +## Two constraints worth knowing + +**A `read`-tier key cannot use this server at all.** The middleware's tier check +gates on HTTP method, and every MCP call is a `POST`, which `read` disallows — +so a read-only key is refused before reaching the view, even though every tool +here is read-only. Use a `read_write` key. Changing that would mean special- +casing MCP paths in the middleware, which is a decision for maintainers rather +than something this app should do unilaterally. + +**These tools see the whole organization.** A platform key resolves to a service +account, and `is_service_account=True` makes `for_user()` managers return +`self.all()` — so listings ignore per-user sharing regardless of the `USER` role +the key was created with. `whoami` and `readMeFirst` both say so, because an +agent that assumed otherwise would draw wrong conclusions from a listing. + +## Adding a tool here + +Same registry mechanics as the deployment server, with `PlatformMCPContext` +(`user`, `platform_key`, `org_name`) as the first argument. Query through the +model's `for_user(context.user)` manager so the platform's own visibility rules +apply, and cap results — `LIST_LIMIT` with a `truncated` note, never silent +truncation. + +Write tools are deliberately absent. `check_tool_allowed` already refuses a +`writes=True` tool to a `read`-tier key, so the guard is in place, but the +read-only promise is pinned by +`test_platform_tier_guard.py::test_platform_registry_is_entirely_read_only` — +adding a write tool will fail that test and force a deliberate decision about +whether an agent should hold that power. + +--- + ## Not implemented OAuth 2.1 with dynamic client registration, which MCP defines for browser-based -one-click connectors. Header/path bearer auth covers Claude Code and -API clients. Adding OAuth is additive — it would mount discovery endpoints -alongside this router without changing the transport. +one-click connectors. Bearer auth covers Claude Code and API clients on both +servers. Adding OAuth is additive — it would mount discovery endpoints +alongside these routers without changing the transport. diff --git a/backend/mcp_server/context.py b/backend/mcp_server/context.py index 13e0d33a26..bff616ed69 100644 --- a/backend/mcp_server/context.py +++ b/backend/mcp_server/context.py @@ -5,10 +5,35 @@ """ from dataclasses import dataclass +from typing import Any from api_v2.models import APIDeployment +@dataclass(frozen=True) +class PlatformMCPContext: + """Everything a tool needs about an organization-scoped caller. + + Built from what ``CustomAuthMiddleware`` already resolved — this server + does not re-authenticate. Deliberately shares no base class with + ``MCPContext``: the two servers expose disjoint tool sets, and a common + base would invite a tool to be written against whichever fields happen to + overlap. + + Attributes: + user: The service-account user the platform key resolves to. Passed to + org-scoped managers so tools inherit the platform's own sharing + rules rather than reimplementing them. + platform_key: The validated key, carried for its permission tier. + org_name: Organization identifier from the URL, already validated + against the key's own organization by the middleware. + """ + + user: Any + platform_key: Any + org_name: str + + @dataclass(frozen=True) class MCPContext: """Everything a tool needs about the caller's deployment. diff --git a/backend/mcp_server/platform_urls.py b/backend/mcp_server/platform_urls.py new file mode 100644 index 0000000000..0e9a33fb4a --- /dev/null +++ b/backend/mcp_server/platform_urls.py @@ -0,0 +1,21 @@ +"""URLs for the organization-scoped MCP server. + +Mounted under the tenant prefix in ``backend/urls_v2.py``, alongside +``platform-api/``: + + POST /api/v1/unstract//mcp/ + +That placement is load-bearing. ``WHITELISTED_PATHS`` is matched with +``startswith``, so the deployment server's ``/mcp/...`` is exempt from +``CustomAuthMiddleware`` while this path is not — which is precisely how this +server inherits platform-key authentication for free. Moving these URLs under +the whitelisted prefix would silently remove all authentication. +""" + +from django.urls import path + +from mcp_server.platform_views import PlatformMCPServerView + +urlpatterns = [ + path("", PlatformMCPServerView.as_view(), name="platform_mcp_server"), +] diff --git a/backend/mcp_server/platform_views.py b/backend/mcp_server/platform_views.py new file mode 100644 index 0000000000..8f713da4c3 --- /dev/null +++ b/backend/mcp_server/platform_views.py @@ -0,0 +1,127 @@ +"""Organization-scoped MCP server. + +Mounted on a tenant URL — deliberately *not* under the whitelisted ``/mcp/`` +prefix — so ``CustomAuthMiddleware`` runs ahead of it and performs the +authentication. By the time this view executes, the middleware has already: + +* resolved the ``Authorization: Bearer `` platform key, +* rejected unknown, inactive, and wrong-organization keys, +* rejected keys whose permission tier disallows the HTTP method, +* set ``request.user`` to the key's service account and pinned + ``ORGANIZATION_ID`` in the state store. + +This view therefore reads that result rather than re-authenticating. Two +consequences worth stating plainly, because both are easy to get wrong: + +1. ``authentication_classes`` must NOT be emptied here. The deployment server + empties it because it owns its own auth; doing the same here is harmless + only by accident, and setting DRF authenticators that return a user would + actively override what the middleware resolved. +2. Because the auth lives in middleware, a test that calls this view through + ``APIRequestFactory`` bypasses it entirely and would pass against a + completely open endpoint. The auth-boundary tests go through the full URL + stack for exactly this reason. +""" + +import logging +from typing import Any + +from platform_api.models import ApiKeyPermission +from rest_framework.request import Request + +from mcp_server.context import PlatformMCPContext +from mcp_server.registry import PLATFORM_TOOLS +from mcp_server.transport import BaseMCPView + +logger = logging.getLogger(__name__) + + +class PlatformMCPServerView(BaseMCPView): + """MCP JSON-RPC endpoint scoped to an organization. + + Exposes read-only discovery tools over the organization's resources, + authenticated by a platform API key. + """ + + registry = PLATFORM_TOOLS + + # Intentionally not overriding authentication_classes: the project sets no + # DEFAULT_AUTHENTICATION_CLASSES that would resolve a user, so the + # middleware's request.user survives into the view. + permission_classes: list = [] + + def initialize_request(self, request: Request, *args: Any, **kwargs: Any) -> Request: + """Skip CSRF. + + The middleware already marks bearer-authenticated requests CSRF-exempt; + this makes the view's behaviour explicit rather than dependent on that. + """ + request.csrf_processing_done = True + return super().initialize_request(request, *args, **kwargs) + + def auth_failure_message(self) -> str: + return "A valid platform API key is required" + + def server_instructions(self, context: PlatformMCPContext) -> str: + return ( + "Unstract runs LLM-driven extraction over unstructured documents " + "and returns structured JSON. This server is scoped to the " + f"organization '{context.org_name}' and is read-only: it discovers " + "deployments, workflows and Prompt Studio projects but cannot run " + "or change anything. Call readMeFirst before any other tool." + ) + + def resolve_context( + self, request: Request, **kwargs: Any + ) -> PlatformMCPContext | None: + """Read the credential the auth middleware already resolved. + + Returns None when the middleware did not attach a platform key. That is + a defence-in-depth check rather than the primary gate: an unauthenticated + request should never reach this view, because the middleware rejects it + first with its own 401. If it does reach here — a whitelist edit, a + routing change — refusing is the safe answer. + """ + platform_key = getattr(request, "platform_api_key", None) + user = getattr(request, "user", None) + + # The organization is not a URL kwarg here: the tenant prefix consumes + # it, and OrganizationMiddleware parses it onto the request. Read it + # from there rather than from the route. + org_name = getattr(request, "organization_id", None) + + if platform_key is None or user is None or not user.is_authenticated: + logger.warning( + "Platform MCP reached without a resolved platform key " + f"(org '{org_name}'). Check that this path is not whitelisted." + ) + return None + + return PlatformMCPContext( + user=user, + platform_key=platform_key, + # Fall back to the key's own organization so the context is never + # built with a null org, whatever the routing did. + org_name=org_name or str(platform_key.organization.organization_id), + ) + + def check_tool_allowed(self, tool: Any, context: PlatformMCPContext) -> str | None: + """Refuse write tools to keys below the read_write tier. + + The middleware's tier check gates on HTTP method, and every MCP call is + a POST — so it cannot distinguish a tool that reads from one that + writes. This is where that distinction is actually enforced. + + Today the platform registry is entirely read-only, so this never fires. + It is here so that adding the first write tool is a one-line change + that is already guarded, rather than a silent privilege widening. + """ + if not tool.writes: + return None + + if context.platform_key.permission == ApiKeyPermission.READ.value: + return ( + f"Tool '{tool.name}' modifies data and requires a platform API " + "key with the read_write or full_access permission." + ) + return None diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index d211c00f96..cf4192e09e 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -68,8 +68,8 @@ def list_schemas(self) -> list[dict[str, Any]]: return [tool.to_mcp_schema() for tool in self._tools.values()] -def build_registry() -> MCPToolRegistry: - """Build the registry of tools exposed by the Unstract MCP server. +def build_deployment_registry() -> MCPToolRegistry: + """Build the tools exposed by the deployment-scoped MCP server. Imported lazily inside the function so that registering a tool cannot trigger Django model imports at module-import time. @@ -151,6 +151,92 @@ def build_registry() -> MCPToolRegistry: return registry -# Single shared registry. The tool set is static, so building it once at import -# time is safe and keeps per-request work down. -TOOL_REGISTRY = build_registry() +def build_platform_registry() -> MCPToolRegistry: + """Build the tools exposed by the organization-scoped MCP server. + + Read-only by design. Every tool here is reachable by any caller holding a + platform key, and an agent will eventually call a destructive tool by + mistake — so writes are deferred until there is a per-tool authorization + story stronger than the key's HTTP-method tier. + """ + from mcp_server.tools.platform import ( + list_api_deployments, + list_prompt_studio_projects, + list_workflows, + no_args_schema, + platform_read_me_first, + whoami, + ) + + registry = MCPToolRegistry() + + registry.register( + MCPTool( + name="readMeFirst", + description=( + "START HERE. Returns a guide to this MCP server: what it can " + "see in the connected Unstract organization, the available " + "tools, and the recommended call sequence. Takes no arguments." + ), + input_schema=no_args_schema(), + handler=platform_read_me_first, + ) + ) + registry.register( + MCPTool( + name="whoami", + description=( + "Describe the credential this session is using: the " + "organization it belongs to, its permission tier, and the " + "scope of what it can see. Call this when a tool returns less " + "or more than you expected. Takes no arguments." + ), + input_schema=no_args_schema(), + handler=whoami, + ) + ) + registry.register( + MCPTool( + name="listApiDeployments", + description=( + "List the organization's API deployments — the deployed " + "extraction endpoints. Use this to discover what can be " + "extracted and to find the api_name needed to connect a " + "deployment-scoped MCP session. Takes no arguments." + ), + input_schema=no_args_schema(), + handler=list_api_deployments, + ) + ) + registry.register( + MCPTool( + name="listWorkflows", + description=( + "List the organization's workflows — the pipelines that API " + "deployments and ETL pipelines run. Takes no arguments." + ), + input_schema=no_args_schema(), + handler=list_workflows, + ) + ) + registry.register( + MCPTool( + name="listPromptStudioProjects", + description=( + "List the organization's Prompt Studio projects, where " + "extraction prompts are authored before being exported as " + "tools and deployed. Takes no arguments." + ), + input_schema=no_args_schema(), + handler=list_prompt_studio_projects, + ) + ) + + return registry + + +# Registries are static, so building them once at import time is safe and keeps +# per-request work down. Kept separate rather than merged with a filter, so a +# tool cannot be exposed on the wrong server by forgetting a flag. +DEPLOYMENT_TOOLS = build_deployment_registry() +PLATFORM_TOOLS = build_platform_registry() diff --git a/backend/mcp_server/tests/test_mcp_auth.py b/backend/mcp_server/tests/test_mcp_auth.py index 5118fb3cf9..e0ff95228b 100644 --- a/backend/mcp_server/tests/test_mcp_auth.py +++ b/backend/mcp_server/tests/test_mcp_auth.py @@ -79,7 +79,7 @@ def _post( return self.view(request, **kwargs) @pytest.mark.critical_path("mcp-server-auth") - @patch("mcp_server.views.TOOL_REGISTRY.get") + @patch("mcp_server.registry.DEPLOYMENT_TOOLS.get") def test_bad_credentials_rejected_before_dispatch(self, registry_get) -> None: cases = [ ("missing header", "live-api", None, ORG_ID), diff --git a/backend/mcp_server/tests/test_mcp_protocol.py b/backend/mcp_server/tests/test_mcp_protocol.py index e6b43df72c..1a58ac4afb 100644 --- a/backend/mcp_server/tests/test_mcp_protocol.py +++ b/backend/mcp_server/tests/test_mcp_protocol.py @@ -17,7 +17,7 @@ from mcp_server.constants import JSONRPC, MCPServer from mcp_server.context import MCPContext from mcp_server.exceptions import MCPToolError -from mcp_server.registry import TOOL_REGISTRY +from mcp_server.registry import DEPLOYMENT_TOOLS from mcp_server.views import MCPServerView ORG_ID = "org-mcp" @@ -61,16 +61,16 @@ def _failing_tool(self, error: Exception, name: str = "getApiInfo"): def boom(*args, **kwargs): raise error - original = TOOL_REGISTRY.get(name) + original = DEPLOYMENT_TOOLS.get(name) return patch.object( - TOOL_REGISTRY, "get", return_value=replace(original, handler=boom) + DEPLOYMENT_TOOLS, "get", return_value=replace(original, handler=boom) ) def _call(self, body): """POST a JSON-RPC body with auth stubbed out.""" request = self.factory.post(f"/mcp/{ORG_ID}/{API_NAME}/", body, format="json") with patch.object( - MCPServerView, "_resolve_context", return_value=self.context + MCPServerView, "resolve_context", return_value=self.context ): response = self.view( request, org_name=ORG_ID, api_name=API_NAME, api_key="test-key" @@ -206,7 +206,7 @@ def test_batch_request_rejected(self) -> None: format="json", ) with patch.object( - MCPServerView, "_resolve_context", return_value=self.context + MCPServerView, "resolve_context", return_value=self.context ): response = self.view( request, org_name=ORG_ID, api_name=API_NAME, api_key="test-key" diff --git a/backend/mcp_server/tests/test_platform_auth.py b/backend/mcp_server/tests/test_platform_auth.py new file mode 100644 index 0000000000..5a281529c5 --- /dev/null +++ b/backend/mcp_server/tests/test_platform_auth.py @@ -0,0 +1,164 @@ +"""Critical path ``mcp-platform-auth``: the organization-scoped MCP endpoint +is authenticated by ``CustomAuthMiddleware``, not by the view. + +That distinction is the whole point of these tests. The deployment MCP server +owns its own auth, so testing it through ``APIRequestFactory`` is sound. This +server does not — its credential is resolved by middleware, which +``APIRequestFactory`` and direct view calls bypass entirely. A test written +that way would pass against a completely unauthenticated endpoint. + +So every test here goes through ``django.test.Client`` against the real tenant +URL, exercising the full middleware stack. Needs a live DB (integration tier). +""" + +from __future__ import annotations + +import json +import uuid + +import pytest +from account_v2.models import Organization, User +from django.test import Client, TestCase +from platform_api.models import PlatformApiKey +from tenant_account_v2.models import OrganizationMember +from utils.user_context import UserContext + +ORG_ID = "org-platform-mcp" +OTHER_ORG_ID = "org-platform-other" + + +def make_org_with_key(org_id: str, permission: str = "read_write"): + """Create an organization and a platform key backed by a service account.""" + org = Organization.objects.create( + name=org_id, display_name=org_id, organization_id=org_id + ) + user = User.objects.create( + username=f"svc-{org_id}", + email=f"svc-{org_id}@platform.internal", + user_id=f"uid-{org_id}", + is_service_account=True, + ) + OrganizationMember.objects.create(user=user, organization=org, role="user") + key = PlatformApiKey.objects.create( + name=f"key-{org_id}", + description="test key", + organization=org, + api_user=user, + permission=permission, + ) + return org, user, key + + +class PlatformMCPAuthTest(TestCase): + def setUp(self) -> None: + self.org, self.user, self.key = make_org_with_key(ORG_ID) + _, _, self.other_key = make_org_with_key(OTHER_ORG_ID) + UserContext.set_organization_identifier(ORG_ID) + self.client = Client() + self.url = f"/api/v1/unstract/{ORG_ID}/mcp/" + + def _post(self, auth: str | None = None, body: dict | None = None, url: str = None): + """POST through the full URL stack so the auth middleware runs.""" + payload = body if body is not None else { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/list", + } + headers = {"HTTP_AUTHORIZATION": auth} if auth else {} + return self.client.post( + url or self.url, + data=json.dumps(payload), + content_type="application/json", + **headers, + ) + + @pytest.mark.critical_path("mcp-platform-auth") + def test_bad_credentials_are_rejected(self) -> None: + """Each of these must be refused before any tool is reachable. + + The middleware answers most of them, which is exactly what is being + verified — that this endpoint sits behind it. + """ + cases = [ + ("no authorization header", None), + ("not a bearer token", f"Token {self.key.key}"), + ("empty bearer", "Bearer "), + ("malformed uuid", "Bearer not-a-uuid"), + ("unknown key", f"Bearer {uuid.uuid4()}"), + ("key from another organization", f"Bearer {self.other_key.key}"), + ] + for label, auth in cases: + with self.subTest(label): + response = self._post(auth) + assert response.status_code in (401, 403), ( + f"{label}: got {response.status_code}, body={response.content[:200]}" + ) + + @pytest.mark.critical_path("mcp-platform-auth") + def test_inactive_key_is_rejected(self) -> None: + self.key.is_active = False + self.key.save() + + response = self._post(f"Bearer {self.key.key}") + + assert response.status_code == 401, response.content + + @pytest.mark.critical_path("mcp-platform-auth") + def test_read_tier_key_cannot_reach_the_server(self) -> None: + """A documented limitation, pinned so it cannot change silently. + + The middleware's tier check gates on HTTP method and every MCP call is + a POST, so a `read` key is refused outright — even though this server + exposes nothing but read tools. + """ + self.key.permission = "read" + self.key.save() + + response = self._post(f"Bearer {self.key.key}") + + assert response.status_code == 403, response.content + + @pytest.mark.critical_path("mcp-platform-auth") + def test_valid_key_reaches_the_tool_listing(self) -> None: + """The inverse guard: a check that refused everything would satisfy + every rejection case above. + """ + response = self._post(f"Bearer {self.key.key}") + + assert response.status_code == 200, response.content + tools = json.loads(response.content)["result"]["tools"] + assert [t["name"] for t in tools] == [ + "readMeFirst", + "whoami", + "listApiDeployments", + "listWorkflows", + "listPromptStudioProjects", + ] + + @pytest.mark.critical_path("mcp-platform-auth") + def test_the_endpoint_is_not_whitelisted(self) -> None: + """Guards the placement this server's security depends on. + + If these URLs were ever moved under the whitelisted ``/mcp/`` prefix, + ``CustomAuthMiddleware`` would skip them and the endpoint would answer + unauthenticated callers. An unauthenticated request must not get a + JSON-RPC result. + """ + from django.conf import settings + + assert not any( + self.url.startswith(path) for path in settings.WHITELISTED_PATHS + ), f"{self.url} is whitelisted — it would bypass authentication entirely" + + response = self._post(auth=None) + assert response.status_code in (401, 403) + assert b'"result"' not in response.content + + @pytest.mark.critical_path("mcp-platform-auth") + def test_org_in_url_must_match_the_key(self) -> None: + """A valid key must not reach another organization's MCP endpoint.""" + response = self._post( + f"Bearer {self.key.key}", url=f"/api/v1/unstract/{OTHER_ORG_ID}/mcp/" + ) + + assert response.status_code in (401, 403), response.content diff --git a/backend/mcp_server/tests/test_platform_tier_guard.py b/backend/mcp_server/tests/test_platform_tier_guard.py new file mode 100644 index 0000000000..ede1936c70 --- /dev/null +++ b/backend/mcp_server/tests/test_platform_tier_guard.py @@ -0,0 +1,76 @@ +"""The per-tool permission guard on the platform server. + +The platform registry is read-only today, so this guard never fires in +practice. It is tested anyway: it exists so that adding the first write tool is +already protected, and an unexercised guard is one that quietly stops working +before anyone depends on it. +""" + +from __future__ import annotations + +from unittest.mock import Mock + +from django.test import SimpleTestCase + +from mcp_server.context import PlatformMCPContext +from mcp_server.platform_views import PlatformMCPServerView +from mcp_server.registry import MCPTool, PLATFORM_TOOLS + + +def a_tool(writes: bool) -> MCPTool: + return MCPTool( + name="doThing", + description="d", + input_schema={"type": "object", "properties": {}}, + handler=lambda ctx: None, + writes=writes, + ) + + +def context_for(tier: str) -> PlatformMCPContext: + return PlatformMCPContext( + user=Mock(is_service_account=True), + platform_key=Mock(permission=tier, name="k"), + org_name="org-tier", + ) + + +class PlatformTierGuardTest(SimpleTestCase): + def setUp(self) -> None: + self.view = PlatformMCPServerView() + + def test_read_tier_is_refused_a_write_tool(self) -> None: + refusal = self.view.check_tool_allowed(a_tool(writes=True), context_for("read")) + + assert refusal is not None + assert "read_write" in refusal + + def test_read_write_tier_is_allowed_a_write_tool(self) -> None: + assert ( + self.view.check_tool_allowed(a_tool(writes=True), context_for("read_write")) + is None + ) + + def test_read_tier_is_allowed_a_read_tool(self) -> None: + """The guard must not block reads — a read key that reaches the server + should still be able to use every read tool on it. + """ + assert ( + self.view.check_tool_allowed(a_tool(writes=False), context_for("read")) + is None + ) + + def test_platform_registry_is_entirely_read_only(self) -> None: + """Pins the documented promise that this server changes nothing. + + If someone registers a write tool here, this fails and forces them to + confirm the tier guard and the README claim still hold. + """ + writers = [ + name for name in PLATFORM_TOOLS.names() if PLATFORM_TOOLS.get(name).writes + ] + + assert writers == [], ( + f"Platform MCP server advertises itself as read-only but exposes " + f"write tools: {writers}" + ) diff --git a/backend/mcp_server/tests/test_platform_tools.py b/backend/mcp_server/tests/test_platform_tools.py new file mode 100644 index 0000000000..bfb08822d6 --- /dev/null +++ b/backend/mcp_server/tests/test_platform_tools.py @@ -0,0 +1,148 @@ +"""Behaviour of the organization-scoped read-only tools. + +These run against a live DB because the whole point of each tool is the +queryset it delegates to — mocking that away would leave nothing worth +asserting. Needs a live DB (integration tier). +""" + +from __future__ import annotations + +from account_v2.models import Organization, User +from api_v2.models import APIDeployment +from django.test import TestCase +from platform_api.models import PlatformApiKey +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from tenant_account_v2.models import OrganizationMember +from utils.user_context import UserContext +from workflow_manager.workflow_v2.models.workflow import Workflow + +from mcp_server.context import PlatformMCPContext +from mcp_server.tools.platform import ( + list_api_deployments, + list_prompt_studio_projects, + list_workflows, + platform_read_me_first, + whoami, +) + +ORG_ID = "org-tools" + + +class PlatformToolsTest(TestCase): + def setUp(self) -> None: + self.org = Organization.objects.create( + name=ORG_ID, display_name="Tools Org", organization_id=ORG_ID + ) + UserContext.set_organization_identifier(ORG_ID) + self.user = User.objects.create( + username="svc-tools", + email="svc-tools@platform.internal", + user_id="uid-tools", + is_service_account=True, + ) + OrganizationMember.objects.create( + user=self.user, organization=self.org, role="user" + ) + self.key = PlatformApiKey.objects.create( + name="tools-key", + description="d", + organization=self.org, + api_user=self.user, + permission="read_write", + ) + + self.workflow = Workflow.objects.create( + workflow_name="invoice-wf", description="Invoices", is_active=True + ) + self.api = APIDeployment.objects.create( + api_name="invoice-api", + display_name="Invoice API", + description="Extracts invoices", + workflow=self.workflow, + ) + self.tool = CustomTool.objects.create( + tool_name="Invoice Prompts", description="Prompt project", author="acme" + ) + + self.context = PlatformMCPContext( + user=self.user, platform_key=self.key, org_name=ORG_ID + ) + + def test_list_api_deployments_exposes_the_api_name(self) -> None: + """api_name is what an agent needs to open a deployment-scoped MCP + session, so it must be in the payload — not just the display name. + """ + result = list_api_deployments(self.context) + + assert result["count"] == 1 + row = result["api_deployments"][0] + assert row["api_name"] == "invoice-api" + assert row["display_name"] == "Invoice API" + assert row["is_active"] is True + assert "truncated" not in result + + def test_list_workflows_returns_the_org_workflow(self) -> None: + result = list_workflows(self.context) + + assert result["count"] == 1 + assert result["workflows"][0]["workflow_name"] == "invoice-wf" + + def test_list_prompt_studio_projects_returns_the_org_project(self) -> None: + result = list_prompt_studio_projects(self.context) + + assert result["count"] == 1 + assert result["prompt_studio_projects"][0]["tool_name"] == "Invoice Prompts" + + def test_listings_do_not_leak_across_organizations(self) -> None: + """The single most important property of these tools. A platform key + must never surface another tenant's resources. + """ + other = Organization.objects.create( + name="org-other", display_name="Other", organization_id="org-other" + ) + UserContext.set_organization_identifier("org-other") + other_wf = Workflow.objects.create(workflow_name="secret-wf", is_active=True) + APIDeployment.objects.create( + api_name="secret-api", display_name="Secret API", workflow=other_wf + ) + CustomTool.objects.create( + tool_name="Secret Prompts", description="d", author="other" + ) + + # Back to the original org: its key must see only its own resources. + UserContext.set_organization_identifier(ORG_ID) + + names = [r["api_name"] for r in list_api_deployments(self.context)["api_deployments"]] + assert names == ["invoice-api"] + + wf_names = [r["workflow_name"] for r in list_workflows(self.context)["workflows"]] + assert "secret-wf" not in wf_names + + tool_names = [ + r["tool_name"] + for r in list_prompt_studio_projects(self.context)["prompt_studio_projects"] + ] + assert "Secret Prompts" not in tool_names + assert other.organization_id == "org-other" # sanity: the other org exists + + def test_whoami_reports_tier_and_service_account_scope(self) -> None: + """An agent that gets more results than a human would should be able to + find out why. + """ + result = whoami(self.context) + + assert result["organization"] == ORG_ID + assert result["permission_tier"] == "read_write" + assert result["is_service_account"] is True + assert "read-only" in result["access"] + + def test_read_me_first_states_the_scope_warning(self) -> None: + """The service-account visibility rule is surprising enough that the + guide must say it outright, not leave it to be discovered. + """ + result = platform_read_me_first(self.context) + + assert "service account" in result["scope_warning"].lower() + assert result["organization"] == ORG_ID + # It must also point at the other server, since this one cannot extract. + assert "extractDocument" in result["to_run_an_extraction"] diff --git a/backend/mcp_server/tools/platform.py b/backend/mcp_server/tools/platform.py new file mode 100644 index 0000000000..a227a9234d --- /dev/null +++ b/backend/mcp_server/tools/platform.py @@ -0,0 +1,169 @@ +"""Read-only tools for the organization-scoped MCP server. + +Every listing goes through the model's own ``for_user`` manager rather than a +raw queryset, so the platform's sharing and visibility rules apply here exactly +as they do in the UI — this app does not get its own opinion about who may see +what. + +One caveat those managers impose, which the tool descriptions state plainly: +a platform key resolves to a service account, and ``for_user`` short-circuits +to "everything in the organization" for service accounts. These tools therefore +see all org resources regardless of the role the key was created with. +""" + +import logging +from typing import Any + +from api_v2.models import APIDeployment +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from workflow_manager.workflow_v2.models.workflow import Workflow + +from mcp_server.context import PlatformMCPContext + +logger = logging.getLogger(__name__) + +# Listings are bounded so a large organization cannot blow up an agent's +# context window. Chosen to comfortably cover real orgs while still being a +# ceiling; tools report when they hit it rather than truncating silently. +LIST_LIMIT = 100 + + +def no_args_schema() -> dict[str, Any]: + return {"type": "object", "properties": {}, "required": []} + + +def _truncation_note(shown: int, total: int) -> dict[str, Any]: + """Describe a capped listing, or nothing when it was complete. + + Silent truncation would read to an agent as "this is everything", which is + exactly the sort of wrong premise it would then build on. + """ + if total <= shown: + return {} + return { + "truncated": True, + "note": ( + f"Showing {shown} of {total}. Narrow the question or use the " + "Unstract UI to see the rest." + ), + } + + +def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: + """Return the orientation guide for an agent on the platform server.""" + return { + "server": "Unstract Platform MCP Server", + "organization": context.org_name, + "what_this_does": ( + "Unstract runs LLM-driven extraction over unstructured documents " + "(PDFs, scans, images) and returns structured JSON. This server is " + "scoped to a whole organization and is READ-ONLY: it lets you " + "discover what exists — deployments, workflows, Prompt Studio " + "projects — but not change or run anything." + ), + "tools": [ + {"name": "whoami", "purpose": "See this credential's org, tier and scope."}, + { + "name": "listApiDeployments", + "purpose": "Find deployed extraction endpoints.", + }, + {"name": "listWorkflows", "purpose": "Find the pipelines behind them."}, + { + "name": "listPromptStudioProjects", + "purpose": "Find where extraction prompts are authored.", + }, + ], + "to_run_an_extraction": ( + "This server cannot extract. Call listApiDeployments to get an " + "api_name, then open a separate MCP session against that " + "deployment's own endpoint — /mcp/// — using that " + "deployment's API key. That session exposes extractDocument." + ), + "scope_warning": ( + "This credential is a service account, which sees ALL resources in " + "the organization regardless of the role it was created with. Do " + "not treat these listings as one person's view." + ), + } + + +def whoami(context: PlatformMCPContext) -> dict[str, Any]: + """Describe the calling credential, so an agent can reason about scope.""" + key = context.platform_key + return { + "organization": context.org_name, + "key_name": key.name, + "permission_tier": key.permission, + "is_service_account": bool(getattr(context.user, "is_service_account", False)), + "access": "read-only (this MCP server exposes no write tools)", + "visibility": ( + "All resources in the organization. Service-account credentials " + "bypass per-user sharing filters." + ), + } + + +def list_api_deployments(context: PlatformMCPContext) -> dict[str, Any]: + """List the organization's API deployments.""" + queryset = APIDeployment.objects.for_user(context.user).order_by("display_name") + total = queryset.count() + rows = queryset[:LIST_LIMIT] + + return { + "count": total, + "api_deployments": [ + { + "id": str(row.id), + "display_name": row.display_name, + # api_name is what an agent needs to open a deployment-scoped + # MCP session, so it is surfaced rather than left to the UI. + "api_name": row.api_name, + "description": row.description or None, + "is_active": row.is_active, + "workflow_id": str(row.workflow_id), + } + for row in rows + ], + **_truncation_note(len(rows), total), + } + + +def list_workflows(context: PlatformMCPContext) -> dict[str, Any]: + """List the organization's workflows.""" + queryset = Workflow.objects.for_user(context.user).order_by("workflow_name") + total = queryset.count() + rows = queryset[:LIST_LIMIT] + + return { + "count": total, + "workflows": [ + { + "id": str(row.id), + "workflow_name": row.workflow_name, + "description": row.description or None, + "is_active": row.is_active, + } + for row in rows + ], + **_truncation_note(len(rows), total), + } + + +def list_prompt_studio_projects(context: PlatformMCPContext) -> dict[str, Any]: + """List the organization's Prompt Studio projects.""" + queryset = CustomTool.objects.for_user(context.user).order_by("tool_name") + total = queryset.count() + rows = queryset[:LIST_LIMIT] + + return { + "count": total, + "prompt_studio_projects": [ + { + "id": str(row.tool_id), + "tool_name": row.tool_name, + "description": row.description or None, + } + for row in rows + ], + **_truncation_note(len(rows), total), + } diff --git a/backend/mcp_server/transport.py b/backend/mcp_server/transport.py new file mode 100644 index 0000000000..ddef3f3e2a --- /dev/null +++ b/backend/mcp_server/transport.py @@ -0,0 +1,280 @@ +"""Shared JSON-RPC 2.0 transport for the hosted MCP servers. + +Two MCP servers are hosted from this app — one scoped to a single API +deployment, one scoped to an organization. They differ only in how a caller is +authenticated and which tools they expose; the wire protocol, dispatch and +error mapping are identical, and live here so the two cannot drift apart on +protocol behaviour. + +Subclasses supply: + registry - the MCPToolRegistry to dispatch into + resolve_context() - authenticate, or return None to reject + server_instructions() - the prose an agent reads on `initialize` +""" + +import json +import logging +from typing import Any + +from django.http import JsonResponse +from rest_framework import status, views +from rest_framework.request import Request +from rest_framework.response import Response + +from mcp_server.constants import JSONRPC, MCPMethod, MCPServer +from mcp_server.exceptions import MCPToolError +from mcp_server.registry import MCPToolRegistry + +logger = logging.getLogger(__name__) + + +def rpc_result(request_id: Any, result: Any) -> JsonResponse: + """Build a JSON-RPC success response. + + JsonResponse rather than DRF's Response so the body is emitted verbatim: + content negotiation must not turn a JSON-RPC envelope into the browsable + API renderer for a client that sent a permissive Accept header. + """ + return JsonResponse({"jsonrpc": JSONRPC.VERSION, "id": request_id, "result": result}) + + +def rpc_error(request_id: Any, code: int, message: str, data: Any = None) -> JsonResponse: + """Build a JSON-RPC error response. + + Always HTTP 200: transport-level success with an application-level error is + exactly what JSON-RPC models, and clients read the envelope, not the status + code. The one exception is authentication, which must be a real 401 so MCP + clients can react to it (see `auth_error`). + """ + error: dict[str, Any] = {"code": code, "message": message} + if data is not None: + error["data"] = data + return JsonResponse({"jsonrpc": JSONRPC.VERSION, "id": request_id, "error": error}) + + +def auth_error( + message: str, http_status: int = status.HTTP_401_UNAUTHORIZED +) -> JsonResponse: + """Build a 401/403 for a rejected credential. + + Unlike other failures this carries a real HTTP status, because a client + that cannot authenticate needs to distinguish "wrong key" from "tool + failed" before it has a usable session. + """ + response = JsonResponse( + { + "jsonrpc": JSONRPC.VERSION, + "id": None, + "error": {"code": JSONRPC.UNAUTHORIZED, "message": message}, + }, + status=http_status, + ) + if http_status == status.HTTP_401_UNAUTHORIZED: + response["WWW-Authenticate"] = 'Bearer realm="unstract-mcp"' + return response + + +def tool_content(text: str, is_error: bool = False) -> dict[str, Any]: + """Wrap tool output in the MCP `tools/call` content envelope.""" + return {"content": [{"type": "text", "text": text}], "isError": is_error} + + +class BaseMCPView(views.APIView): + """JSON-RPC transport shared by the deployment and platform MCP servers.""" + + #: Tools this server exposes. Set by each subclass. + registry: MCPToolRegistry = None + + def get_registry(self) -> MCPToolRegistry: + """Return the registry to dispatch into. + + Indirection exists so tests can patch a single seam rather than a + module-level global shared by both servers. + """ + return self.registry + + def resolve_context(self, **kwargs: Any) -> Any | None: + """Authenticate the caller and build the per-request context. + + Return None to reject. Subclasses must answer every failure mode + identically so the endpoint cannot be used to probe what exists. + """ + raise NotImplementedError + + def server_instructions(self, context: Any) -> str: + """Prose returned on `initialize`, read by the agent as a prompt.""" + raise NotImplementedError + + def auth_failure_message(self) -> str: + """Message used for every rejected credential on this server.""" + return "Invalid or missing credentials" + + def server_info(self) -> dict[str, Any]: + """Identity advertised on GET and on `initialize`.""" + return { + "name": MCPServer.NAME, + "version": MCPServer.VERSION, + "protocolVersion": MCPServer.PROTOCOL_VERSION, + "transport": "http", + "authMethods": ["bearer"], + } + + def get(self, request: Request, **kwargs: Any) -> Response: + """Advertise server identity. + + Clients probe with GET to check connectivity before opening a session. + Deliberately free of tenant detail — it reveals only that an MCP server + is mounted here. + """ + return Response(self.server_info()) + + def post(self, request: Request, **kwargs: Any) -> JsonResponse: + """Handle a single JSON-RPC request.""" + context = self.resolve_context(request=request, **kwargs) + if context is None: + return auth_error(self.auth_failure_message()) + + body = request.data + if not isinstance(body, dict): + # Batch requests are valid JSON-RPC 2.0 but are not used by MCP + # clients; rejecting them explicitly beats a confusing downstream + # AttributeError. + return rpc_error( + None, + JSONRPC.INVALID_REQUEST, + "Invalid Request", + "Expected a single JSON-RPC object", + ) + + request_id = body.get("id") + method = body.get("method") + + if body.get("jsonrpc") != JSONRPC.VERSION: + return rpc_error( + request_id, + JSONRPC.INVALID_REQUEST, + "Invalid Request", + "Only JSON-RPC 2.0 is supported", + ) + if not method: + return rpc_error( + request_id, JSONRPC.INVALID_REQUEST, "Invalid Request", "Missing method" + ) + + return self._dispatch( + method=method, + request_id=request_id, + params=body.get("params") or {}, + context=context, + ) + + def _dispatch( + self, method: str, request_id: Any, params: dict[str, Any], context: Any + ) -> JsonResponse: + """Route a JSON-RPC method to its handler.""" + if method == MCPMethod.INITIALIZE: + return rpc_result( + request_id, + { + "protocolVersion": MCPServer.PROTOCOL_VERSION, + "capabilities": {"tools": {"listChanged": False}}, + "serverInfo": { + "name": MCPServer.NAME, + "version": MCPServer.VERSION, + }, + "instructions": self.server_instructions(context), + }, + ) + + if method.startswith(MCPMethod.NOTIFICATION_PREFIX): + # Notifications (e.g. notifications/initialized) carry no id and + # must not receive a result body. + return JsonResponse({}, status=status.HTTP_202_ACCEPTED) + + if method == MCPMethod.PING: + return rpc_result(request_id, {}) + + if method == MCPMethod.TOOLS_LIST: + return rpc_result(request_id, {"tools": self.get_registry().list_schemas()}) + + if method == MCPMethod.TOOLS_CALL: + return self._call_tool(request_id=request_id, params=params, context=context) + + return rpc_error( + request_id, + JSONRPC.METHOD_NOT_FOUND, + "Method not found", + f"Unsupported method '{method}'", + ) + + def check_tool_allowed(self, tool: Any, context: Any) -> str | None: + """Authorize a specific tool for this caller. + + Return None to allow, or a message explaining the refusal. The default + allows everything the registry exposes; servers with permission tiers + override this. + """ + return None + + def _call_tool( + self, request_id: Any, params: dict[str, Any], context: Any + ) -> JsonResponse: + """Invoke a registered tool and wrap its result in MCP content format.""" + registry = self.get_registry() + tool_name = params.get("name") + if not tool_name: + return rpc_error( + request_id, JSONRPC.INVALID_PARAMS, "Invalid params", "Missing tool name" + ) + + tool = registry.get(tool_name) + if tool is None: + return rpc_error( + request_id, + JSONRPC.METHOD_NOT_FOUND, + "Method not found", + f"Tool '{tool_name}' not found. Available tools: {registry.names()}", + ) + + refusal = self.check_tool_allowed(tool, context) + if refusal is not None: + return rpc_error( + request_id, JSONRPC.UNAUTHORIZED, "Permission denied", refusal + ) + + arguments = params.get("arguments") or {} + if not isinstance(arguments, dict): + return rpc_error( + request_id, + JSONRPC.INVALID_PARAMS, + "Invalid params", + "'arguments' must be an object", + ) + + try: + result = tool.handler(context, **arguments) + except MCPToolError as error: + # Expected, agent-actionable failure: the message is written for + # the agent, so pass it through as an error result rather than a + # protocol error. isError=True lets the agent see and retry it. + return rpc_result(request_id, tool_content(str(error), is_error=True)) + except TypeError as error: + # Almost always the agent passing arguments the tool does not + # accept; report it as bad params rather than a server fault. + logger.warning(f"MCP tool '{tool_name}' called with bad arguments: {error}") + return rpc_error( + request_id, JSONRPC.INVALID_PARAMS, "Invalid params", str(error) + ) + except Exception as error: + logger.exception(f"MCP tool '{tool_name}' failed: {error}") + return rpc_error( + request_id, + JSONRPC.TOOL_EXECUTION_ERROR, + "Tool execution failed", + f"Tool '{tool_name}' failed unexpectedly. " + "Contact your Unstract administrator if this persists.", + ) + + return rpc_result( + request_id, tool_content(json.dumps(result, indent=2, default=str)) + ) diff --git a/backend/mcp_server/views.py b/backend/mcp_server/views.py index b4599ad9ad..7f0552d67d 100644 --- a/backend/mcp_server/views.py +++ b/backend/mcp_server/views.py @@ -1,75 +1,29 @@ -"""JSON-RPC 2.0 transport for the hosted Unstract MCP server. +"""Deployment-scoped MCP server. Mirrors the shape of the API deployment execution endpoint: the URL carries the organization and the deployment name, and the deployment's own API key authenticates the caller. An MCP session is therefore scoped to exactly one API deployment, and reuses that deployment's existing key management. + +This endpoint is in ``WHITELISTED_PATHS``, so ``CustomAuthMiddleware`` skips it +and the view owns its own authentication — deliberately, because a deployment +key is not a platform key and must not resolve to a user. """ -import json import logging from typing import Any from api_v2.deployment_helper import DeploymentHelper -from django.http import JsonResponse -from rest_framework import status, views from rest_framework.request import Request -from rest_framework.response import Response -from mcp_server.constants import JSONRPC, MCPMethod, MCPServer from mcp_server.context import MCPContext -from mcp_server.exceptions import MCPToolError -from mcp_server.registry import TOOL_REGISTRY +from mcp_server.registry import DEPLOYMENT_TOOLS +from mcp_server.transport import BaseMCPView logger = logging.getLogger(__name__) -def _rpc_result(request_id: Any, result: Any) -> JsonResponse: - """Build a JSON-RPC success response. - - JsonResponse rather than DRF's Response so the body is emitted verbatim: - content negotiation must not turn a JSON-RPC envelope into the browsable - API renderer for a client that sent a permissive Accept header. - """ - return JsonResponse({"jsonrpc": JSONRPC.VERSION, "id": request_id, "result": result}) - - -def _rpc_error( - request_id: Any, code: int, message: str, data: Any = None -) -> JsonResponse: - """Build a JSON-RPC error response. - - Always HTTP 200: transport-level success with an application-level error is - exactly what JSON-RPC models, and clients read the envelope, not the status - code. The one exception is authentication, which must be a real 401 so MCP - clients can react to it (see `_auth_error`). - """ - error: dict[str, Any] = {"code": code, "message": message} - if data is not None: - error["data"] = data - return JsonResponse({"jsonrpc": JSONRPC.VERSION, "id": request_id, "error": error}) - - -def _auth_error(message: str) -> JsonResponse: - """Build a 401 for a failed or missing credential. - - Unlike other failures this carries a real HTTP status, because a client - that cannot authenticate needs to distinguish "wrong key" from "tool - failed" before it has a usable session. - """ - response = JsonResponse( - { - "jsonrpc": JSONRPC.VERSION, - "id": None, - "error": {"code": JSONRPC.UNAUTHORIZED, "message": message}, - }, - status=status.HTTP_401_UNAUTHORIZED, - ) - response["WWW-Authenticate"] = 'Bearer realm="unstract-mcp"' - return response - - -class MCPServerView(views.APIView): +class MCPServerView(BaseMCPView): """MCP JSON-RPC endpoint for a single API deployment. Authentication is the deployment's API key, accepted either as a bearer @@ -78,6 +32,8 @@ class MCPServerView(views.APIView): execution endpoint it sits beside, so session auth does not apply. """ + registry = DEPLOYMENT_TOOLS + authentication_classes: list = [] permission_classes: list = [] @@ -86,29 +42,31 @@ def initialize_request(self, request: Request, *args: Any, **kwargs: Any) -> Req request.csrf_processing_done = True return super().initialize_request(request, *args, **kwargs) - def get( - self, request: Request, org_name: str, api_name: str, api_key: str | None = None - ) -> Response: - """Advertise server identity. + def auth_failure_message(self) -> str: + return "Invalid API key or unknown API deployment" - Clients probe with GET to check connectivity before opening a session. - Deliberately unauthenticated and free of deployment detail — it reveals - only that an MCP server is mounted here. - """ - return Response( - { - "name": MCPServer.NAME, - "version": MCPServer.VERSION, - "protocolVersion": MCPServer.PROTOCOL_VERSION, - "transport": "http", - "authMethods": ["bearer"], - } + def server_instructions(self, context: MCPContext) -> str: + return ( + "Unstract runs LLM-driven extraction over unstructured documents " + "and returns structured JSON. This server is scoped to the API " + f"deployment '{context.api.display_name}'. Call readMeFirst before " + "any other tool." ) - def post( - self, request: Request, org_name: str, api_name: str, api_key: str | None = None - ) -> JsonResponse: - """Handle a single JSON-RPC request.""" + def resolve_context( + self, + request: Request, + org_name: str, + api_name: str, + api_key: str | None = None, + **kwargs: Any, + ) -> MCPContext | None: + """Authenticate the key against the named deployment. + + Returns None for every failure mode — unknown deployment, wrong key, + malformed key — so the caller answers all of them identically and the + endpoint cannot be used to probe which deployment names exist. + """ # `api_key` in the path takes precedence when present; otherwise fall # back to the Authorization header. if not api_key: @@ -117,61 +75,12 @@ def post( api_key = header.split(" ", 1)[1].strip() if not api_key: - return _auth_error("Missing API key") - - context = self._resolve_context( - org_name=org_name, api_name=api_name, api_key=api_key - ) - if context is None: - return _auth_error("Invalid API key or unknown API deployment") - - body = request.data - if not isinstance(body, dict): - # Batch requests are a valid part of JSON-RPC 2.0 but are not used - # by MCP clients; rejecting them explicitly beats a confusing - # downstream AttributeError. - return _rpc_error( - None, - JSONRPC.INVALID_REQUEST, - "Invalid Request", - "Expected a single JSON-RPC object", - ) - - request_id = body.get("id") - method = body.get("method") - - if body.get("jsonrpc") != JSONRPC.VERSION: - return _rpc_error( - request_id, - JSONRPC.INVALID_REQUEST, - "Invalid Request", - "Only JSON-RPC 2.0 is supported", - ) - if not method: - return _rpc_error( - request_id, JSONRPC.INVALID_REQUEST, "Invalid Request", "Missing method" - ) - - return self._dispatch( - method=method, - request_id=request_id, - params=body.get("params") or {}, - context=context, - ) - - def _resolve_context( - self, org_name: str, api_name: str, api_key: str - ) -> MCPContext | None: - """Authenticate the key against the named deployment. + return None - Returns None for every failure mode — unknown deployment, wrong key, - malformed key — so the caller answers all of them identically and the - endpoint cannot be used to probe which deployment names exist. - """ # Pin the organization before touching org-scoped managers; the # deployment lookup below filters on it. DeploymentHelper.validate_parameters( - self.request, api_name=api_name, org_name=org_name + request, api_name=api_name, org_name=org_name ) api_deployment = DeploymentHelper.get_deployment_by_api_name(api_name=api_name) @@ -184,109 +93,3 @@ def _resolve_context( return None return MCPContext(api=api_deployment, api_key=api_key, org_name=org_name) - - def _dispatch( - self, method: str, request_id: Any, params: dict[str, Any], context: MCPContext - ) -> JsonResponse: - """Route a JSON-RPC method to its handler.""" - if method == MCPMethod.INITIALIZE: - return _rpc_result( - request_id, - { - "protocolVersion": MCPServer.PROTOCOL_VERSION, - "capabilities": {"tools": {"listChanged": False}}, - "serverInfo": { - "name": MCPServer.NAME, - "version": MCPServer.VERSION, - }, - "instructions": ( - "Unstract runs LLM-driven extraction over unstructured " - "documents and returns structured JSON. This server is " - "scoped to the API deployment " - f"'{context.api.display_name}'. Call readMeFirst before " - "any other tool." - ), - }, - ) - - if method.startswith(MCPMethod.NOTIFICATION_PREFIX): - # Notifications (e.g. notifications/initialized) carry no id and - # must not receive a result body. - return JsonResponse({}, status=status.HTTP_202_ACCEPTED) - - if method == MCPMethod.PING: - return _rpc_result(request_id, {}) - - if method == MCPMethod.TOOLS_LIST: - return _rpc_result(request_id, {"tools": TOOL_REGISTRY.list_schemas()}) - - if method == MCPMethod.TOOLS_CALL: - return self._call_tool(request_id=request_id, params=params, context=context) - - return _rpc_error( - request_id, - JSONRPC.METHOD_NOT_FOUND, - "Method not found", - f"Unsupported method '{method}'", - ) - - def _call_tool( - self, request_id: Any, params: dict[str, Any], context: MCPContext - ) -> JsonResponse: - """Invoke a registered tool and wrap its result in MCP content format.""" - tool_name = params.get("name") - if not tool_name: - return _rpc_error( - request_id, JSONRPC.INVALID_PARAMS, "Invalid params", "Missing tool name" - ) - - tool = TOOL_REGISTRY.get(tool_name) - if tool is None: - return _rpc_error( - request_id, - JSONRPC.METHOD_NOT_FOUND, - "Method not found", - f"Tool '{tool_name}' not found. Available tools: {TOOL_REGISTRY.names()}", - ) - - arguments = params.get("arguments") or {} - if not isinstance(arguments, dict): - return _rpc_error( - request_id, - JSONRPC.INVALID_PARAMS, - "Invalid params", - "'arguments' must be an object", - ) - - try: - result = tool.handler(context, **arguments) - except MCPToolError as error: - # Expected, agent-actionable failure: the message is written for - # the agent, so pass it through as an error result rather than a - # protocol error. isError=True lets the agent see and retry it. - return _rpc_result(request_id, _tool_content(str(error), is_error=True)) - except TypeError as error: - # Almost always the agent passing arguments the tool does not - # accept; report it as bad params rather than a server fault. - logger.warning(f"MCP tool '{tool_name}' called with bad arguments: {error}") - return _rpc_error( - request_id, JSONRPC.INVALID_PARAMS, "Invalid params", str(error) - ) - except Exception as error: - logger.exception(f"MCP tool '{tool_name}' failed: {error}") - return _rpc_error( - request_id, - JSONRPC.TOOL_EXECUTION_ERROR, - "Tool execution failed", - f"Tool '{tool_name}' failed unexpectedly. " - "Contact your Unstract administrator if this persists.", - ) - - return _rpc_result( - request_id, _tool_content(json.dumps(result, indent=2, default=str)) - ) - - -def _tool_content(text: str, is_error: bool = False) -> dict[str, Any]: - """Wrap tool output in the MCP `tools/call` content envelope.""" - return {"content": [{"type": "text", "text": text}], "isError": is_error} diff --git a/tests/critical_paths.yaml b/tests/critical_paths.yaml index 7e8a7a97d1..94aa29f006 100644 --- a/tests/critical_paths.yaml +++ b/tests/critical_paths.yaml @@ -71,6 +71,11 @@ paths: covered_by: [integration-backend] proof: marker + - id: mcp-platform-auth + description: "The org-scoped MCP endpoint stays behind the platform-API-key middleware; unauthenticated or mis-scoped calls reach no tool." + covered_by: [integration-backend] + proof: marker + - id: prompt-studio-author description: "Create a Prompt Studio project and add a prompt to it." covered_by: [integration-backend] From cd5cd56df10042feb23d8699c380b72fe804f1d0 Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sat, 25 Jul 2026 20:25:13 +0530 Subject: [PATCH 05/16] feat(mcp): move deployment MCP onto the execution URL, add platform write tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three changes. 1. The deployment MCP endpoint moves to a suffix on the deployment's own execution URL: POST /deployment/api/// (REST) POST /deployment/api///mcp (MCP) Its URLs now live in api_v2/execution_urls.py, next to the endpoint they extend, which also makes them inherit the deployment prefix's existing WHITELISTED_PATHS entry instead of needing one of their own. 2. MCP_PATH_PREFIX is removed. It was introduced by this branch's first commit and nothing outside it ever depended on it; the deployment MCP path is now determined by API_DEPLOYMENT_PATH_PREFIX, and a second knob for a prefix that should never change was just a way to break clients. 3. The platform server gains write tools: setApiDeploymentActive, setPipelineActive and executePipeline, alongside a new listPipelines. Authorization for the write tools reuses machinery that already exists rather than inventing a scheme. Platform key tiers are defined over HTTP methods (ApiKeyPermission.allows), but every MCP call is a POST — so the middleware's tier check cannot tell listWorkflows from executePipeline. Each tool now declares the method its REST equivalent would use, and check_tool_allowed re-applies the key's tier against that. A tool marked DELETE is therefore full_access-only for exactly the same reason a REST DELETE is. Two categories stay out, deliberately: * Credential operations. Key creation and rotation return the secret in their response, so exposing them would give an agent processing untrusted document content a way to mint or exfiltrate credentials. The codebase already reasons this way — see CanRotatePlatformApiKey's docstring. A test asserts no tool name suggests credential handling. * Deletions. The write tools here are reversible by construction; removing a workflow or project is not. The previous read-only invariant test is replaced rather than dropped: it now asserts every writes=True tool declares a non-GET required_method, so a write tool cannot silently slip past the guard by keeping the default. Write tools resolve their target through for_user, so a key cannot touch another organization's resources — asserted directly, since the consequence there is mutation rather than mere disclosure. The service-account caveat now appears in every write tool's description, not just whoami: for reads it was a disclosure note, for writes it means the key can modify any org resource. Tests: 59 in this app (was 50); full backend unit tier 311 passing. Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/api_v2/execution_urls.py | 12 +- backend/api_v2/serializers.py | 2 +- backend/backend/base_urls.py | 2 - backend/backend/exceptions.py | 1 - backend/backend/settings/base.py | 11 +- backend/backend/urls_v2.py | 6 +- backend/mcp_server/README.md | 127 +++++--- backend/mcp_server/context.py | 6 + backend/mcp_server/platform_views.py | 55 ++-- backend/mcp_server/registry.py | 100 +++++- backend/mcp_server/tests/test_mcp_auth.py | 39 ++- .../mcp_server/tests/test_platform_auth.py | 4 + .../tests/test_platform_tier_guard.py | 110 +++++-- .../mcp_server/tests/test_platform_tools.py | 129 +++++++- backend/mcp_server/tools/platform.py | 304 ++++++++++++++++-- backend/mcp_server/urls.py | 19 +- 16 files changed, 782 insertions(+), 145 deletions(-) diff --git a/backend/api_v2/execution_urls.py b/backend/api_v2/execution_urls.py index 738340013b..8fb3fd46f4 100644 --- a/backend/api_v2/execution_urls.py +++ b/backend/api_v2/execution_urls.py @@ -1,4 +1,4 @@ -from django.urls import re_path +from django.urls import include, re_path from rest_framework.urlpatterns import format_suffix_patterns from api_v2.api_deployment_views import DeploymentExecution @@ -6,7 +6,15 @@ execute = DeploymentExecution.as_view() -urlpatterns = format_suffix_patterns( +# The deployment MCP server hangs off the execution URL +# (``…//mcp``), so it is matched first. The execution pattern below +# anchors with ``$`` and cannot swallow the ``/mcp`` suffix on its own, but +# ordering it first keeps that independent of the other pattern's exact shape. +# Kept outside ``format_suffix_patterns`` because that helper rewrites leaf +# patterns and is not meant to wrap an ``include()``. +urlpatterns = [re_path(r"^", include("mcp_server.urls"))] + +urlpatterns += format_suffix_patterns( [ re_path( r"^api/(?P[\w-]+)/(?P[\w-]+)/?$", diff --git a/backend/api_v2/serializers.py b/backend/api_v2/serializers.py index 0eb99f1bce..6c13d8e49c 100644 --- a/backend/api_v2/serializers.py +++ b/backend/api_v2/serializers.py @@ -4,6 +4,7 @@ from typing import Any from urllib.parse import urlparse +from backend.serializers import AuditSerializer from django.apps import apps from django.core.validators import RegexValidator from pipeline_v2.models import Pipeline @@ -35,7 +36,6 @@ from api_v2.constants import ApiExecution from api_v2.models import APIDeployment, APIKey -from backend.serializers import AuditSerializer class APIDeploymentSerializer(IntegrityErrorMixin, AuditSerializer): diff --git a/backend/backend/base_urls.py b/backend/backend/base_urls.py index ee539e39de..8450add04d 100644 --- a/backend/backend/base_urls.py +++ b/backend/backend/base_urls.py @@ -22,8 +22,6 @@ f"{settings.API_DEPLOYMENT_PATH_PREFIX}/pipeline/", include("pipeline_v2.public_api_urls"), ), - # Hosted MCP server, authenticated by the API deployment's own key - path(f"{settings.MCP_PATH_PREFIX}/", include("mcp_server.urls")), path("", include("health.urls")), # Internal API for worker communication path("internal/", include("backend.internal_base_urls")), diff --git a/backend/backend/exceptions.py b/backend/backend/exceptions.py index caf4df575b..c7ca940752 100644 --- a/backend/backend/exceptions.py +++ b/backend/backend/exceptions.py @@ -3,7 +3,6 @@ from rest_framework.exceptions import APIException from rest_framework.response import Response from rest_framework.views import exception_handler - from unstract.connectors.exceptions import ConnectorBaseException, ConnectorError diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 85947ad807..b21c15df04 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -118,9 +118,6 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | API_DEPLOYMENT_PATH_PREFIX = os.environ.get( "API_DEPLOYMENT_PATH_PREFIX", "deployment" ).strip("/") -# Mount point for the hosted MCP server. Changing this invalidates the MCP -# endpoint URLs already configured in users' MCP clients. -MCP_PATH_PREFIX = os.environ.get("MCP_PATH_PREFIX", "mcp").strip("/") # Maximum file size for presigned URLs in API deployments (in MB) API_DEPL_PRESIGNED_URL_MAX_FILE_SIZE_MB = int( @@ -651,13 +648,11 @@ def filter(self, record): "/static", ] WHITELISTED_PATHS = [f"/{PATH_PREFIX}{PATH}" for PATH in WHITELISTED_PATHS_LIST] -# White lists workflow-api-deployment path +# White lists workflow-api-deployment path. This also covers the deployment MCP +# server, which hangs off the same URL and authenticates with the deployment's +# own API key rather than a session. WHITELISTED_PATHS.append(f"/{API_DEPLOYMENT_PATH_PREFIX}") -# White lists the hosted MCP server, which authenticates with the API -# deployment's own key rather than a session. -WHITELISTED_PATHS.append(f"/{MCP_PATH_PREFIX}") - # Whitelisting health check API WHITELISTED_PATHS.append("/health") diff --git a/backend/backend/urls_v2.py b/backend/backend/urls_v2.py index be12a14d1a..5e68db6b61 100644 --- a/backend/backend/urls_v2.py +++ b/backend/backend/urls_v2.py @@ -65,8 +65,8 @@ path("execution/", include("workflow_manager.file_execution.urls")), path("metrics/", include("dashboard_metrics.urls")), path("platform-api/", include("platform_api.urls")), - # Organization-scoped MCP server. Mounted here rather than under - # MCP_PATH_PREFIX so CustomAuthMiddleware authenticates the platform API - # key — that prefix is whitelisted and would bypass auth entirely. + # Organization-scoped MCP server. Mounted on this tenant path rather than + # under the deployment prefix so CustomAuthMiddleware authenticates the + # platform API key — that prefix is whitelisted and would bypass auth. path("mcp/", include("mcp_server.platform_urls")), ] diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index 4d8e862d4a..a73ee92211 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -15,8 +15,8 @@ authenticated and which tools they expose. | Scope | one API deployment | one organization | | Credential | that deployment's API key | a platform API key | | Authenticated by | the view itself | `CustomAuthMiddleware` | -| URL | `/mcp///` | `/api/v1/unstract//mcp/` | -| Tools | extract, poll status | read-only discovery | +| URL | `/deployment/api///mcp` | `/api/v1/unstract//mcp/` | +| Tools | extract, poll status | discovery + state changes | The split is not cosmetic. A deployment key grants exactly one workflow and resolves to no user, so it cannot authorize anything organization-wide; a @@ -34,7 +34,7 @@ shape of that deployment's REST endpoint: ``` POST /deployment/api/// # REST -POST /mcp/// # MCP +POST /deployment/api///mcp # MCP ``` Authentication is the deployment's **existing API key** — the same key used for @@ -49,7 +49,7 @@ For MCP clients that cannot attach custom headers, the key may instead be given as a path segment: ``` -POST /mcp//// +POST /deployment/api///mcp// ``` The path key takes precedence over the header when both are present. @@ -63,7 +63,7 @@ With Claude Code: ```bash claude mcp add --transport http unstract \ - https:///mcp/// \ + https:///deployment/api///mcp \ --header "Authorization: Bearer " ``` @@ -126,9 +126,9 @@ internal detail does not leak to the client. # Platform server -Organization-scoped and **read-only**: it lets an agent discover what exists in -an Unstract organization — deployments, workflows, Prompt Studio projects — but -cannot run or change anything. +Organization-scoped: an agent can discover what exists in an Unstract +organization — deployments, workflows, pipelines, Prompt Studio projects — and +change the running state of those resources. ## Endpoint @@ -147,24 +147,58 @@ Platform API keys are managed at `/api/v1/unstract//platform-api/keys/`. ## Tools -| Tool | Purpose | -| --- | --- | -| `readMeFirst` | Orientation, including what this credential can see. | -| `whoami` | The key's organization, permission tier and visibility scope. | -| `listApiDeployments` | Deployed extraction endpoints, with their `api_name`. | -| `listWorkflows` | The pipelines behind them. | -| `listPromptStudioProjects` | Where extraction prompts are authored. | - -To actually extract, an agent calls `listApiDeployments` to find an `api_name`, -then opens a **separate** session against the deployment server above. +| Tool | Tier | Purpose | +| --- | --- | --- | +| `readMeFirst` | read | Orientation, including what this credential can reach. | +| `whoami` | read | The key's organization, tier and what it may do. | +| `listApiDeployments` | read | Deployed extraction endpoints, with their `api_name`. | +| `listWorkflows` | read | The workflows behind them. | +| `listPipelines` | read | ETL and task pipelines, with schedule and last-run state. | +| `listPromptStudioProjects` | read | Where extraction prompts are authored. | +| `setApiDeploymentActive` | `read_write` | Take a deployment offline, or bring it back. | +| `setPipelineActive` | `read_write` | Pause or resume a pipeline's schedule. | +| `executePipeline` | `read_write` | Trigger a real pipeline run. **Consumes quota.** | + +To actually extract a document, an agent calls `listApiDeployments` to find an +`api_name`, then opens a **separate** session against the deployment server +above. + +## How write tools are authorized + +Platform API key tiers are defined in terms of HTTP methods +(`ApiKeyPermission.allows`), but every MCP call is a `POST` — so the auth +middleware's tier check cannot tell `listWorkflows` from `executePipeline`. + +Each tool therefore declares the method its REST equivalent would use +(`required_method`), and `check_tool_allowed` re-applies the key's tier against +*that*. This reuses the semantics the codebase already enforces instead of +inventing a parallel scheme, and it means a tool marked `DELETE` is +`full_access`-only for the same reason a REST `DELETE` is. + +A `writes=True` tool left at the default `required_method="GET"` would slip past +that guard, so a test asserts none exists. + +## What is deliberately not exposed + +**Credential operations.** Creating or rotating an API key returns the secret in +its response. An MCP tool for it would hand an agent — one that may be +processing untrusted document content — a way to mint or exfiltrate credentials. +The codebase already reasons this way: see `CanRotatePlatformApiKey`'s docstring +on why rotation is `full_access`-only. A test asserts no tool name suggests +credential handling. + +**Deletions.** Removing a workflow, deployment or Prompt Studio project destroys +work that no inverse call restores. Every write tool here is reversible by +construction — the opposite call puts things back. ## Why the URL matters This server is mounted under the tenant prefix, next to `platform-api/`, and -**not** under `MCP_PATH_PREFIX`. `WHITELISTED_PATHS` is matched with -`startswith`, so `/mcp/...` is exempt from `CustomAuthMiddleware` while -`/api/v1/unstract//mcp/` is not — which is exactly how this server -inherits platform-key authentication. +**not** under the deployment prefix. `WHITELISTED_PATHS` is matched with +`startswith`, so everything under `/deployment/...` — including the deployment +MCP server — is exempt from `CustomAuthMiddleware`, while +`/api/v1/unstract//mcp/` is not. That is exactly how this server inherits +platform-key authentication. Moving these URLs under the whitelisted prefix would silently remove all authentication. `test_platform_auth.py::test_the_endpoint_is_not_whitelisted` @@ -178,31 +212,38 @@ bypasses the middleware and would pass against a completely open endpoint. **A `read`-tier key cannot use this server at all.** The middleware's tier check gates on HTTP method, and every MCP call is a `POST`, which `read` disallows — -so a read-only key is refused before reaching the view, even though every tool -here is read-only. Use a `read_write` key. Changing that would mean special- -casing MCP paths in the middleware, which is a decision for maintainers rather -than something this app should do unilaterally. - -**These tools see the whole organization.** A platform key resolves to a service -account, and `is_service_account=True` makes `for_user()` managers return -`self.all()` — so listings ignore per-user sharing regardless of the `USER` role -the key was created with. `whoami` and `readMeFirst` both say so, because an -agent that assumed otherwise would draw wrong conclusions from a listing. +so a read-only key is refused before reaching the view, even for the read tools. +Use a `read_write` key. Changing that would mean special-casing MCP paths in the +middleware, which is a decision for maintainers rather than something this app +should do unilaterally. + +**These tools reach the whole organization.** A platform key resolves to a +service account, and `is_service_account=True` makes `for_user()` managers +return `self.all()` — so it ignores per-user sharing regardless of the `USER` +role the key was created with. For reads that is a disclosure caveat; for +writes it means the key can modify *any* org resource. `whoami`, `readMeFirst` +and every write tool's description say so. ## Adding a tool here Same registry mechanics as the deployment server, with `PlatformMCPContext` -(`user`, `platform_key`, `org_name`) as the first argument. Query through the -model's `for_user(context.user)` manager so the platform's own visibility rules -apply, and cap results — `LIST_LIMIT` with a `truncated` note, never silent -truncation. - -Write tools are deliberately absent. `check_tool_allowed` already refuses a -`writes=True` tool to a `read`-tier key, so the guard is in place, but the -read-only promise is pinned by -`test_platform_tier_guard.py::test_platform_registry_is_entirely_read_only` — -adding a write tool will fail that test and force a deliberate decision about -whether an agent should hold that power. +(`user`, `platform_key`, `org_name`, `request`) as the first argument. Query +through the model's `for_user(context.user)` manager so the platform's own +visibility rules apply, and cap listings — `LIST_LIMIT` with a `truncated` note, +never silent truncation. + +For a write tool, also: + +- set `writes=True` and a `required_method` (`"POST"` for a mutation, `"DELETE"` + for anything destructive — that makes it `full_access`-only); +- resolve the target through `for_user` so a caller cannot reach another + organization's resource, and raise `MCPToolError` naming the list tool when it + is not found; +- include `_ORG_WIDE_WARNING` in the description — the description is the prompt + the agent reads before acting, and the service-account caveat matters more + here than anywhere else; +- return `changed: False` rather than silently succeeding on a no-op, so an + agent can tell "I did this" from "this was already so". --- diff --git a/backend/mcp_server/context.py b/backend/mcp_server/context.py index bff616ed69..ab1f8058b7 100644 --- a/backend/mcp_server/context.py +++ b/backend/mcp_server/context.py @@ -27,11 +27,17 @@ class PlatformMCPContext: platform_key: The validated key, carried for its permission tier. org_name: Organization identifier from the URL, already validated against the key's own organization by the middleware. + request: The DRF request. Carried because several platform helpers + (notably ``PipelineManager.execute_pipeline``) take a request + rather than a user and mutate ``request.data`` on the way through. + Passing the real request is safer than fabricating one, since it + keeps the org and auth state those helpers read. """ user: Any platform_key: Any org_name: str + request: Any = None @dataclass(frozen=True) diff --git a/backend/mcp_server/platform_views.py b/backend/mcp_server/platform_views.py index 8f713da4c3..37444ae072 100644 --- a/backend/mcp_server/platform_views.py +++ b/backend/mcp_server/platform_views.py @@ -66,9 +66,12 @@ def server_instructions(self, context: PlatformMCPContext) -> str: return ( "Unstract runs LLM-driven extraction over unstructured documents " "and returns structured JSON. This server is scoped to the " - f"organization '{context.org_name}' and is read-only: it discovers " - "deployments, workflows and Prompt Studio projects but cannot run " - "or change anything. Call readMeFirst before any other tool." + f"organization '{context.org_name}'. It can discover deployments, " + "workflows, pipelines and Prompt Studio projects, and can change " + "their running state — activating deployments, pausing schedules, " + "triggering pipeline runs. Those writes affect the whole " + "organization and consume quota. Call readMeFirst before any " + "other tool." ) def resolve_context( @@ -103,25 +106,39 @@ def resolve_context( # Fall back to the key's own organization so the context is never # built with a null org, whatever the routing did. org_name=org_name or str(platform_key.organization.organization_id), + request=request, ) def check_tool_allowed(self, tool: Any, context: PlatformMCPContext) -> str | None: - """Refuse write tools to keys below the read_write tier. - - The middleware's tier check gates on HTTP method, and every MCP call is - a POST — so it cannot distinguish a tool that reads from one that - writes. This is where that distinction is actually enforced. - - Today the platform registry is entirely read-only, so this never fires. - It is here so that adding the first write tool is a one-line change - that is already guarded, rather than a silent privilege widening. + """Apply the platform key's permission tier to this specific tool. + + The middleware already checked the tier, but it can only check it + against the HTTP method of the request — and every MCP call is a POST. + That makes its verdict uselessly coarse here: it cannot tell + ``listWorkflows`` from ``executePipeline``. + + So each tool declares the method its REST equivalent would use, and the + tier is re-evaluated against that. This deliberately reuses + ``ApiKeyPermission.allows`` rather than inventing a second permission + scheme, so a tool marked DELETE is full_access-only for exactly the + same reason a REST DELETE is. """ - if not tool.writes: + try: + permission = ApiKeyPermission(context.platform_key.permission) + except ValueError: + # The middleware rejects unrecognised tiers, so this is + # unreachable in practice; refusing is the safe answer if it ever + # becomes reachable. + logger.error( + f"Unrecognised platform key tier " + f"'{context.platform_key.permission}'; refusing tool access." + ) + return "This API key has an unrecognised permission tier." + + if permission.allows(tool.required_method): return None - if context.platform_key.permission == ApiKeyPermission.READ.value: - return ( - f"Tool '{tool.name}' modifies data and requires a platform API " - "key with the read_write or full_access permission." - ) - return None + return ( + f"Tool '{tool.name}' requires a platform API key permitting " + f"{tool.required_method}; this key is '{permission.value}'." + ) diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index cf4192e09e..3649831c9f 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -25,6 +25,14 @@ class MCPTool: handler: Callable invoked as ``handler(context, **arguments)``. writes: True when the tool has side effects (consumes quota, starts an execution). Read-only tools are safe to retry; write tools are not. + required_method: The HTTP method this tool's REST equivalent would use. + Platform API key tiers are defined in terms of HTTP methods + (``ApiKeyPermission.allows``), and every MCP call is a POST — so + declaring the *equivalent* method is what lets the existing tier + semantics apply per tool instead of per request. "GET" for reads, + "POST" for mutations, "DELETE" for destructive operations (which + only ``full_access`` permits). Unused by the deployment server, + whose key has no tiers. """ name: str @@ -32,6 +40,7 @@ class MCPTool: input_schema: dict[str, Any] handler: Callable[..., Any] writes: bool = False + required_method: str = "GET" def to_mcp_schema(self) -> dict[str, Any]: """Serialize to the shape returned by `tools/list`.""" @@ -154,17 +163,35 @@ def build_deployment_registry() -> MCPToolRegistry: def build_platform_registry() -> MCPToolRegistry: """Build the tools exposed by the organization-scoped MCP server. - Read-only by design. Every tool here is reachable by any caller holding a - platform key, and an agent will eventually call a destructive tool by - mistake — so writes are deferred until there is a per-tool authorization - story stronger than the key's HTTP-method tier. + Each tool declares the HTTP method its REST equivalent would use, which is + how the platform key's existing permission tier applies per tool rather + than per request — see ``MCPTool.required_method``. + + Two categories are deliberately absent: + + * **Credential operations.** Creating or rotating an API key returns the + secret in its response, so an MCP tool for it would be an exfiltration + path for an agent processing untrusted document content. The codebase + already reasons this way — see ``CanRotatePlatformApiKey``'s docstring on + why rotation is ``full_access``-only. + * **Deletions.** Removing a workflow, deployment or Prompt Studio project + destroys work that no inverse call restores. The write tools here are + reversible by construction. """ from mcp_server.tools.platform import ( + _ORG_WIDE_WARNING, + execute_pipeline, + execute_pipeline_schema, list_api_deployments, + list_pipelines, list_prompt_studio_projects, list_workflows, no_args_schema, platform_read_me_first, + set_api_deployment_active, + set_api_deployment_active_schema, + set_pipeline_active, + set_pipeline_active_schema, whoami, ) @@ -231,6 +258,71 @@ def build_platform_registry() -> MCPToolRegistry: handler=list_prompt_studio_projects, ) ) + registry.register( + MCPTool( + name="listPipelines", + description=( + "List the organization's ETL and task pipelines, with their " + "schedule state and the status of their last run. Takes no " + "arguments." + ), + input_schema=no_args_schema(), + handler=list_pipelines, + ) + ) + + registry.register( + MCPTool( + name="setApiDeploymentActive", + description=( + "Activate or deactivate an API deployment. A deactivated " + "deployment rejects extraction requests; activating it again " + "restores service.\n\n" + "Use this to take a misbehaving deployment offline, or to " + "bring one back. The change is immediate and affects every " + "caller of that deployment, not just this session.\n\n" + f"{_ORG_WIDE_WARNING}" + ), + input_schema=set_api_deployment_active_schema(), + handler=set_api_deployment_active, + writes=True, + required_method="POST", + ) + ) + registry.register( + MCPTool( + name="setPipelineActive", + description=( + "Enable or pause a pipeline's schedule. A paused pipeline " + "stops running on its schedule but is not deleted, and " + "enabling it resumes the existing schedule.\n\n" + f"{_ORG_WIDE_WARNING}" + ), + input_schema=set_pipeline_active_schema(), + handler=set_pipeline_active, + writes=True, + required_method="POST", + ) + ) + registry.register( + MCPTool( + name="executePipeline", + description=( + "Trigger an immediate run of an ETL or task pipeline, " + "independently of its schedule.\n\n" + "This does real work: it processes whatever documents the " + "pipeline's source currently holds and writes to its " + "destination, consuming the organization's quota. It is not a " + "dry run and cannot be undone — do not call it to 'test' a " + "pipeline.\n\n" + f"{_ORG_WIDE_WARNING}" + ), + input_schema=execute_pipeline_schema(), + handler=execute_pipeline, + writes=True, + required_method="POST", + ) + ) return registry diff --git a/backend/mcp_server/tests/test_mcp_auth.py b/backend/mcp_server/tests/test_mcp_auth.py index e0ff95228b..7c3391d7fe 100644 --- a/backend/mcp_server/tests/test_mcp_auth.py +++ b/backend/mcp_server/tests/test_mcp_auth.py @@ -71,7 +71,7 @@ def _post( "method": "tools/list", } request = self.factory.post( - f"/mcp/{org}/{api_name}/", body, format="json", **headers + f"/deployment/api/{org}/{api_name}/mcp", body, format="json", **headers ) kwargs = {"org_name": org, "api_name": api_name} if path_key is not None: @@ -146,6 +146,43 @@ def test_path_key_wins_over_valid_header(self) -> None: ) assert response.status_code == 401, response.content + @pytest.mark.critical_path("mcp-server-auth") + def test_mcp_url_hangs_off_the_execution_url(self) -> None: + """Pins the endpoint's address and its relationship to the REST one. + + The MCP endpoint is a suffix on the deployment's own execution URL, so + the two are visibly the same resource. It must also stay inside the + whitelisted deployment prefix: this server authenticates a deployment + key itself, and routing it somewhere the session middleware applies + would break every client. + """ + from django.conf import settings + from django.test import Client + from django.urls import resolve + + base = f"/deployment/api/{ORG_ID}/live-api" + + assert resolve(f"{base}/mcp").url_name == "mcp_server" + assert resolve(f"{base}/mcp/{self.key.api_key}/").url_name == ( + "mcp_server_with_key" + ) + # The execution endpoint must still resolve to execution, not MCP. + assert resolve(f"{base}/").url_name == "api_deployment_execution" + + assert any( + f"{base}/mcp".startswith(path) for path in settings.WHITELISTED_PATHS + ), "deployment MCP path must stay inside the whitelisted deployment prefix" + + # And it must really answer there, not merely resolve. + response = Client().post( + f"{base}/mcp", + data=json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}), + content_type="application/json", + HTTP_AUTHORIZATION=f"Bearer {self.key.api_key}", + ) + assert response.status_code == 200, response.content + assert "readMeFirst" in response.content.decode() + @pytest.mark.critical_path("mcp-server-auth") def test_get_probe_does_not_leak_deployment_details(self) -> None: """The unauthenticated GET probe advertises the server, not the diff --git a/backend/mcp_server/tests/test_platform_auth.py b/backend/mcp_server/tests/test_platform_auth.py index 5a281529c5..bae3bcaa9f 100644 --- a/backend/mcp_server/tests/test_platform_auth.py +++ b/backend/mcp_server/tests/test_platform_auth.py @@ -133,6 +133,10 @@ def test_valid_key_reaches_the_tool_listing(self) -> None: "listApiDeployments", "listWorkflows", "listPromptStudioProjects", + "listPipelines", + "setApiDeploymentActive", + "setPipelineActive", + "executePipeline", ] @pytest.mark.critical_path("mcp-platform-auth") diff --git a/backend/mcp_server/tests/test_platform_tier_guard.py b/backend/mcp_server/tests/test_platform_tier_guard.py index ede1936c70..1e3765082d 100644 --- a/backend/mcp_server/tests/test_platform_tier_guard.py +++ b/backend/mcp_server/tests/test_platform_tier_guard.py @@ -1,9 +1,10 @@ -"""The per-tool permission guard on the platform server. +"""Per-tool authorization on the platform server. -The platform registry is read-only today, so this guard never fires in -practice. It is tested anyway: it exists so that adding the first write tool is -already protected, and an unexercised guard is one that quietly stops working -before anyone depends on it. +The auth middleware checks the key's permission tier against the request's HTTP +method — but every MCP call is a POST, so its verdict cannot distinguish +``listWorkflows`` from ``executePipeline``. This guard re-applies the tier +against the method each tool *declares*, and is the only thing standing between +a low-tier key and a write tool. """ from __future__ import annotations @@ -14,23 +15,24 @@ from mcp_server.context import PlatformMCPContext from mcp_server.platform_views import PlatformMCPServerView -from mcp_server.registry import MCPTool, PLATFORM_TOOLS +from mcp_server.registry import PLATFORM_TOOLS, MCPTool -def a_tool(writes: bool) -> MCPTool: +def a_tool(required_method: str, writes: bool = True) -> MCPTool: return MCPTool( name="doThing", description="d", input_schema={"type": "object", "properties": {}}, handler=lambda ctx: None, writes=writes, + required_method=required_method, ) def context_for(tier: str) -> PlatformMCPContext: return PlatformMCPContext( user=Mock(is_service_account=True), - platform_key=Mock(permission=tier, name="k"), + platform_key=Mock(permission=tier), org_name="org-tier", ) @@ -39,38 +41,84 @@ class PlatformTierGuardTest(SimpleTestCase): def setUp(self) -> None: self.view = PlatformMCPServerView() - def test_read_tier_is_refused_a_write_tool(self) -> None: - refusal = self.view.check_tool_allowed(a_tool(writes=True), context_for("read")) + def test_tier_matrix(self) -> None: + """The whole authorization model in one table. - assert refusal is not None + Mirrors ``ApiKeyPermission.allows`` deliberately: if that mapping ever + changes, this fails and forces the MCP surface to be reconsidered + rather than silently inheriting a wider grant. + """ + cases = [ + # tier, method, allowed + ("read", "GET", True), + ("read", "POST", False), + ("read", "DELETE", False), + ("read_write", "GET", True), + ("read_write", "POST", True), + # read_write must NOT reach destructive tools. + ("read_write", "DELETE", False), + ("full_access", "GET", True), + ("full_access", "POST", True), + ("full_access", "DELETE", True), + ] + for tier, method, allowed in cases: + with self.subTest(f"{tier} -> {method}"): + refusal = self.view.check_tool_allowed( + a_tool(method), context_for(tier) + ) + assert (refusal is None) == allowed, ( + f"{tier} {method}: expected allowed={allowed}, got {refusal!r}" + ) + + def test_refusal_names_the_tool_and_the_tier(self) -> None: + """The refusal is read by an agent, which should be able to tell its + operator what to change rather than just retrying. + """ + refusal = self.view.check_tool_allowed(a_tool("DELETE"), context_for("read_write")) + + assert "doThing" in refusal + assert "DELETE" in refusal assert "read_write" in refusal - def test_read_write_tier_is_allowed_a_write_tool(self) -> None: - assert ( - self.view.check_tool_allowed(a_tool(writes=True), context_for("read_write")) - is None - ) + def test_unrecognised_tier_is_refused(self) -> None: + refusal = self.view.check_tool_allowed(a_tool("GET"), context_for("wat")) - def test_read_tier_is_allowed_a_read_tool(self) -> None: - """The guard must not block reads — a read key that reaches the server - should still be able to use every read tool on it. + assert refusal is not None + + def test_every_write_tool_declares_a_non_get_method(self) -> None: + """The invariant that keeps the guard meaningful. + + A ``writes=True`` tool left at the default ``required_method="GET"`` + would be reachable by any key that can reach the server at all — the + guard would pass it silently. This is what makes registering a write + tool safe by default. """ - assert ( - self.view.check_tool_allowed(a_tool(writes=False), context_for("read")) - is None + unguarded = [ + name + for name in PLATFORM_TOOLS.names() + if PLATFORM_TOOLS.get(name).writes + and PLATFORM_TOOLS.get(name).required_method == "GET" + ] + + assert unguarded == [], ( + f"These platform tools mutate state but declare required_method=" + f"'GET', so the tier guard will not protect them: {unguarded}" ) - def test_platform_registry_is_entirely_read_only(self) -> None: - """Pins the documented promise that this server changes nothing. + def test_no_credential_tools_are_exposed(self) -> None: + """API key creation and rotation return the secret in their response. - If someone registers a write tool here, this fails and forces them to - confirm the tier guard and the README claim still hold. + Exposing either as an MCP tool would hand an agent — one that may be + processing untrusted document content — a way to mint or exfiltrate + credentials. Named explicitly so adding one is a deliberate act. """ - writers = [ - name for name in PLATFORM_TOOLS.names() if PLATFORM_TOOLS.get(name).writes + forbidden = ("key", "credential", "secret", "token", "rotate") + offenders = [ + name + for name in PLATFORM_TOOLS.names() + if any(word in name.lower() for word in forbidden) ] - assert writers == [], ( - f"Platform MCP server advertises itself as read-only but exposes " - f"write tools: {writers}" + assert offenders == [], ( + f"Platform MCP server must not expose credential tools: {offenders}" ) diff --git a/backend/mcp_server/tests/test_platform_tools.py b/backend/mcp_server/tests/test_platform_tools.py index bfb08822d6..142874c37b 100644 --- a/backend/mcp_server/tests/test_platform_tools.py +++ b/backend/mcp_server/tests/test_platform_tools.py @@ -7,9 +7,12 @@ from __future__ import annotations +from unittest.mock import Mock, patch + from account_v2.models import Organization, User from api_v2.models import APIDeployment from django.test import TestCase +from pipeline_v2.models import Pipeline from platform_api.models import PlatformApiKey from prompt_studio.prompt_studio_core_v2.models import CustomTool from tenant_account_v2.models import OrganizationMember @@ -17,11 +20,15 @@ from workflow_manager.workflow_v2.models.workflow import Workflow from mcp_server.context import PlatformMCPContext +from mcp_server.exceptions import MCPToolError from mcp_server.tools.platform import ( + execute_pipeline, list_api_deployments, list_prompt_studio_projects, list_workflows, platform_read_me_first, + set_api_deployment_active, + set_pipeline_active, whoami, ) @@ -63,9 +70,15 @@ def setUp(self) -> None: self.tool = CustomTool.objects.create( tool_name="Invoice Prompts", description="Prompt project", author="acme" ) + self.pipeline = Pipeline.objects.create( + pipeline_name="etl-pipeline", workflow=self.workflow + ) self.context = PlatformMCPContext( - user=self.user, platform_key=self.key, org_name=ORG_ID + user=self.user, + platform_key=self.key, + org_name=ORG_ID, + request=Mock(data={}), ) def test_list_api_deployments_exposes_the_api_name(self) -> None: @@ -134,7 +147,119 @@ def test_whoami_reports_tier_and_service_account_scope(self) -> None: assert result["organization"] == ORG_ID assert result["permission_tier"] == "read_write" assert result["is_service_account"] is True - assert "read-only" in result["access"] + # read_write may write but must not reach destructive tools. + assert result["can_use_write_tools"] is True + assert result["can_use_destructive_tools"] is False + assert "modify" in result["visibility"] + + def test_set_api_deployment_active_toggles_and_reports_the_change(self) -> None: + result = set_api_deployment_active( + self.context, api_name="invoice-api", is_active=False + ) + + assert result["changed"] is True + assert result["previous_state"] is True + assert result["is_active"] is False + self.api.refresh_from_db() + assert self.api.is_active is False + + def test_set_api_deployment_active_is_a_noop_when_already_in_state(self) -> None: + """Reporting `changed: False` rather than silently succeeding lets an + agent tell "I did this" from "this was already so". + """ + result = set_api_deployment_active( + self.context, api_name="invoice-api", is_active=True + ) + + assert result["changed"] is False + + def test_set_api_deployment_active_rejects_an_unknown_name(self) -> None: + with self.assertRaises(MCPToolError) as caught: + set_api_deployment_active( + self.context, api_name="no-such-api", is_active=False + ) + + # The message must point at the tool that lists valid names. + assert "listApiDeployments" in str(caught.exception) + + def test_set_pipeline_active_toggles_the_schedule(self) -> None: + result = set_pipeline_active( + self.context, pipeline_id=str(self.pipeline.id), active=True + ) + + assert result["changed"] is True + self.pipeline.refresh_from_db() + assert self.pipeline.active is True + + def test_execute_pipeline_delegates_to_the_platform_manager(self) -> None: + """Guards the tool-kwarg -> PipelineManager mapping. + + Nothing else executes this call, so a renamed kwarg would pass every + other test and fail on the first real pipeline run. + """ + fake_response = Mock(status_code=200, data={"execution": "started"}) + with patch( + "pipeline_v2.manager.PipelineManager.execute_pipeline", + return_value=fake_response, + ) as execute: + result = execute_pipeline(self.context, pipeline_id=str(self.pipeline.id)) + + execute.assert_called_once_with( + request=self.context.request, pipeline_id=str(self.pipeline.id) + ) + assert result["status"] == 200 + assert result["pipeline_name"] == "etl-pipeline" + + def test_execute_pipeline_rejects_an_unknown_id(self) -> None: + with patch("pipeline_v2.manager.PipelineManager.execute_pipeline") as execute: + with self.assertRaises(MCPToolError): + execute_pipeline( + self.context, pipeline_id="99999999-9999-9999-9999-999999999999" + ) + + execute.assert_not_called() + + def test_writes_do_not_cross_organizations(self) -> None: + """The most important property of the write tools. + + A platform key must not be able to deactivate, pause or run another + tenant's resources — the read-side isolation test has a write-side + equivalent because the consequences here are not merely disclosure. + """ + UserContext.set_organization_identifier("org-write-other") + Organization.objects.create( + name="org-write-other", + display_name="Other", + organization_id="org-write-other", + ) + other_wf = Workflow.objects.create(workflow_name="other-wf", is_active=True) + other_api = APIDeployment.objects.create( + api_name="other-api", display_name="Other API", workflow=other_wf + ) + other_pipeline = Pipeline.objects.create( + pipeline_name="other-pipeline", workflow=other_wf + ) + + UserContext.set_organization_identifier(ORG_ID) + + with self.assertRaises(MCPToolError): + set_api_deployment_active( + self.context, api_name="other-api", is_active=False + ) + with self.assertRaises(MCPToolError): + set_pipeline_active( + self.context, pipeline_id=str(other_pipeline.id), active=True + ) + with patch("pipeline_v2.manager.PipelineManager.execute_pipeline") as execute: + with self.assertRaises(MCPToolError): + execute_pipeline(self.context, pipeline_id=str(other_pipeline.id)) + execute.assert_not_called() + + # And nothing was touched. + other_api.refresh_from_db() + other_pipeline.refresh_from_db() + assert other_api.is_active is True + assert other_pipeline.active is False def test_read_me_first_states_the_scope_warning(self) -> None: """The service-account visibility rule is surprising enough that the diff --git a/backend/mcp_server/tools/platform.py b/backend/mcp_server/tools/platform.py index a227a9234d..c1b82f13af 100644 --- a/backend/mcp_server/tools/platform.py +++ b/backend/mcp_server/tools/platform.py @@ -1,24 +1,31 @@ -"""Read-only tools for the organization-scoped MCP server. +"""Tools for the organization-scoped MCP server. -Every listing goes through the model's own ``for_user`` manager rather than a -raw queryset, so the platform's sharing and visibility rules apply here exactly -as they do in the UI — this app does not get its own opinion about who may see -what. +Every query goes through the model's own ``for_user`` manager rather than a raw +queryset, so the platform's sharing and visibility rules apply here exactly as +they do in the UI — this app does not get its own opinion about who may see or +touch what. -One caveat those managers impose, which the tool descriptions state plainly: -a platform key resolves to a service account, and ``for_user`` short-circuits -to "everything in the organization" for service accounts. These tools therefore -see all org resources regardless of the role the key was created with. +One caveat those managers impose, which the tool descriptions state plainly: a +platform key resolves to a service account, and ``for_user`` short-circuits to +"everything in the organization" for service accounts. These tools therefore +reach all org resources regardless of the role the key was created with — for +the write tools below that means they can modify anything in the organization, +not merely what one user could. + +Write tools are authorized per tool by the view, using the platform key's +existing permission tier — see ``MCPTool.required_method``. """ import logging from typing import Any from api_v2.models import APIDeployment +from pipeline_v2.models import Pipeline from prompt_studio.prompt_studio_core_v2.models import CustomTool from workflow_manager.workflow_v2.models.workflow import Workflow from mcp_server.context import PlatformMCPContext +from mcp_server.exceptions import MCPToolError logger = logging.getLogger(__name__) @@ -32,6 +39,16 @@ def no_args_schema() -> dict[str, Any]: return {"type": "object", "properties": {}, "required": []} +# Every write tool repeats this, because the tool description is the prompt the +# agent reads before deciding to act — and the blast radius of a write is wider +# than the caveat's usual home in whoami. +_ORG_WIDE_WARNING = ( + "This credential is a service account and can modify ANY resource in the " + "organization, not only those shared with a particular user. Confirm you " + "have the right target before calling." +) + + def _truncation_note(shown: int, total: int) -> dict[str, Any]: """Describe a capped listing, or nothing when it was complete. @@ -57,48 +74,86 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "what_this_does": ( "Unstract runs LLM-driven extraction over unstructured documents " "(PDFs, scans, images) and returns structured JSON. This server is " - "scoped to a whole organization and is READ-ONLY: it lets you " - "discover what exists — deployments, workflows, Prompt Studio " - "projects — but not change or run anything." + "scoped to a whole organization: it discovers what exists and can " + "change the running state of those resources." ), - "tools": [ + "read_tools": [ {"name": "whoami", "purpose": "See this credential's org, tier and scope."}, { "name": "listApiDeployments", "purpose": "Find deployed extraction endpoints.", }, - {"name": "listWorkflows", "purpose": "Find the pipelines behind them."}, + {"name": "listWorkflows", "purpose": "Find the workflows behind them."}, + {"name": "listPipelines", "purpose": "Find ETL and task pipelines."}, { "name": "listPromptStudioProjects", "purpose": "Find where extraction prompts are authored.", }, ], + "write_tools": [ + { + "name": "setApiDeploymentActive", + "purpose": "Take a deployment offline, or bring it back.", + }, + { + "name": "setPipelineActive", + "purpose": "Pause or resume a pipeline's schedule.", + }, + { + "name": "executePipeline", + "purpose": "Trigger a real pipeline run. Consumes quota.", + }, + ], + "before_writing": ( + "Write tools change state for the whole organization and take " + "effect immediately — there is no staging or dry-run mode. " + "executePipeline in particular does real work and consumes quota; " + "it is not a way to test a pipeline. Confirm the target with a " + "list tool first." + ), + "not_available": ( + "Creating or rotating API keys, and deleting workflows, " + "deployments or projects, are deliberately not exposed here." + ), "to_run_an_extraction": ( "This server cannot extract. Call listApiDeployments to get an " "api_name, then open a separate MCP session against that " - "deployment's own endpoint — /mcp/// — using that " - "deployment's API key. That session exposes extractDocument." + "deployment's own endpoint — /deployment/api///mcp " + "— using that deployment's API key. That session exposes " + "extractDocument." ), "scope_warning": ( - "This credential is a service account, which sees ALL resources in " - "the organization regardless of the role it was created with. Do " - "not treat these listings as one person's view." + "This credential is a service account, which reaches ALL resources " + "in the organization regardless of the role it was created with. " + "Do not treat these listings as one person's view, and remember " + "the write tools have the same reach." ), } def whoami(context: PlatformMCPContext) -> dict[str, Any]: """Describe the calling credential, so an agent can reason about scope.""" + from platform_api.models import ApiKeyPermission + key = context.platform_key + try: + permission = ApiKeyPermission(key.permission) + can_write = permission.allows("POST") + can_delete = permission.allows("DELETE") + except ValueError: + can_write = can_delete = False + return { "organization": context.org_name, "key_name": key.name, "permission_tier": key.permission, "is_service_account": bool(getattr(context.user, "is_service_account", False)), - "access": "read-only (this MCP server exposes no write tools)", + "can_use_write_tools": can_write, + "can_use_destructive_tools": can_delete, "visibility": ( "All resources in the organization. Service-account credentials " - "bypass per-user sharing filters." + "bypass per-user sharing filters, so this key can read AND modify " + "any resource in the organization." ), } @@ -167,3 +222,210 @@ def list_prompt_studio_projects(context: PlatformMCPContext) -> dict[str, Any]: ], **_truncation_note(len(rows), total), } + + +def list_pipelines(context: PlatformMCPContext) -> dict[str, Any]: + """List the organization's ETL and task pipelines.""" + queryset = Pipeline.objects.for_user(context.user).order_by("pipeline_name") + total = queryset.count() + rows = queryset[:LIST_LIMIT] + + return { + "count": total, + "pipelines": [ + { + "id": str(row.id), + "pipeline_name": row.pipeline_name, + "pipeline_type": row.pipeline_type, + "active": row.active, + "last_run_status": row.last_run_status, + } + for row in rows + ], + **_truncation_note(len(rows), total), + } + + +# ============================ write tools ============================ +# +# Authorization for these is not decided here. Each is registered with a +# `required_method`, and the view checks it against the platform key's tier +# using ``ApiKeyPermission.allows`` — the same method-based semantics the auth +# middleware already enforces for REST. Destructive operations therefore fall +# out as full_access-only without a second, parallel permission scheme. + + +def set_api_deployment_active_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "api_name": { + "type": "string", + "description": ( + "The api_name of the deployment, as returned by listApiDeployments." + ), + }, + "is_active": { + "type": "boolean", + "description": ( + "True to activate the deployment so it accepts extraction " + "requests, False to take it offline." + ), + }, + }, + "required": ["api_name", "is_active"], + } + + +def set_api_deployment_active( + context: PlatformMCPContext, api_name: str, is_active: bool +) -> dict[str, Any]: + """Activate or deactivate an API deployment. + + Reversible by construction — the inverse call restores the previous state — + which is why this is the write operation exposed rather than anything that + destroys configuration. + """ + deployment = ( + APIDeployment.objects.for_user(context.user).filter(api_name=api_name).first() + ) + if deployment is None: + raise MCPToolError( + f"No API deployment named '{api_name}' in organization " + f"'{context.org_name}'. Call listApiDeployments to see valid names." + ) + + was_active = deployment.is_active + if was_active == is_active: + return { + "api_name": api_name, + "is_active": is_active, + "changed": False, + "message": f"Deployment '{api_name}' was already " + f"{'active' if is_active else 'inactive'}.", + } + + deployment.is_active = is_active + deployment.save(update_fields=["is_active"]) + logger.info( + f"Platform MCP set deployment '{api_name}' active={is_active} " + f"(org '{context.org_name}', key '{context.platform_key.name}')" + ) + + return { + "api_name": api_name, + "is_active": is_active, + "changed": True, + "previous_state": was_active, + } + + +def execute_pipeline_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "pipeline_id": { + "type": "string", + "description": ( + "UUID of the pipeline to run, as returned by listPipelines." + ), + }, + }, + "required": ["pipeline_id"], + } + + +def execute_pipeline(context: PlatformMCPContext, pipeline_id: str) -> dict[str, Any]: + """Trigger a run of an ETL or task pipeline. + + Delegates to ``PipelineManager``, which is what the REST endpoint calls, so + scheduling, execution records and notifications behave identically to a run + started from the UI. + """ + from pipeline_v2.manager import PipelineManager + + pipeline = Pipeline.objects.for_user(context.user).filter(id=pipeline_id).first() + if pipeline is None: + raise MCPToolError( + f"No pipeline with id '{pipeline_id}' in organization " + f"'{context.org_name}'. Call listPipelines to see valid ids." + ) + + if context.request is None: + # The manager mutates request.data on the way to the workflow viewset, + # so there is no meaningful fallback — fail loudly rather than + # fabricate a request and produce a half-configured execution. + raise MCPToolError( + "Pipeline execution is unavailable in this context. " + "Contact your Unstract administrator." + ) + + logger.info( + f"Platform MCP executing pipeline '{pipeline_id}' " + f"(org '{context.org_name}', key '{context.platform_key.name}')" + ) + response = PipelineManager.execute_pipeline( + request=context.request, pipeline_id=str(pipeline.id) + ) + + return { + "pipeline_id": str(pipeline.id), + "pipeline_name": pipeline.pipeline_name, + "status": getattr(response, "status_code", None), + "result": getattr(response, "data", None), + } + + +def set_pipeline_active_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "pipeline_id": { + "type": "string", + "description": ("UUID of the pipeline, as returned by listPipelines."), + }, + "active": { + "type": "boolean", + "description": ( + "True to enable the pipeline's schedule, False to pause it." + ), + }, + }, + "required": ["pipeline_id", "active"], + } + + +def set_pipeline_active( + context: PlatformMCPContext, pipeline_id: str, active: bool +) -> dict[str, Any]: + """Enable or pause a pipeline's schedule.""" + pipeline = Pipeline.objects.for_user(context.user).filter(id=pipeline_id).first() + if pipeline is None: + raise MCPToolError( + f"No pipeline with id '{pipeline_id}' in organization " + f"'{context.org_name}'. Call listPipelines to see valid ids." + ) + + was_active = pipeline.active + if was_active == active: + return { + "pipeline_id": str(pipeline.id), + "active": active, + "changed": False, + "message": f"Pipeline was already {'active' if active else 'paused'}.", + } + + pipeline.active = active + pipeline.save(update_fields=["active"]) + logger.info( + f"Platform MCP set pipeline '{pipeline_id}' active={active} " + f"(org '{context.org_name}', key '{context.platform_key.name}')" + ) + + return { + "pipeline_id": str(pipeline.id), + "pipeline_name": pipeline.pipeline_name, + "active": active, + "changed": True, + "previous_state": was_active, + } diff --git a/backend/mcp_server/urls.py b/backend/mcp_server/urls.py index 89d42b3892..4be26f670e 100644 --- a/backend/mcp_server/urls.py +++ b/backend/mcp_server/urls.py @@ -1,10 +1,15 @@ -"""URLs for the hosted MCP server. +"""URLs for the deployment-scoped MCP server. -Mounted alongside the API deployment execution endpoint so an MCP endpoint is -the same shape as the REST endpoint for the same deployment: +Included from ``api_v2/execution_urls.py`` so the MCP endpoint hangs directly +off the deployment's own execution URL — the same resource, reached two ways: - POST /deployment/api/// (REST) - POST /mcp/// (MCP) + POST /deployment/api/// (REST) + POST /deployment/api///mcp (MCP) + +Sitting under ``API_DEPLOYMENT_PATH_PREFIX`` means these paths are already +covered by that prefix's entry in ``WHITELISTED_PATHS``, so +``CustomAuthMiddleware`` skips them and the view authenticates the deployment +key itself — exactly as it does for the execution endpoint next door. """ from django.urls import re_path @@ -15,7 +20,7 @@ urlpatterns = [ re_path( - r"^(?P[\w-]+)/(?P[\w-]+)/?$", + r"^api/(?P[\w-]+)/(?P[\w-]+)/mcp/?$", mcp_server, name="mcp_server", ), @@ -23,7 +28,7 @@ # Authorization header to the request. The key is a UUID, so the pattern # cannot collide with the header-authenticated route above. re_path( - r"^(?P[\w-]+)/(?P[\w-]+)/" + r"^api/(?P[\w-]+)/(?P[\w-]+)/mcp/" r"(?P[0-9a-fA-F-]{36})/?$", mcp_server, name="mcp_server_with_key", From 5073754600930682293df39655a8d00e9e929326 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 14:57:18 +0000 Subject: [PATCH 06/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- backend/api_v2/serializers.py | 2 +- backend/backend/exceptions.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/backend/api_v2/serializers.py b/backend/api_v2/serializers.py index 6c13d8e49c..0eb99f1bce 100644 --- a/backend/api_v2/serializers.py +++ b/backend/api_v2/serializers.py @@ -4,7 +4,6 @@ from typing import Any from urllib.parse import urlparse -from backend.serializers import AuditSerializer from django.apps import apps from django.core.validators import RegexValidator from pipeline_v2.models import Pipeline @@ -36,6 +35,7 @@ from api_v2.constants import ApiExecution from api_v2.models import APIDeployment, APIKey +from backend.serializers import AuditSerializer class APIDeploymentSerializer(IntegrityErrorMixin, AuditSerializer): diff --git a/backend/backend/exceptions.py b/backend/backend/exceptions.py index c7ca940752..caf4df575b 100644 --- a/backend/backend/exceptions.py +++ b/backend/backend/exceptions.py @@ -3,6 +3,7 @@ from rest_framework.exceptions import APIException from rest_framework.response import Response from rest_framework.views import exception_handler + from unstract.connectors.exceptions import ConnectorBaseException, ConnectorError From 8d128fa2812ad9df827108dd92dfa1ec9a29770a Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sat, 25 Jul 2026 23:21:27 +0530 Subject: [PATCH 07/16] feat(mcp): billable Prompt Studio tools with a spend guard, plus the deferred read tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four changes. 1. Documents the platform server's exclusions in its README — the credential-bearing surfaces and destructive operations no tool wraps, and why. Framed as what this server does not do rather than as an inventory of which endpoints return secrets, since this repository is public. 2. Adds the billable Prompt Studio operations: indexDocument, fetchResponse, bulkFetchResponse and singlePassExtraction. Each delegates to PromptStudioCoreView rather than reimplementing run-id generation, profile resolution, lookup gating and Celery dispatch — a parallel implementation would drift from the UI on the first change to either. 3. Adds a spend guard for them. There is no OSS token budget to gate on (usage_v2 reports after the fact; subscription enforcement is enterprise), so this is a per-organization *call* budget in Redis. Enforced at the transport via a `billable` flag, so a new costly tool cannot be added without being budgeted. Three deliberate behaviours: budget is consumed on invocation and never refunded — the opposite of the rate-limit slot in tools/execution.py, because that models concurrency while this models money already spent, and refunding would let an agent spend without limit by failing in a loop; exhaustion returns a retryable isError result rather than a permission error, so the agent waits instead of giving up; and it fails open, because bounding runaway loops is not worth taking the server offline when Redis blips. whoami reports remaining budget so an agent can pace itself. 4. Adds the read tools previously skipped for no good reason: listExecutions, getExecutionDetail, getUsageSummary, getWorkflowEndpoints, listToolInstances, listTags and listPipelines. Execution history in particular closes the gap where an agent could trigger work but not see how it went. Two things surfaced while building this, both now fixed: * A registry-wide test seeds an org with fake secrets, invokes every read tool and fails if any appears in the output. It caught getExecutionDetail being uncovered, then caught a real leak — an execution's error_message can embed the connection string a failed connector tried, which listExecutions returned verbatim. Error text is now passed through redact_secrets. * spend_guard.peek() did not fail open, so whoami broke when the cache was unavailable. Now matches consume(). getWorkflowEndpoints is the case to look at for the pattern: a workflow endpoint references a connector instance whose metadata decrypts on access, so the tool returns the shape of the connection and never the configuration. All responses here are built field-by-field from named attributes; never serializer.data. Tests: 82 in this app (was 59); full backend unit tier 331 passing. Co-Authored-By: Claude Opus 4.8 (1M context) --- CAPTURE3.md | 110 +++++ backend/backend/settings/base.py | 8 + backend/mcp_server/README.md | 148 ++++++- backend/mcp_server/platform_views.py | 24 +- backend/mcp_server/registry.py | 185 ++++++++ backend/mcp_server/spend_guard.py | 153 +++++++ .../tests/test_no_credential_leak.py | 241 +++++++++++ .../mcp_server/tests/test_platform_auth.py | 21 +- backend/mcp_server/tests/test_redaction.py | 81 ++++ backend/mcp_server/tests/test_spend_guard.py | 275 ++++++++++++ backend/mcp_server/tools/observability.py | 400 ++++++++++++++++++ backend/mcp_server/tools/platform.py | 79 +++- backend/mcp_server/tools/prompt_studio.py | 261 ++++++++++++ backend/mcp_server/transport.py | 21 + 14 files changed, 1974 insertions(+), 33 deletions(-) create mode 100644 CAPTURE3.md create mode 100644 backend/mcp_server/spend_guard.py create mode 100644 backend/mcp_server/tests/test_no_credential_leak.py create mode 100644 backend/mcp_server/tests/test_redaction.py create mode 100644 backend/mcp_server/tests/test_spend_guard.py create mode 100644 backend/mcp_server/tools/observability.py create mode 100644 backend/mcp_server/tools/prompt_studio.py diff --git a/CAPTURE3.md b/CAPTURE3.md new file mode 100644 index 0000000000..e97f777648 --- /dev/null +++ b/CAPTURE3.md @@ -0,0 +1,110 @@ +# CAPTURE3 — security findings in the existing REST API + +Two findings surfaced while surveying the tenant API surface to decide what the +platform MCP server should expose (PR #2207). **Neither is caused by that PR, +and neither is fixed by it** — both are pre-existing properties of the REST API, +captured here to be addressed separately. + +--- + +## 1. Platform API keys are readable repeatedly, despite being documented "shown once" + +**Severity:** medium — turns a write-once secret into a read-anytime one. + +`PlatformApiKeyDetailSerializer` includes the full `key` field, and its own +docstring states the intent: + +```python +# backend/platform_api/serializers.py:89-95 +class PlatformApiKeyDetailSerializer(serializers.ModelSerializer): + """Used for create/rotate responses where the full key is shown once.""" + + class Meta: + model = PlatformApiKey + fields = ["id", "name", "key", "is_active"] +``` + +`get_serializer_class` routes `list` → masked, `create` → create serializer, and +`partial_update` → update serializer — but `retrieve` is not in that mapping and +instead selects the detail serializer explicitly: + +```python +# backend/platform_api/views.py:59-62 +def retrieve(self, request, *args, **kwargs): + instance = self.get_object() + serializer = PlatformApiKeyDetailSerializer(instance) + return Response(serializer.data) +``` + +So `GET /api/v1/unstract//platform-api/keys//` returns the full +plaintext key on **every** call, not once at creation. + +**Why it matters.** `list` deliberately masks the key (`****-last4` via +`PlatformApiKeyListSerializer.get_key`), which shows the intent was for the +secret not to be freely readable. `retrieve` silently defeats that: anyone who +can reach the detail endpoint can recover any key in the organization at any +time, so a key leaked once cannot be reasoned about as "shown at creation only", +and rotation is the only remediation. + +**Options to consider.** +- Mask `key` in `retrieve` (reuse the list serializer), and return the full key + only from `create` and `rotate`, which is what the docstring already promises. +- If a re-read capability is genuinely wanted, make it an explicit, separately + authorized action rather than the default detail representation. + +**Related, not the same finding:** `GET /api/keys/api//` in `api_v2` +returns the full deployment API key and is gated at `IsOwnerOrSharedUser` rather +than owner-only — so any user a deployment is *shared with* can read its live +key. Worth reviewing alongside this one. + +--- + +## 2. Two destructive endpoints are exposed over `GET` + +**Severity:** medium — destructive actions reachable by link, prefetch or CSRF. + +```python +# backend/file_management/urls.py:37-47 +file_delete = FileManagementViewSet.as_view({"get": "delete"}) +... +path("file/delete", file_delete, name="delete"), +``` + +```python +# backend/workflow_manager/workflow_v2/urls/workflow.py:25 +workflow_clear_file_marker = WorkflowViewSet.as_view({"get": "clear_file_marker"}) +``` + +`GET /api/v1/unstract//file/delete` deletes files, and +`GET /api/v1/unstract//workflow//clear-file-marker/` clears file +markers. Both mutate state behind a verb that the entire web stack treats as +safe and idempotent. + +**Why it matters.** +- **CSRF.** Django's CSRF protection does not apply to `GET`, so a cross-origin + link, image tag or redirect can trigger either action using the victim's + session. The usual "unsafe methods only" assumption does not hold here. +- **Prefetch and crawlers.** Browsers, link previews, and security scanners + follow `GET` URLs speculatively. Any of these can fire a deletion nobody + requested. +- **Caching and logs.** Intermediaries may cache or replay `GET`, and the URL + (with parameters identifying what is deleted) lands in access logs and browser + history. + +Clearing file markers additionally forces re-processing of the affected files, +so it has a **cost** consequence as well as a data one. + +**Options to consider.** Move both to `POST`/`DELETE`. This is a breaking change +for any existing caller, so it likely needs a deprecation window: accept both +verbs, log `GET` usage, then remove `GET` once callers have migrated. + +--- + +## Not in scope here + +Several other endpoints return decrypted credentials in their normal responses +(connector metadata, adapter provider keys, notification webhook tokens, +`platform_settings_v2` keys, the Postman collection download). Those are +intentional today insofar as the UI relies on them, and they are a broader +design question than these two — the platform MCP server simply does not wrap +any of them. See `backend/mcp_server/README.md` for that exclusion list. diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index b21c15df04..ea8693abe4 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -119,6 +119,14 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | "API_DEPLOYMENT_PATH_PREFIX", "deployment" ).strip("/") +# Budget for billable MCP tool calls (LLM inference, indexing, pipeline runs), +# counted per organization over a rolling window. Bounds how often an agent can +# trigger paid work; it counts calls, not tokens — see mcp_server/spend_guard.py. +MCP_BILLABLE_CALL_LIMIT = int(os.environ.get("MCP_BILLABLE_CALL_LIMIT", 50)) +MCP_BILLABLE_WINDOW_SECONDS = int( + os.environ.get("MCP_BILLABLE_WINDOW_SECONDS", 3600) +) + # Maximum file size for presigned URLs in API deployments (in MB) API_DEPL_PRESIGNED_URL_MAX_FILE_SIZE_MB = int( os.environ.get("API_DEPL_PRESIGNED_URL_MAX_FILE_SIZE_MB", 20) diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index a73ee92211..273267063a 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -147,21 +147,81 @@ Platform API keys are managed at `/api/v1/unstract//platform-api/keys/`. ## Tools +**Discovery** — what exists: + | Tool | Tier | Purpose | | --- | --- | --- | | `readMeFirst` | read | Orientation, including what this credential can reach. | -| `whoami` | read | The key's organization, tier and what it may do. | +| `whoami` | read | The key's organization, tier, and remaining spend budget. | | `listApiDeployments` | read | Deployed extraction endpoints, with their `api_name`. | | `listWorkflows` | read | The workflows behind them. | | `listPipelines` | read | ETL and task pipelines, with schedule and last-run state. | | `listPromptStudioProjects` | read | Where extraction prompts are authored. | +| `getWorkflowEndpoints` | read | How a workflow is wired — shape only, never config. | +| `listToolInstances` | read | The tool steps inside a workflow. | +| `listTags` | read | Execution tags. | + +**Observability** — what happened: + +| Tool | Tier | Purpose | +| --- | --- | --- | +| `listExecutions` | read | Recent runs and their status. Start here to debug. | +| `getExecutionDetail` | read | Per-file results and errors for one run. | +| `getUsageSummary` | read | Tokens and cost recorded so far. | + +**State changes** — cheap, reversible: + +| Tool | Tier | Purpose | +| --- | --- | --- | | `setApiDeploymentActive` | `read_write` | Take a deployment offline, or bring it back. | | `setPipelineActive` | `read_write` | Pause or resume a pipeline's schedule. | -| `executePipeline` | `read_write` | Trigger a real pipeline run. **Consumes quota.** | -To actually extract a document, an agent calls `listApiDeployments` to find an -`api_name`, then opens a **separate** session against the deployment server -above. +**Billable** — each call costs real money, and is budgeted: + +| Tool | Tier | Purpose | +| --- | --- | --- | +| `executePipeline` | `read_write` | Trigger a real pipeline run. | +| `indexDocument` | `read_write` | Embed a document so prompts can run against it. | +| `fetchResponse` | `read_write` | Run one prompt against an indexed document. | +| `bulkFetchResponse` | `read_write` | Run several prompts in one pass. | +| `singlePassExtraction` | `read_write` | Run a project's whole prompt set. | + +To extract a document through a *deployed* API, an agent calls +`listApiDeployments` to find an `api_name`, then opens a **separate** session +against the deployment server above. The Prompt Studio tools here are for +working with prompts before they are deployed. + +## The spend guard + +The billable tools drive LLM inference, embedding and vector-store writes. An +agent looping over a list, or retrying on a misread error, can spend a great +deal without anyone watching — so they are budgeted per organization over a +rolling window (`MCP_BILLABLE_CALL_LIMIT`, default 50 per +`MCP_BILLABLE_WINDOW_SECONDS`, default one hour). + +**It counts calls, not tokens or currency.** Unstract's open-source backend +records usage after the fact but has no pre-flight allowance to check against — +subscription and quota enforcement live in the enterprise overlay, which this +app cannot depend on. A call counter is the strongest guard implementable here, +and it is a blunt one: it bounds how *often* an agent triggers paid work, not +how expensive each call is. + +Three behaviours are deliberate: + +- **Budget is consumed on invocation and never refunded**, including when the + tool then fails. A Prompt Studio call that fails partway may already have + spent tokens upstream, and refunding on failure would let an agent burn + unbounded spend by failing in a loop. This is the opposite of the rate-limit + slot in `tools/execution.py`, which models concurrency rather than cost. +- **Exhaustion is temporary, not a permission error.** It comes back as an + `isError` result naming when to retry, so the agent waits rather than + concluding the tool is forbidden. +- **It fails open.** If the cache is unavailable the call is allowed and the + event logged. This guard bounds runaway loops; it is not a licence check, and + taking the whole MCP surface offline because Redis blipped would be worse. + +`whoami` reports the remaining budget so an agent can pace itself instead of +discovering the limit by hitting it. ## How write tools are authorized @@ -180,16 +240,64 @@ that guard, so a test asserts none exists. ## What is deliberately not exposed -**Credential operations.** Creating or rotating an API key returns the secret in -its response. An MCP tool for it would hand an agent — one that may be -processing untrusted document content — a way to mint or exfiltrate credentials. -The codebase already reasons this way: see `CanRotatePlatformApiKey`'s docstring -on why rotation is `full_access`-only. A test asserts no tool name suggests -credential handling. +The platform has a much larger API surface than this server wraps. What is left +out is left out on purpose, and falls into three groups. + +### Anything whose response carries a credential + +Several subsystems return decrypted secrets as part of their ordinary +responses — connector configuration includes the credentials used to reach the +source system, adapter configuration includes the provider API key, notification +configuration includes the webhook's authorization token, and the key-management +endpoints return key material by design. -**Deletions.** Removing a workflow, deployment or Prompt Studio project destroys -work that no inverse call restores. Every write tool here is reversible by -construction — the opposite call puts things back. +**No tool wraps any of them.** An agent's context is not a safe place for a +credential: it is logged, it may be replayed to a model provider, and an agent +processing an untrusted document can be induced to repeat what it has seen. + +Concretely, this server has no tool for: + +- connector or adapter configuration (`listConnectors`, `getAdapter`, …) +- platform, deployment, or platform-API key management — including **creating + or rotating** a key, which returns the new secret in its response +- notification / webhook configuration +- the Postman collection export, which embeds a live API key + +Two mechanisms keep this true as tools are added: + +- `test_no_credential_leak.py` seeds an organization with recognizable fake + secrets, invokes **every** read tool in the registry, and fails if any of them + appears in the output. New tools are covered the moment they are registered. +- A test asserts that no tool name suggests credential handling, and none + suggests connector or adapter access. + +Where a tool must touch something adjacent to a credential, it names the fields +it returns rather than serializing a model. `getWorkflowEndpoints` is the case +to look at: a workflow endpoint points at a connector instance, so the tool +returns the *shape* of the connection — endpoint type, connection type, +connector name — and never the configuration. Free-form error text from failed +executions is additionally passed through `redact_secrets`, because a failing +connector reports the connection string it tried. + +### Destructive operations + +Deleting a workflow, deployment, Prompt Studio project, tag or connector +destroys work that no inverse call restores; so does removing organization +members. Every write tool here is reversible by construction — the opposite call +puts things back — and `required_method="DELETE"` exists so that if a +destructive tool is ever added it is `full_access`-only from the start. + +Also excluded: password resets, role assignment and revocation, and member +removal. These change who can access the organization, which is not a decision +to delegate to an agent. + +### Everything else, for now + +Not excluded on principle, simply not built: file upload/download, +tool-instance and endpoint *configuration* (as opposed to reading their shape), +Prompt Studio project/prompt authoring and import/export, and the connector and +adapter *test* endpoints (which make live outbound calls). These are candidates +for later, with the same rules applied. ## Why the URL matters @@ -232,6 +340,11 @@ through the model's `for_user(context.user)` manager so the platform's own visibility rules apply, and cap listings — `LIST_LIMIT` with a `truncated` note, never silent truncation. +**Never** build a response with `serializer.data`, `model_to_dict`, or a `**` +splat. Name every field. Several models decrypt credentials on attribute +access, so wholesale serialization is how a secret escapes — see the exclusion +list above. + For a write tool, also: - set `writes=True` and a `required_method` (`"POST"` for a mutation, `"DELETE"` @@ -245,6 +358,13 @@ For a write tool, also: - return `changed: False` rather than silently succeeding on a no-op, so an agent can tell "I did this" from "this was already so". +For a tool that costs money, also set `billable=True`. That flag is the only +thing wiring it into the spend guard, so forgetting it leaves the tool +unbudgeted — `test_spend_guard.py` names the tools that must carry it. + +If a tool returns free-form text that originates outside this app (an error +message, a log line), pass it through `redact_secrets` first. + --- ## Not implemented diff --git a/backend/mcp_server/platform_views.py b/backend/mcp_server/platform_views.py index 37444ae072..2ae7dd554b 100644 --- a/backend/mcp_server/platform_views.py +++ b/backend/mcp_server/platform_views.py @@ -29,6 +29,7 @@ from platform_api.models import ApiKeyPermission from rest_framework.request import Request +from mcp_server import spend_guard from mcp_server.context import PlatformMCPContext from mcp_server.registry import PLATFORM_TOOLS from mcp_server.transport import BaseMCPView @@ -39,8 +40,11 @@ class PlatformMCPServerView(BaseMCPView): """MCP JSON-RPC endpoint scoped to an organization. - Exposes read-only discovery tools over the organization's resources, - authenticated by a platform API key. + Exposes discovery, observability, state-change and billable tools over the + organization's resources, authenticated by a platform API key. Tool access + is gated twice: by the key's permission tier (``check_tool_allowed``) and, + for tools that cost money, by a per-organization budget + (``check_spend_allowed``). """ registry = PLATFORM_TOOLS @@ -142,3 +146,19 @@ def check_tool_allowed(self, tool: Any, context: PlatformMCPContext) -> str | No f"Tool '{tool.name}' requires a platform API key permitting " f"{tool.required_method}; this key is '{permission.value}'." ) + + def check_spend_allowed(self, tool: Any, context: PlatformMCPContext) -> str | None: + """Claim one unit of the organization's billable-call budget. + + Only billable tools consume budget; everything else is free. The claim + is made here rather than inside each handler so a new billable tool + cannot be added without being budgeted — the flag on the registry entry + is the only thing it has to get right. + """ + if not tool.billable: + return None + + state = spend_guard.consume(context.org_name) + if state.allowed: + return None + return state.message() diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index 3649831c9f..87ba37c2bb 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -33,6 +33,11 @@ class MCPTool: "POST" for mutations, "DELETE" for destructive operations (which only ``full_access`` permits). Unused by the deployment server, whose key has no tiers. + billable: True when invoking the tool costs money — LLM inference, + embedding, indexing, or a pipeline run. Budgeted per organization + by ``mcp_server.spend_guard``. Deliberately distinct from + ``writes``: a tool can mutate cheaply (pausing a schedule) or cost + money without changing configuration (running an extraction). """ name: str @@ -41,6 +46,7 @@ class MCPTool: handler: Callable[..., Any] writes: bool = False required_method: str = "GET" + billable: bool = False def to_mcp_schema(self) -> dict[str, Any]: """Serialize to the shape returned by `tools/list`.""" @@ -178,6 +184,19 @@ def build_platform_registry() -> MCPToolRegistry: destroys work that no inverse call restores. The write tools here are reversible by construction. """ + from mcp_server.tools.observability import ( + get_execution_detail, + get_execution_detail_schema, + get_usage_summary, + get_usage_summary_schema, + get_workflow_endpoints, + get_workflow_endpoints_schema, + list_executions, + list_executions_schema, + list_tags, + list_tool_instances, + list_tool_instances_schema, + ) from mcp_server.tools.platform import ( _ORG_WIDE_WARNING, execute_pipeline, @@ -194,6 +213,16 @@ def build_platform_registry() -> MCPToolRegistry: set_pipeline_active_schema, whoami, ) + from mcp_server.tools.prompt_studio import ( + bulk_fetch_response, + bulk_fetch_response_schema, + fetch_response, + fetch_response_schema, + index_document, + index_document_schema, + single_pass_extraction, + single_pass_extraction_schema, + ) registry = MCPToolRegistry() @@ -321,6 +350,162 @@ def build_platform_registry() -> MCPToolRegistry: handler=execute_pipeline, writes=True, required_method="POST", + billable=True, + ) + ) + + # ---- observability: what happened, and what it cost ---- + + registry.register( + MCPTool( + name="listExecutions", + description=( + "List recent workflow executions, newest first, with per-run " + "status and file counts.\n\n" + "This is the tool for 'did that run?' and 'why did it fail?'. " + "Optionally filter by workflow_id or status." + ), + input_schema=list_executions_schema(), + handler=list_executions, + ) + ) + registry.register( + MCPTool( + name="getExecutionDetail", + description=( + "Get per-file detail for one execution — which files " + "succeeded, which failed, and the error for each. Use after " + "listExecutions has identified a run worth investigating." + ), + input_schema=get_execution_detail_schema(), + handler=get_execution_detail, + ) + ) + registry.register( + MCPTool( + name="getUsageSummary", + description=( + "Aggregate token and cost usage recorded for this " + "organization. Historical accounting of what has already been " + "spent — not a budget or a limit. Takes no arguments." + ), + input_schema=get_usage_summary_schema(), + handler=get_usage_summary, + ) + ) + registry.register( + MCPTool( + name="getWorkflowEndpoints", + description=( + "Describe how a workflow is wired: its source and destination " + "endpoint types, and the name of the connector each uses.\n\n" + "Connector credentials and endpoint configuration are " + "deliberately not returned — only the shape of the connection." + ), + input_schema=get_workflow_endpoints_schema(), + handler=get_workflow_endpoints, + ) + ) + registry.register( + MCPTool( + name="listToolInstances", + description=( + "List the tool steps configured inside a workflow, in " + "execution order. Use this to understand what a workflow " + "actually does before running it." + ), + input_schema=list_tool_instances_schema(), + handler=list_tool_instances, + ) + ) + registry.register( + MCPTool( + name="listTags", + description=( + "List the organization's tags, which label workflow " + "executions for grouping and filtering. Takes no arguments." + ), + input_schema=no_args_schema(), + handler=list_tags, + ) + ) + + # ---- billable Prompt Studio operations ---- + # + # Each drives real LLM inference or embedding work, so each is billable and + # budgeted per organization by mcp_server.spend_guard. + + registry.register( + MCPTool( + name="indexDocument", + description=( + "Index a document in a Prompt Studio project so prompts can be " + "run against it.\n\n" + "**Costs money**: this embeds the document and writes to the " + "vector store. Index once, then run as many prompts as you " + "need against it — do not re-index between prompts.\n\n" + f"{_ORG_WIDE_WARNING}" + ), + input_schema=index_document_schema(), + handler=index_document, + writes=True, + required_method="POST", + billable=True, + ) + ) + registry.register( + MCPTool( + name="fetchResponse", + description=( + "Run one Prompt Studio prompt against an indexed document and " + "return the extracted response.\n\n" + "**Costs money**: this is a live LLM call. If you need several " + "prompts against the same document, use bulkFetchResponse " + "instead — it is cheaper and avoids an indexing race.\n\n" + f"{_ORG_WIDE_WARNING}" + ), + input_schema=fetch_response_schema(), + handler=fetch_response, + writes=True, + required_method="POST", + billable=True, + ) + ) + registry.register( + MCPTool( + name="bulkFetchResponse", + description=( + "Run several Prompt Studio prompts against one document in a " + "single pass.\n\n" + "**Costs money.** Prefer this over repeated fetchResponse " + "calls: it indexes once and dispatches one task, which is both " + "cheaper and avoids the 'document being indexed' race that " + "concurrent single calls provoke.\n\n" + f"{_ORG_WIDE_WARNING}" + ), + input_schema=bulk_fetch_response_schema(), + handler=bulk_fetch_response, + writes=True, + required_method="POST", + billable=True, + ) + ) + registry.register( + MCPTool( + name="singlePassExtraction", + description=( + "Run a Prompt Studio project's entire prompt set against a " + "document in one LLM pass.\n\n" + "**Costs money**, and is the most expensive tool here. Use it " + "when you want the project's full output; use fetchResponse " + "when you want one field.\n\n" + f"{_ORG_WIDE_WARNING}" + ), + input_schema=single_pass_extraction_schema(), + handler=single_pass_extraction, + writes=True, + required_method="POST", + billable=True, ) ) diff --git a/backend/mcp_server/spend_guard.py b/backend/mcp_server/spend_guard.py new file mode 100644 index 0000000000..a6c6322f6c --- /dev/null +++ b/backend/mcp_server/spend_guard.py @@ -0,0 +1,153 @@ +"""Per-organization budget for billable MCP tool calls. + +Some platform tools cost real money each time they run: they drive LLM +inference, embedding and vector-store writes. An agent looping over a list, or +retrying on a misread error, can spend a great deal without anyone watching. So +billable tools are budgeted per organization over a rolling window. + +**This is a call budget, not a spend budget.** It counts invocations, not +tokens or currency. Unstract's open-source backend records usage after the fact +(``usage_v2`` aggregates what was consumed) but has no pre-flight token +allowance to check against — subscription and quota enforcement live in the +enterprise overlay, which this app cannot depend on. A call counter is the +strongest guard implementable here, and it is a blunt one: it bounds *how many +times* an agent can trigger billable work, not how expensive each call is. + +Two properties are deliberate and easy to get backwards: + +* **Consumption happens on invocation and is never refunded**, including when + the tool then fails. A failed Prompt Studio call may already have spent + tokens upstream, and refunding on failure would let an agent burn unbounded + spend by triggering failures in a loop. This is the opposite of the + rate-limit slot handling in ``tools/execution.py``, where the slot models + concurrency rather than cost — see ``test_spend_guard`` for the pin. +* **Exhaustion is temporary, not a permission error.** The caller is told when + the window resets so it can retry later, rather than concluding the tool is + forbidden and giving up. +""" + +import logging +from dataclasses import dataclass + +from django.conf import settings +from django.core.cache import cache + +logger = logging.getLogger(__name__) + +# Key namespace. The window start is not encoded in the key: the TTL is what +# expires the counter, giving a fixed window that begins at the first billable +# call rather than at a wall-clock boundary. +_KEY_PREFIX = "mcp:billable" + + +def _key(org_id: str) -> str: + return f"{_KEY_PREFIX}:{org_id}" + + +@dataclass(frozen=True) +class BudgetState: + """The outcome of a budget check.""" + + allowed: bool + used: int + limit: int + window_seconds: int + retry_after_seconds: int | None = None + + def message(self) -> str: + """Explain exhaustion to the calling agent. + + Written so the agent understands this is a wait, not a refusal — and + knows roughly how long to wait. + """ + minutes = max(1, round((self.retry_after_seconds or self.window_seconds) / 60)) + return ( + f"This organization has used its budget for billable MCP " + f"operations ({self.used}/{self.limit} in the last " + f"{round(self.window_seconds / 60)} minutes). This is a temporary " + f"limit, not a permission problem — retry in about {minutes} " + f"minute(s), or ask an administrator to raise " + f"MCP_BILLABLE_CALL_LIMIT." + ) + + +def get_limit() -> int: + return int(getattr(settings, "MCP_BILLABLE_CALL_LIMIT", 50)) + + +def get_window_seconds() -> int: + return int(getattr(settings, "MCP_BILLABLE_WINDOW_SECONDS", 3600)) + + +def peek(org_id: str) -> BudgetState: + """Report budget state without consuming any of it. + + Used by ``whoami`` so an agent can pace itself instead of discovering the + limit by hitting it. + """ + limit = get_limit() + window = get_window_seconds() + try: + used = cache.get(_key(org_id)) or 0 + except Exception as error: + # Reporting the budget must never be the thing that breaks whoami — + # `consume` fails open for the same reason, and a read is even less + # load-bearing than a claim. + logger.warning(f"MCP spend guard unreadable for org '{org_id}': {error}") + used = 0 + return BudgetState( + allowed=used < limit, used=used, limit=limit, window_seconds=window + ) + + +def consume(org_id: str) -> BudgetState: + """Claim one billable call for this organization. + + Returns a state with ``allowed=False`` when the budget is already spent, in + which case nothing is consumed and the caller must not run the tool. + + Fails **open** if the cache is unavailable: this guard exists to stop an + agent looping, not to be a hard licence check, and taking the whole MCP + surface offline because Redis blipped would be a worse failure than briefly + unbudgeted calls. The event is logged so it is visible. + """ + limit = get_limit() + window = get_window_seconds() + key = _key(org_id) + + try: + # `add` only sets when absent, so it both initialises the counter and + # starts the window's TTL without clobbering an in-flight window. + cache.add(key, 0, timeout=window) + used = cache.incr(key) + except ValueError: + # The key expired between `add` and `incr`. Re-seed and count this call + # as the first of a fresh window. + cache.set(key, 1, timeout=window) + used = 1 + except Exception as error: + logger.error( + f"MCP spend guard unavailable for org '{org_id}', allowing call: {error}" + ) + return BudgetState(allowed=True, used=0, limit=limit, window_seconds=window) + + if used > limit: + # Over budget. The counter is deliberately left incremented: decrementing + # here would let a caller probe the limit repeatedly at no cost. + ttl = None + try: + ttl = cache.ttl(key) + except Exception: + pass + logger.warning( + f"MCP billable budget exhausted for org '{org_id}': {used}/{limit}" + ) + return BudgetState( + allowed=False, + used=used, + limit=limit, + window_seconds=window, + retry_after_seconds=ttl, + ) + + return BudgetState(allowed=True, used=used, limit=limit, window_seconds=window) diff --git a/backend/mcp_server/tests/test_no_credential_leak.py b/backend/mcp_server/tests/test_no_credential_leak.py new file mode 100644 index 0000000000..0ca3b96e71 --- /dev/null +++ b/backend/mcp_server/tests/test_no_credential_leak.py @@ -0,0 +1,241 @@ +"""Registry-wide invariant: no tool output may contain a credential. + +This is the test that keeps the README's exclusion list true as tools get +added. Several of this project's serializers return decrypted secrets in their +normal responses — ``ConnectorInstance.connector_metadata`` carries database +passwords and object-store keys, ``AdapterInstance.metadata`` carries provider +API keys — so a tool that reaches a model through ``serializer.data``, +``model_to_dict`` or a ``**`` splat would pipe those into an LLM's context +without anyone noticing at review time. + +Rather than trusting each tool to be written carefully, this seeds an +organization with recognizable fake secrets, invokes **every** read tool in the +platform registry, and asserts none of those strings comes back. A new tool is +covered the moment it is registered. + +Needs a live DB (integration tier). +""" + +from __future__ import annotations + +import json +from unittest.mock import Mock, patch + +import pytest +from account_v2.models import Organization, User +from adapter_processor_v2.models import AdapterInstance +from api_v2.models import APIDeployment +from connector_v2.models import ConnectorInstance +from django.test import TestCase, override_settings +from platform_api.models import PlatformApiKey +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from tenant_account_v2.models import OrganizationMember +from utils.user_context import UserContext +from workflow_manager.endpoint_v2.models import WorkflowEndpoint +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow + +from mcp_server.context import PlatformMCPContext +from mcp_server.registry import PLATFORM_TOOLS + +ORG_ID = "org-leak" + +# `whoami` reports the spend budget, which reads the cache. Use local memory so +# this test exercises tool output rather than Redis availability. +LOCMEM = { + "default": { + "BACKEND": "django.core.cache.backends.locmem.LocMemCache", + "LOCATION": "mcp-leak-tests", + } +} + +# Distinctive sentinels: if any of these appears in tool output, a secret +# escaped. Chosen not to collide with anything a tool would legitimately emit. +DB_PASSWORD = "LEAKCANARY-db-password-9f3a" +S3_SECRET = "LEAKCANARY-s3-secret-key-7c1b" +LLM_API_KEY = "LEAKCANARY-openai-api-key-4e8d" + +# Only secrets are canaries. A platform key's *name* is deliberately surfaced +# by whoami — it is a label the operator chose, not credential material, and +# treating it as one would make this test object to correct behaviour. +CANARIES = (DB_PASSWORD, S3_SECRET, LLM_API_KEY) + + +@override_settings(CACHES=LOCMEM) +class NoCredentialLeakTest(TestCase): + def setUp(self) -> None: + self.org = Organization.objects.create( + name=ORG_ID, display_name="Leak Org", organization_id=ORG_ID + ) + UserContext.set_organization_identifier(ORG_ID) + self.user = User.objects.create( + username="svc-leak", + email="svc-leak@platform.internal", + user_id="uid-leak", + is_service_account=True, + ) + OrganizationMember.objects.create( + user=self.user, organization=self.org, role="user" + ) + self.key = PlatformApiKey.objects.create( + name="leak-test-key", + description="d", + organization=self.org, + api_user=self.user, + permission="read_write", + ) + + # A connector holding credentials, wired into a workflow endpoint — + # the transitive path that makes getWorkflowEndpoints risky. + self.connector = ConnectorInstance.objects.create( + connector_name="Leak Postgres", + connector_id="postgresql|4d3a1b0e-1f2c-4a3d-9e8f-0a1b2c3d4e5f", + connector_metadata={ + "user": "admin", + "password": DB_PASSWORD, + "secret_access_key": S3_SECRET, + }, + ) + # `adapter_metadata` is the writable field; `metadata` is the property + # that decrypts it on read — which is exactly why a tool must never + # serialize this model wholesale. + self.adapter = AdapterInstance.objects.create( + adapter_name="Leak LLM", + adapter_id="openai|1a2b3c4d-5e6f-7a8b-9c0d-1e2f3a4b5c6d", + adapter_type="LLM", + adapter_metadata={"api_key": LLM_API_KEY}, + ) + + self.workflow = Workflow.objects.create( + workflow_name="leak-wf", description="d", is_active=True + ) + WorkflowEndpoint.objects.create( + workflow=self.workflow, + endpoint_type=WorkflowEndpoint.EndpointType.SOURCE, + connection_type=WorkflowEndpoint.ConnectionType.DATABASE, + connector_instance=self.connector, + configuration={"table": "docs", "password_hint": DB_PASSWORD}, + ) + APIDeployment.objects.create( + api_name="leak-api", display_name="Leak API", workflow=self.workflow + ) + CustomTool.objects.create( + tool_name="Leak Prompts", description="d", author="acme" + ) + # WorkflowExecution.save() mirrors into a Redis-backed execution cache + # via its own client, which override_settings(CACHES=...) does not + # cover. That mirroring is irrelevant here — the tool reads the DB — + # so it is stubbed to keep this test off live infra. + with patch( + "workflow_manager.workflow_v2.models.execution." + "WorkflowExecution._handle_execution_cache" + ): + self.execution = WorkflowExecution.objects.create( + workflow_id=self.workflow.id, + status="ERROR", + # Error text is echoed back to the agent, so seed a canary + # here too: a stack trace is where a connection string leaks. + error_message=f"connect failed: password={DB_PASSWORD}", + ) + + self.context = PlatformMCPContext( + user=self.user, + platform_key=self.key, + org_name=ORG_ID, + request=Mock(data={}), + ) + + def _read_tools(self): + """Every non-billable, non-writing tool — the ones safe to just call.""" + for name in PLATFORM_TOOLS.names(): + tool = PLATFORM_TOOLS.get(name) + if tool.writes or tool.billable: + continue + yield name, tool + + @pytest.mark.critical_path("mcp-platform-auth") + def test_no_read_tool_returns_a_credential(self) -> None: + """Invoke every read tool and scan its output for the canaries. + + Tools needing an argument get the seeded ids; anything that still + cannot be called is reported rather than silently skipped, so this + cannot quietly stop covering the registry. + """ + arguments = { + "getWorkflowEndpoints": {"workflow_id": str(self.workflow.id)}, + "listToolInstances": {"workflow_id": str(self.workflow.id)}, + "listExecutions": {}, + "getUsageSummary": {}, + "getExecutionDetail": {"execution_id": str(self.execution.id)}, + } + + checked = [] + for name, tool in self._read_tools(): + kwargs = arguments.get(name, {}) + required = tool.input_schema.get("required", []) + if any(field not in kwargs for field in required): + # A read tool needing arguments this test does not know how to + # supply. Fail loudly: an uncovered tool is the whole risk. + raise AssertionError( + f"Read tool '{name}' requires {required} but this test has " + f"no fixture for it — add one so it stays covered." + ) + + result = tool.handler(self.context, **kwargs) + blob = json.dumps(result, default=str) + checked.append(name) + + for canary in CANARIES: + assert canary not in blob, ( + f"Tool '{name}' leaked a credential ({canary}). Build its " + f"response from named fields — never serializer.data, " + f"model_to_dict, or a ** splat." + ) + + # Guard the guard: if the filter above ever matches nothing, this test + # would pass vacuously. + assert len(checked) >= 5, f"Expected to check several tools, got {checked}" + + @pytest.mark.critical_path("mcp-platform-auth") + def test_workflow_endpoints_omits_connector_configuration(self) -> None: + """The specific transitive path worth naming. + + A workflow endpoint points at a connector instance whose metadata is + decrypted on access. Returning the endpoint's own `configuration`, or + following the relation into the connector, would leak it. + """ + from mcp_server.tools.observability import get_workflow_endpoints + + result = get_workflow_endpoints(self.context, workflow_id=str(self.workflow.id)) + blob = json.dumps(result, default=str) + + assert result["endpoints"], "fixture should produce one endpoint" + endpoint = result["endpoints"][0] + # The shape is exposed... + assert endpoint["connection_type"] == "DATABASE" + assert endpoint["connector_name"] == "Leak Postgres" + assert endpoint["is_configured"] is True + # ...but never the contents. + assert "configuration" not in endpoint + assert "connector_metadata" not in blob + assert DB_PASSWORD not in blob + + @pytest.mark.critical_path("mcp-platform-auth") + def test_no_tool_wraps_connectors_or_adapters(self) -> None: + """These subsystems are excluded wholesale, not filtered. + + Their serializers return decrypted credentials by design, so the safe + boundary is to expose no tool for them at all rather than to trust a + field-level filter to stay correct. + """ + forbidden = ("connector", "adapter") + offenders = [ + name + for name in PLATFORM_TOOLS.names() + if any(word in name.lower() for word in forbidden) + ] + + assert offenders == [], ( + f"Connector and adapter responses carry decrypted credentials; " + f"no MCP tool should wrap them: {offenders}" + ) diff --git a/backend/mcp_server/tests/test_platform_auth.py b/backend/mcp_server/tests/test_platform_auth.py index bae3bcaa9f..bfc036e4e3 100644 --- a/backend/mcp_server/tests/test_platform_auth.py +++ b/backend/mcp_server/tests/test_platform_auth.py @@ -126,18 +126,23 @@ def test_valid_key_reaches_the_tool_listing(self) -> None: response = self._post(f"Bearer {self.key.key}") assert response.status_code == 200, response.content - tools = json.loads(response.content)["result"]["tools"] - assert [t["name"] for t in tools] == [ - "readMeFirst", + tools = [t["name"] for t in json.loads(response.content)["result"]["tools"]] + + # readMeFirst must lead: agents weight earlier tools more heavily, and + # it is what explains the budget and the org-wide reach of the rest. + assert tools[0] == "readMeFirst" + # A representative tool from each group, rather than the full ordered + # list — pinning all 19 makes this fail on every addition without + # telling anyone anything useful. + assert { "whoami", "listApiDeployments", - "listWorkflows", - "listPromptStudioProjects", - "listPipelines", + "listExecutions", + "getUsageSummary", "setApiDeploymentActive", - "setPipelineActive", "executePipeline", - ] + "bulkFetchResponse", + } <= set(tools) @pytest.mark.critical_path("mcp-platform-auth") def test_the_endpoint_is_not_whitelisted(self) -> None: diff --git a/backend/mcp_server/tests/test_redaction.py b/backend/mcp_server/tests/test_redaction.py new file mode 100644 index 0000000000..bac618b2c1 --- /dev/null +++ b/backend/mcp_server/tests/test_redaction.py @@ -0,0 +1,81 @@ +"""Redaction of credential-shaped text in execution errors. + +Structured fields are safe because they are named explicitly. Free-form error +text is not: a connector that fails to connect reports the connection string it +tried, and a provider client echoes the key it authenticated with. That text +goes straight to the agent, so it is scrubbed first. + +This was found by ``test_no_credential_leak`` rather than by review — the +canary seeded in an execution's ``error_message`` came back through +``listExecutions``. +""" + +from __future__ import annotations + +from django.test import SimpleTestCase + +from mcp_server.tools.observability import redact_secrets + + +class RedactSecretsTest(SimpleTestCase): + def test_key_value_secrets_are_masked(self) -> None: + cases = [ + "connect failed: password=hunter2", + "connect failed: PASSWORD: hunter2", + "auth error api_key=sk-abc123", + "auth error api-key = sk-abc123", + "bad token=eyJhbGciOi", + "secret_access_key=AKIAIOSFODNN7EXAMPLE", + ] + for text in cases: + with self.subTest(text): + out = redact_secrets(text) + assert "[REDACTED]" in out + for leaked in ("hunter2", "sk-abc123", "eyJhbGciOi", "AKIA"): + assert leaked not in out + + def test_connection_string_password_is_masked(self) -> None: + out = redact_secrets( + "could not connect to postgresql://admin:s3cr3tpw@db.internal:5432/x" + ) + + assert "s3cr3tpw" not in out + assert "[REDACTED]" in out + # The rest of the URL survives, or the message stops being diagnosable. + assert "db.internal:5432" in out + assert "admin" in out + + def test_bearer_token_is_masked(self) -> None: + out = redact_secrets("401 from provider, sent Authorization: Bearer abcdef123456") + + assert "abcdef123456" not in out + assert "[REDACTED]" in out + + def test_surrounding_message_is_preserved(self) -> None: + """Redaction that eats the whole message would trade one problem for + another — the agent still has to be able to diagnose the failure. + """ + out = redact_secrets( + "Failed after 3 attempts on table `invoices`: password=hunter2 (code 28P01)" + ) + + assert "Failed after 3 attempts" in out + assert "invoices" in out + assert "28P01" in out + assert "hunter2" not in out + + def test_ordinary_error_text_is_untouched(self) -> None: + text = "File not found: invoice-2024-03.pdf (no such object in bucket)" + + assert redact_secrets(text) == text + + def test_empty_and_none_are_safe(self) -> None: + assert redact_secrets(None) is None + assert redact_secrets("") == "" + + def test_multiple_secrets_in_one_message_are_all_masked(self) -> None: + out = redact_secrets("password=one and api_key=two and token=three") + + for leaked in ("one", "two", "three"): + assert f"={leaked}" not in out + assert out.count("[REDACTED]") == 3 diff --git a/backend/mcp_server/tests/test_spend_guard.py b/backend/mcp_server/tests/test_spend_guard.py new file mode 100644 index 0000000000..3e1172803a --- /dev/null +++ b/backend/mcp_server/tests/test_spend_guard.py @@ -0,0 +1,275 @@ +"""The billable-call budget. + +The guard is what stands between an agent in a retry loop and an unbounded LLM +bill, so its edges are pinned here — particularly the one that contradicts the +neighbouring rate-limiter pattern. +""" + +from __future__ import annotations + +import json +from unittest.mock import Mock, patch + +from django.core.cache import cache +from django.test import SimpleTestCase, override_settings + +from mcp_server import spend_guard +from mcp_server.context import PlatformMCPContext +from mcp_server.platform_views import PlatformMCPServerView +from mcp_server.registry import PLATFORM_TOOLS, MCPTool + +ORG = "org-budget" + + +def a_tool(billable: bool) -> MCPTool: + return MCPTool( + name="costly" if billable else "cheap", + description="d", + input_schema={"type": "object", "properties": {}}, + handler=lambda ctx: None, + billable=billable, + required_method="POST" if billable else "GET", + ) + + +def context() -> PlatformMCPContext: + return PlatformMCPContext( + user=Mock(is_service_account=True), + platform_key=Mock(permission="read_write"), + org_name=ORG, + ) + + +# The guard's behaviour is what is under test, not Redis's. A local-memory +# cache exercises the same django.core.cache API deterministically and without +# requiring live infra, keeping these in the unit tier. +LOCMEM = { + "default": { + "BACKEND": "django.core.cache.backends.locmem.LocMemCache", + "LOCATION": "mcp-spend-guard-tests", + } +} + + +@override_settings( + MCP_BILLABLE_CALL_LIMIT=3, MCP_BILLABLE_WINDOW_SECONDS=60, CACHES=LOCMEM +) +class SpendGuardTest(SimpleTestCase): + def setUp(self) -> None: + cache.delete(f"mcp:billable:{ORG}") + + def tearDown(self) -> None: + cache.delete(f"mcp:billable:{ORG}") + + def test_budget_allows_up_to_the_limit_then_refuses(self) -> None: + for expected_used in (1, 2, 3): + state = spend_guard.consume(ORG) + assert state.allowed is True + assert state.used == expected_used + + exhausted = spend_guard.consume(ORG) + assert exhausted.allowed is False + + def test_exhaustion_message_reads_as_temporary_not_forbidden(self) -> None: + """An agent that reads this as a permission error stops retrying; one + that reads it as a wait comes back later. The wording is the only thing + that distinguishes them. + """ + for _ in range(4): + state = spend_guard.consume(ORG) + + message = state.message() + assert "temporary" in message.lower() + assert "retry" in message.lower() + assert "3" in message # the limit is named + + def test_peek_does_not_consume(self) -> None: + """whoami calls peek; surfacing the budget must not spend it.""" + spend_guard.consume(ORG) + + before = spend_guard.peek(ORG) + after = spend_guard.peek(ORG) + + assert before.used == 1 + assert after.used == 1 + + def test_budget_is_per_organization(self) -> None: + for _ in range(4): + spend_guard.consume(ORG) + + other = spend_guard.consume("org-other-budget") + try: + assert other.allowed is True, "one org's spend must not block another's" + finally: + cache.delete("mcp:billable:org-other-budget") + + def test_cache_failure_fails_open(self) -> None: + """This guard bounds runaway loops; it is not a licence check. Taking + the whole MCP surface offline because Redis blipped would be the worse + failure, so it allows and logs. + """ + with patch( + "mcp_server.spend_guard.cache.add", side_effect=RuntimeError("redis down") + ): + state = spend_guard.consume(ORG) + + assert state.allowed is True + + +@override_settings( + MCP_BILLABLE_CALL_LIMIT=2, MCP_BILLABLE_WINDOW_SECONDS=60, CACHES=LOCMEM +) +class SpendGuardEnforcementTest(SimpleTestCase): + def setUp(self) -> None: + cache.delete(f"mcp:billable:{ORG}") + self.view = PlatformMCPServerView() + + def tearDown(self) -> None: + cache.delete(f"mcp:billable:{ORG}") + + def test_only_billable_tools_consume_budget(self) -> None: + for _ in range(5): + assert self.view.check_spend_allowed(a_tool(billable=False), context()) is None + + assert spend_guard.peek(ORG).used == 0 + + def test_billable_tool_is_refused_once_the_budget_is_spent(self) -> None: + assert self.view.check_spend_allowed(a_tool(billable=True), context()) is None + assert self.view.check_spend_allowed(a_tool(billable=True), context()) is None + + refusal = self.view.check_spend_allowed(a_tool(billable=True), context()) + + assert refusal is not None + assert "budget" in refusal.lower() + + def test_budget_is_not_refunded_when_a_tool_then_fails(self) -> None: + """Deliberately the opposite of the rate-limit slot in tools/execution. + + That slot models concurrency, so releasing it on failure is right. This + counter models money already spent: a Prompt Studio call that fails + partway may have burned tokens upstream, and refunding would let an + agent spend without limit by failing in a loop. + """ + self.view.check_spend_allowed(a_tool(billable=True), context()) + used_before = spend_guard.peek(ORG).used + + # Simulate the handler blowing up after the budget was claimed. + try: + raise RuntimeError("tool exploded") + except RuntimeError: + pass + + assert spend_guard.peek(ORG).used == used_before, ( + "budget must not be refunded on failure" + ) + + def test_probing_after_exhaustion_does_not_reset_the_counter(self) -> None: + for _ in range(2): + self.view.check_spend_allowed(a_tool(billable=True), context()) + + for _ in range(3): + refusal = self.view.check_spend_allowed(a_tool(billable=True), context()) + assert refusal is not None + + +@override_settings( + MCP_BILLABLE_CALL_LIMIT=1, MCP_BILLABLE_WINDOW_SECONDS=60, CACHES=LOCMEM +) +class SpendGuardDispatchTest(SimpleTestCase): + """The budget as an MCP client actually experiences it. + + The enforcement tests above call the hook directly; this drives the real + JSON-RPC dispatch, which is where the distinction between "refused + permanently" and "retry later" is actually made. + """ + + def setUp(self) -> None: + cache.delete(f"mcp:billable:{ORG}") + self.view = PlatformMCPServerView() + self.context = context() + + def tearDown(self) -> None: + cache.delete(f"mcp:billable:{ORG}") + + def _call_billable(self): + from dataclasses import replace + + tool = replace( + PLATFORM_TOOLS.get("executePipeline"), + handler=lambda ctx, **kw: {"ran": True}, + ) + with patch.object(PLATFORM_TOOLS, "get", return_value=tool): + response = self.view._call_tool( + request_id=1, + params={"name": "executePipeline", "arguments": {}}, + context=self.context, + ) + return json.loads(response.content) + + def test_first_call_runs_and_second_is_refused_retryably(self) -> None: + first = self._call_billable() + assert first["result"]["isError"] is False + + second = self._call_billable() + + # Crucially a *result*, not a JSON-RPC error: clients treat protocol + # errors as unrecoverable, and this condition clears on its own. + assert "error" not in second + assert second["result"]["isError"] is True + text = second["result"]["content"][0]["text"] + assert "temporary" in text.lower() + assert "retry" in text.lower() + + def test_non_billable_tools_still_work_once_the_budget_is_gone(self) -> None: + """Spending the budget must not take the whole server down with it — + an agent should still be able to look around and report what happened. + """ + self._call_billable() + self._call_billable() # exhausts + + response = self.view._call_tool( + request_id=2, + params={"name": "whoami", "arguments": {}}, + context=self.context, + ) + body = json.loads(response.content) + + assert body["result"]["isError"] is False + + +class BillableRegistryInvariantTest(SimpleTestCase): + def test_every_costly_platform_tool_is_marked_billable(self) -> None: + """The flag is the only thing wiring a tool into the budget, so a + costly tool that forgets it is unguarded. Named explicitly rather than + inferred, so adding one to this list is a deliberate act. + """ + must_be_billable = { + "executePipeline", + "indexDocument", + "fetchResponse", + "bulkFetchResponse", + "singlePassExtraction", + } + + unguarded = [ + name + for name in must_be_billable + if PLATFORM_TOOLS.get(name) is None or not PLATFORM_TOOLS.get(name).billable + ] + + assert unguarded == [], ( + f"These tools cost money but are not budgeted: {unguarded}" + ) + + def test_billable_tools_are_also_write_gated(self) -> None: + """Spending money is a write in every sense that matters, so a billable + tool must never be reachable by a key that cannot write. + """ + wrong = [ + name + for name in PLATFORM_TOOLS.names() + if PLATFORM_TOOLS.get(name).billable + and PLATFORM_TOOLS.get(name).required_method == "GET" + ] + + assert wrong == [], f"Billable tools must not be GET-tier: {wrong}" diff --git a/backend/mcp_server/tools/observability.py b/backend/mcp_server/tools/observability.py new file mode 100644 index 0000000000..c7a162a95f --- /dev/null +++ b/backend/mcp_server/tools/observability.py @@ -0,0 +1,400 @@ +"""Read-only tools for execution history, usage and workflow structure. + +Every response here is built **field by field from named model attributes**. +That is a security boundary, not a style preference: several of this project's +serializers return decrypted credentials in their normal output — connector +metadata carries database passwords and object-store keys, adapter metadata +carries provider API keys — and a workflow endpoint references a connector +instance. Passing a model through ``serializer.data``, ``model_to_dict`` or a +``**`` splat would pipe those secrets into an LLM's context and defeat the +whole exclusion list in this app's README. + +So: name every field you return. ``test_no_credentials_in_tool_output`` +enforces this across the whole registry. +""" + +import logging +import re +from typing import Any + +from usage_v2.models import Usage +from workflow_manager.endpoint_v2.models import WorkflowEndpoint +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow + +from mcp_server.context import PlatformMCPContext +from mcp_server.exceptions import MCPToolError + +logger = logging.getLogger(__name__) + +# Executions are the highest-cardinality thing here, and an agent asking "what +# happened" wants the recent tail rather than the archive. +EXECUTION_LIMIT = 25 +LOG_LIMIT = 100 + + +# Error text from a failed execution is one of the few places a secret can +# reach an agent without any tool asking for it: a connector that fails to +# connect may report the connection string it tried, and a provider client may +# echo the key it authenticated with. These patterns catch the common shapes. +_REDACTED = "[REDACTED]" + +# Each entry is (pattern, replacement). The replacement keeps the surrounding +# context — the key name, the URL prefix — so the message stays diagnosable +# while the value itself is gone. +_SECRET_PATTERNS = ( + # Bearer tokens first: "Authorization: Bearer " would otherwise be + # consumed by the key=value rule below, which would mask the word "Bearer" + # and leave the token itself in place. + ( + re.compile(r"(?i)\b(bearer\s+)([A-Za-z0-9._\-]{8,})"), + rf"\1{_REDACTED}", + ), + # key=value / key: value, where the key looks credential-ish. Stops at + # whitespace, quote or delimiter so the rest of the message survives. + ( + re.compile( + r"(?i)\b(password|passwd|pwd|secret|secret_access_key|access_key|" + r"api[_-]?key|token|authorization|credential)" + r"(\s*[=:]\s*)" + r"([^\s,;'\"&)}\]]+)" + ), + rf"\1\2{_REDACTED}", + ), + # Credentials embedded in a URL: scheme://user:secret@host + ( + re.compile(r"(?i)\b([a-z][a-z0-9+.-]*://[^\s:/@]+:)([^\s@]+)(@)"), + rf"\1{_REDACTED}\3", + ), +) + + +def redact_secrets(text: str | None) -> str | None: + """Mask credential-shaped substrings in free-form text. + + Applied to every error message these tools return. Execution errors are one + of the few places a secret reaches an agent without any tool asking for it: + a connector that fails to connect may report the connection string it + tried, and a provider client may echo the key it authenticated with. + + This is a safety net, not a guarantee — it cannot recognise a secret that + looks like ordinary prose — so it complements the rule that structured + fields are named explicitly rather than replacing it. + """ + if not text: + return text + + redacted = text + for pattern, replacement in _SECRET_PATTERNS: + redacted = pattern.sub(replacement, redacted) + return redacted + + +def _org_workflow_ids(context: PlatformMCPContext) -> list[Any]: + """Workflow ids the caller may see, for scoping execution queries. + + WorkflowExecution has no ``for_user`` manager, so visibility is derived + from the workflows the caller can reach rather than assumed. + """ + return list(Workflow.objects.for_user(context.user).values_list("id", flat=True)) + + +def list_executions_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "workflow_id": { + "type": "string", + "description": ( + "Optional UUID to show executions of one workflow only. " + "Omit for the organization's most recent executions." + ), + }, + "status": { + "type": "string", + "description": ( + "Optional status filter, e.g. COMPLETED, ERROR, EXECUTING, PENDING." + ), + }, + }, + "required": [], + } + + +def list_executions( + context: PlatformMCPContext, + workflow_id: str | None = None, + status: str | None = None, +) -> dict[str, Any]: + """List recent workflow executions, newest first. + + This is the tool for "why did that fail" and "did it run" — the questions + an agent actually has after triggering work. + """ + visible = _org_workflow_ids(context) + queryset = WorkflowExecution.objects.filter(workflow_id__in=visible) + + if workflow_id: + if not any(str(wid) == str(workflow_id) for wid in visible): + raise MCPToolError( + f"No workflow with id '{workflow_id}' in organization " + f"'{context.org_name}'. Call listWorkflows for valid ids." + ) + queryset = queryset.filter(workflow_id=workflow_id) + if status: + queryset = queryset.filter(status=status.upper()) + + queryset = queryset.order_by("-created_at") + total = queryset.count() + rows = queryset[:EXECUTION_LIMIT] + + return { + "count": total, + "showing": min(total, EXECUTION_LIMIT), + "executions": [ + { + "id": str(row.id), + "workflow_id": str(row.workflow_id), + "pipeline_id": str(row.pipeline_id) if row.pipeline_id else None, + "status": row.status, + "execution_mode": row.execution_mode, + "total_files": row.total_files, + "successful_files": row.successful_files, + "failed_files": row.failed_files, + # Truncated: an error message can be a full stack trace, and a + # page of them would crowd out everything else. + "error_message": redact_secrets((row.error_message or "")[:500]) or None, + "attempts": row.attempts, + "execution_time_seconds": row.execution_time, + "created_at": row.created_at, + } + for row in rows + ], + **( + { + "truncated": True, + "note": ( + f"Showing the {EXECUTION_LIMIT} most recent of {total}. " + "Filter by workflow_id or status to narrow." + ), + } + if total > EXECUTION_LIMIT + else {} + ), + } + + +def get_execution_detail_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "execution_id": { + "type": "string", + "description": "UUID of the execution, from listExecutions.", + }, + }, + "required": ["execution_id"], + } + + +def get_execution_detail( + context: PlatformMCPContext, execution_id: str +) -> dict[str, Any]: + """Per-file detail for one execution — which files failed, and why.""" + visible = _org_workflow_ids(context) + execution = WorkflowExecution.objects.filter( + id=execution_id, workflow_id__in=visible + ).first() + if execution is None: + raise MCPToolError( + f"No execution with id '{execution_id}' in organization " + f"'{context.org_name}'. Call listExecutions for valid ids." + ) + + files = [] + try: + for fe in execution.file_executions.all()[:LOG_LIMIT]: + files.append( + { + "id": str(fe.id), + "file_name": fe.file_name, + "status": fe.status, + "error": redact_secrets((fe.execution_error or "")[:500]) or None, + } + ) + except Exception as error: # pragma: no cover - defensive + logger.warning(f"Could not load file executions for '{execution_id}': {error}") + + return { + "id": str(execution.id), + "workflow_id": str(execution.workflow_id), + "status": execution.status, + "total_files": execution.total_files, + "successful_files": execution.successful_files, + "failed_files": execution.failed_files, + "error_message": redact_secrets((execution.error_message or "")[:1000]) or None, + "execution_time_seconds": execution.execution_time, + "created_at": execution.created_at, + "files": files, + } + + +def get_usage_summary_schema() -> dict[str, Any]: + return {"type": "object", "properties": {}, "required": []} + + +def get_usage_summary(context: PlatformMCPContext) -> dict[str, Any]: + """Aggregate token and cost usage for the organization. + + Aggregated in the database rather than by pulling rows: a busy + organization has far more usage records than belong in an agent's context, + and the totals are what the question is actually about. + """ + from django.db.models import Count, Sum + + totals = Usage.objects.aggregate( + records=Count("id"), + prompt_tokens=Sum("prompt_tokens"), + completion_tokens=Sum("completion_tokens"), + total_tokens=Sum("total_tokens"), + embedding_tokens=Sum("embedding_tokens"), + cost_in_dollars=Sum("cost_in_dollars"), + ) + + return { + "organization": context.org_name, + "records": totals["records"] or 0, + "prompt_tokens": totals["prompt_tokens"] or 0, + "completion_tokens": totals["completion_tokens"] or 0, + "total_tokens": totals["total_tokens"] or 0, + "embedding_tokens": totals["embedding_tokens"] or 0, + "cost_in_dollars": float(totals["cost_in_dollars"] or 0), + "note": ( + "Totals across all recorded usage for this organization. This is " + "historical accounting, not a spending limit." + ), + } + + +def list_tool_instances_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "workflow_id": { + "type": "string", + "description": "UUID of the workflow, from listWorkflows.", + }, + }, + "required": ["workflow_id"], + } + + +def list_tool_instances(context: PlatformMCPContext, workflow_id: str) -> dict[str, Any]: + """List the tool steps inside a workflow, in execution order. + + Tool settings are omitted deliberately: they are free-form JSON that can + carry adapter ids and, in some tools, credential-shaped values. The step + order and identity answer "what does this workflow do" without that risk. + """ + from tool_instance_v2.models import ToolInstance + + workflow = Workflow.objects.for_user(context.user).filter(id=workflow_id).first() + if workflow is None: + raise MCPToolError( + f"No workflow with id '{workflow_id}' in organization " + f"'{context.org_name}'. Call listWorkflows for valid ids." + ) + + rows = ToolInstance.objects.filter(workflow=workflow).order_by("step") + + return { + "workflow_id": str(workflow.id), + "workflow_name": workflow.workflow_name, + "tool_instances": [ + { + "id": str(row.id), + "tool_id": row.tool_id, + "step": row.step, + } + for row in rows + ], + "note": ( + "Tool settings are not exposed — they can reference adapters and " + "credential-shaped values. Use the Unstract UI to inspect them." + ), + } + + +def list_tags(context: PlatformMCPContext) -> dict[str, Any]: + """List the organization's execution tags.""" + from tags.models import Tag + + queryset = Tag.objects.all().order_by("name") + total = queryset.count() + rows = queryset[:100] + + return { + "count": total, + "tags": [ + {"id": str(row.id), "name": row.name, "description": row.description or None} + for row in rows + ], + } + + +def get_workflow_endpoints_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "workflow_id": { + "type": "string", + "description": "UUID of the workflow, from listWorkflows.", + }, + }, + "required": ["workflow_id"], + } + + +def get_workflow_endpoints( + context: PlatformMCPContext, workflow_id: str +) -> dict[str, Any]: + """Describe a workflow's source and destination wiring. + + Returns the *shape* of the connection — endpoint type, connection type, and + the connector's name — and deliberately never its configuration. A + connector instance holds decrypted credentials in ``connector_metadata``, + so only named, non-secret attributes are read here. Do not extend this to + return ``configuration`` or ``connector_metadata``. + """ + workflow = Workflow.objects.for_user(context.user).filter(id=workflow_id).first() + if workflow is None: + raise MCPToolError( + f"No workflow with id '{workflow_id}' in organization " + f"'{context.org_name}'. Call listWorkflows for valid ids." + ) + + endpoints = [] + for endpoint in WorkflowEndpoint.objects.filter(workflow=workflow): + connector = getattr(endpoint, "connector_instance", None) + endpoints.append( + { + "id": str(endpoint.id), + "endpoint_type": endpoint.endpoint_type, + "connection_type": endpoint.connection_type, + # Name and id only. The connector's metadata is credential + # material and must not appear here. + "connector_name": getattr(connector, "connector_name", None), + "connector_id": str(connector.id) if connector else None, + "is_configured": bool(endpoint.configuration), + } + ) + + return { + "workflow_id": str(workflow.id), + "workflow_name": workflow.workflow_name, + "endpoints": endpoints, + "note": ( + "Connector credentials and endpoint configuration are deliberately " + "not exposed. Use the Unstract UI to inspect or change them." + ), + } diff --git a/backend/mcp_server/tools/platform.py b/backend/mcp_server/tools/platform.py index c1b82f13af..87edbaf73e 100644 --- a/backend/mcp_server/tools/platform.py +++ b/backend/mcp_server/tools/platform.py @@ -24,6 +24,7 @@ from prompt_studio.prompt_studio_core_v2.models import CustomTool from workflow_manager.workflow_v2.models.workflow import Workflow +from mcp_server import spend_guard from mcp_server.context import PlatformMCPContext from mcp_server.exceptions import MCPToolError @@ -77,8 +78,8 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "scoped to a whole organization: it discovers what exists and can " "change the running state of those resources." ), - "read_tools": [ - {"name": "whoami", "purpose": "See this credential's org, tier and scope."}, + "discovery_tools": [ + {"name": "whoami", "purpose": "See this credential's org, tier and budget."}, { "name": "listApiDeployments", "purpose": "Find deployed extraction endpoints.", @@ -89,8 +90,31 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "name": "listPromptStudioProjects", "purpose": "Find where extraction prompts are authored.", }, + { + "name": "getWorkflowEndpoints", + "purpose": "See how a workflow is wired (shape only, no config).", + }, + { + "name": "listToolInstances", + "purpose": "See the tool steps inside a workflow.", + }, + {"name": "listTags", "purpose": "List execution tags."}, + ], + "observability_tools": [ + { + "name": "listExecutions", + "purpose": "Recent runs and their status. Start here to debug.", + }, + { + "name": "getExecutionDetail", + "purpose": "Per-file results and errors for one run.", + }, + { + "name": "getUsageSummary", + "purpose": "Tokens and cost recorded so far.", + }, ], - "write_tools": [ + "state_change_tools": [ { "name": "setApiDeploymentActive", "purpose": "Take a deployment offline, or bring it back.", @@ -99,17 +123,41 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "name": "setPipelineActive", "purpose": "Pause or resume a pipeline's schedule.", }, + ], + "billable_tools": [ { "name": "executePipeline", - "purpose": "Trigger a real pipeline run. Consumes quota.", + "purpose": "Trigger a real pipeline run.", + }, + { + "name": "indexDocument", + "purpose": "Embed a document so prompts can run against it.", + }, + { + "name": "fetchResponse", + "purpose": "Run one prompt against an indexed document.", + }, + { + "name": "bulkFetchResponse", + "purpose": "Run several prompts in one pass. Prefer over looping.", + }, + { + "name": "singlePassExtraction", + "purpose": "Run a project's whole prompt set. Most expensive.", }, ], "before_writing": ( - "Write tools change state for the whole organization and take " - "effect immediately — there is no staging or dry-run mode. " - "executePipeline in particular does real work and consumes quota; " - "it is not a way to test a pipeline. Confirm the target with a " - "list tool first." + "State-change tools take effect immediately for the whole " + "organization — there is no staging or dry-run mode. Confirm the " + "target with a discovery tool first." + ), + "before_spending": ( + "Billable tools cost real money each call: they drive LLM " + "inference, embedding and vector-store writes. They are budgeted " + "per organization — call whoami to see how much of that budget " + "remains. Index a document once and then run prompts against it; " + "prefer bulkFetchResponse over repeated fetchResponse calls; and " + "never call a billable tool to 'test' whether something works." ), "not_available": ( "Creating or rotating API keys, and deleting workflows, " @@ -143,6 +191,8 @@ def whoami(context: PlatformMCPContext) -> dict[str, Any]: except ValueError: can_write = can_delete = False + budget = spend_guard.peek(context.org_name) + return { "organization": context.org_name, "key_name": key.name, @@ -150,6 +200,17 @@ def whoami(context: PlatformMCPContext) -> dict[str, Any]: "is_service_account": bool(getattr(context.user, "is_service_account", False)), "can_use_write_tools": can_write, "can_use_destructive_tools": can_delete, + "billable_budget": { + "used": budget.used, + "limit": budget.limit, + "remaining": max(0, budget.limit - budget.used), + "window_minutes": round(budget.window_seconds / 60), + "note": ( + "Counts calls to billable tools, not tokens. Pace yourself " + "against `remaining`; when it reaches zero, billable tools are " + "refused until the window rolls over." + ), + }, "visibility": ( "All resources in the organization. Service-account credentials " "bypass per-user sharing filters, so this key can read AND modify " diff --git a/backend/mcp_server/tools/prompt_studio.py b/backend/mcp_server/tools/prompt_studio.py new file mode 100644 index 0000000000..f4906a360f --- /dev/null +++ b/backend/mcp_server/tools/prompt_studio.py @@ -0,0 +1,261 @@ +"""Billable Prompt Studio tools for the organization-scoped MCP server. + +These drive real LLM inference, embedding and vector-store writes, so each is +registered ``billable=True`` and budgeted by ``mcp_server.spend_guard``. + +Every tool here delegates to ``PromptStudioCoreView`` rather than +reimplementing its logic. Those actions carry a great deal that is easy to get +subtly wrong — run-id generation, profile resolution, lookup gating, indexing +race avoidance, Celery dispatch — and a parallel implementation would drift +from the UI's behaviour on the first change to either. Dispatching the real +view keeps one code path, and keeps its permission and sharing checks. +""" + +import logging +from typing import Any + +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from prompt_studio.prompt_studio_core_v2.views import PromptStudioCoreView + +from mcp_server.context import PlatformMCPContext +from mcp_server.exceptions import MCPToolError + +logger = logging.getLogger(__name__) + + +def _resolve_project(context: PlatformMCPContext, project_id: str) -> CustomTool: + """Find a Prompt Studio project the caller may reach, or refuse. + + Resolving through ``for_user`` is what keeps a key inside its own + organization — the view's own ``get_object`` would apply the same rules, + but failing here produces a message the agent can act on rather than a + 404 surfaced as an unexpected error. + """ + project = CustomTool.objects.for_user(context.user).filter(tool_id=project_id).first() + if project is None: + raise MCPToolError( + f"No Prompt Studio project with id '{project_id}' in organization " + f"'{context.org_name}'. Call listPromptStudioProjects for valid ids." + ) + return project + + +def _dispatch(context: PlatformMCPContext, action: str, project_id: str, payload: dict): + """Invoke a PromptStudioCoreView action with the MCP request. + + The underlying view reads its arguments from ``request.data`` and resolves + the organization from the request, both of which the platform MCP request + already carries — the middleware set the org, and this server's context + holds the real DRF request. + """ + if context.request is None: + raise MCPToolError( + "Prompt Studio operations are unavailable in this context. " + "Contact your Unstract administrator." + ) + + request = context.request + # The view reads its inputs from request.data. Replacing it wholesale keeps + # the JSON-RPC envelope (which carries the MCP method, not the tool's + # arguments) from reaching the view. + original_data = request.data + try: + request._full_data = payload + view = PromptStudioCoreView.as_view({"post": action}) + response = view(request, pk=project_id) + finally: + request._full_data = original_data + + return response + + +def _result(response, project: CustomTool) -> dict[str, Any]: + """Normalise a DRF response into a tool result. + + A non-2xx status is returned as data rather than raised: the view's own + error bodies (a missing prompt id, an unindexed document) are precisely + what the agent needs to correct its next call. + """ + status_code = getattr(response, "status_code", None) + data = getattr(response, "data", None) + ok = status_code is not None and 200 <= status_code < 300 + + result: dict[str, Any] = { + "project_id": str(project.tool_id), + "project_name": project.tool_name, + "ok": ok, + "status": status_code, + "result": data, + } + if not ok: + result["note"] = ( + "The operation was rejected. Read `result` for the reason — it is " + "usually a missing or invalid argument, or a document that has not " + "been indexed yet." + ) + return result + + +def index_document_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "project_id": { + "type": "string", + "description": ( + "UUID of the Prompt Studio project, from listPromptStudioProjects." + ), + }, + "document_id": { + "type": "string", + "description": "UUID of the document within that project to index.", + }, + }, + "required": ["project_id", "document_id"], + } + + +def index_document( + context: PlatformMCPContext, project_id: str, document_id: str +) -> dict[str, Any]: + """Index a document so prompts can be run against it.""" + project = _resolve_project(context, project_id) + logger.info( + f"MCP index_document project='{project_id}' document='{document_id}' " + f"(org '{context.org_name}', key '{context.platform_key.name}')" + ) + response = _dispatch( + context, "index_document", project_id, {"document_id": document_id} + ) + return _result(response, project) + + +def fetch_response_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "project_id": { + "type": "string", + "description": "UUID of the Prompt Studio project.", + }, + "document_id": { + "type": "string", + "description": "UUID of the document to run the prompt against.", + }, + "prompt_id": { + "type": "string", + "description": "UUID of the prompt to run.", + }, + "profile_manager_id": { + "type": "string", + "description": ( + "Optional UUID of an LLM profile to override the project's default." + ), + }, + }, + "required": ["project_id", "document_id", "prompt_id"], + } + + +def fetch_response( + context: PlatformMCPContext, + project_id: str, + document_id: str, + prompt_id: str, + profile_manager_id: str | None = None, +) -> dict[str, Any]: + """Run a single prompt against an indexed document.""" + project = _resolve_project(context, project_id) + payload: dict[str, Any] = {"document_id": document_id, "id": prompt_id} + if profile_manager_id: + payload["profile_manager"] = profile_manager_id + + logger.info( + f"MCP fetch_response project='{project_id}' prompt='{prompt_id}' " + f"(org '{context.org_name}', key '{context.platform_key.name}')" + ) + response = _dispatch(context, "fetch_response", project_id, payload) + return _result(response, project) + + +def bulk_fetch_response_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "project_id": { + "type": "string", + "description": "UUID of the Prompt Studio project.", + }, + "document_id": { + "type": "string", + "description": "UUID of the document to run the prompts against.", + }, + "prompt_ids": { + "type": "array", + "items": {"type": "string"}, + "minItems": 1, + "description": "UUIDs of the prompts to run.", + }, + }, + "required": ["project_id", "document_id", "prompt_ids"], + } + + +def bulk_fetch_response( + context: PlatformMCPContext, + project_id: str, + document_id: str, + prompt_ids: list[str], +) -> dict[str, Any]: + """Run several prompts against one document in a single pass. + + Preferred over repeated ``fetchResponse`` calls: it indexes once and + dispatches one task, which is both cheaper and avoids the + document-being-indexed race that concurrent single calls provoke. + """ + project = _resolve_project(context, project_id) + logger.info( + f"MCP bulk_fetch_response project='{project_id}' " + f"prompts={len(prompt_ids)} (org '{context.org_name}', " + f"key '{context.platform_key.name}')" + ) + response = _dispatch( + context, + "bulk_fetch_response", + project_id, + {"document_id": document_id, "prompt_ids": prompt_ids}, + ) + return _result(response, project) + + +def single_pass_extraction_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "project_id": { + "type": "string", + "description": "UUID of the Prompt Studio project.", + }, + "document_id": { + "type": "string", + "description": "UUID of the document to extract from.", + }, + }, + "required": ["project_id", "document_id"], + } + + +def single_pass_extraction( + context: PlatformMCPContext, project_id: str, document_id: str +) -> dict[str, Any]: + """Run the project's whole prompt set in a single LLM pass.""" + project = _resolve_project(context, project_id) + logger.info( + f"MCP single_pass_extraction project='{project_id}' " + f"document='{document_id}' (org '{context.org_name}', " + f"key '{context.platform_key.name}')" + ) + response = _dispatch( + context, "single_pass_extraction", project_id, {"document_id": document_id} + ) + return _result(response, project) diff --git a/backend/mcp_server/transport.py b/backend/mcp_server/transport.py index ddef3f3e2a..8998878162 100644 --- a/backend/mcp_server/transport.py +++ b/backend/mcp_server/transport.py @@ -216,6 +216,18 @@ def check_tool_allowed(self, tool: Any, context: Any) -> str | None: """ return None + def check_spend_allowed(self, tool: Any, context: Any) -> str | None: + """Claim budget for a billable tool. + + Return None to allow, or a message explaining that the budget is spent. + Called once per invocation immediately before the handler runs, so + overriding servers should treat it as consuming, not merely checking. + + The default allows everything — the deployment server's costs are + already bounded by the API deployment rate limiter. + """ + return None + def _call_tool( self, request_id: Any, params: dict[str, Any], context: Any ) -> JsonResponse: @@ -242,6 +254,15 @@ def _call_tool( request_id, JSONRPC.UNAUTHORIZED, "Permission denied", refusal ) + # Budget is checked after permission and before the handler, so an + # unauthorized call never consumes budget. Unlike the permission + # refusal above this is temporal, so it comes back as an isError + # *result* the agent can read and retry — a protocol error would read + # as "this tool does not work" and stop it retrying later. + over_budget = self.check_spend_allowed(tool, context) + if over_budget is not None: + return rpc_result(request_id, tool_content(over_budget, is_error=True)) + arguments = params.get("arguments") or {} if not isinstance(arguments, dict): return rpc_error( From 8e73615d28063f6912a96c3ffde0274c8aadc038 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 17:53:23 +0000 Subject: [PATCH 08/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- backend/backend/settings/base.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index ea8693abe4..9353108bfd 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -123,9 +123,7 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | # counted per organization over a rolling window. Bounds how often an agent can # trigger paid work; it counts calls, not tokens — see mcp_server/spend_guard.py. MCP_BILLABLE_CALL_LIMIT = int(os.environ.get("MCP_BILLABLE_CALL_LIMIT", 50)) -MCP_BILLABLE_WINDOW_SECONDS = int( - os.environ.get("MCP_BILLABLE_WINDOW_SECONDS", 3600) -) +MCP_BILLABLE_WINDOW_SECONDS = int(os.environ.get("MCP_BILLABLE_WINDOW_SECONDS", 3600)) # Maximum file size for presigned URLs in API deployments (in MB) API_DEPL_PRESIGNED_URL_MAX_FILE_SIZE_MB = int( From ce191f6ed76196cdac7e1e2b34b97df5bbfccffe Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sun, 26 Jul 2026 11:41:20 +0530 Subject: [PATCH 09/16] feat(mcp): add MCP_REDIS_DB, drop CAPTURE3.md from the branch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds MCP_REDIS_DB, defaulting to REDIS_DB. Nothing about the current deployment changes: with both unset the budget counter stays in the shared cache DB and the guard keeps using django.core.cache, which is also what makes override_settings(CACHES=...) work in tests. The knob exists so MCP state can be moved to its own Redis DB later without that being a breaking change. When the two differ the guard builds a client for the configured DB, preferring a CACHES["mcp"] alias if the operator configured one so they control pooling and auth rather than having them inferred. This follows the shape of CacheService.clear_cache_optimized: Django cache by default, raw client only when a specific DB is required. Worth recording why the default matters. Sharing DB 0 means a FLUSHDB or a cache-wide eviction resets every organization's window — a fail-open outcome, consistent with a guard that already allows calls when the cache is unreachable. It is a loop bound, not an audited ledger, and the README now says so alongside the key layout and TTL-as-window mechanics. Also adds MCP_BILLABLE_CALL_LIMIT and MCP_BILLABLE_WINDOW_SECONDS to sample.env, which were introduced without being documented there. Removes CAPTURE3.md from the branch. The findings it captured are unfixed issues in the existing REST API and are being tracked elsewhere; a public PR was the wrong place for them. Tests: 87 in this app (was 82) — the new ones pin that a matching DB uses the shared cache, that REDIS_DB's empty-string default coerces to 0 the same way downstream code coerces it, and that a differing DB really does get a different client. peek() now has its own fail-open test. Co-Authored-By: Claude Opus 4.8 (1M context) --- CAPTURE3.md | 110 ------------------- backend/backend/settings/base.py | 6 + backend/mcp_server/README.md | 20 ++++ backend/mcp_server/spend_guard.py | 61 +++++++++- backend/mcp_server/tests/test_spend_guard.py | 79 ++++++++++++- backend/sample.env | 9 ++ 6 files changed, 170 insertions(+), 115 deletions(-) delete mode 100644 CAPTURE3.md diff --git a/CAPTURE3.md b/CAPTURE3.md deleted file mode 100644 index e97f777648..0000000000 --- a/CAPTURE3.md +++ /dev/null @@ -1,110 +0,0 @@ -# CAPTURE3 — security findings in the existing REST API - -Two findings surfaced while surveying the tenant API surface to decide what the -platform MCP server should expose (PR #2207). **Neither is caused by that PR, -and neither is fixed by it** — both are pre-existing properties of the REST API, -captured here to be addressed separately. - ---- - -## 1. Platform API keys are readable repeatedly, despite being documented "shown once" - -**Severity:** medium — turns a write-once secret into a read-anytime one. - -`PlatformApiKeyDetailSerializer` includes the full `key` field, and its own -docstring states the intent: - -```python -# backend/platform_api/serializers.py:89-95 -class PlatformApiKeyDetailSerializer(serializers.ModelSerializer): - """Used for create/rotate responses where the full key is shown once.""" - - class Meta: - model = PlatformApiKey - fields = ["id", "name", "key", "is_active"] -``` - -`get_serializer_class` routes `list` → masked, `create` → create serializer, and -`partial_update` → update serializer — but `retrieve` is not in that mapping and -instead selects the detail serializer explicitly: - -```python -# backend/platform_api/views.py:59-62 -def retrieve(self, request, *args, **kwargs): - instance = self.get_object() - serializer = PlatformApiKeyDetailSerializer(instance) - return Response(serializer.data) -``` - -So `GET /api/v1/unstract//platform-api/keys//` returns the full -plaintext key on **every** call, not once at creation. - -**Why it matters.** `list` deliberately masks the key (`****-last4` via -`PlatformApiKeyListSerializer.get_key`), which shows the intent was for the -secret not to be freely readable. `retrieve` silently defeats that: anyone who -can reach the detail endpoint can recover any key in the organization at any -time, so a key leaked once cannot be reasoned about as "shown at creation only", -and rotation is the only remediation. - -**Options to consider.** -- Mask `key` in `retrieve` (reuse the list serializer), and return the full key - only from `create` and `rotate`, which is what the docstring already promises. -- If a re-read capability is genuinely wanted, make it an explicit, separately - authorized action rather than the default detail representation. - -**Related, not the same finding:** `GET /api/keys/api//` in `api_v2` -returns the full deployment API key and is gated at `IsOwnerOrSharedUser` rather -than owner-only — so any user a deployment is *shared with* can read its live -key. Worth reviewing alongside this one. - ---- - -## 2. Two destructive endpoints are exposed over `GET` - -**Severity:** medium — destructive actions reachable by link, prefetch or CSRF. - -```python -# backend/file_management/urls.py:37-47 -file_delete = FileManagementViewSet.as_view({"get": "delete"}) -... -path("file/delete", file_delete, name="delete"), -``` - -```python -# backend/workflow_manager/workflow_v2/urls/workflow.py:25 -workflow_clear_file_marker = WorkflowViewSet.as_view({"get": "clear_file_marker"}) -``` - -`GET /api/v1/unstract//file/delete` deletes files, and -`GET /api/v1/unstract//workflow//clear-file-marker/` clears file -markers. Both mutate state behind a verb that the entire web stack treats as -safe and idempotent. - -**Why it matters.** -- **CSRF.** Django's CSRF protection does not apply to `GET`, so a cross-origin - link, image tag or redirect can trigger either action using the victim's - session. The usual "unsafe methods only" assumption does not hold here. -- **Prefetch and crawlers.** Browsers, link previews, and security scanners - follow `GET` URLs speculatively. Any of these can fire a deletion nobody - requested. -- **Caching and logs.** Intermediaries may cache or replay `GET`, and the URL - (with parameters identifying what is deleted) lands in access logs and browser - history. - -Clearing file markers additionally forces re-processing of the affected files, -so it has a **cost** consequence as well as a data one. - -**Options to consider.** Move both to `POST`/`DELETE`. This is a breaking change -for any existing caller, so it likely needs a deprecation window: accept both -verbs, log `GET` usage, then remove `GET` once callers have migrated. - ---- - -## Not in scope here - -Several other endpoints return decrypted credentials in their normal responses -(connector metadata, adapter provider keys, notification webhook tokens, -`platform_settings_v2` keys, the Postman collection download). Those are -intentional today insofar as the UI relies on them, and they are a broader -design question than these two — the platform MCP server simply does not wrap -any of them. See `backend/mcp_server/README.md` for that exclusion list. diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 9353108bfd..ed383e4210 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -103,6 +103,12 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | REDIS_HOST = os.environ.get("REDIS_HOST", "localhost") REDIS_PORT = os.environ.get("REDIS_PORT", "6379") REDIS_DB = os.environ.get("REDIS_DB", "") +# Redis DB for MCP server state (currently the billable-call budget). Defaults +# to REDIS_DB, so out of the box MCP state shares the general cache DB and +# nothing changes. It exists as a separate knob so MCP state can later be moved +# to its own DB — for isolation, or so a cache flush does not reset in-flight +# budget windows — without that becoming a breaking change. +MCP_REDIS_DB = int(os.environ.get("MCP_REDIS_DB") or REDIS_DB or 0) SESSION_EXPIRATION_TIME_IN_SECOND = os.environ.get( "SESSION_EXPIRATION_TIME_IN_SECOND", 3600 ) diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index 273267063a..759db6d5be 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -223,6 +223,26 @@ Three behaviours are deliberate: `whoami` reports the remaining budget so an agent can pace itself instead of discovering the limit by hitting it. +### Where the counter lives + +One key per organization, `mcp:billable:`, expired by TTL — there is no +window-start timestamp, so the window begins at the first billable call rather +than at a wall-clock boundary. `cache.add` initialises the counter without +disturbing an in-flight window, and `cache.incr` is atomic, so concurrent +requests cannot race past the limit. + +`MCP_REDIS_DB` selects the Redis DB and **defaults to `REDIS_DB`**, so out of +the box MCP state shares the general cache DB and behaves like any other cache +user — including honouring `override_settings(CACHES=...)` in tests. Set it only +to move MCP state onto its own DB; the guard then builds a client for that DB, +preferring a `CACHES["mcp"]` alias if one is configured. The knob exists so that +move is available later without being a breaking change. + +Note the consequence of the default: sharing DB 0 means a `FLUSHDB` or a +cache-wide eviction resets every organization's window. That is a fail-open +outcome, consistent with the rest of the guard's design — this is a loop +bound, not an audited ledger. + ## How write tools are authorized Platform API key tiers are defined in terms of HTTP methods diff --git a/backend/mcp_server/spend_guard.py b/backend/mcp_server/spend_guard.py index a6c6322f6c..96ce0ce3a2 100644 --- a/backend/mcp_server/spend_guard.py +++ b/backend/mcp_server/spend_guard.py @@ -30,7 +30,7 @@ from dataclasses import dataclass from django.conf import settings -from django.core.cache import cache +from django.core.cache import cache as default_cache logger = logging.getLogger(__name__) @@ -44,6 +44,62 @@ def _key(org_id: str) -> str: return f"{_KEY_PREFIX}:{org_id}" +def _default_redis_db() -> int: + """The DB the project's shared cache is configured to use.""" + return int(getattr(settings, "REDIS_DB", "") or 0) + + +def get_cache(): + """Return the cache client for MCP budget state. + + ``MCP_REDIS_DB`` defaults to ``REDIS_DB``, so ordinarily this is just the + project's shared cache and everything behaves as any other cache user does + — including honouring ``override_settings(CACHES=...)`` in tests. + + When the two differ, an operator has deliberately moved MCP state to its + own Redis DB. ``django.core.cache`` cannot express that (the DB is fixed by + ``CACHES``), so a dedicated client is built for that DB. This mirrors how + ``CacheService.clear_cache_optimized`` reaches the workers' DB: use the + Django cache by default, drop to a raw client only when a specific DB is + required. + """ + configured_db = int(getattr(settings, "MCP_REDIS_DB", _default_redis_db())) + if configured_db == _default_redis_db(): + return default_cache + + return _dedicated_cache(configured_db) + + +def _dedicated_cache(db: int): + """Build a cache client pinned to ``db``. + + Deliberately constructed per call rather than cached at import: settings can + change under ``override_settings`` in tests, and this path is only taken by + installations that opted into a separate DB, where one extra client + construction per billable call is immaterial next to the LLM call it guards. + """ + from django.core.cache import caches + from django.core.cache.backends.base import InvalidCacheBackendError + + # Prefer an explicitly configured alias if the operator added one; it lets + # them control pooling and auth the same way as the default cache. + try: + return caches["mcp"] + except InvalidCacheBackendError: + pass + + from django_redis.cache import RedisCache + + base = settings.CACHES["default"] + options = dict(base.get("OPTIONS", {})) + options["DB"] = db + location = base["LOCATION"] + if isinstance(location, str) and location.rsplit("/", 1)[-1].isdigit(): + location = f"{location.rsplit('/', 1)[0]}/{db}" + + return RedisCache(location, {**base, "OPTIONS": options}) + + @dataclass(frozen=True) class BudgetState: """The outcome of a budget check.""" @@ -88,7 +144,7 @@ def peek(org_id: str) -> BudgetState: limit = get_limit() window = get_window_seconds() try: - used = cache.get(_key(org_id)) or 0 + used = get_cache().get(_key(org_id)) or 0 except Exception as error: # Reporting the budget must never be the thing that breaks whoami — # `consume` fails open for the same reason, and a read is even less @@ -114,6 +170,7 @@ def consume(org_id: str) -> BudgetState: limit = get_limit() window = get_window_seconds() key = _key(org_id) + cache = get_cache() try: # `add` only sets when absent, so it both initialises the counter and diff --git a/backend/mcp_server/tests/test_spend_guard.py b/backend/mcp_server/tests/test_spend_guard.py index 3e1172803a..621cf730a8 100644 --- a/backend/mcp_server/tests/test_spend_guard.py +++ b/backend/mcp_server/tests/test_spend_guard.py @@ -108,13 +108,23 @@ def test_cache_failure_fails_open(self) -> None: the whole MCP surface offline because Redis blipped would be the worse failure, so it allows and logs. """ - with patch( - "mcp_server.spend_guard.cache.add", side_effect=RuntimeError("redis down") - ): + broken = Mock() + broken.add.side_effect = RuntimeError("redis down") + with patch("mcp_server.spend_guard.get_cache", return_value=broken): state = spend_guard.consume(ORG) assert state.allowed is True + def test_peek_also_fails_open(self) -> None: + """whoami reports the budget, so an unreachable cache must not break it.""" + broken = Mock() + broken.get.side_effect = RuntimeError("redis down") + with patch("mcp_server.spend_guard.get_cache", return_value=broken): + state = spend_guard.peek(ORG) + + assert state.allowed is True + assert state.used == 0 + @override_settings( MCP_BILLABLE_CALL_LIMIT=2, MCP_BILLABLE_WINDOW_SECONDS=60, CACHES=LOCMEM @@ -237,6 +247,69 @@ def test_non_billable_tools_still_work_once_the_budget_is_gone(self) -> None: assert body["result"]["isError"] is False +class RedisDbSelectionTest(SimpleTestCase): + """Which Redis DB the budget lives in. + + ``MCP_REDIS_DB`` defaults to ``REDIS_DB``, so the shipped behaviour is + "same DB as everything else" and nothing about the existing deployment + changes. The knob exists so MCP state can be moved later without that being + a breaking change. + """ + + @override_settings(REDIS_DB="0", MCP_REDIS_DB=0) + def test_matching_db_uses_the_shared_django_cache(self) -> None: + """The default path. Using the shared cache is what keeps + ``override_settings(CACHES=...)`` working in tests and keeps the guard + consistent with every other cache user in the backend. + """ + from django.core.cache import cache as django_cache + + assert spend_guard.get_cache() is django_cache + + @override_settings(REDIS_DB="", MCP_REDIS_DB=0) + def test_unset_redis_db_is_treated_as_zero(self) -> None: + """REDIS_DB ships as an empty string, which downstream code coerces to + 0. The comparison here must coerce the same way or the guard would + wrongly conclude the DBs differ and build a needless client. + """ + from django.core.cache import cache as django_cache + + assert spend_guard.get_cache() is django_cache + + @override_settings(REDIS_DB="0", MCP_REDIS_DB=7) + def test_differing_db_does_not_use_the_shared_cache(self) -> None: + """An operator who set a different DB must actually get a different + client — silently continuing on DB 0 would make the setting a lie. + """ + from django.core.cache import cache as django_cache + + with patch("mcp_server.spend_guard._dedicated_cache") as dedicated: + client = spend_guard.get_cache() + + dedicated.assert_called_once_with(7) + assert client is not django_cache + + @override_settings( + REDIS_DB="0", + MCP_REDIS_DB=7, + CACHES={ + **LOCMEM, + "mcp": { + "BACKEND": "django.core.cache.backends.locmem.LocMemCache", + "LOCATION": "mcp-alias-test", + }, + }, + ) + def test_an_explicit_mcp_cache_alias_wins(self) -> None: + """If an operator configures a CACHES['mcp'] alias they control pooling + and auth themselves, which is better than this module inferring them + from the default cache's settings. + """ + from django.core.cache import caches + + assert spend_guard.get_cache() is caches["mcp"] + + class BillableRegistryInvariantTest(SimpleTestCase): def test_every_costly_platform_tool_is_marked_billable(self) -> None: """The flag is the only thing wiring a tool into the budget, so a diff --git a/backend/sample.env b/backend/sample.env index e04445822d..65871c8a22 100644 --- a/backend/sample.env +++ b/backend/sample.env @@ -259,3 +259,12 @@ HITL_FILES_FILE_STORAGE_CREDENTIALS='{"provider": "minio", "credentials": {"endp # File active cache redis db FILE_ACTIVE_CACHE_REDIS_DB=0 + +# Hosted MCP servers +# Redis DB for MCP state (the billable-call budget). Leave unset to share +# REDIS_DB with the general cache; set it to move MCP state to its own DB. +# MCP_REDIS_DB=0 +# Budget for billable MCP tool calls (LLM inference, indexing, pipeline runs), +# counted per organization over a rolling window. Counts calls, not tokens. +MCP_BILLABLE_CALL_LIMIT=50 +MCP_BILLABLE_WINDOW_SECONDS=3600 From 41679676bc0379b8c9fda5a204cfe32a30ce6989 Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sun, 26 Jul 2026 15:06:27 +0530 Subject: [PATCH 10/16] feat(mcp): make the billable Prompt Studio tools reachable, and answer GET per spec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The four billable Prompt Studio tools each require a document_id, and fetchResponse also a prompt_id, but nothing in the registry returned either. An MCP client sees only tools/list — it has no database and cannot invent a UUID — so the most expensive tools on the server were listed and uncallable. Adds listPromptStudioDocuments and listPrompts as the producers. Both resolve through the project rather than by a bare id lookup: the platform key authenticates as a service account, for which for_user() returns everything, so the project join is the scoping that actually holds. listPrompts builds its response field by field instead of via a serializer, because ToolStudioPrompt carries a profile_manager FK (the LLM, embedding and vector-store adapters) and a webhook URL that may embed a token — both of which the README promises never reach an agent, and neither of which was pinned by a test until now. The durable fix is test_registry_reachability: every id a tool requires must have a declared producer. The bug class was invisible per tool and only wrong in aggregate, so the invariant is asserted over the registry. Also, per the Streamable HTTP spec (rev 2025-06-18), GET opens a server-to-client SSE stream and a server offering none must answer 405. This server pushes nothing, so GET now returns 405 with Allow: POST, keeping the identity body for uptime probes. The "every MCP call is a POST" comments were imprecise in the same way and now say that every JSON-RPC *message* arrives as a POST — which is what makes required_method necessary. Co-Authored-By: Claude Opus 5 (1M context) --- backend/mcp_server/README.md | 50 +++++-- backend/mcp_server/platform_views.py | 3 +- backend/mcp_server/registry.py | 38 +++++- backend/mcp_server/tests/test_mcp_auth.py | 7 +- .../tests/test_platform_tier_guard.py | 3 +- .../mcp_server/tests/test_platform_tools.py | 126 +++++++++++++++++ .../tests/test_registry_reachability.py | 129 ++++++++++++++++++ backend/mcp_server/tools/prompt_studio.py | 90 ++++++++++++ backend/mcp_server/transport.py | 23 +++- 9 files changed, 451 insertions(+), 18 deletions(-) create mode 100644 backend/mcp_server/tests/test_registry_reachability.py diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index 759db6d5be..db2684d39a 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -157,6 +157,8 @@ Platform API keys are managed at `/api/v1/unstract//platform-api/keys/`. | `listWorkflows` | read | The workflows behind them. | | `listPipelines` | read | ETL and task pipelines, with schedule and last-run state. | | `listPromptStudioProjects` | read | Where extraction prompts are authored. | +| `listPromptStudioDocuments` | read | Documents in a project. The only source of `document_id`. | +| `listPrompts` | read | A project's prompts. The only source of `prompt_id`. | | `getWorkflowEndpoints` | read | How a workflow is wired — shape only, never config. | | `listToolInstances` | read | The tool steps inside a workflow. | | `listTags` | read | Execution tags. | @@ -246,8 +248,9 @@ bound, not an audited ledger. ## How write tools are authorized Platform API key tiers are defined in terms of HTTP methods -(`ApiKeyPermission.allows`), but every MCP call is a `POST` — so the auth -middleware's tier check cannot tell `listWorkflows` from `executePipeline`. +(`ApiKeyPermission.allows`), but every JSON-RPC message arrives as an HTTP +`POST` whatever the tool inside it does — so the auth middleware's tier check +cannot tell `listWorkflows` from `executePipeline`. Each tool therefore declares the method its REST equivalent would use (`required_method`), and `check_tool_allowed` re-applies the key's tier against @@ -311,13 +314,43 @@ Also excluded: password resets, role assignment and revocation, and member removal. These change who can access the organization, which is not a decision to delegate to an agent. +### Endpoints that reach out to a caller-influenced host + +The connector and adapter *test* endpoints exist to verify a configuration by +making a live outbound call to the host it names. Exposing them would hand any +platform API key a request-forgery primitive: the agent chooses the +destination, the server makes the call from inside the deployment's network, +and the response comes back. That the same endpoints also take credentials as +*input* is a second reason, but the outbound call is the disqualifying one. + +Prompt Studio project **import** is excluded on the same grounds — it ingests a +bundle the caller supplies. **Export** is excluded because a project bundle may +carry adapter references, which is the credential rule above applied to a file +rather than a JSON field. + +### Workflow endpoint and tool-instance *configuration* + +Reading the shape of a workflow's endpoints is exposed +(`getWorkflowEndpoints`, `listToolInstances`); writing it is not. Endpoint +configuration is where connector instances are bound to a workflow, so a write +there redirects where a workflow reads its input from and writes its output to +— without changing anything that a later read would show as unusual. That is a +data-exfiltration path dressed as a settings change, so it stays out. + ### Everything else, for now -Not excluded on principle, simply not built: file upload/download, -tool-instance and endpoint *configuration* (as opposed to reading their shape), -Prompt Studio project/prompt authoring and import/export, and the connector and -adapter *test* endpoints (which make live outbound calls). These are candidates -for later, with the same rules applied. +Not excluded on principle, simply not built: file upload/download, and Prompt +Studio project and prompt *authoring*. The coherent scope of this server is +that an agent operates on projects and documents a human has already set up — +it can discover them, run them, and read the results, but it does not create +the raw material. These are candidates for later, with the same rules applied. + +Document and prompt *listing* used to sit in this section. It has since been +built (`listPromptStudioDocuments`, `listPrompts`) because it was not a nicety: +those two tools are the only source of the `document_id` and `prompt_id` that +every billable Prompt Studio tool requires, so without them those tools were +listed but uncallable. `tests/test_registry_reachability.py` now fails if any +tool requires an id that no tool is declared to produce. ## Why the URL matters @@ -339,7 +372,8 @@ bypasses the middleware and would pass against a completely open endpoint. ## Two constraints worth knowing **A `read`-tier key cannot use this server at all.** The middleware's tier check -gates on HTTP method, and every MCP call is a `POST`, which `read` disallows — +gates on HTTP method, and every JSON-RPC message is a `POST`, which `read` +disallows — so a read-only key is refused before reaching the view, even for the read tools. Use a `read_write` key. Changing that would mean special-casing MCP paths in the middleware, which is a decision for maintainers rather than something this app diff --git a/backend/mcp_server/platform_views.py b/backend/mcp_server/platform_views.py index 2ae7dd554b..f20f21f56b 100644 --- a/backend/mcp_server/platform_views.py +++ b/backend/mcp_server/platform_views.py @@ -117,7 +117,8 @@ def check_tool_allowed(self, tool: Any, context: PlatformMCPContext) -> str | No """Apply the platform key's permission tier to this specific tool. The middleware already checked the tier, but it can only check it - against the HTTP method of the request — and every MCP call is a POST. + against the HTTP method of the request — and every JSON-RPC message + arrives as an HTTP POST, whatever the tool inside it does. That makes its verdict uselessly coarse here: it cannot tell ``listWorkflows`` from ``executePipeline``. diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index 87ba37c2bb..d5a6ab4544 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -27,7 +27,8 @@ class MCPTool: execution). Read-only tools are safe to retry; write tools are not. required_method: The HTTP method this tool's REST equivalent would use. Platform API key tiers are defined in terms of HTTP methods - (``ApiKeyPermission.allows``), and every MCP call is a POST — so + (``ApiKeyPermission.allows``), and every JSON-RPC message arrives + as an HTTP POST regardless of the tool it carries — so declaring the *equivalent* method is what lets the existing tier semantics apply per tool instead of per request. "GET" for reads, "POST" for mutations, "DELETE" for destructive operations (which @@ -220,6 +221,10 @@ def build_platform_registry() -> MCPToolRegistry: fetch_response_schema, index_document, index_document_schema, + list_prompt_studio_documents, + list_prompt_studio_documents_schema, + list_prompts, + list_prompts_schema, single_pass_extraction, single_pass_extraction_schema, ) @@ -287,6 +292,37 @@ def build_platform_registry() -> MCPToolRegistry: handler=list_prompt_studio_projects, ) ) + # Registered next to the project listing rather than beside the billable + # tools they feed: an agent reads this list top-down, and the id-producing + # step belongs immediately after the project it drills into. + registry.register( + MCPTool( + name="listPromptStudioDocuments", + description=( + "List the documents uploaded to a Prompt Studio project.\n\n" + "Call this to obtain the `document_id` that indexDocument, " + "fetchResponse, bulkFetchResponse and singlePassExtraction " + "require. Free to call — it reads metadata only and runs no " + "inference." + ), + input_schema=list_prompt_studio_documents_schema(), + handler=list_prompt_studio_documents, + ) + ) + registry.register( + MCPTool( + name="listPrompts", + description=( + "List a Prompt Studio project's prompts: their keys, text, " + "type and order.\n\n" + "Call this to obtain the `prompt_id` that fetchResponse and " + "bulkFetchResponse require, or to see what a project extracts " + "before spending money running it. Free to call." + ), + input_schema=list_prompts_schema(), + handler=list_prompts, + ) + ) registry.register( MCPTool( name="listPipelines", diff --git a/backend/mcp_server/tests/test_mcp_auth.py b/backend/mcp_server/tests/test_mcp_auth.py index 7c3391d7fe..2b7784aefa 100644 --- a/backend/mcp_server/tests/test_mcp_auth.py +++ b/backend/mcp_server/tests/test_mcp_auth.py @@ -187,10 +187,15 @@ def test_mcp_url_hangs_off_the_execution_url(self) -> None: def test_get_probe_does_not_leak_deployment_details(self) -> None: """The unauthenticated GET probe advertises the server, not the deployment behind it. + + It answers 405: a Streamable HTTP client issues GET to open an SSE + stream, and this server pushes nothing, so declining is the conformant + response. The identity body rides along for uptime probes. """ request = self.factory.get(f"/mcp/{ORG_ID}/live-api/") response = self.view(request, org_name=ORG_ID, api_name="live-api") - assert response.status_code == 200 + assert response.status_code == 405 + assert response["Allow"] == "POST" assert response.data["name"] == "unstract" assert "live-api" not in json.dumps(response.data) diff --git a/backend/mcp_server/tests/test_platform_tier_guard.py b/backend/mcp_server/tests/test_platform_tier_guard.py index 1e3765082d..bb68381941 100644 --- a/backend/mcp_server/tests/test_platform_tier_guard.py +++ b/backend/mcp_server/tests/test_platform_tier_guard.py @@ -1,7 +1,8 @@ """Per-tool authorization on the platform server. The auth middleware checks the key's permission tier against the request's HTTP -method — but every MCP call is a POST, so its verdict cannot distinguish +method — but every JSON-RPC message arrives as an HTTP POST whatever the tool +inside it does, so the middleware's verdict cannot distinguish ``listWorkflows`` from ``executePipeline``. This guard re-applies the tier against the method each tool *declares*, and is the only thing standing between a low-tier key and a write tool. diff --git a/backend/mcp_server/tests/test_platform_tools.py b/backend/mcp_server/tests/test_platform_tools.py index 142874c37b..ca70e084e7 100644 --- a/backend/mcp_server/tests/test_platform_tools.py +++ b/backend/mcp_server/tests/test_platform_tools.py @@ -7,6 +7,7 @@ from __future__ import annotations +import json from unittest.mock import Mock, patch from account_v2.models import Organization, User @@ -15,6 +16,8 @@ from pipeline_v2.models import Pipeline from platform_api.models import PlatformApiKey from prompt_studio.prompt_studio_core_v2.models import CustomTool +from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager +from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt from tenant_account_v2.models import OrganizationMember from utils.user_context import UserContext from workflow_manager.workflow_v2.models.workflow import Workflow @@ -31,6 +34,10 @@ set_pipeline_active, whoami, ) +from mcp_server.tools.prompt_studio import ( + list_prompt_studio_documents, + list_prompts, +) ORG_ID = "org-tools" @@ -271,3 +278,122 @@ def test_read_me_first_states_the_scope_warning(self) -> None: assert result["organization"] == ORG_ID # It must also point at the other server, since this one cannot extract. assert "extractDocument" in result["to_run_an_extraction"] + + +class PromptStudioProducerToolsTest(TestCase): + """The two tools that make the billable Prompt Studio tools reachable. + + ``indexDocument``, ``fetchResponse``, ``bulkFetchResponse`` and + ``singlePassExtraction`` all require ids an agent has no way to invent. + These tools are the only source of them, so their payload shape is load + bearing in a way a listing's usually is not — a missing ``document_id`` + field here silently disables every billable tool on the server. + """ + + def setUp(self) -> None: + self.org = Organization.objects.create( + name="org-producers", + display_name="Producers Org", + organization_id="org-producers", + ) + UserContext.set_organization_identifier("org-producers") + self.user = User.objects.create( + username="svc-producers", + email="svc-producers@platform.internal", + user_id="uid-producers", + is_service_account=True, + ) + OrganizationMember.objects.create( + user=self.user, organization=self.org, role="user" + ) + self.key = PlatformApiKey.objects.create( + name="producers-key", + description="d", + organization=self.org, + api_user=self.user, + permission="read_write", + ) + self.project = CustomTool.objects.create( + tool_name="Invoice Prompts", description="Prompt project", author="acme" + ) + self.document = DocumentManager.objects.create( + document_name="invoice-001.pdf", tool=self.project + ) + self.prompt = ToolStudioPrompt.objects.create( + prompt_key="invoice_number", + prompt="What is the invoice number?", + tool_id=self.project, + sequence_number=1, + prompt_type="Text", + enforce_type="text", + # Set deliberately: the redaction assertion below is only meaningful + # if there is a webhook URL that *could* have leaked. + postprocessing_webhook_url="https://hooks.internal/acme?token=s3cr3t", + ) + self.context = PlatformMCPContext( + user=self.user, + platform_key=self.key, + org_name="org-producers", + request=Mock(data={}), + ) + + def test_documents_are_listed_with_the_id_billable_tools_need(self) -> None: + result = list_prompt_studio_documents( + self.context, project_id=str(self.project.tool_id) + ) + + assert result["documents"] == [ + { + "document_id": str(self.document.document_id), + "document_name": "invoice-001.pdf", + } + ] + + def test_prompts_are_listed_with_the_id_fetch_response_needs(self) -> None: + result = list_prompts(self.context, project_id=str(self.project.tool_id)) + + assert len(result["prompts"]) == 1 + row = result["prompts"][0] + assert row["prompt_id"] == str(self.prompt.prompt_id) + assert row["prompt_key"] == "invoice_number" + assert row["prompt"] == "What is the invoice number?" + + def test_listing_prompts_does_not_leak_adapters_or_webhooks(self) -> None: + """The promise the README makes, pinned. + + ``ToolStudioPrompt`` carries a ``profile_manager`` FK — the LLM, + embedding and vector-store adapters behind the prompt — and a + webhook URL that may embed a token. A serializer would have carried + both out to the agent, which is why the handler builds its dict by + hand. This test is what stops someone swapping it back. + """ + result = list_prompts(self.context, project_id=str(self.project.tool_id)) + + payload = json.dumps(result) + assert "profile_manager" not in payload + assert "postprocessing_webhook_url" not in payload + assert "s3cr3t" not in payload, "a webhook token reached the agent" + assert "hooks.internal" not in payload + + def test_an_unknown_project_is_refused_with_a_usable_message(self) -> None: + with self.assertRaises(MCPToolError) as caught: + list_prompts(self.context, project_id="00000000-0000-0000-0000-000000000000") + + assert "listPromptStudioProjects" in str(caught.exception) + + def test_documents_do_not_leak_across_projects(self) -> None: + """Documents are resolved through the project, never by a bare id + lookup. The platform key is a service account, for which ``for_user`` + returns everything, so the project join is the scoping that holds. + """ + other = CustomTool.objects.create( + tool_name="Other Project", description="d", author="acme" + ) + DocumentManager.objects.create(document_name="secret.pdf", tool=other) + + result = list_prompt_studio_documents( + self.context, project_id=str(self.project.tool_id) + ) + + names = [row["document_name"] for row in result["documents"]] + assert names == ["invoice-001.pdf"], "a sibling project's document leaked" diff --git a/backend/mcp_server/tests/test_registry_reachability.py b/backend/mcp_server/tests/test_registry_reachability.py new file mode 100644 index 0000000000..71d65bec10 --- /dev/null +++ b/backend/mcp_server/tests/test_registry_reachability.py @@ -0,0 +1,129 @@ +"""Every id a tool asks for must be obtainable from another tool. + +An MCP client sees only ``tools/list``. It has no database, no UI and no way to +guess a UUID, so a tool whose required argument no other tool ever returns is +dead on arrival — it will be listed, attempted, and fail every time. + +This shipped once already: the four billable Prompt Studio tools all required a +``document_id`` while nothing in the registry produced one, which made the most +expensive tools on the server the only unreachable ones. The failure was +invisible because each tool was correct in isolation; only the registry as a +whole was wrong. So the invariant is asserted over the registry, not per tool. +""" + +from __future__ import annotations + +from django.test import SimpleTestCase + +from mcp_server.registry import DEPLOYMENT_TOOLS, PLATFORM_TOOLS + +# Which tool an agent is expected to call to obtain each opaque id. +# +# Written out rather than inferred from response shapes: handlers return plain +# dicts built at runtime, so inference would mean executing them, and a map +# that guesses would quietly bless the next unreachable id. Adding an entry +# here is a deliberate statement that the named tool really does return it. +PLATFORM_ID_PRODUCERS = { + "project_id": "listPromptStudioProjects", + "document_id": "listPromptStudioDocuments", + "prompt_id": "listPrompts", + "prompt_ids": "listPrompts", + "pipeline_id": "listPipelines", + "workflow_id": "listWorkflows", + "api_id": "listApiDeployments", + "execution_id": "listExecutions", +} + +DEPLOYMENT_ID_PRODUCERS = { + # The deployment server is scoped to a single API, and its execution id is + # handed back by the extraction call itself rather than by a listing. + "execution_id": "extractDocument", +} + + +def _required_ids(registry) -> dict[str, list[str]]: + """Map each id-shaped argument to the tools that require it.""" + consumed: dict[str, list[str]] = {} + for name in registry.names(): + schema = registry.get(name).input_schema + required = set(schema.get("required", [])) + for argument in schema.get("properties", {}): + if argument in required and argument.endswith(("_id", "_ids")): + consumed.setdefault(argument, []).append(name) + return consumed + + +class PlatformRegistryReachabilityTest(SimpleTestCase): + def test_every_required_id_has_a_declared_producer(self) -> None: + consumed = _required_ids(PLATFORM_TOOLS) + + unproducible = { + argument: tools + for argument, tools in consumed.items() + if argument not in PLATFORM_ID_PRODUCERS + } + + assert unproducible == {}, ( + "These tools require an id no tool is declared to produce, so an " + f"agent can never call them: {unproducible}. Either add the tool " + "that returns the id, or record its producer in " + "PLATFORM_ID_PRODUCERS." + ) + + def test_every_declared_producer_is_actually_registered(self) -> None: + """The map is only as good as its referents. + + A producer that was renamed or dropped would leave the test above + passing while the id became unobtainable again. + """ + missing = sorted( + { + producer + for producer in PLATFORM_ID_PRODUCERS.values() + if PLATFORM_TOOLS.get(producer) is None + } + ) + + assert missing == [], f"Declared id producers are not registered: {missing}" + + def test_the_billable_prompt_studio_chain_is_walkable(self) -> None: + """The specific path that was broken, pinned end to end. + + Named separately from the general invariant because this is the + sequence an agent must actually perform to spend money: project, then + document, then prompt, then the billable call. + """ + for step in ( + "listPromptStudioProjects", + "listPromptStudioDocuments", + "listPrompts", + "fetchResponse", + ): + assert PLATFORM_TOOLS.get(step) is not None, f"{step} is missing" + + for producer in ("listPromptStudioDocuments", "listPrompts"): + tool = PLATFORM_TOOLS.get(producer) + assert tool.billable is False, ( + f"{producer} exists so an agent can find ids before spending; " + "making it billable would put the discovery step behind the " + "budget it is meant to help the agent respect." + ) + assert tool.required_method == "GET", ( + f"{producer} only reads, so a read-tier key must reach it — " + "otherwise the billable tools stay unreachable for that key." + ) + + +class DeploymentRegistryReachabilityTest(SimpleTestCase): + def test_every_required_id_has_a_declared_producer(self) -> None: + consumed = _required_ids(DEPLOYMENT_TOOLS) + + unproducible = { + argument: tools + for argument, tools in consumed.items() + if argument not in DEPLOYMENT_ID_PRODUCERS + } + + assert unproducible == {}, ( + f"Unreachable on the deployment server: {unproducible}" + ) diff --git a/backend/mcp_server/tools/prompt_studio.py b/backend/mcp_server/tools/prompt_studio.py index f4906a360f..a844f9f94a 100644 --- a/backend/mcp_server/tools/prompt_studio.py +++ b/backend/mcp_server/tools/prompt_studio.py @@ -259,3 +259,93 @@ def single_pass_extraction( context, "single_pass_extraction", project_id, {"document_id": document_id} ) return _result(response, project) + + +def list_prompt_studio_documents_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "project_id": { + "type": "string", + "description": ( + "UUID of the Prompt Studio project, from listPromptStudioProjects." + ), + }, + }, + "required": ["project_id"], + } + + +def list_prompt_studio_documents( + context: PlatformMCPContext, project_id: str +) -> dict[str, Any]: + """List the documents already uploaded to a Prompt Studio project. + + This is the only producer of the ``document_id`` that every billable tool + in this module requires, so without it those tools are unreachable. + + Documents are resolved *through* the project rather than by a direct + ``DocumentManager`` lookup. The platform key authenticates as a service + account, for which ``for_user`` returns everything, so the project join in + ``_resolve_project`` is the scoping that actually holds. + """ + project = _resolve_project(context, project_id) + documents = project.document_managers.all().order_by("document_name") + + return { + "project_id": str(project.tool_id), + "project_name": project.tool_name, + "documents": [ + { + "document_id": str(document.document_id), + "document_name": document.document_name, + } + for document in documents + ], + } + + +def list_prompts_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "project_id": { + "type": "string", + "description": ( + "UUID of the Prompt Studio project, from listPromptStudioProjects." + ), + }, + }, + "required": ["project_id"], + } + + +def list_prompts(context: PlatformMCPContext, project_id: str) -> dict[str, Any]: + """List a project's prompts, which is where ``prompt_id`` comes from. + + The response is built field by field rather than through a serializer. + ``ToolStudioPrompt`` carries a ``profile_manager`` FK — the LLM, embedding + and vector-store adapters backing the prompt — and a + ``postprocessing_webhook_url``. A serializer would carry both out to the + agent, which is exactly what this server's README promises it does not do. + Adding a field here is therefore a deliberate act. + """ + project = _resolve_project(context, project_id) + prompts = project.mapped_prompt.all().order_by("sequence_number") + + return { + "project_id": str(project.tool_id), + "project_name": project.tool_name, + "prompts": [ + { + "prompt_id": str(prompt.prompt_id), + "prompt_key": prompt.prompt_key, + "prompt": prompt.prompt, + "prompt_type": prompt.prompt_type, + "enforce_type": prompt.enforce_type, + "sequence_number": prompt.sequence_number, + "active": prompt.active, + } + for prompt in prompts + ], + } diff --git a/backend/mcp_server/transport.py b/backend/mcp_server/transport.py index 8998878162..cd16bb4693 100644 --- a/backend/mcp_server/transport.py +++ b/backend/mcp_server/transport.py @@ -120,13 +120,24 @@ def server_info(self) -> dict[str, Any]: } def get(self, request: Request, **kwargs: Any) -> Response: - """Advertise server identity. - - Clients probe with GET to check connectivity before opening a session. - Deliberately free of tenant detail — it reveals only that an MCP server - is mounted here. + """Refuse the SSE stream, but say who is here. + + Under Streamable HTTP a client issues GET to open a server-to-client + SSE stream, and a server that offers none must answer 405 (spec rev + 2025-06-18). Nothing here pushes messages — every tool call is + request/response — so 405 is the honest answer, and returning + ``200 application/json`` instead would leave a conformant client + parsing an identity document as an event stream. + + The body is kept anyway: uptime checks and humans with curl probe this + path, and a 405 may carry one. It stays deliberately free of tenant + detail — it reveals only that an MCP server is mounted here. """ - return Response(self.server_info()) + response = Response(self.server_info(), status=405) + # RFC 9110 requires Allow on a 405, and it tells a client which method + # this endpoint actually speaks. + response["Allow"] = "POST" + return response def post(self, request: Request, **kwargs: Any) -> JsonResponse: """Handle a single JSON-RPC request.""" From 3279a4529367ac1ff187e2c4bcd791790c382c2d Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sun, 26 Jul 2026 15:11:30 +0530 Subject: [PATCH 11/16] fix(mcp): cover the new read tools in the leak sweep, correct the GET docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_no_credential_leak drives every read tool through a hardcoded argument map and raises if a tool requires an argument the map lacks. Both new tools require project_id, so without an entry the sweep would have failed outright — and the whole reason listPrompts builds its response by hand is that ToolStudioPrompt carries credential-bearing fields. Seeds a project, a document and a prompt whose webhook URL embeds a canary, so the sweep now actually exercises the tool whose leak risk motivated the hand-built dict. Also updates the README's GET description, which still promised a 200 identity response, and the last "every MCP call is a POST" comment. Co-Authored-By: Claude Opus 5 (1M context) --- backend/mcp_server/README.md | 8 +++++-- .../tests/test_no_credential_leak.py | 22 ++++++++++++++++++- .../mcp_server/tests/test_platform_auth.py | 7 +++--- 3 files changed, 31 insertions(+), 6 deletions(-) diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index db2684d39a..4416be84c8 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -54,8 +54,12 @@ POST /deployment/api///mcp// The path key takes precedence over the header when both are present. -`GET` on either URL returns server identity for clients that probe before -connecting. It is unauthenticated and reveals nothing about the deployment. +`GET` on either URL answers `405 Method Not Allowed` with `Allow: POST`. Under +Streamable HTTP a client issues `GET` to open a server-to-client SSE stream, +and a server that offers none must decline — every tool call here is +request/response, so there is nothing to stream. The 405 still carries a small +server-identity body for uptime probes and humans with `curl`; it is +unauthenticated and reveals nothing about the deployment behind it. ## Connecting diff --git a/backend/mcp_server/tests/test_no_credential_leak.py b/backend/mcp_server/tests/test_no_credential_leak.py index 0ca3b96e71..8be0bca052 100644 --- a/backend/mcp_server/tests/test_no_credential_leak.py +++ b/backend/mcp_server/tests/test_no_credential_leak.py @@ -29,6 +29,8 @@ from django.test import TestCase, override_settings from platform_api.models import PlatformApiKey from prompt_studio.prompt_studio_core_v2.models import CustomTool +from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager +from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt from tenant_account_v2.models import OrganizationMember from utils.user_context import UserContext from workflow_manager.endpoint_v2.models import WorkflowEndpoint @@ -119,9 +121,25 @@ def setUp(self) -> None: APIDeployment.objects.create( api_name="leak-api", display_name="Leak API", workflow=self.workflow ) - CustomTool.objects.create( + self.project = CustomTool.objects.create( tool_name="Leak Prompts", description="d", author="acme" ) + DocumentManager.objects.create( + document_name="leak-doc.pdf", tool=self.project + ) + # The webhook URL carries a canary because ToolStudioPrompt is the one + # model in this sweep with a credential-bearing field that is not a + # connector or adapter — a serializer-built response would carry it + # straight out, and this is what catches that. + ToolStudioPrompt.objects.create( + prompt_key="leak_check", + prompt="What is the total?", + tool_id=self.project, + sequence_number=1, + prompt_type="Text", + enforce_type="text", + postprocessing_webhook_url=f"https://hooks.internal/x?token={LLM_API_KEY}", + ) # WorkflowExecution.save() mirrors into a Redis-backed execution cache # via its own client, which override_settings(CACHES=...) does not # cover. That mirroring is irrelevant here — the tool reads the DB — @@ -167,6 +185,8 @@ def test_no_read_tool_returns_a_credential(self) -> None: "listExecutions": {}, "getUsageSummary": {}, "getExecutionDetail": {"execution_id": str(self.execution.id)}, + "listPromptStudioDocuments": {"project_id": str(self.project.tool_id)}, + "listPrompts": {"project_id": str(self.project.tool_id)}, } checked = [] diff --git a/backend/mcp_server/tests/test_platform_auth.py b/backend/mcp_server/tests/test_platform_auth.py index bfc036e4e3..4c170ac89c 100644 --- a/backend/mcp_server/tests/test_platform_auth.py +++ b/backend/mcp_server/tests/test_platform_auth.py @@ -107,9 +107,10 @@ def test_inactive_key_is_rejected(self) -> None: def test_read_tier_key_cannot_reach_the_server(self) -> None: """A documented limitation, pinned so it cannot change silently. - The middleware's tier check gates on HTTP method and every MCP call is - a POST, so a `read` key is refused outright — even though this server - exposes nothing but read tools. + The middleware's tier check gates on HTTP method and every JSON-RPC + message arrives as a POST whatever the tool inside it does, so a `read` + key is refused outright — even though this server exposes plenty of + read tools. """ self.key.permission = "read" self.key.save() From afbbe64b5521349c2fdb63da4761b744ea31f55f Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sun, 26 Jul 2026 20:46:20 +0530 Subject: [PATCH 12/16] fix(mcp): list the id-producing tools in readMeFirst, and pin the guide to the registry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit readMeFirst's tool listing is hand-maintained prose, so it drifted the moment listPromptStudioDocuments and listPrompts were added to the registry without being added to the guide. That is a softer form of the unreachability bug those two tools were built to fix: an agent trusting the guide would find the four billable tools and no way to obtain the document_id and prompt_id they require. The reachability invariant did not catch it because it checks the registry's schemas, not this prose. Extends it to call both readMeFirst handlers and compare their listings against the registry — every registered tool must be named, no phantom tools may be advertised, and the billable grouping must match the registry's billable flags exactly, since that grouping is the guide's main signal about which calls cost money. Co-Authored-By: Claude Opus 5 (1M context) --- .../tests/test_registry_reachability.py | 112 ++++++++++++++++++ backend/mcp_server/tools/platform.py | 14 +++ 2 files changed, 126 insertions(+) diff --git a/backend/mcp_server/tests/test_registry_reachability.py b/backend/mcp_server/tests/test_registry_reachability.py index 71d65bec10..e5f58cfbfb 100644 --- a/backend/mcp_server/tests/test_registry_reachability.py +++ b/backend/mcp_server/tests/test_registry_reachability.py @@ -13,6 +13,8 @@ from __future__ import annotations +from unittest.mock import Mock + from django.test import SimpleTestCase from mcp_server.registry import DEPLOYMENT_TOOLS, PLATFORM_TOOLS @@ -127,3 +129,113 @@ def test_every_required_id_has_a_declared_producer(self) -> None: assert unproducible == {}, ( f"Unreachable on the deployment server: {unproducible}" ) + + +class ReadMeFirstMatchesTheRegistryTest(SimpleTestCase): + """``readMeFirst`` must name every tool the server actually exposes. + + This is the tool whose entire job is telling an agent how to reach the + others, and its listing is hand-maintained prose rather than a projection + of the registry — so it drifts silently. It already did: the two tools that + produce ``document_id`` and ``prompt_id`` were added to the registry and + not to the guide, leaving an agent that trusts the guide able to see the + billable tools with no way to call them. + + Kept as a listing rather than generated from the registry on purpose. The + ``purpose`` text is written for an agent deciding *whether* to call a tool + — which is worth more than the tool's own description — so the value is in + a human writing it. This test only enforces that the set is complete. + """ + + # Keys in the platform guide that enumerate tools. Named explicitly so + # adding a new grouping is a deliberate act rather than a silent hole. + PLATFORM_LISTING_KEYS = ( + "discovery_tools", + "observability_tools", + "state_change_tools", + "billable_tools", + ) + + def _named_in(self, guide: dict, keys) -> set[str]: + return { + entry["name"] + for key in keys + for entry in guide[key] + if isinstance(entry, dict) and "name" in entry + } + + def test_platform_guide_names_every_platform_tool(self) -> None: + from mcp_server.tools.platform import platform_read_me_first + + guide = platform_read_me_first( + Mock(org_name="org-guide", platform_key=Mock(permission="read_write")) + ) + listed = self._named_in(guide, self.PLATFORM_LISTING_KEYS) + + # readMeFirst does not list itself; an agent has already called it. + registered = set(PLATFORM_TOOLS.names()) - {"readMeFirst"} + + assert registered - listed == set(), ( + f"readMeFirst does not mention {sorted(registered - listed)}. An " + "agent that trusts the guide will never call them — add each to " + "the appropriate listing in platform_read_me_first." + ) + assert listed - registered == set(), ( + f"readMeFirst advertises tools that do not exist: " + f"{sorted(listed - registered)}." + ) + + def test_billable_tools_are_listed_as_billable(self) -> None: + """A costly tool listed under discovery reads as free to call. + + The grouping is the guide's main signal about consequence, so a tool + in the wrong group actively misleads rather than merely omitting. + """ + from mcp_server.tools.platform import platform_read_me_first + + guide = platform_read_me_first( + Mock(org_name="org-guide", platform_key=Mock(permission="read_write")) + ) + listed_billable = self._named_in(guide, ("billable_tools",)) + actually_billable = { + name for name in PLATFORM_TOOLS.names() if PLATFORM_TOOLS.get(name).billable + } + + assert listed_billable == actually_billable, ( + "readMeFirst's billable_tools must match the registry's billable " + f"flags exactly. Listed: {sorted(listed_billable)}, actual: " + f"{sorted(actually_billable)}." + ) + + def test_free_tools_are_not_listed_as_billable(self) -> None: + """The producers exist so an agent can find ids *before* spending. + + Listing them as billable would push an agent to avoid the very step + that makes the billable tools callable. + """ + from mcp_server.tools.platform import platform_read_me_first + + guide = platform_read_me_first( + Mock(org_name="org-guide", platform_key=Mock(permission="read_write")) + ) + listed_discovery = self._named_in(guide, ("discovery_tools",)) + + for producer in ("listPromptStudioDocuments", "listPrompts"): + assert producer in listed_discovery, ( + f"{producer} produces the ids the billable tools need, so it " + "belongs in discovery_tools where an agent looks first." + ) + + def test_deployment_guide_names_every_deployment_tool(self) -> None: + from mcp_server.tools.info import read_me_first + + guide = read_me_first( + Mock(api=Mock(display_name="d", description="x", is_active=True)) + ) + listed = self._named_in(guide, ("tools",)) + registered = set(DEPLOYMENT_TOOLS.names()) - {"readMeFirst"} + + assert listed == registered, ( + f"Deployment readMeFirst is out of sync. Missing: " + f"{sorted(registered - listed)}, phantom: {sorted(listed - registered)}." + ) diff --git a/backend/mcp_server/tools/platform.py b/backend/mcp_server/tools/platform.py index 87edbaf73e..92a84f0deb 100644 --- a/backend/mcp_server/tools/platform.py +++ b/backend/mcp_server/tools/platform.py @@ -90,6 +90,20 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "name": "listPromptStudioProjects", "purpose": "Find where extraction prompts are authored.", }, + { + "name": "listPromptStudioDocuments", + "purpose": ( + "List a project's documents. The only source of the " + "document_id every billable tool below requires." + ), + }, + { + "name": "listPrompts", + "purpose": ( + "List a project's prompts. The only source of the " + "prompt_id fetchResponse and bulkFetchResponse require." + ), + }, { "name": "getWorkflowEndpoints", "purpose": "See how a workflow is wired (shape only, no config).", From 06b4df8a9029383634f5149c72174fbf833b1525 Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sun, 26 Jul 2026 21:36:14 +0530 Subject: [PATCH 13/16] feat(mcp): run extractions from the platform server too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Requiring a deployment key to extract bought no safety: the platform server already spends money through executePipeline and the Prompt Studio tools, and it deliberately cannot hand over a deployment key (key retrieval is excluded), so an agent was told to go get a credential it had no way to obtain. Both servers now extract, and the caller picks the blast radius — one deployment, or one credential across the organization. The work is delegated to tools.execution rather than reimplemented, so both servers share one path through validation, rate limiting and the execution helper. Only the target lookup differs: there it arrives with the credential, here it is resolved from api_name through for_user. Two divergences, both consequences of the credential: llm_profile_id is not offered. It is validated against the API key's owner, and a platform key resolves to a service account with no meaningful owner. Passing some deployment's key to satisfy that check would assert a principal the caller is not. Omitting the field is also what keeps the serializer from ever reading the api_key this context lacks — an invisible coupling, so it is pinned. getExecutionStatus re-checks the execution against the organization first. DeploymentHelper.get_execution_status does a bare lookup by id with no tenant filter, and TENANT_APPS is empty so there is no per-tenant schema to fall back on. The deployment server's API key made this moot; nothing does here. Also corrects the _is_service_account docstring, which still claimed DELETE was blocked for all API keys — untrue since full_access was added, and it contradicted platform_api/permissions.py. The conclusion it draws was right; the stated reason was not. Co-Authored-By: Claude Opus 5 (1M context) --- backend/mcp_server/README.md | 39 ++- backend/mcp_server/context.py | 17 +- backend/mcp_server/registry.py | 43 ++++ .../tests/test_no_credential_leak.py | 3 + .../tests/test_registry_reachability.py | 68 +++++- backend/mcp_server/tests/test_spend_guard.py | 3 + backend/mcp_server/tools/platform.py | 28 ++- .../mcp_server/tools/platform_execution.py | 225 ++++++++++++++++++ backend/permissions/permission.py | 25 +- 9 files changed, 428 insertions(+), 23 deletions(-) create mode 100644 backend/mcp_server/tools/platform_execution.py diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index 4416be84c8..b07b0822a0 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -16,13 +16,22 @@ authenticated and which tools they expose. | Credential | that deployment's API key | a platform API key | | Authenticated by | the view itself | `CustomAuthMiddleware` | | URL | `/deployment/api///mcp` | `/api/v1/unstract//mcp/` | -| Tools | extract, poll status | discovery + state changes | +| Tools | extract, poll status | extract, plus discovery, observability, state changes and Prompt Studio | The split is not cosmetic. A deployment key grants exactly one workflow and resolves to no user, so it cannot authorize anything organization-wide; a platform key resolves to a service-account user and is checked by the shared auth middleware. Neither key works on the other server. +**Both servers run extraction, and that is deliberate.** They are two blast +radii for the same operation. A deployment key reaches one deployment, so a +leaked one costs one workflow — but obtaining it is a separate step, and the +platform server cannot hand it over (key retrieval is excluded, see below). A +platform key reaches every deployment in the organization and needs no second +credential. Requiring the narrow key bought no safety when the platform server +already spends money through `executePipeline` and Prompt Studio; it only made +extraction awkward to reach. Pick whichever tradeoff you want. + --- # Deployment server @@ -174,6 +183,7 @@ Platform API keys are managed at `/api/v1/unstract//platform-api/keys/`. | `listExecutions` | read | Recent runs and their status. Start here to debug. | | `getExecutionDetail` | read | Per-file results and errors for one run. | | `getUsageSummary` | read | Tokens and cost recorded so far. | +| `getExecutionStatus` | read | Poll an extraction started by `extractDocument`. | **State changes** — cheap, reversible: @@ -191,11 +201,32 @@ Platform API keys are managed at `/api/v1/unstract//platform-api/keys/`. | `fetchResponse` | `read_write` | Run one prompt against an indexed document. | | `bulkFetchResponse` | `read_write` | Run several prompts in one pass. | | `singlePassExtraction` | `read_write` | Run a project's whole prompt set. | +| `extractDocument` | `read_write` | Run a named deployment's extraction workflow. | To extract a document through a *deployed* API, an agent calls -`listApiDeployments` to find an `api_name`, then opens a **separate** session -against the deployment server above. The Prompt Studio tools here are for -working with prompts before they are deployed. +`listApiDeployments` for an `api_name`, then `extractDocument` with that name — +no second credential and no second session. Polling is free: `getExecutionStatus` +is a read tool, so an agent waiting on a result never spends budget to check on +it. The Prompt Studio tools are for working with prompts *before* they are +deployed; `extractDocument` runs what is already deployed. + +Two things differ from the same tool on the deployment server, both because the +caller is org-scoped rather than deployment-scoped: + +- **`api_name` is a required argument.** The deployment server learns its target + from the credential; here the agent names it. +- **`llm_profile_id` is not offered.** That argument is validated against the + API key's owner, and a platform key resolves to a service account with no + meaningful owner to check against. Supplying some deployment's key to satisfy + the check would assert a principal the caller is not, so the argument is + dropped instead. Pinned by `test_registry_reachability.py`, because the + coupling is otherwise invisible: the field is what would make the serializer + read an `api_key` this context does not have. + +`getExecutionStatus` re-checks the execution against the organization before +returning it. `DeploymentHelper.get_execution_status` does a bare lookup by id +with no tenant filter, and this deployment runs one shared schema rather than a +schema per tenant, so the check is made here rather than assumed. ## The spend guard diff --git a/backend/mcp_server/context.py b/backend/mcp_server/context.py index ab1f8058b7..2ed88fa7b0 100644 --- a/backend/mcp_server/context.py +++ b/backend/mcp_server/context.py @@ -46,13 +46,22 @@ class MCPContext: Attributes: api: The API deployment this MCP server session is scoped to. - api_key: The validated API key used to authenticate the request. - Carried through because downstream execution helpers record it - against the execution. + api_key: The validated API key used to authenticate the request, or + ``None`` when the caller reached this deployment with an + organization-scoped platform key instead. + + It is read in exactly one place: validating that an + ``llm_profile_id`` belongs to the key's owner + (``ExecutionRequestSerializer.validate_llm_profile_id``). The + platform server therefore does not offer that argument, so the + validator never runs and the ``None`` is never dereferenced — see + ``tools/platform_execution``. Supplying some other deployment's key + to fill the field would make that ownership check pass against a + principal the caller is not. org_name: Organization identifier taken from the URL, matching the value already placed in the state store by the auth layer. """ api: APIDeployment - api_key: str + api_key: str | None org_name: str diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index d5a6ab4544..58df20d5e8 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -214,6 +214,12 @@ def build_platform_registry() -> MCPToolRegistry: set_pipeline_active_schema, whoami, ) + from mcp_server.tools.platform_execution import ( + get_platform_execution_status, + platform_execution_status_schema, + platform_extract_document, + platform_extract_document_schema, + ) from mcp_server.tools.prompt_studio import ( bulk_fetch_response, bulk_fetch_response_schema, @@ -544,6 +550,43 @@ def build_platform_registry() -> MCPToolRegistry: billable=True, ) ) + registry.register( + MCPTool( + name="extractDocument", + description=( + "Run a named API deployment's extraction workflow over one or " + "more documents.\n\n" + "**Costs money** and consumes the organization's extraction " + "quota. Call listApiDeployments first for a valid `api_name`; " + "the deployment fixes the prompts and output schema, so you " + "supply documents, not instructions.\n\n" + "The same operation is available on a deployment-scoped MCP " + "server, where the credential reaches only one deployment. " + "Use that one if you want a narrower blast radius.\n\n" + f"{_ORG_WIDE_WARNING}" + ), + input_schema=platform_extract_document_schema(), + handler=platform_extract_document, + writes=True, + required_method="POST", + billable=True, + ) + ) + registry.register( + MCPTool( + name="getExecutionStatus", + description=( + "Poll for the result of an extraction started by " + "extractDocument. Free to call.\n\n" + "If extractDocument already returned execution_status " + "COMPLETED, the result is in that response and no polling is " + "needed. Otherwise poll this until COMPLETED or ERROR, " + "pausing a few seconds between calls." + ), + input_schema=platform_execution_status_schema(), + handler=get_platform_execution_status, + ) + ) return registry diff --git a/backend/mcp_server/tests/test_no_credential_leak.py b/backend/mcp_server/tests/test_no_credential_leak.py index 8be0bca052..5ec5398137 100644 --- a/backend/mcp_server/tests/test_no_credential_leak.py +++ b/backend/mcp_server/tests/test_no_credential_leak.py @@ -187,6 +187,9 @@ def test_no_read_tool_returns_a_credential(self) -> None: "getExecutionDetail": {"execution_id": str(self.execution.id)}, "listPromptStudioDocuments": {"project_id": str(self.project.tool_id)}, "listPrompts": {"project_id": str(self.project.tool_id)}, + # Seeded against this org's workflow, so the tool gets past its + # org-scope check and actually renders a response to scan. + "getExecutionStatus": {"execution_id": str(self.execution.id)}, } checked = [] diff --git a/backend/mcp_server/tests/test_registry_reachability.py b/backend/mcp_server/tests/test_registry_reachability.py index e5f58cfbfb..ffe165fdd5 100644 --- a/backend/mcp_server/tests/test_registry_reachability.py +++ b/backend/mcp_server/tests/test_registry_reachability.py @@ -33,9 +33,19 @@ "pipeline_id": "listPipelines", "workflow_id": "listWorkflows", "api_id": "listApiDeployments", + # listExecutions produces ids for past runs; extractDocument returns the id + # of the run it just started. Either is a valid handle for polling. "execution_id": "listExecutions", + # Not an id by name, but just as unguessable: an agent cannot invent a + # deployment's api_name any more than it can invent a UUID. + "api_name": "listApiDeployments", } +# Argument names that identify a specific resource the caller must already +# know. Extends past the ``_id`` suffix because ``api_name`` is an opaque +# handle in every way that matters here. +_OPAQUE_SUFFIXES = ("_id", "_ids", "_name") + DEPLOYMENT_ID_PRODUCERS = { # The deployment server is scoped to a single API, and its execution id is # handed back by the extraction call itself rather than by a listing. @@ -50,7 +60,7 @@ def _required_ids(registry) -> dict[str, list[str]]: schema = registry.get(name).input_schema required = set(schema.get("required", [])) for argument in schema.get("properties", {}): - if argument in required and argument.endswith(("_id", "_ids")): + if argument in required and argument.endswith(_OPAQUE_SUFFIXES): consumed.setdefault(argument, []).append(name) return consumed @@ -239,3 +249,59 @@ def test_deployment_guide_names_every_deployment_tool(self) -> None: f"Deployment readMeFirst is out of sync. Missing: " f"{sorted(registered - listed)}, phantom: {sorted(listed - registered)}." ) + + +class PlatformExtractionShapeTest(SimpleTestCase): + """Invariants of extraction on the organization-scoped server. + + Both servers run the same extraction. What differs is the credential, and + two consequences of that are enforced here because neither is visible at + the call site. + """ + + def test_platform_extraction_does_not_offer_an_llm_profile(self) -> None: + """The coupling that keeps a missing api_key from ever being read. + + ``ExecutionRequestSerializer.validate_llm_profile_id`` resolves the + profile against the API key's owner. A platform caller has no + deployment key — ``MCPContext.api_key`` is None — and the only reason + that is safe is that this argument is never offered, so DRF never runs + that validator. Adding the field back would surface as "Unable to + validate LLM profile ownership" from a code path with no obvious link + to this schema. + """ + schema = PLATFORM_TOOLS.get("extractDocument").input_schema + + assert "llm_profile_id" not in schema["properties"], ( + "llm_profile_id is validated against the API key's owner, which a " + "platform key does not have. Offering it here would either break " + "at runtime or require borrowing another principal's key." + ) + + def test_platform_extraction_still_mirrors_the_deployment_limits(self) -> None: + """The platform schema is derived from the deployment one, so the file + cap and tag rules cannot drift between the two servers. + """ + platform = PLATFORM_TOOLS.get("extractDocument").input_schema["properties"] + deployment = DEPLOYMENT_TOOLS.get("extractDocument").input_schema["properties"] + + assert platform["document_urls"] == deployment["document_urls"] + assert platform["timeout"] == deployment["timeout"] + assert platform["tags"] == deployment["tags"] + + def test_platform_extraction_requires_naming_a_deployment(self) -> None: + """The deployment server gets its target from the credential; here the + agent must name one, so it is required rather than optional. + """ + schema = PLATFORM_TOOLS.get("extractDocument").input_schema + + assert "api_name" in schema["required"] + assert "document_urls" in schema["required"] + + def test_polling_is_free_but_extracting_is_not(self) -> None: + """An agent told that polling costs money will avoid polling, which is + the one thing it should do freely while waiting. + """ + assert PLATFORM_TOOLS.get("extractDocument").billable is True + assert PLATFORM_TOOLS.get("getExecutionStatus").billable is False + assert PLATFORM_TOOLS.get("getExecutionStatus").required_method == "GET" diff --git a/backend/mcp_server/tests/test_spend_guard.py b/backend/mcp_server/tests/test_spend_guard.py index 621cf730a8..16392eeac8 100644 --- a/backend/mcp_server/tests/test_spend_guard.py +++ b/backend/mcp_server/tests/test_spend_guard.py @@ -322,6 +322,9 @@ def test_every_costly_platform_tool_is_marked_billable(self) -> None: "fetchResponse", "bulkFetchResponse", "singlePassExtraction", + # The most expensive tool on the server: it runs a whole + # extraction workflow and consumes the org's extraction quota. + "extractDocument", } unguarded = [ diff --git a/backend/mcp_server/tools/platform.py b/backend/mcp_server/tools/platform.py index 92a84f0deb..1063e92c5f 100644 --- a/backend/mcp_server/tools/platform.py +++ b/backend/mcp_server/tools/platform.py @@ -127,6 +127,13 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "name": "getUsageSummary", "purpose": "Tokens and cost recorded so far.", }, + { + "name": "getExecutionStatus", + "purpose": ( + "Poll an extraction started by extractDocument. Free to " + "call, unlike the extraction itself." + ), + }, ], "state_change_tools": [ { @@ -159,6 +166,13 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "name": "singlePassExtraction", "purpose": "Run a project's whole prompt set. Most expensive.", }, + { + "name": "extractDocument", + "purpose": ( + "Run a deployment's extraction workflow over documents. " + "Needs an api_name from listApiDeployments." + ), + }, ], "before_writing": ( "State-change tools take effect immediately for the whole " @@ -178,11 +192,15 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "deployments or projects, are deliberately not exposed here." ), "to_run_an_extraction": ( - "This server cannot extract. Call listApiDeployments to get an " - "api_name, then open a separate MCP session against that " - "deployment's own endpoint — /deployment/api///mcp " - "— using that deployment's API key. That session exposes " - "extractDocument." + "Call listApiDeployments for an api_name, then extractDocument " + "with that name and your document URLs. If the response is not " + "already COMPLETED, poll getExecutionStatus with the returned " + "execution_id.\n\n" + "The same extraction is also available on a deployment-scoped " + "server at /deployment/api///mcp, which takes that " + "one deployment's own API key and can reach nothing else. Both " + "run the identical workflow — pick this server for one credential " + "across the organization, or that one for a narrower blast radius." ), "scope_warning": ( "This credential is a service account, which reaches ALL resources " diff --git a/backend/mcp_server/tools/platform_execution.py b/backend/mcp_server/tools/platform_execution.py new file mode 100644 index 0000000000..c7f4f254af --- /dev/null +++ b/backend/mcp_server/tools/platform_execution.py @@ -0,0 +1,225 @@ +"""Extraction from the organization-scoped server. + +The deployment server and this one now both run extractions. They are not +redundant: they are two blast radii for the same operation, and the caller +picks. A deployment key reaches exactly one deployment, so a leaked one costs +one workflow; the platform key reaches every deployment in the organization, +and buys the caller not having to obtain a second credential to extract at all. + +The work itself is delegated to ``tools.execution`` rather than reimplemented, +so both servers share one code path through validation, rate limiting and the +execution helper. What differs is only how the target deployment is found: +there it arrives with the credential, here it is looked up by ``api_name``. + +Two deliberate divergences from the deployment server: + +* **No ``llm_profile_id``.** On the deployment server that argument is checked + against the API key's owner (``ExecutionRequestSerializer.validate_llm_profile_id``). + A platform caller is a service account with org-wide reach, so there is no + meaningful owner to check it against, and borrowing some deployment key to + satisfy the check would be asserting a principal this caller is not. The + argument is therefore not offered here. This also keeps the serializer from + ever needing the ``api_key`` that this context does not have — omitting the + field means DRF never runs that validator. +* **Execution status is scoped to the organization.** See + ``get_platform_execution_status``. +""" + +import logging +from typing import Any + +from api_v2.models import APIDeployment + +from mcp_server.context import MCPContext, PlatformMCPContext +from mcp_server.exceptions import MCPToolError +from mcp_server.tools.execution import ( + DEFAULT_TIMEOUT_SEC, + extract_document, + extract_document_schema, +) + +logger = logging.getLogger(__name__) + + +def _resolve_deployment(context: PlatformMCPContext, api_name: str) -> APIDeployment: + """Find an API deployment the caller may reach, or refuse. + + ``for_user`` is what keeps this inside the caller's organization. The + platform key authenticates as a service account, for which that manager + returns every deployment in the org — which is the intended reach here — + but it is still the boundary that stops another organization's deployment + being named. + """ + deployment = ( + APIDeployment.objects.for_user(context.user).filter(api_name=api_name).first() + ) + if deployment is None: + raise MCPToolError( + f"No API deployment named '{api_name}' in organization " + f"'{context.org_name}'. Call listApiDeployments for valid names." + ) + return deployment + + +def _as_deployment_context( + context: PlatformMCPContext, deployment: APIDeployment +) -> MCPContext: + """Adapt a platform context to the one ``tools.execution`` expects. + + ``api_key`` is deliberately ``None``. It is read in exactly one place — + validating ``llm_profile_id`` ownership — and this tool does not accept + that argument, so the serializer never invokes that validator. Passing a + real deployment key instead would make the ownership check pass against + whichever principal happens to own it, which is worse than not offering + the feature. + """ + return MCPContext(api=deployment, api_key=None, org_name=context.org_name) + + +def platform_extract_document_schema() -> dict[str, Any]: + """The deployment schema plus ``api_name``, minus ``llm_profile_id``. + + Derived from the deployment schema rather than restated, so the two cannot + drift on file caps, tag rules or the timeout range. + """ + schema = extract_document_schema() + properties = dict(schema["properties"]) + # See the module docstring: not offerable without a deployment key to check + # ownership against. + properties.pop("llm_profile_id", None) + + return { + "type": "object", + "properties": { + "api_name": { + "type": "string", + "description": ( + "Name of the API deployment to run, from " + "listApiDeployments. This fixes the prompts and output " + "schema — you supply documents, not instructions." + ), + }, + **properties, + }, + "required": ["api_name", *schema["required"]], + } + + +def platform_extract_document( + context: PlatformMCPContext, + api_name: str, + document_urls: list[str], + timeout: int = DEFAULT_TIMEOUT_SEC, + include_metadata: bool = False, + include_metrics: bool = False, + include_extracted_text: bool = False, + tags: list[str] | None = None, +) -> dict[str, Any]: + """Run a named deployment's workflow over the supplied document URLs.""" + deployment = _resolve_deployment(context, api_name) + logger.info( + f"MCP platform extract_document api='{api_name}' " + f"(org '{context.org_name}', key '{context.platform_key.name}')" + ) + + result = extract_document( + _as_deployment_context(context, deployment), + document_urls=document_urls, + timeout=timeout, + include_metadata=include_metadata, + include_metrics=include_metrics, + include_extracted_text=include_extracted_text, + tags=tags, + ) + # The agent named a deployment rather than being scoped to one, so echo + # back which one actually ran. + return {"api_name": api_name, **result} + + +def platform_execution_status_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "execution_id": { + "type": "string", + "description": ( + "The execution_id returned by a previous extractDocument call." + ), + }, + "include_metadata": { + "type": "boolean", + "default": False, + "description": "Include extraction metadata in the result.", + }, + "include_metrics": { + "type": "boolean", + "default": False, + "description": "Include token and cost metrics in the result.", + }, + "include_extracted_text": { + "type": "boolean", + "default": False, + "description": ( + "Include the full raw text extracted from each document. " + "This can be very large." + ), + }, + }, + "required": ["execution_id"], + } + + +def get_platform_execution_status( + context: PlatformMCPContext, + execution_id: str, + include_metadata: bool = False, + include_metrics: bool = False, + include_extracted_text: bool = False, +) -> dict[str, Any]: + """Return the status, and once ready the result, of an extraction. + + The execution id is checked against the organization *before* the status + is fetched. ``DeploymentHelper.get_execution_status`` resolves a bare + ``WorkflowExecution.objects.get(id=...)`` with no tenant filter, and this + deployment runs a single shared schema rather than one per tenant, so + without this check a caller could poll an execution id belonging to + another organization. On the deployment server the API key narrowed the + caller to one deployment; here nothing does, so the check is made + explicitly. + """ + from workflow_manager.workflow_v2.models.execution import WorkflowExecution + + from mcp_server.tools.execution import get_execution_status + from mcp_server.tools.observability import _org_workflow_ids + + visible = _org_workflow_ids(context) + execution = ( + WorkflowExecution.objects.filter(id=execution_id, workflow_id__in=visible) + .only("id", "workflow_id") + .first() + ) + if execution is None: + raise MCPToolError( + f"No execution with id '{execution_id}' in organization " + f"'{context.org_name}'. Poll only ids returned by extractDocument, " + "or call listExecutions." + ) + + deployment = ( + APIDeployment.objects.for_user(context.user) + .filter(workflow_id=execution.workflow_id) + .first() + ) + if deployment is None: + raise MCPToolError( + f"Execution '{execution_id}' did not come from an API deployment. " + "Use getExecutionDetail for pipeline and ETL runs." + ) + + return get_execution_status( + _as_deployment_context(context, deployment), + execution_id=execution_id, + include_metadata=include_metadata, + include_metrics=include_metrics, + include_extracted_text=include_extracted_text, + ) diff --git a/backend/permissions/permission.py b/backend/permissions/permission.py index 0c0a3b88eb..7c7557791c 100644 --- a/backend/permissions/permission.py +++ b/backend/permissions/permission.py @@ -11,15 +11,22 @@ def _is_service_account(request: Request) -> bool: - """Allow service accounts through for all non-DELETE methods. - - Two constraints are enforced upstream in the authentication middleware - before this check is ever reached: - 1. DELETE is blocked for all API keys. - 2. Write methods (POST/PUT/PATCH) are blocked for READ-only API keys. - - Therefore any service-account request that arrives here is permitted to - proceed regardless of HTTP method. + """Allow service accounts through regardless of HTTP method. + + The authentication middleware has already checked the key's permission + tier against the request method (``ApiKeyPermission.allows``) before this + is reached, so the method itself needs no second look here: + 1. ``read`` keys reach only GET/HEAD/OPTIONS. + 2. ``read_write`` keys add POST/PUT/PATCH. + 3. ``full_access`` keys add DELETE. + + Note the third case: an earlier version of this docstring said DELETE was + blocked for all API keys, which stopped being true when ``full_access`` + was introduced. A ``full_access`` service-account request does reach this + check on DELETE and is allowed through — object-level ownership is not + consulted, because a service account is org-wide by design. That tier is + the only thing standing between an API key and deleting any resource in + the organization, so grant it deliberately. """ return getattr(request.user, "is_service_account", False) From 5345c2eebda2299d3bd3b6a25e5046b088ca45fd Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sun, 26 Jul 2026 23:05:57 +0530 Subject: [PATCH 14/16] test(mcp): pin the org boundary on platform extraction, keep the leak sweep off Redis MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The org-scope check in get_platform_execution_status is the one genuinely new security boundary this change creates — on the deployment server the API key made it moot — and nothing exercised it. Adds the refusal cases: another tenant's execution, an unknown id, a pipeline run with no deployment behind it, and extraction naming a deployment outside the org. The cross-tenant test also asserts the refusal does not name the owning organization, since confirming an id exists elsewhere is itself a disclosure. The leak sweep now invokes getExecutionStatus, and the seeded execution has status ERROR — which ExecutionStatus counts as terminal, so the tool takes its is_completed branch and reads results through a raw Redis client that override_settings(CACHES=...) does not reach. Stubbed the same way _handle_execution_cache already is, so the sweep stays in the unit tier's reach rather than failing in CI for a reason unrelated to credentials. Co-Authored-By: Claude Opus 5 (1M context) --- .../tests/test_no_credential_leak.py | 14 +- .../mcp_server/tests/test_platform_tools.py | 132 ++++++++++++++++++ 2 files changed, 145 insertions(+), 1 deletion(-) diff --git a/backend/mcp_server/tests/test_no_credential_leak.py b/backend/mcp_server/tests/test_no_credential_leak.py index 5ec5398137..23698668b9 100644 --- a/backend/mcp_server/tests/test_no_credential_leak.py +++ b/backend/mcp_server/tests/test_no_credential_leak.py @@ -193,6 +193,17 @@ def test_no_read_tool_returns_a_credential(self) -> None: } checked = [] + # The seeded execution has status ERROR, which ExecutionStatus counts + # as terminal — so getExecutionStatus takes its is_completed branch and + # reads results through a raw Redis client that + # override_settings(CACHES=...) does not reach, exactly like + # _handle_execution_cache above. Stubbed to keep this test off live + # infra; the cached results are not what is being scanned here. + result_cache = patch( + "workflow_manager.endpoint_v2.result_cache_utils." + "ResultCacheUtils.get_api_results", + return_value=[], + ) for name, tool in self._read_tools(): kwargs = arguments.get(name, {}) required = tool.input_schema.get("required", []) @@ -204,7 +215,8 @@ def test_no_read_tool_returns_a_credential(self) -> None: f"no fixture for it — add one so it stays covered." ) - result = tool.handler(self.context, **kwargs) + with result_cache: + result = tool.handler(self.context, **kwargs) blob = json.dumps(result, default=str) checked.append(name) diff --git a/backend/mcp_server/tests/test_platform_tools.py b/backend/mcp_server/tests/test_platform_tools.py index ca70e084e7..6ec102f8be 100644 --- a/backend/mcp_server/tests/test_platform_tools.py +++ b/backend/mcp_server/tests/test_platform_tools.py @@ -8,6 +8,7 @@ from __future__ import annotations import json +from typing import Any from unittest.mock import Mock, patch from account_v2.models import Organization, User @@ -20,10 +21,15 @@ from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt from tenant_account_v2.models import OrganizationMember from utils.user_context import UserContext +from workflow_manager.workflow_v2.models.execution import WorkflowExecution from workflow_manager.workflow_v2.models.workflow import Workflow from mcp_server.context import PlatformMCPContext from mcp_server.exceptions import MCPToolError +from mcp_server.tools.platform_execution import ( + get_platform_execution_status, + platform_extract_document, +) from mcp_server.tools.platform import ( execute_pipeline, list_api_deployments, @@ -397,3 +403,129 @@ def test_documents_do_not_leak_across_projects(self) -> None: names = [row["document_name"] for row in result["documents"]] assert names == ["invoice-001.pdf"], "a sibling project's document leaked" + + +class PlatformExtractionScopeTest(TestCase): + """The organization boundary on platform-side extraction. + + On the deployment server the API key narrowed a caller to one deployment, + so an execution id from elsewhere was unreachable by construction. Nothing + narrows a platform caller, and ``get_status_of_async_task`` resolves a bare + ``WorkflowExecution.objects.get(id=...)`` with no tenant filter — this + deployment runs one shared schema (``TENANT_APPS`` is empty), so there is + no per-tenant database to fall back on. These tests pin the check that + replaces the one the credential used to provide. + """ + + def setUp(self) -> None: + self.org = Organization.objects.create( + name="org-extract", + display_name="Extract Org", + organization_id="org-extract", + ) + UserContext.set_organization_identifier("org-extract") + self.user = User.objects.create( + username="svc-extract", + email="svc-extract@platform.internal", + user_id="uid-extract", + is_service_account=True, + ) + OrganizationMember.objects.create( + user=self.user, organization=self.org, role="user" + ) + self.key = PlatformApiKey.objects.create( + name="extract-key", + description="d", + organization=self.org, + api_user=self.user, + permission="read_write", + ) + self.workflow = Workflow.objects.create( + workflow_name="extract-wf", is_active=True + ) + self.api = APIDeployment.objects.create( + api_name="extract-api", display_name="Extract API", workflow=self.workflow + ) + self.context = PlatformMCPContext( + user=self.user, + platform_key=self.key, + org_name="org-extract", + request=Mock(data={}), + ) + + def _execution_for(self, workflow) -> Any: + with patch( + "workflow_manager.workflow_v2.models.execution." + "WorkflowExecution._handle_execution_cache" + ): + return WorkflowExecution.objects.create( + workflow_id=workflow.id, status="PENDING" + ) + + def test_another_organizations_execution_is_not_readable(self) -> None: + """The boundary this change actually creates. A platform key must not + be able to poll an execution belonging to another tenant. + """ + UserContext.set_organization_identifier("org-extract-other") + Organization.objects.create( + name="org-extract-other", + display_name="Other", + organization_id="org-extract-other", + ) + other_wf = Workflow.objects.create(workflow_name="other-wf", is_active=True) + APIDeployment.objects.create( + api_name="other-api", display_name="Other API", workflow=other_wf + ) + foreign = self._execution_for(other_wf) + + UserContext.set_organization_identifier("org-extract") + + with self.assertRaises(MCPToolError) as caught: + get_platform_execution_status(self.context, execution_id=str(foreign.id)) + + # The message must not confirm the id exists elsewhere. + assert "org-extract-other" not in str(caught.exception) + + def test_an_unknown_execution_is_refused(self) -> None: + with self.assertRaises(MCPToolError): + get_platform_execution_status( + self.context, execution_id="00000000-0000-0000-0000-000000000000" + ) + + def test_a_pipeline_run_is_refused_with_a_pointer_to_the_right_tool(self) -> None: + """The second gate. A workflow in this org that is not behind an API + deployment has no deployment to build a context from, so it must be + refused rather than fail deeper with an unhelpful error. + """ + pipeline_wf = Workflow.objects.create( + workflow_name="pipeline-only-wf", is_active=True + ) + execution = self._execution_for(pipeline_wf) + + with self.assertRaises(MCPToolError) as caught: + get_platform_execution_status(self.context, execution_id=str(execution.id)) + + assert "getExecutionDetail" in str(caught.exception) + + def test_extraction_refuses_a_deployment_outside_the_organization(self) -> None: + UserContext.set_organization_identifier("org-extract-far") + Organization.objects.create( + name="org-extract-far", + display_name="Far", + organization_id="org-extract-far", + ) + far_wf = Workflow.objects.create(workflow_name="far-wf", is_active=True) + APIDeployment.objects.create( + api_name="far-api", display_name="Far API", workflow=far_wf + ) + + UserContext.set_organization_identifier("org-extract") + + with self.assertRaises(MCPToolError) as caught: + platform_extract_document( + self.context, + api_name="far-api", + document_urls=["https://s3.example.com/doc.pdf?X-Amz-Signature=x"], + ) + + assert "listApiDeployments" in str(caught.exception) From dbd85e34642f342c16ce5a1129aa119d2c630ba5 Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Sun, 26 Jul 2026 23:25:32 +0530 Subject: [PATCH 15/16] fix(mcp): validate a billable tool's target before claiming its budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The budget is consumed on invocation and never refunded — right for a call that may have spent tokens upstream, wrong for one refused on an argument the server could check for free. extractDocument made this visible: api_name is caller-supplied prose rather than an id copied from a listing, so an agent guessing at deployment names could exhaust the organization's window without ever reaching an LLM. Adds an optional preflight to MCPTool, run after the permission check and before the budget is claimed, and gives every billable tool one that resolves what it names. The handlers already did these lookups; they just did them on the far side of the counter. executePipeline was found by the new invariant rather than by inspection — it had the same shape and no preflight — so its resolution is now shared between the preflight and the handler. Pinned four ways: a failed preflight leaves the counter at zero, repeated bad names never exhaust the window, a passing preflight still consumes budget (the hook must not become a way around the guard), and an exhausted budget still reports the bad name rather than masking it — that last one is what distinguishes running preflight *before* the claim from merely reporting after it. Reversing the two in transport.py fails three of them. Co-Authored-By: Claude Opus 5 (1M context) --- backend/mcp_server/README.md | 11 ++ backend/mcp_server/registry.py | 23 ++++ backend/mcp_server/tests/test_spend_guard.py | 122 ++++++++++++++++++ backend/mcp_server/tools/platform.py | 30 ++++- .../mcp_server/tools/platform_execution.py | 19 +++ backend/mcp_server/tools/prompt_studio.py | 16 +++ backend/mcp_server/transport.py | 38 ++++-- 7 files changed, 244 insertions(+), 15 deletions(-) diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index b07b0822a0..de48a96f40 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -236,6 +236,17 @@ deal without anyone watching — so they are budgeted per organization over a rolling window (`MCP_BILLABLE_CALL_LIMIT`, default 50 per `MCP_BILLABLE_WINDOW_SECONDS`, default one hour). +**A call that could never have run does not cost a slot.** The budget is +consumed on invocation and never refunded, which is right for a call that may +have spent tokens upstream — but wrong for one refused on an argument the +server could check for free. Every billable tool therefore declares a +`preflight` that resolves its target (the deployment, project or pipeline it +names) *before* the budget is claimed. This matters most for `extractDocument`, +whose `api_name` is caller-supplied prose rather than an id copied from a +listing: without it, an agent guessing at names could exhaust the window +without ever reaching an LLM. `test_spend_guard.py` fails if a billable tool is +registered without one. + **It counts calls, not tokens or currency.** Unstract's open-source backend records usage after the fact but has no pre-flight allowance to check against — subscription and quota enforcement live in the enterprise overlay, which this diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index 58df20d5e8..287f0a217d 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -39,6 +39,17 @@ class MCPTool: by ``mcp_server.spend_guard``. Deliberately distinct from ``writes``: a tool can mutate cheaply (pausing a schedule) or cost money without changing configuration (running an extraction). + preflight: Optional ``preflight(context, **arguments)`` run *before* the + billable budget is claimed. It exists because the budget is + consumed on invocation and never refunded, so a call refused for a + reason the caller could have been told about up front — a + deployment name that does not resolve — would otherwise burn a slot + having spent nothing. Raise ``MCPToolError`` to refuse. + + Only for checks that are cheap and certain: resolving a named + resource, not anything that could itself fail transiently. A tool + whose arguments are all opaque ids the caller got from a listing + does not need one. """ name: str @@ -48,6 +59,7 @@ class MCPTool: writes: bool = False required_method: str = "GET" billable: bool = False + preflight: Callable[..., Any] | None = None def to_mcp_schema(self) -> dict[str, Any]: """Serialize to the shape returned by `tools/list`.""" @@ -208,6 +220,7 @@ def build_platform_registry() -> MCPToolRegistry: list_workflows, no_args_schema, platform_read_me_first, + preflight_pipeline, set_api_deployment_active, set_api_deployment_active_schema, set_pipeline_active, @@ -219,6 +232,7 @@ def build_platform_registry() -> MCPToolRegistry: platform_execution_status_schema, platform_extract_document, platform_extract_document_schema, + preflight_extract_document, ) from mcp_server.tools.prompt_studio import ( bulk_fetch_response, @@ -231,6 +245,7 @@ def build_platform_registry() -> MCPToolRegistry: list_prompt_studio_documents_schema, list_prompts, list_prompts_schema, + preflight_project, single_pass_extraction, single_pass_extraction_schema, ) @@ -393,6 +408,7 @@ def build_platform_registry() -> MCPToolRegistry: writes=True, required_method="POST", billable=True, + preflight=preflight_pipeline, ) ) @@ -493,6 +509,7 @@ def build_platform_registry() -> MCPToolRegistry: writes=True, required_method="POST", billable=True, + preflight=preflight_project, ) ) registry.register( @@ -511,6 +528,7 @@ def build_platform_registry() -> MCPToolRegistry: writes=True, required_method="POST", billable=True, + preflight=preflight_project, ) ) registry.register( @@ -530,6 +548,7 @@ def build_platform_registry() -> MCPToolRegistry: writes=True, required_method="POST", billable=True, + preflight=preflight_project, ) ) registry.register( @@ -548,6 +567,7 @@ def build_platform_registry() -> MCPToolRegistry: writes=True, required_method="POST", billable=True, + preflight=preflight_project, ) ) registry.register( @@ -570,6 +590,9 @@ def build_platform_registry() -> MCPToolRegistry: writes=True, required_method="POST", billable=True, + # api_name is caller-supplied prose rather than an id from a + # listing, so resolve it before the budget is claimed. + preflight=preflight_extract_document, ) ) registry.register( diff --git a/backend/mcp_server/tests/test_spend_guard.py b/backend/mcp_server/tests/test_spend_guard.py index 16392eeac8..79e424a0d2 100644 --- a/backend/mcp_server/tests/test_spend_guard.py +++ b/backend/mcp_server/tests/test_spend_guard.py @@ -15,6 +15,7 @@ from mcp_server import spend_guard from mcp_server.context import PlatformMCPContext +from mcp_server.exceptions import MCPToolError from mcp_server.platform_views import PlatformMCPServerView from mcp_server.registry import PLATFORM_TOOLS, MCPTool @@ -207,6 +208,10 @@ def _call_billable(self): tool = replace( PLATFORM_TOOLS.get("executePipeline"), handler=lambda ctx, **kw: {"ran": True}, + # The real preflight resolves a pipeline_id against the DB, which + # this unit-tier test has no fixture for. Budget behaviour is what + # is under test here; preflight has its own tests below. + preflight=None, ) with patch.object(PLATFORM_TOOLS, "get", return_value=tool): response = self.view._call_tool( @@ -349,3 +354,120 @@ def test_billable_tools_are_also_write_gated(self) -> None: ] assert wrong == [], f"Billable tools must not be GET-tier: {wrong}" + + +@override_settings( + MCP_BILLABLE_CALL_LIMIT=2, MCP_BILLABLE_WINDOW_SECONDS=60, CACHES=LOCMEM +) +class PreflightProtectsTheBudgetTest(SimpleTestCase): + """A call that could never have run must not cost a billable slot. + + The budget is consumed on invocation and never refunded, which is right + for a call that may have spent tokens upstream — but wrong for one refused + on an argument the server could check for free. ``extractDocument`` is the + case that motivated this: ``api_name`` is caller-supplied prose rather than + an id copied from a listing, so an agent guessing at names could exhaust + the window without ever reaching an LLM. + """ + + def setUp(self) -> None: + cache.delete(f"mcp:billable:{ORG}") + self.view = PlatformMCPServerView() + self.context = context() + + def tearDown(self) -> None: + cache.delete(f"mcp:billable:{ORG}") + + def _call(self, tool: MCPTool, arguments: dict): + with patch.object(PLATFORM_TOOLS, "get", return_value=tool): + response = self.view._call_tool( + request_id=1, + params={"name": tool.name, "arguments": arguments}, + context=self.context, + ) + return json.loads(response.content) + + def _billable_tool(self, preflight) -> MCPTool: + from dataclasses import replace + + return replace( + PLATFORM_TOOLS.get("extractDocument"), + handler=lambda ctx, **kw: {"ran": True}, + preflight=preflight, + ) + + def test_a_failed_preflight_does_not_consume_budget(self) -> None: + def refuse(ctx, **kwargs): + raise MCPToolError("No API deployment named 'typo'.") + + body = self._call(self._billable_tool(refuse), {"api_name": "typo"}) + + assert body["result"]["isError"] is True + assert "typo" in body["result"]["content"][0]["text"] + assert spend_guard.peek(ORG).used == 0, ( + "a call refused before it could run must not cost a billable slot" + ) + + def test_repeated_bad_names_never_exhaust_the_window(self) -> None: + """The failure mode this exists to prevent: an agent guessing at + deployment names locking the organization out of real extractions. + """ + + def refuse(ctx, **kwargs): + raise MCPToolError("No such deployment.") + + for _ in range(5): + self._call(self._billable_tool(refuse), {"api_name": "guess"}) + + assert spend_guard.peek(ORG).used == 0 + + # The budget is still intact for a call that does resolve. + body = self._call(self._billable_tool(lambda ctx, **kw: None), {}) + assert body["result"]["isError"] is False + assert spend_guard.peek(ORG).used == 1 + + def test_a_passing_preflight_still_consumes_budget(self) -> None: + """Preflight guards the budget; it must not become a way around it.""" + body = self._call(self._billable_tool(lambda ctx, **kw: None), {}) + + assert body["result"]["isError"] is False + assert spend_guard.peek(ORG).used == 1 + + def test_preflight_runs_before_the_budget_not_merely_instead_of_it(self) -> None: + """Ordering, pinned directly. + + If the budget were claimed first and the preflight merely reported the + refusal, the counter would still have moved. Exhausting the budget and + then calling with a bad name distinguishes the two: with correct + ordering the caller is told the name is wrong, not that it is broke. + """ + + def refuse(ctx, **kwargs): + raise MCPToolError("No API deployment named 'typo'.") + + for _ in range(2): + self._call(self._billable_tool(lambda ctx, **kw: None), {}) + assert spend_guard.peek(ORG).used == 2 # exhausted + + body = self._call(self._billable_tool(refuse), {"api_name": "typo"}) + + text = body["result"]["content"][0]["text"] + assert "typo" in text, f"budget refusal masked the real problem: {text}" + + def test_every_billable_tool_has_a_preflight(self) -> None: + """Each billable tool resolves something before it spends, so each can + refuse for free. A new one without a preflight silently reintroduces + the burn-a-slot-on-a-typo behaviour. + """ + unguarded = [ + name + for name in PLATFORM_TOOLS.names() + if PLATFORM_TOOLS.get(name).billable + and PLATFORM_TOOLS.get(name).preflight is None + ] + + assert unguarded == [], ( + "These billable tools consume budget before validating their " + f"arguments: {unguarded}. Give each a preflight that resolves its " + "target, or state here why it cannot fail cheaply." + ) diff --git a/backend/mcp_server/tools/platform.py b/backend/mcp_server/tools/platform.py index 1063e92c5f..221f8fc2d0 100644 --- a/backend/mcp_server/tools/platform.py +++ b/backend/mcp_server/tools/platform.py @@ -428,6 +428,29 @@ def execute_pipeline_schema() -> dict[str, Any]: } +def _resolve_pipeline(context: PlatformMCPContext, pipeline_id: str) -> Pipeline: + """Find a pipeline the caller may reach, or refuse.""" + pipeline = Pipeline.objects.for_user(context.user).filter(id=pipeline_id).first() + if pipeline is None: + raise MCPToolError( + f"No pipeline with id '{pipeline_id}' in organization " + f"'{context.org_name}'. Call listPipelines to see valid ids." + ) + return pipeline + + +def preflight_pipeline( + context: PlatformMCPContext, pipeline_id: str, **_ignored: Any +) -> None: + """Resolve the pipeline before any budget is claimed. + + ``execute_pipeline`` resolves it again inside the handler, but that runs + after the budget is consumed and the budget is never refunded — so a stale + id would cost a billable slot for a run that never started. + """ + _resolve_pipeline(context, pipeline_id) + + def execute_pipeline(context: PlatformMCPContext, pipeline_id: str) -> dict[str, Any]: """Trigger a run of an ETL or task pipeline. @@ -437,12 +460,7 @@ def execute_pipeline(context: PlatformMCPContext, pipeline_id: str) -> dict[str, """ from pipeline_v2.manager import PipelineManager - pipeline = Pipeline.objects.for_user(context.user).filter(id=pipeline_id).first() - if pipeline is None: - raise MCPToolError( - f"No pipeline with id '{pipeline_id}' in organization " - f"'{context.org_name}'. Call listPipelines to see valid ids." - ) + pipeline = _resolve_pipeline(context, pipeline_id) if context.request is None: # The manager mutates request.data on the way to the workflow viewset, diff --git a/backend/mcp_server/tools/platform_execution.py b/backend/mcp_server/tools/platform_execution.py index c7f4f254af..b2a3432277 100644 --- a/backend/mcp_server/tools/platform_execution.py +++ b/backend/mcp_server/tools/platform_execution.py @@ -105,6 +105,25 @@ def platform_extract_document_schema() -> dict[str, Any]: } +def preflight_extract_document( + context: PlatformMCPContext, api_name: str, **_ignored: Any +) -> None: + """Resolve the named deployment before any budget is claimed. + + ``api_name`` is the one argument here an agent can plausibly get wrong: it + is a human-chosen string rather than an id copied from a listing, so a + guess or a typo is likely. The budget is never refunded, so without this + check each wrong guess would cost a billable slot for a call that reached + no LLM at all — and an agent guessing repeatedly could exhaust the window + without ever running an extraction. + + Extra keyword arguments are accepted and ignored so that adding an argument + to the tool does not require touching this function; the handler is what + validates the rest. + """ + _resolve_deployment(context, api_name) + + def platform_extract_document( context: PlatformMCPContext, api_name: str, diff --git a/backend/mcp_server/tools/prompt_studio.py b/backend/mcp_server/tools/prompt_studio.py index a844f9f94a..3e8740c53d 100644 --- a/backend/mcp_server/tools/prompt_studio.py +++ b/backend/mcp_server/tools/prompt_studio.py @@ -349,3 +349,19 @@ def list_prompts(context: PlatformMCPContext, project_id: str) -> dict[str, Any] for prompt in prompts ], } + + +def preflight_project( + context: PlatformMCPContext, project_id: str, **_ignored: Any +) -> None: + """Resolve the project before any budget is claimed. + + Every billable tool in this module already calls ``_resolve_project`` — but + inside the handler, which runs *after* the budget is consumed. A stale or + mistyped ``project_id`` would therefore cost a slot for a call that never + reached an LLM. Running the same resolution first makes that refusal free. + + Cheap and deterministic by design: a single indexed lookup the handler + repeats anyway, not work that could itself fail transiently. + """ + _resolve_project(context, project_id) diff --git a/backend/mcp_server/transport.py b/backend/mcp_server/transport.py index cd16bb4693..79cf3fa55e 100644 --- a/backend/mcp_server/transport.py +++ b/backend/mcp_server/transport.py @@ -265,15 +265,6 @@ def _call_tool( request_id, JSONRPC.UNAUTHORIZED, "Permission denied", refusal ) - # Budget is checked after permission and before the handler, so an - # unauthorized call never consumes budget. Unlike the permission - # refusal above this is temporal, so it comes back as an isError - # *result* the agent can read and retry — a protocol error would read - # as "this tool does not work" and stop it retrying later. - over_budget = self.check_spend_allowed(tool, context) - if over_budget is not None: - return rpc_result(request_id, tool_content(over_budget, is_error=True)) - arguments = params.get("arguments") or {} if not isinstance(arguments, dict): return rpc_error( @@ -283,6 +274,35 @@ def _call_tool( "'arguments' must be an object", ) + # Validate the arguments a billable tool was given *before* claiming + # budget. The budget is consumed on invocation and never refunded, so + # without this a caller naming a deployment that does not exist would + # pay a slot for a call that spent nothing upstream — and an agent + # guessing at a name could exhaust the window without ever running a + # single extraction. + if tool.preflight is not None: + try: + tool.preflight(context, **arguments) + except MCPToolError as error: + return rpc_result(request_id, tool_content(str(error), is_error=True)) + except TypeError as error: + logger.warning( + f"MCP tool '{tool_name}' preflight rejected arguments: {error}" + ) + return rpc_error( + request_id, JSONRPC.INVALID_PARAMS, "Invalid params", str(error) + ) + + # Budget is checked after permission and preflight, and before the + # handler, so neither an unauthorized call nor one that could never + # have run consumes budget. Unlike the permission refusal above this is + # temporal, so it comes back as an isError *result* the agent can read + # and retry — a protocol error would read as "this tool does not work" + # and stop it retrying later. + over_budget = self.check_spend_allowed(tool, context) + if over_budget is not None: + return rpc_result(request_id, tool_content(over_budget, is_error=True)) + try: result = tool.handler(context, **arguments) except MCPToolError as error: From 8e031e2103fab2fc17ca749b318ab592c7fc40be Mon Sep 17 00:00:00 2001 From: Hari John Kuriakose Date: Mon, 27 Jul 2026 00:43:47 +0530 Subject: [PATCH 16/16] fix(mcp): advertise 2025-06-18, and actually negotiate the version MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The server advertised protocol revision 2024-11-05 while the transport followed 2025-06-18 rules — Streamable HTTP, and answering GET with 405 rather than an SSE stream. The version string was describing a server this is not. Bumping the constant alone would have been worse than the inconsistency: the initialize handler ignored the client's requested version entirely and always returned the server's own, so every client pinned to 2024-11-05 would have been told to disconnect. The spec's rule is an echo, not an announcement — respond with the requested version when supported, and offer one of ours otherwise — so that is now implemented. 2024-11-05 stays in the supported set deliberately. The subset implemented here (initialize, tools/list, tools/call, ping) is identical across the two revisions; what differs is server-initiated streaming and session ids, neither of which this server does. Accepting the older revision therefore costs nothing and keeps working clients working. Pinned through the view as well as the pure function, so the handler is proven wired to the negotiation rather than to the constant: removing the echo fails three tests. Co-Authored-By: Claude Opus 5 (1M context) --- backend/mcp_server/README.md | 8 ++ backend/mcp_server/constants.py | 23 +++++- backend/mcp_server/tests/test_mcp_protocol.py | 73 +++++++++++++++++++ backend/mcp_server/transport.py | 28 ++++++- 4 files changed, 127 insertions(+), 5 deletions(-) diff --git a/backend/mcp_server/README.md b/backend/mcp_server/README.md index de48a96f40..052ceaa590 100644 --- a/backend/mcp_server/README.md +++ b/backend/mcp_server/README.md @@ -63,6 +63,14 @@ POST /deployment/api///mcp// The path key takes precedence over the header when both are present. +Both servers speak MCP revision **2025-06-18** (Streamable HTTP), and also +accept a client that asks for `2024-11-05`: the subset implemented here — +`initialize`, `tools/list`, `tools/call`, `ping` — is identical across the two, +so refusing the older revision would disconnect working clients for no +behavioural reason. Per the spec's negotiation rule `initialize` *echoes* the +client's revision when it is one of these and offers `2025-06-18` otherwise, +rather than announcing its own regardless. + `GET` on either URL answers `405 Method Not Allowed` with `Allow: POST`. Under Streamable HTTP a client issues `GET` to open a server-to-client SSE stream, and a server that offers none must decline — every tool call here is diff --git a/backend/mcp_server/constants.py b/backend/mcp_server/constants.py index baa838651f..049e420697 100644 --- a/backend/mcp_server/constants.py +++ b/backend/mcp_server/constants.py @@ -6,10 +6,25 @@ class MCPServer: NAME: str = "unstract" VERSION: str = "0.1.0" - # Protocol revision this transport implements. Clients that request a - # different revision are answered with this one, per the MCP spec's - # version-negotiation rules. - PROTOCOL_VERSION: str = "2024-11-05" + + # Protocol revision this transport implements, and the one offered to a + # client that asks for something unsupported. This is 2025-06-18 because + # that is the revision whose rules the transport actually follows — + # Streamable HTTP, and answering GET with 405 rather than an SSE stream. + PROTOCOL_VERSION: str = "2025-06-18" + + # Revisions this server will speak if a client asks for one by name. + # + # 2024-11-05 is kept because the request/response subset used here — the + # `initialize` handshake, `tools/list`, `tools/call` and `ping` — is + # unchanged between the two, and a client pinned to it works against this + # server today. What differs is the parts this server does not implement + # anyway (server-initiated streaming, session ids), so accepting the older + # revision costs nothing and disconnects no one. + # + # Ordered newest first: `initialize` echoes the client's revision when it + # appears here, and otherwise offers the first entry. + SUPPORTED_PROTOCOL_VERSIONS: tuple[str, ...] = ("2025-06-18", "2024-11-05") class JSONRPC: diff --git a/backend/mcp_server/tests/test_mcp_protocol.py b/backend/mcp_server/tests/test_mcp_protocol.py index 1a58ac4afb..7d55f6efd6 100644 --- a/backend/mcp_server/tests/test_mcp_protocol.py +++ b/backend/mcp_server/tests/test_mcp_protocol.py @@ -18,6 +18,7 @@ from mcp_server.context import MCPContext from mcp_server.exceptions import MCPToolError from mcp_server.registry import DEPLOYMENT_TOOLS +from mcp_server.transport import negotiate_protocol_version from mcp_server.views import MCPServerView ORG_ID = "org-mcp" @@ -91,6 +92,33 @@ def test_initialize_returns_protocol_and_server_info(self) -> None: # have to guess what the tools will operate on. assert "Invoice Extractor" in result["instructions"] + def test_initialize_echoes_the_clients_requested_revision(self) -> None: + """Through the real view, so the handler is proven wired to the + negotiation and not merely to the constant. + """ + _, body = self._call( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2024-11-05"}, + } + ) + + assert body["result"]["protocolVersion"] == "2024-11-05" + + def test_initialize_offers_our_revision_when_the_client_asks_for_junk(self) -> None: + _, body = self._call( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "1.0.0"}, + } + ) + + assert body["result"]["protocolVersion"] == MCPServer.PROTOCOL_VERSION + def test_notification_gets_no_result_body(self) -> None: """Notifications carry no id; answering one with a result would be a protocol violation that some clients treat as a fatal error. @@ -214,3 +242,48 @@ def test_batch_request_rejected(self) -> None: body = json.loads(response.content) assert body["error"]["code"] == JSONRPC.INVALID_REQUEST + + +class ProtocolVersionNegotiationTest(SimpleTestCase): + """`initialize` echoes the client's revision when it is one we speak. + + The spec's rule is an echo, not an announcement — a server that always + replies with its own preferred revision silently tells every client on an + older one to disconnect. Tested against the pure function rather than the + view so the cases stay readable; the view path is covered above. + """ + + def test_a_supported_revision_is_echoed_back(self) -> None: + for revision in MCPServer.SUPPORTED_PROTOCOL_VERSIONS: + assert negotiate_protocol_version(revision) == revision + + def test_the_older_revision_is_still_accepted(self) -> None: + """Named explicitly rather than left to the loop above. + + The request/response subset this server implements is identical across + the two revisions, so dropping the older one would disconnect working + clients for no behavioural reason. + """ + assert negotiate_protocol_version("2024-11-05") == "2024-11-05" + + def test_an_unknown_revision_is_answered_with_our_preferred_one(self) -> None: + """Not an error: the spec has the server offer what it does support and + lets the client decide whether to continue. + """ + assert negotiate_protocol_version("1.0.0") == MCPServer.PROTOCOL_VERSION + assert negotiate_protocol_version("2099-01-01") == MCPServer.PROTOCOL_VERSION + + def test_a_missing_or_malformed_version_falls_back(self) -> None: + """Clients in the wild omit the field or send the wrong type; neither + should fail the handshake before any tool is reachable. + """ + for bad in (None, 20250618, {"version": "2025-06-18"}, []): + assert negotiate_protocol_version(bad) == MCPServer.PROTOCOL_VERSION + + def test_the_preferred_revision_is_the_one_the_transport_implements(self) -> None: + """The bug this replaced: the server advertised 2024-11-05 while the + transport followed 2025-06-18 rules (405 on GET rather than an SSE + stream). The version string must track the behaviour. + """ + assert MCPServer.PROTOCOL_VERSION == "2025-06-18" + assert MCPServer.SUPPORTED_PROTOCOL_VERSIONS[0] == MCPServer.PROTOCOL_VERSION diff --git a/backend/mcp_server/transport.py b/backend/mcp_server/transport.py index 79cf3fa55e..62d44c6497 100644 --- a/backend/mcp_server/transport.py +++ b/backend/mcp_server/transport.py @@ -28,6 +28,30 @@ logger = logging.getLogger(__name__) +def negotiate_protocol_version(requested: Any) -> str: + """Choose the protocol revision to answer `initialize` with. + + The spec's rule is an echo, not an announcement: "If the server supports + the requested protocol version, it MUST respond with the same version. + Otherwise, the server MUST respond with another protocol version it + supports" (rev 2025-06-18, Lifecycle → Version Negotiation). + + So a client asking for a revision this server speaks gets that revision + back, and anything else — an unknown revision, a malformed value, or an + omitted field — is answered with this server's preferred one. The client + then decides whether it can live with that and disconnects if not, which + is why an unsupported request is not an error response here. + """ + if isinstance(requested, str) and requested in MCPServer.SUPPORTED_PROTOCOL_VERSIONS: + return requested + if requested is not None: + logger.info( + f"MCP client requested unsupported protocol version {requested!r}; " + f"offering {MCPServer.PROTOCOL_VERSION}" + ) + return MCPServer.PROTOCOL_VERSION + + def rpc_result(request_id: Any, result: Any) -> JsonResponse: """Build a JSON-RPC success response. @@ -187,7 +211,9 @@ def _dispatch( return rpc_result( request_id, { - "protocolVersion": MCPServer.PROTOCOL_VERSION, + "protocolVersion": negotiate_protocol_version( + params.get("protocolVersion") + ), "capabilities": {"tools": {"listChanged": False}}, "serverInfo": { "name": MCPServer.NAME,