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/backend/settings/base.py b/backend/backend/settings/base.py index 94071ac3b8..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 ) @@ -119,6 +125,12 @@ 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) @@ -379,6 +391,7 @@ def filter(self, record): "pipeline_v2", "platform_settings_v2", "api_v2", + "mcp_server", "usage_v2", "notification_v2", "prompt_studio.prompt_profile_manager_v2", @@ -647,7 +660,9 @@ 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}") # Whitelisting health check API diff --git a/backend/backend/urls_v2.py b/backend/backend/urls_v2.py index fd91d3cc2f..5e68db6b61 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 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 new file mode 100644 index 0000000000..052ceaa590 --- /dev/null +++ b/backend/mcp_server/README.md @@ -0,0 +1,483 @@ +# Hosted MCP servers + +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. + +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 | `/deployment/api///mcp` | `/api/v1/unstract//mcp/` | +| 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 + +## 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 /deployment/api///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 /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 +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 + +With Claude Code: + +```bash +claude mcp add --transport http unstract \ + https:///deployment/api///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 **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. + +## 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. + +--- + +# Platform server + +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 + +``` +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 + +**Discovery** — what exists: + +| Tool | Tier | Purpose | +| --- | --- | --- | +| `readMeFirst` | read | Orientation, including what this credential can reach. | +| `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. | +| `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. | + +**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. | +| `getExecutionStatus` | read | Poll an extraction started by `extractDocument`. | + +**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. | + +**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. | +| `extractDocument` | `read_write` | Run a named deployment's extraction workflow. | + +To extract a document through a *deployed* API, an agent calls +`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 + +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). + +**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 +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. + +### 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 +(`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 +*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 + +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. + +**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. + +### 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, 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 + +This server is mounted under the tenant prefix, next to `platform-api/`, and +**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` +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 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 +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`, `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. + +**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"` + 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". + +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 + +OAuth 2.1 with dynamic client registration, which MCP defines for browser-based +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/__init__.py b/backend/mcp_server/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/mcp_server/apps.py b/backend/mcp_server/apps.py new file mode 100644 index 0000000000..3447b53558 --- /dev/null +++ b/backend/mcp_server/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class MCPConfig(AppConfig): + name = "mcp_server" diff --git a/backend/mcp_server/constants.py b/backend/mcp_server/constants.py new file mode 100644 index 0000000000..049e420697 --- /dev/null +++ b/backend/mcp_server/constants.py @@ -0,0 +1,56 @@ +"""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, 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: + """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_server/context.py b/backend/mcp_server/context.py new file mode 100644 index 0000000000..2ed88fa7b0 --- /dev/null +++ b/backend/mcp_server/context.py @@ -0,0 +1,67 @@ +"""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 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. + 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) +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, 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 | None + org_name: str diff --git a/backend/mcp_server/exceptions.py b/backend/mcp_server/exceptions.py new file mode 100644 index 0000000000..3200cba49d --- /dev/null +++ b/backend/mcp_server/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_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..f20f21f56b --- /dev/null +++ b/backend/mcp_server/platform_views.py @@ -0,0 +1,165 @@ +"""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 import spend_guard +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 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 + + # 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}'. 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( + 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), + request=request, + ) + + def check_tool_allowed(self, tool: Any, context: PlatformMCPContext) -> str | None: + """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 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``. + + 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. + """ + 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 + + return ( + 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 new file mode 100644 index 0000000000..287f0a217d --- /dev/null +++ b/backend/mcp_server/registry.py @@ -0,0 +1,621 @@ +"""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. + 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 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 + 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). + 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 + description: str + input_schema: dict[str, Any] + handler: Callable[..., Any] + 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`.""" + 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_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. + """ + from mcp_server.tools.execution import ( + extract_document, + extract_document_schema, + get_execution_status, + get_execution_status_schema, + ) + from mcp_server.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 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" + "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 + + +def build_platform_registry() -> MCPToolRegistry: + """Build the tools exposed by the organization-scoped MCP server. + + 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.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, + execute_pipeline_schema, + list_api_deployments, + list_pipelines, + list_prompt_studio_projects, + list_workflows, + no_args_schema, + platform_read_me_first, + preflight_pipeline, + set_api_deployment_active, + set_api_deployment_active_schema, + set_pipeline_active, + 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, + preflight_extract_document, + ) + from mcp_server.tools.prompt_studio import ( + bulk_fetch_response, + bulk_fetch_response_schema, + fetch_response, + fetch_response_schema, + index_document, + index_document_schema, + list_prompt_studio_documents, + list_prompt_studio_documents_schema, + list_prompts, + list_prompts_schema, + preflight_project, + single_pass_extraction, + single_pass_extraction_schema, + ) + + 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, + ) + ) + # 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", + 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", + billable=True, + preflight=preflight_pipeline, + ) + ) + + # ---- 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, + preflight=preflight_project, + ) + ) + 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, + preflight=preflight_project, + ) + ) + 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, + preflight=preflight_project, + ) + ) + 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, + preflight=preflight_project, + ) + ) + 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, + # 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( + 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 + + +# 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/spend_guard.py b/backend/mcp_server/spend_guard.py new file mode 100644 index 0000000000..96ce0ce3a2 --- /dev/null +++ b/backend/mcp_server/spend_guard.py @@ -0,0 +1,210 @@ +"""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 as default_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}" + + +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.""" + + 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 = 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 + # 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) + cache = get_cache() + + 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/__init__.py b/backend/mcp_server/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/mcp_server/tests/test_mcp_auth.py b/backend/mcp_server/tests/test_mcp_auth.py new file mode 100644 index 0000000000..2b7784aefa --- /dev/null +++ b/backend/mcp_server/tests/test_mcp_auth.py @@ -0,0 +1,201 @@ +"""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_server.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"/deployment/api/{org}/{api_name}/mcp", 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_server.registry.DEPLOYMENT_TOOLS.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_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 + 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 == 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_mcp_protocol.py b/backend/mcp_server/tests/test_mcp_protocol.py new file mode 100644 index 0000000000..7d55f6efd6 --- /dev/null +++ b/backend/mcp_server/tests/test_mcp_protocol.py @@ -0,0 +1,289 @@ +"""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_server.constants import JSONRPC, MCPServer +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" +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 = DEPLOYMENT_TOOLS.get(name) + return patch.object( + 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 + ): + 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_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. + """ + 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 + + +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/tests/test_no_credential_leak.py b/backend/mcp_server/tests/test_no_credential_leak.py new file mode 100644 index 0000000000..23698668b9 --- /dev/null +++ b/backend/mcp_server/tests/test_no_credential_leak.py @@ -0,0 +1,276 @@ +"""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 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 +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 + ) + 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 — + # 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)}, + "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 = [] + # 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", []) + 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." + ) + + with result_cache: + 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 new file mode 100644 index 0000000000..4c170ac89c --- /dev/null +++ b/backend/mcp_server/tests/test_platform_auth.py @@ -0,0 +1,174 @@ +"""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 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() + + 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 = [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", + "listExecutions", + "getUsageSummary", + "setApiDeploymentActive", + "executePipeline", + "bulkFetchResponse", + } <= set(tools) + + @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..bb68381941 --- /dev/null +++ b/backend/mcp_server/tests/test_platform_tier_guard.py @@ -0,0 +1,125 @@ +"""Per-tool authorization on the platform server. + +The auth middleware checks the key's permission tier against the request's HTTP +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. +""" + +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 PLATFORM_TOOLS, 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), + org_name="org-tier", + ) + + +class PlatformTierGuardTest(SimpleTestCase): + def setUp(self) -> None: + self.view = PlatformMCPServerView() + + def test_tier_matrix(self) -> None: + """The whole authorization model in one table. + + 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_unrecognised_tier_is_refused(self) -> None: + refusal = self.view.check_tool_allowed(a_tool("GET"), context_for("wat")) + + 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. + """ + 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_no_credential_tools_are_exposed(self) -> None: + """API key creation and rotation return the secret in their response. + + 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. + """ + 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 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 new file mode 100644 index 0000000000..6ec102f8be --- /dev/null +++ b/backend/mcp_server/tests/test_platform_tools.py @@ -0,0 +1,531 @@ +"""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 + +import json +from typing import Any +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 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.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, + list_prompt_studio_projects, + list_workflows, + platform_read_me_first, + set_api_deployment_active, + set_pipeline_active, + whoami, +) +from mcp_server.tools.prompt_studio import ( + list_prompt_studio_documents, + list_prompts, +) + +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.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, + request=Mock(data={}), + ) + + 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 + # 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 + 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"] + + +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" + + +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) 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_registry_reachability.py b/backend/mcp_server/tests/test_registry_reachability.py new file mode 100644 index 0000000000..ffe165fdd5 --- /dev/null +++ b/backend/mcp_server/tests/test_registry_reachability.py @@ -0,0 +1,307 @@ +"""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 unittest.mock import Mock + +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", + # 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. + "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(_OPAQUE_SUFFIXES): + 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}" + ) + + +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)}." + ) + + +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 new file mode 100644 index 0000000000..79e424a0d2 --- /dev/null +++ b/backend/mcp_server/tests/test_spend_guard.py @@ -0,0 +1,473 @@ +"""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.exceptions import MCPToolError +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. + """ + 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 +) +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}, + # 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( + 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 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 + 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", + # The most expensive tool on the server: it runs a whole + # extraction workflow and consumes the org's extraction quota. + "extractDocument", + } + + 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}" + + +@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/tests/test_tool_errors.py b/backend/mcp_server/tests/test_tool_errors.py new file mode 100644 index 0000000000..f8a70823e2 --- /dev/null +++ b/backend/mcp_server/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_server.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_server/tests/test_tool_execution.py b/backend/mcp_server/tests/test_tool_execution.py new file mode 100644 index 0000000000..869fc5f2b3 --- /dev/null +++ b/backend/mcp_server/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_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. +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_server.tools.execution.DeploymentHelper.load_presigned_files" + ) as load_files, + patch( + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + return_value=(True, {}), + ), + patch( + "mcp_server.tools.execution.APIDeploymentRateLimiter.release_slot" + ) as release, + patch( + "mcp_server.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_server.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_server.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_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire" + ) as acquire, + patch( + "mcp_server.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_server.tools.execution.DeploymentHelper.load_presigned_files" + ) as load_files, + patch( + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + return_value=( + False, + {"current_usage": 5, "limit": 5, "limit_type": "organization"}, + ), + ), + patch( + "mcp_server.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_server.tools.execution.DeploymentHelper.load_presigned_files", + side_effect=RuntimeError("403 Forbidden"), + ), + patch( + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + return_value=(True, {}), + ), + patch( + "mcp_server.tools.execution.APIDeploymentRateLimiter.release_slot" + ) as release, + patch( + "mcp_server.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_server.tools.execution.DeploymentHelper.load_presigned_files"), + patch( + "mcp_server.tools.execution.APIDeploymentRateLimiter.check_and_acquire", + return_value=(True, {}), + ), + patch( + "mcp_server.tools.execution.APIDeploymentRateLimiter.release_slot" + ) as release, + patch( + "mcp_server.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_server.tools.execution.ExecutionQuerySerializer.is_valid", + return_value=True, + ), + patch( + "mcp_server.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_server.tools.execution.DeploymentHelper.get_execution_status", + return_value=response, + ), + patch( + "mcp_server.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_server.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_server/tools/__init__.py b/backend/mcp_server/tools/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/mcp_server/tools/execution.py b/backend/mcp_server/tools/execution.py new file mode 100644 index 0000000000..40f00bcadc --- /dev/null +++ b/backend/mcp_server/tools/execution.py @@ -0,0 +1,321 @@ +"""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 tags.serializers import TagParamsSerializer +from utils.enums import CeleryTaskState +from workflow_manager.workflow_v2.dto import ExecutionResponse + +from mcp_server.context import MCPContext +from mcp_server.exceptions import MCPToolError + +logger = logging.getLogger(__name__) + +# 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 +# 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": ( + "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": { + "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"}, + "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", + "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, []) + + 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 + ) + 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: + file_objs: list[Any] = [] + DeploymentHelper.load_presigned_files(presigned_urls, file_objs) + 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 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 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( + 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_server/tools/info.py b/backend/mcp_server/tools/info.py new file mode 100644 index 0000000000..41fc08928b --- /dev/null +++ b/backend/mcp_server/tools/info.py @@ -0,0 +1,84 @@ +"""Read-only tools describing the connected API deployment.""" + +import logging +from typing import Any + +from mcp_server.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, 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.", + ], + } + + +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_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 new file mode 100644 index 0000000000..221f8fc2d0 --- /dev/null +++ b/backend/mcp_server/tools/platform.py @@ -0,0 +1,542 @@ +"""Tools for the organization-scoped MCP server. + +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 +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 import spend_guard +from mcp_server.context import PlatformMCPContext +from mcp_server.exceptions import MCPToolError + +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": []} + + +# 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. + + 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: it discovers what exists and can " + "change the running state of those resources." + ), + "discovery_tools": [ + {"name": "whoami", "purpose": "See this credential's org, tier and budget."}, + { + "name": "listApiDeployments", + "purpose": "Find deployed extraction endpoints.", + }, + {"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.", + }, + { + "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).", + }, + { + "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.", + }, + { + "name": "getExecutionStatus", + "purpose": ( + "Poll an extraction started by extractDocument. Free to " + "call, unlike the extraction itself." + ), + }, + ], + "state_change_tools": [ + { + "name": "setApiDeploymentActive", + "purpose": "Take a deployment offline, or bring it back.", + }, + { + "name": "setPipelineActive", + "purpose": "Pause or resume a pipeline's schedule.", + }, + ], + "billable_tools": [ + { + "name": "executePipeline", + "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.", + }, + { + "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 " + "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, " + "deployments or projects, are deliberately not exposed here." + ), + "to_run_an_extraction": ( + "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 " + "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 + + budget = spend_guard.peek(context.org_name) + + return { + "organization": context.org_name, + "key_name": key.name, + "permission_tier": key.permission, + "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 " + "any resource in the organization." + ), + } + + +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), + } + + +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 _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. + + 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 = _resolve_pipeline(context, pipeline_id) + + 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/tools/platform_execution.py b/backend/mcp_server/tools/platform_execution.py new file mode 100644 index 0000000000..b2a3432277 --- /dev/null +++ b/backend/mcp_server/tools/platform_execution.py @@ -0,0 +1,244 @@ +"""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 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, + 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/mcp_server/tools/prompt_studio.py b/backend/mcp_server/tools/prompt_studio.py new file mode 100644 index 0000000000..3e8740c53d --- /dev/null +++ b/backend/mcp_server/tools/prompt_studio.py @@ -0,0 +1,367 @@ +"""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) + + +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 + ], + } + + +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 new file mode 100644 index 0000000000..62d44c6497 --- /dev/null +++ b/backend/mcp_server/transport.py @@ -0,0 +1,358 @@ +"""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 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. + + 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: + """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. + """ + 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.""" + 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": negotiate_protocol_version( + params.get("protocolVersion") + ), + "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 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: + """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", + ) + + # 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: + # 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/urls.py b/backend/mcp_server/urls.py new file mode 100644 index 0000000000..4be26f670e --- /dev/null +++ b/backend/mcp_server/urls.py @@ -0,0 +1,36 @@ +"""URLs for the deployment-scoped MCP server. + +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 /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 + +from mcp_server.views import MCPServerView + +mcp_server = MCPServerView.as_view() + +urlpatterns = [ + re_path( + r"^api/(?P[\w-]+)/(?P[\w-]+)/mcp/?$", + 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"^api/(?P[\w-]+)/(?P[\w-]+)/mcp/" + r"(?P[0-9a-fA-F-]{36})/?$", + mcp_server, + name="mcp_server_with_key", + ), +] diff --git a/backend/mcp_server/views.py b/backend/mcp_server/views.py new file mode 100644 index 0000000000..7f0552d67d --- /dev/null +++ b/backend/mcp_server/views.py @@ -0,0 +1,95 @@ +"""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 logging +from typing import Any + +from api_v2.deployment_helper import DeploymentHelper +from rest_framework.request import Request + +from mcp_server.context import MCPContext +from mcp_server.registry import DEPLOYMENT_TOOLS +from mcp_server.transport import BaseMCPView + +logger = logging.getLogger(__name__) + + +class MCPServerView(BaseMCPView): + """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. + """ + + registry = DEPLOYMENT_TOOLS + + 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 auth_failure_message(self) -> str: + return "Invalid API key or unknown API deployment" + + 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 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: + header = request.headers.get("Authorization", "") + if header.startswith("Bearer "): + api_key = header.split(" ", 1)[1].strip() + + if not api_key: + return None + + # Pin the organization before touching org-scoped managers; the + # deployment lookup below filters on it. + DeploymentHelper.validate_parameters( + 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) 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) 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 diff --git a/tests/critical_paths.yaml b/tests/critical_paths.yaml index 00429280e6..94aa29f006 100644 --- a/tests/critical_paths.yaml +++ b/tests/critical_paths.yaml @@ -66,6 +66,16 @@ 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: 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]