From 90513b939802f0ca31410be5e3101f7023fdc0b3 Mon Sep 17 00:00:00 2001 From: Ciaran Sweet Date: Fri, 12 Jun 2026 11:55:27 +0100 Subject: [PATCH 1/2] feat(security)!: refactor clients and runtime to support per-user credentials --- README.md | 76 +++++++- packages/mcp-agent/pyproject.toml | 1 + packages/mcp-agent/src/mcp_agent/main.py | 172 ++++++++++++++++-- packages/mcp-agent/src/mcp_agent/web.py | 69 ++++++- packages/mcp-agent/tests/test_agent.py | 93 +++++++++- packages/mcp-cli/src/mcp_cli/main.py | 49 ++++- packages/mcp-cli/tests/test_cli.py | 28 +++ .../src/mcp_runtime/credentials.py | 80 ++++++++ packages/mcp-runtime/src/mcp_runtime/index.py | 18 +- .../mcp-runtime/src/mcp_runtime/server.py | 30 ++- packages/mcp-runtime/tests/test_contract.py | 6 +- .../mcp-runtime/tests/test_credentials.py | 31 ++++ packages/mcp-runtime/tests/test_index.py | 2 + packages/mcp-runtime/tests/test_server.py | 17 ++ pyproject.toml | 2 + toolsets/credential-demo/pyproject.toml | 16 ++ .../src/credential_demo/__init__.py | 1 + .../src/credential_demo/tools.py | 49 +++++ .../tests/test_credential_demo.py | 31 ++++ toolsets/credential-demo/toolset.yaml | 2 + uv.lock | 20 ++ 21 files changed, 763 insertions(+), 30 deletions(-) create mode 100644 packages/mcp-cli/tests/test_cli.py create mode 100644 packages/mcp-runtime/src/mcp_runtime/credentials.py create mode 100644 packages/mcp-runtime/tests/test_credentials.py create mode 100644 toolsets/credential-demo/pyproject.toml create mode 100644 toolsets/credential-demo/src/credential_demo/__init__.py create mode 100644 toolsets/credential-demo/src/credential_demo/tools.py create mode 100644 toolsets/credential-demo/tests/test_credential_demo.py create mode 100644 toolsets/credential-demo/toolset.yaml diff --git a/README.md b/README.md index 15d0fbf..7099561 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,11 @@ tools and merge. ``` `TOOLS` is the only required export. Non-empty docstrings are enforced by a - contract test. + contract test. If a tool does I/O (HTTP, database), write it as + `async def` — `@tool` supports coroutines natively; sync tools are fine + for pure computation (the runtime runs them in a thread pool). If a tool + needs the *user's* credentials, read them from the request headers — see + [Per-user credentials](#per-user-credentials). 3. Add tests in `toolsets/my-toolset/tests/test_my_toolset.py` and run `./scripts/test`. @@ -285,6 +289,76 @@ Each Helm release owns its own Ingress for the same host and the controller merges them, so the domain's routing table tracks deploys with no central config to edit; the index's `/` path only catches what no toolset claims. +### Per-user credentials + +Tools that act on a user's behalf (with credentials that differ per calling +user) must not bake secrets into the deployment — and must not take them as +tool arguments either, or the model sees them and they land in chat history +and traces. Instead the client sends them as HTTP headers on every MCP +call, and the tool reads them at call time: + +```python +from mcp_runtime.credentials import credential_from_header + +@tool +def whoami() -> dict[str, Any]: + """Report which account the calling user's credential belongs to.""" + token = credential_from_header("x-demo-token") + ... + +TOOLS = [whoami] +CREDENTIAL_HEADERS = ["x-demo-token"] # advertised; validated by the contract test +``` + +The `CREDENTIAL_HEADERS` export is advertised in the toolset's `/health` and +in the index's `toolsets` entries, so clients know which toolset needs which +credential — and send each one *only* to the connections that declare it, +never to unrelated toolsets. `toolsets/credential-demo` is a working +(stubbed) example. Clients attach the header per connection — agents by +decorating the index's `connections` map, `mcp-cli` with `-H`: + +```python +connections = httpx.get("https:///").json()["connections"] +connections["credential-demo"]["headers"] = {"X-Demo-Token": user_token} +tools = await MultiServerMCPClient(connections).get_tools() +``` + +`mcp-agent` goes further, in the shape a multi-user deployment needs: the +agent is built **once** and credentials are supplied per call. Each +connection gets an httpx client factory that, at request time, injects the +calling user's headers — only those the toolset's advertised declaration +names (for a direct single-server URL the agent asks the endpoint's sibling +`/health` for its declaration): + +```python +from mcp_agent.main import user_credentials + +with user_credentials({"x-demo-token": the_users_token}): + result = await agent.ainvoke(...) +``` + +The Chainlit UI builds a settings field (⚙ by the message box) for every +credential header the connected toolsets advertise and applies the values +per message — so one long-lived agent process serves many users, each with +their own credentials. + +```sh +uv run mcp-cli call whoami \ + --url https:///credential-demo/mcp -H "X-Demo-Token: $TOKEN" +``` + +The secret rides the transport (TLS-encrypted at the ingress), never the +conversation, and the service stays stateless: every call carries its own +credential, so one pod serves all users. A missing header raises a +`MissingCredentialError` whose message tells the caller how to supply it. +Test credential-using tools without a server via +`mcp_runtime.credentials.header_context`: + +```python +with header_context({"x-demo-token": "secret"}): + whoami.invoke({}) +``` + ## Development ```sh diff --git a/packages/mcp-agent/pyproject.toml b/packages/mcp-agent/pyproject.toml index f881d33..62c5f45 100644 --- a/packages/mcp-agent/pyproject.toml +++ b/packages/mcp-agent/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "langchain<2.0.0,>=1.3.7", "langchain-mcp-adapters<0.4.0,>=0.3.0", "langchain-mistralai<2.0.0,>=1.1.5", + "mcp>=1.27.2,<2.0.0", "pydantic-settings<3.0.0,>=2.12.0", "rich<16.0.0,>=15.0.0", "typer<0.27.0,>=0.26.7", diff --git a/packages/mcp-agent/src/mcp_agent/main.py b/packages/mcp-agent/src/mcp_agent/main.py index 0af5c32..9426ad9 100644 --- a/packages/mcp-agent/src/mcp_agent/main.py +++ b/packages/mcp-agent/src/mcp_agent/main.py @@ -7,11 +7,16 @@ """ import asyncio +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar from typing import Annotated, Any, cast import httpx import typer from langchain.agents import create_agent +from mcp.client.streamable_http import create_mcp_http_client +from mcp.shared.exceptions import McpError from langchain_core.messages import BaseMessage, HumanMessage from langchain_core.tools import BaseTool from langchain_mcp_adapters.client import MultiServerMCPClient @@ -53,21 +58,164 @@ def connections_from(url: str, payload: Any) -> dict[str, Any]: return {"server": {"transport": "streamable_http", "url": url}} -def fetch_connections(url: str) -> dict[str, Any]: - """Resolve an index or single-server URL to a MultiServerMCPClient config.""" +def credential_headers_from(payload: Any) -> dict[str, list[str]] | None: + """Per-toolset credential header names from an index payload. + + ``None`` means the payload was not an index (a direct single-server URL), + so no declarations are available. + """ + if not (isinstance(payload, dict) and isinstance(payload.get("toolsets"), list)): + return None + return { + entry["name"]: [ + header.lower() for header in entry.get("credential_headers", []) + ] + for entry in payload["toolsets"] + if isinstance(entry, dict) and entry.get("name") + } + + +# Connection failures an agent should report rather than crash on. +CONNECT_ERRORS = (httpx.HTTPError, OSError, McpError) + + +def first_leaf(error: BaseException) -> BaseException: + """Unwrap (possibly nested) ExceptionGroups to the first real exception.""" + while isinstance(error, BaseExceptionGroup): + error = error.exceptions[0] + return error + + +def connect_error_hint(url: str) -> str: + """A nudge for the most common misconfiguration: a missing /mcp path.""" + if url.rstrip("/").endswith("/mcp"): + return "" + return ( + " Hint: single-toolset servers serve MCP under /mcp " + "(e.g. http://localhost:8000/mcp); only an index is served at the root." + ) + + +def health_url_for(url: str) -> str | None: + """Derive a direct MCP endpoint's sibling /health URL, if there is one.""" + base = url.rstrip("/") + return base.removesuffix("/mcp") + "/health" if base.endswith("/mcp") else None + + +async def single_server_credential_headers( + client: httpx.AsyncClient, url: str +) -> dict[str, list[str]] | None: + """Ask a direct MCP endpoint's /health which credential headers it reads. + + Returns ``None`` when there is no health route or it doesn't advertise + credentials (e.g. a non-mcp-toolsets server). + """ + health_url = health_url_for(url) + if health_url is None: + return None try: - payload = httpx.get(url, follow_redirects=True, timeout=10.0).json() - except (httpx.HTTPError, ValueError): - payload = None - return connections_from(url, payload) + health = (await client.get(health_url)).json() + headers = health.get("credential_headers") + except (httpx.HTTPError, ValueError, AttributeError): + return None + if not isinstance(headers, list): + return None + return {"server": [str(header).lower() for header in headers]} + + +async def fetch_connections( + url: str, +) -> tuple[dict[str, Any], dict[str, list[str]] | None]: + """Resolve a URL to a MultiServerMCPClient config plus credential needs.""" + async with httpx.AsyncClient(follow_redirects=True, timeout=10.0) as client: + try: + payload = (await client.get(url)).json() + except (httpx.HTTPError, ValueError): + payload = None + connections = connections_from(url, payload) + required = credential_headers_from(payload) + if required is None: + required = await single_server_credential_headers(client, url) + return connections, required + + +_credentials: ContextVar[dict[str, str] | None] = ContextVar( + "user_credentials", default=None +) + + +@contextmanager +def user_credentials(headers: dict[str, str] | None) -> Iterator[None]: + """Provide the calling user's credential headers for the duration. + + This is how an agent passes a user's secrets to the tools without the + model ever seeing them: they ride the MCP transport, not the conversation. + The agent is built once; wrap each turn (``run_turn``) in this and the + tool calls made inside read the values at request time, so one long-lived + agent serves many users with different credentials. + """ + token = _credentials.set(headers) + try: + yield + finally: + _credentials.reset(token) + + +def credential_client_factory(allowed: list[str] | None) -> Any: + """Build an httpx client factory injecting the current user's credentials. + + Only headers named in ``allowed`` (the toolset's advertised declaration) + are injected, so unrelated toolsets never receive them; ``None`` means no + declaration was discoverable (a server the user pointed at directly) and + every provided header is sent. + """ + wanted = None if allowed is None else {header.lower() for header in allowed} + + def factory( + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + provided = _credentials.get() or {} + send = { + header: value + for header, value in provided.items() + if wanted is None or header.lower() in wanted + } + return create_mcp_http_client( + headers={**(headers or {}), **send}, timeout=timeout, auth=auth + ) + + return factory + + +def with_credential_support( + connections: dict[str, Any], required: dict[str, list[str]] | None +) -> dict[str, Any]: + """Wire each connection to inject per-user credentials at call time.""" + return { + name: { + **connection, + "httpx_client_factory": credential_client_factory( + None if required is None else required.get(name, []) + ), + } + for name, connection in connections.items() + } async def build_agent( url: str, model: str, api_key: SecretStr ) -> tuple[Any, dict[str, Any], list[BaseTool]]: - """Discover the servers behind ``url`` and build a tool-calling agent.""" - connections = fetch_connections(url) - tools = await MultiServerMCPClient(connections).get_tools() + """Discover the servers behind ``url`` and build a tool-calling agent. + + Built once per process/session: per-user credentials are not baked in but + read from :func:`user_credentials` on every tool call. + """ + connections, required = await fetch_connections(url) + tools = await MultiServerMCPClient( + with_credential_support(connections, required) + ).get_tools() agent = create_agent( ChatMistralAI(model_name=model, api_key=api_key), tools, @@ -89,11 +237,13 @@ async def run_turn( async def chat_loop(url: str, model: str, api_key: SecretStr) -> None: try: agent, connections, tools = await build_agent(url, model, api_key) - except* (httpx.HTTPError, OSError) as group: + except* CONNECT_ERRORS as group: console.print( f"[red]Could not reach the MCP server(s) behind {url}: " - f"{group.exceptions[0]}[/red]" + f"{first_leaf(group)}[/red]" ) + if hint := connect_error_hint(url): + console.print(f"[yellow]{hint.strip()}[/yellow]") raise typer.Exit(1) from None console.print( diff --git a/packages/mcp-agent/src/mcp_agent/web.py b/packages/mcp-agent/src/mcp_agent/web.py index aaf8088..8a319b1 100644 --- a/packages/mcp-agent/src/mcp_agent/web.py +++ b/packages/mcp-agent/src/mcp_agent/web.py @@ -5,6 +5,13 @@ ``MCP_URL`` (index root or single MCP endpoint, default ``http://localhost:8000/mcp``), ``MISTRAL_MODEL`` and ``CHAINLIT_PORT`` (default 8080). + +Per-user credentials: every credential header the connected toolsets +advertise gets a field in the chat's settings panel; values are sent as HTTP +headers on the MCP calls — only to the toolsets that declared them — so the +model and the chat history never see them. The agent is built once per +session; credentials apply per message via ``user_credentials``, the same +mechanism a public multi-user API would use with one shared agent. """ import os @@ -13,18 +20,43 @@ from typing import Any import chainlit as cl +from chainlit.input_widget import InputWidget, TextInput from langchain_core.messages import BaseMessage, ToolMessage from pydantic import ValidationError -from mcp_agent.main import AgentSettings, build_agent, run_turn +from mcp_agent.main import ( + AgentSettings, + build_agent, + connect_error_hint, + fetch_connections, + first_leaf, + run_turn, + user_credentials, +) @cl.on_chat_start async def start() -> None: settings = AgentSettings() - agent, connections, tools = await build_agent( - settings.mcp_url, settings.mistral_model, settings.mistral_api_key + _, required = await fetch_connections(settings.mcp_url) + header_names = sorted( + {name for names in (required or {}).values() for name in names} ) + if header_names: + fields: list[InputWidget] = [ + TextInput(id=name, label=name) for name in header_names + ] + await cl.ChatSettings(fields).send() + try: + agent, connections, tools = await build_agent( + settings.mcp_url, settings.mistral_model, settings.mistral_api_key + ) + except Exception as error: # noqa: BLE001 - surface in the UI, not the logs + await cl.Message( + f"Could not reach the MCP server(s) behind {settings.mcp_url}: " + f"{first_leaf(error)}.{connect_error_hint(settings.mcp_url)}" + ).send() + return cl.user_session.set("agent", agent) cl.user_session.set("messages", []) await cl.Message( @@ -32,14 +64,43 @@ async def start() -> None: f"({', '.join(connections)}) with **{len(tools)}** tools: " f"{', '.join(tool.name for tool in tools)}." ).send() + if header_names: + needing = ", ".join( + f"{toolset} ({', '.join(names)})" + for toolset, names in sorted((required or {}).items()) + if names + ) + await cl.Message( + f"Some tools act on your behalf and need credentials: {needing}. " + "Set them in the settings panel (⚙ by the message box); each is " + "sent only to the toolset that declares it, never to the model." + ).send() + + +@cl.on_settings_update +async def apply_credentials(values: dict[str, Any]) -> None: + headers = { + name: value.strip() + for name, value in values.items() + if isinstance(value, str) and value.strip() + } + cl.user_session.set("credentials", headers or None) + await cl.Message("Credentials updated — your next tool calls will use them.").send() @cl.on_message async def on_message(message: cl.Message) -> None: agent = cl.user_session.get("agent") + if agent is None: + await cl.Message( + "Not connected to any MCP server — fix MCP_URL and reload the page." + ).send() + return messages: list[BaseMessage] = cl.user_session.get("messages") or [] + credentials: dict[str, str] | None = cl.user_session.get("credentials") try: - history, new_messages = await run_turn(agent, messages, message.content) + with user_credentials(credentials): + history, new_messages = await run_turn(agent, messages, message.content) except Exception as error: # noqa: BLE001 - surface in the UI, keep chatting await cl.Message(f"Error: {error}").send() return diff --git a/packages/mcp-agent/tests/test_agent.py b/packages/mcp-agent/tests/test_agent.py index fba5a61..11d9cb4 100644 --- a/packages/mcp-agent/tests/test_agent.py +++ b/packages/mcp-agent/tests/test_agent.py @@ -1,7 +1,17 @@ import pytest from pydantic import ValidationError -from mcp_agent.main import AgentSettings, connections_from +from mcp_agent.main import ( + AgentSettings, + connect_error_hint, + connections_from, + credential_client_factory, + credential_headers_from, + first_leaf, + health_url_for, + user_credentials, + with_credential_support, +) def test_settings_from_env(monkeypatch): @@ -46,3 +56,84 @@ def test_connections_from_non_index_payload_wraps_url(): } assert connections_from("http://localhost:8000/mcp", None) == expected assert connections_from("http://localhost:8000/mcp", {"status": "ok"}) == expected + + +def test_credential_headers_from_index_payload(): + payload = { + "connections": {}, + "toolsets": [ + {"name": "credential-demo", "credential_headers": ["X-Demo-Token"]}, + {"name": "dataset-search", "credential_headers": []}, + {"name": "aoi-generator"}, + ], + } + assert credential_headers_from(payload) == { + "credential-demo": ["x-demo-token"], + "dataset-search": [], + "aoi-generator": [], + } + + +def test_credential_headers_from_non_index_payload(): + assert credential_headers_from(None) is None + assert credential_headers_from({"status": "ok"}) is None + + +async def test_credential_factory_injects_only_declared_headers(): + factory = credential_client_factory(["x-demo-token"]) + with user_credentials({"X-Demo-Token": "secret", "x-other-cred": "nope"}): + client = factory(headers={"existing": "kept"}) + async with client: + assert client.headers["x-demo-token"] == "secret" + assert client.headers["existing"] == "kept" + assert "x-other-cred" not in client.headers + + +async def test_credential_factory_without_declaration_sends_all(): + factory = credential_client_factory(None) + with user_credentials({"x-demo-token": "secret", "x-other-cred": "yes"}): + client = factory() + async with client: + assert client.headers["x-demo-token"] == "secret" + assert client.headers["x-other-cred"] == "yes" + + +async def test_credential_factory_outside_context_injects_nothing(): + factory = credential_client_factory(["x-demo-token"]) + with user_credentials({"x-demo-token": "secret"}): + pass # context exited: credentials no longer available + async with factory() as client: + assert "x-demo-token" not in client.headers + + +def test_with_credential_support_wires_every_connection(): + connections = { + "credential-demo": {"transport": "streamable_http", "url": "http://a/mcp"}, + "dataset-search": {"transport": "streamable_http", "url": "http://b/mcp"}, + } + wired = with_credential_support(connections, {"credential-demo": ["x-demo-token"]}) + assert all(callable(c["httpx_client_factory"]) for c in wired.values()) + assert "httpx_client_factory" not in connections["credential-demo"] # untouched + + +def test_connect_error_hint_only_for_urls_missing_mcp_path(): + assert "under /mcp" in connect_error_hint("http://localhost:8000") + assert "under /mcp" in connect_error_hint("http://localhost:8000/") + assert connect_error_hint("http://localhost:8000/mcp") == "" + assert connect_error_hint("https://mcp.example.com/credential-demo/mcp/") == "" + + +def test_first_leaf_unwraps_nested_groups(): + error = ValueError("inner") + group = ExceptionGroup("outer", [ExceptionGroup("nested", [error])]) + assert first_leaf(group) is error + assert first_leaf(error) is error + + +def test_health_url_for(): + assert health_url_for("http://localhost:8000/mcp") == "http://localhost:8000/health" + assert ( + health_url_for("https://mcp.example.com/credential-demo/mcp/") + == "https://mcp.example.com/credential-demo/health" + ) + assert health_url_for("https://mcp.example.com/") is None diff --git a/packages/mcp-cli/src/mcp_cli/main.py b/packages/mcp-cli/src/mcp_cli/main.py index 40ea3e5..9cb23d1 100644 --- a/packages/mcp-cli/src/mcp_cli/main.py +++ b/packages/mcp-cli/src/mcp_cli/main.py @@ -22,6 +22,28 @@ console = Console() UrlOption = Annotated[str, typer.Option("--url", help="MCP server URL.")] +HeaderOption = Annotated[ + list[str] | None, + typer.Option( + "--header", + "-H", + help='HTTP header sent with every request, as "Name: value" (repeatable) — ' + "e.g. per-user credentials a tool reads server-side.", + ), +] + + +def parse_headers(pairs: list[str] | None) -> dict[str, str] | None: + """Parse ``Name: value`` header pairs (as curl's ``-H``).""" + if not pairs: + return None + headers = {} + for pair in pairs: + name, sep, value = pair.partition(":") + if not sep or not name.strip() or not value.strip(): + raise ValueError(f'expected "Name: value", got {pair!r}') + headers[name.strip()] = value.strip() + return headers def parse_tool_args(pairs: list[str]) -> dict[str, Any]: @@ -39,9 +61,11 @@ def parse_tool_args(pairs: list[str]) -> dict[str, Any]: async def _with_session[T]( - url: str, action: Callable[[ClientSession], Awaitable[T]] + url: str, + action: Callable[[ClientSession], Awaitable[T]], + headers: dict[str, str] | None = None, ) -> T: - async with streamablehttp_client(url) as (read, write, _): + async with streamablehttp_client(url, headers=headers) as (read, write, _): async with ClientSession(read, write) as session: await session.initialize() return await action(session) @@ -71,14 +95,22 @@ def _print_result(result: types.CallToolResult) -> None: console.print_json(json.dumps(result.structuredContent)) +def _parse_headers_or_exit(header: list[str] | None) -> dict[str, str] | None: + try: + return parse_headers(header) + except ValueError as error: + raise typer.BadParameter(str(error)) from error + + @app.command("list") -def list_tools(url: UrlOption = DEFAULT_URL) -> None: +def list_tools(url: UrlOption = DEFAULT_URL, header: HeaderOption = None) -> None: """List the server's tools with their argument schemas.""" + headers = _parse_headers_or_exit(header) async def action(session: ClientSession) -> list[types.Tool]: return (await session.list_tools()).tools - console.print(_tools_table(asyncio.run(_with_session(url, action)))) + console.print(_tools_table(asyncio.run(_with_session(url, action, headers)))) @app.command() @@ -88,8 +120,10 @@ def call( list[str] | None, typer.Argument(help="Arguments as key=value pairs.") ] = None, url: UrlOption = DEFAULT_URL, + header: HeaderOption = None, ) -> None: """Call a tool once and print its result.""" + headers = _parse_headers_or_exit(header) try: arguments = parse_tool_args(args or []) except ValueError as error: @@ -98,12 +132,13 @@ def call( async def action(session: ClientSession) -> types.CallToolResult: return await session.call_tool(tool, arguments) - _print_result(asyncio.run(_with_session(url, action))) + _print_result(asyncio.run(_with_session(url, action, headers))) @app.command() -def repl(url: UrlOption = DEFAULT_URL) -> None: +def repl(url: UrlOption = DEFAULT_URL, header: HeaderOption = None) -> None: """Interactive loop on a single session: `list`, ` key=value ...`, `quit`.""" + headers = _parse_headers_or_exit(header) async def action(session: ClientSession) -> None: tools = (await session.list_tools()).tools @@ -132,4 +167,4 @@ async def action(session: ClientSession) -> None: except Exception as error: # noqa: BLE001 - keep the repl alive console.print(f"[red]{error}[/red]") - asyncio.run(_with_session(url, action)) + asyncio.run(_with_session(url, action, headers)) diff --git a/packages/mcp-cli/tests/test_cli.py b/packages/mcp-cli/tests/test_cli.py new file mode 100644 index 0000000..44827df --- /dev/null +++ b/packages/mcp-cli/tests/test_cli.py @@ -0,0 +1,28 @@ +import pytest + +from mcp_cli.main import parse_headers, parse_tool_args + + +def test_parse_headers(): + assert parse_headers(None) is None + assert parse_headers([]) is None + assert parse_headers(["X-Demo-Token: secret", "Other: spaced "]) == { + "X-Demo-Token": "secret", + "Other": "spaced", + } + + +@pytest.mark.parametrize("bad", ["no-separator", ": empty-name", "empty-value:"]) +def test_parse_headers_rejects_malformed(bad): + with pytest.raises(ValueError, match="Name: value"): + parse_headers([bad]) + + +def test_parse_tool_args(): + assert parse_tool_args(["query=era5", "limit=3", "flag=true"]) == { + "query": "era5", + "limit": 3, + "flag": True, + } + with pytest.raises(ValueError, match="key=value"): + parse_tool_args(["nope"]) diff --git a/packages/mcp-runtime/src/mcp_runtime/credentials.py b/packages/mcp-runtime/src/mcp_runtime/credentials.py new file mode 100644 index 0000000..00bbbec --- /dev/null +++ b/packages/mcp-runtime/src/mcp_runtime/credentials.py @@ -0,0 +1,80 @@ +"""Read per-user credentials from the calling MCP request's HTTP headers. + +Tools that act on a user's behalf (e.g. downloading from a source that needs +the user's API key) must not bake secrets into the deployment. Instead the +MCP client sends them as HTTP headers on every call — `MultiServerMCPClient` +takes a ``headers`` dict per connection, ``mcp-cli`` takes ``--header`` — and +the tool reads them at call time with :func:`credential_from_header`. + +The credential rides the transport, so it never appears in the model context, +tool schemas, chat history or traces. Works with the stateless streamable +HTTP runtime (every tool call is its own request, carrying its own headers) +and from both sync and async tools. +""" + +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any, cast + +from mcp.server.lowlevel.server import request_ctx +from mcp.shared.context import RequestContext +from starlette.requests import Request + + +class MissingCredentialError(Exception): + """A required credential header was absent from the calling request.""" + + def __init__(self, header: str) -> None: + super().__init__( + f"missing credential: send the {header!r} HTTP header with your " + f"MCP requests (MultiServerMCPClient connections take a 'headers' " + f"dict; mcp-cli takes --header)" + ) + self.header = header + + +def credential_from_header(header: str) -> str: + """Return the named header from the MCP request that invoked this tool. + + Raises :class:`MissingCredentialError` if the header is absent or the + tool was not invoked over HTTP (e.g. called directly in tests). + """ + try: + request = request_ctx.get().request + except LookupError: + request = None + value = request.headers.get(header) if isinstance(request, Request) else None + if not value: + raise MissingCredentialError(header) + return value + + +@contextmanager +def header_context(headers: dict[str, str]) -> Iterator[None]: + """Run as if inside an MCP request carrying ``headers`` (test support). + + Lets toolset tests exercise credential-reading tools without a server: + + with header_context({"x-demo-token": "secret"}): + my_tool.invoke({...}) + """ + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (name.lower().encode(), value.encode()) for name, value in headers.items() + ], + } + context: RequestContext[Any, Any, Any] = RequestContext( + request_id=0, + meta=None, + session=cast(Any, None), + lifespan_context=None, + request=Request(scope), + ) + token = request_ctx.set(context) + try: + yield + finally: + request_ctx.reset(token) diff --git a/packages/mcp-runtime/src/mcp_runtime/index.py b/packages/mcp-runtime/src/mcp_runtime/index.py index c51c0b3..56fa3d4 100644 --- a/packages/mcp-runtime/src/mcp_runtime/index.py +++ b/packages/mcp-runtime/src/mcp_runtime/index.py @@ -44,12 +44,18 @@ class ToolsetService(NamedTuple): class ToolsetEntry(BaseModel): - """One deployed toolset in the directory.""" + """One deployed toolset in the directory. + + ``credential_headers`` names the per-user HTTP headers the toolset's + tools read; clients should send those credentials only to this toolset's + connection. + """ name: str url: str status: Literal["ok", "unreachable"] tools: list[str] + credential_headers: list[str] = [] class Connection(BaseModel): @@ -111,12 +117,18 @@ async def describe( try: response = await client.get(f"{service.base_url}/health") response.raise_for_status() - tools = response.json().get("tools", []) + health = response.json() except httpx.HTTPError: return ToolsetEntry( name=service.toolset, url=url, status="unreachable", tools=[] ) - return ToolsetEntry(name=service.toolset, url=url, status="ok", tools=tools) + return ToolsetEntry( + name=service.toolset, + url=url, + status="ok", + tools=health.get("tools", []), + credential_headers=health.get("credential_headers", []), + ) def build_app(public_url: str) -> FastAPI: diff --git a/packages/mcp-runtime/src/mcp_runtime/server.py b/packages/mcp-runtime/src/mcp_runtime/server.py index 2afdc2a..fa5f868 100644 --- a/packages/mcp-runtime/src/mcp_runtime/server.py +++ b/packages/mcp-runtime/src/mcp_runtime/server.py @@ -51,6 +51,24 @@ def load_tools(module_name: str) -> list[BaseTool]: return list(tools) +def load_credential_headers(module_name: str) -> list[str]: + """Return a tools module's optional ``CREDENTIAL_HEADERS`` export. + + Names of per-user HTTP headers the toolset's tools read (via + ``mcp_runtime.credentials``). Advertised in ``/health`` and the index so + clients attach each credential only to the toolsets that declare it. + """ + module = importlib.import_module(module_name) + headers = getattr(module, "CREDENTIAL_HEADERS", []) + if not isinstance(headers, list) or not all( + isinstance(header, str) and header for header in headers + ): + raise RuntimeError( + f"{module_name}.CREDENTIAL_HEADERS must be a list of header names" + ) + return sorted(header.lower() for header in headers) + + def build_server( toolset: str, module_name: str | None = None, @@ -58,7 +76,9 @@ def build_server( port: int = 8000, ) -> FastMCP: """Build a stateless FastMCP server exposing the toolset's TOOLS.""" - tools = load_tools(module_name or toolset_module_name(toolset)) + module_name = module_name or toolset_module_name(toolset) + tools = load_tools(module_name) + credential_headers = load_credential_headers(module_name) server = FastMCP( name=f"mcp-{toolset}", tools=[to_fastmcp(tool) for tool in tools], @@ -71,7 +91,13 @@ def build_server( @server.custom_route("/health", methods=["GET"]) async def health(request: Request) -> Response: - return JSONResponse({"status": "ok", "tools": tool_names}) + return JSONResponse( + { + "status": "ok", + "tools": tool_names, + "credential_headers": credential_headers, + } + ) return server diff --git a/packages/mcp-runtime/tests/test_contract.py b/packages/mcp-runtime/tests/test_contract.py index 0b9826f..85ba955 100644 --- a/packages/mcp-runtime/tests/test_contract.py +++ b/packages/mcp-runtime/tests/test_contract.py @@ -11,6 +11,8 @@ import pytest from langchain_core.tools import BaseTool +from mcp_runtime.server import load_credential_headers + TOOLSETS_DIR = Path(__file__).resolve().parents[3] / "toolsets" TOOLSET_NAMES = sorted( path.name for path in TOOLSETS_DIR.iterdir() if (path / "pyproject.toml").is_file() @@ -23,7 +25,9 @@ def test_toolsets_discovered(): @pytest.mark.parametrize("toolset", TOOLSET_NAMES) def test_toolset_contract(toolset): - module = importlib.import_module(toolset.replace("-", "_") + ".tools") + module_name = toolset.replace("-", "_") + ".tools" + module = importlib.import_module(module_name) + load_credential_headers(module_name) # validates the optional export's shape tools = module.TOOLS assert isinstance(tools, list) and tools, f"{toolset}: TOOLS must be non-empty" for tool in tools: diff --git a/packages/mcp-runtime/tests/test_credentials.py b/packages/mcp-runtime/tests/test_credentials.py new file mode 100644 index 0000000..1522da6 --- /dev/null +++ b/packages/mcp-runtime/tests/test_credentials.py @@ -0,0 +1,31 @@ +import pytest + +from mcp_runtime.credentials import ( + MissingCredentialError, + credential_from_header, + header_context, +) + + +def test_reads_header_case_insensitively(): + with header_context({"X-Demo-Token": "secret"}): + assert credential_from_header("x-demo-token") == "secret" + assert credential_from_header("X-Demo-Token") == "secret" + + +def test_missing_header_names_it_in_the_error(): + with header_context({"other": "value"}): + with pytest.raises(MissingCredentialError, match="x-demo-token"): + credential_from_header("x-demo-token") + + +def test_outside_a_request_raises(): + with pytest.raises(MissingCredentialError, match="x-demo-token"): + credential_from_header("x-demo-token") + + +def test_header_context_resets(): + with header_context({"x-demo-token": "secret"}): + credential_from_header("x-demo-token") + with pytest.raises(MissingCredentialError): + credential_from_header("x-demo-token") diff --git a/packages/mcp-runtime/tests/test_index.py b/packages/mcp-runtime/tests/test_index.py index 122be38..467bf15 100644 --- a/packages/mcp-runtime/tests/test_index.py +++ b/packages/mcp-runtime/tests/test_index.py @@ -70,6 +70,7 @@ async def fake_describe(client, service, public_url): url=f"{public_url}/{service.toolset}/mcp", status="ok", tools=["search_datasets"], + credential_headers=["x-demo-token"], ) monkeypatch.setattr(mcp_runtime.index, "fetch_toolset_services", fake_fetch) @@ -91,6 +92,7 @@ async def fake_describe(client, service, public_url): "url": "https://mcp.example.com/dataset-search/mcp", "status": "ok", "tools": ["search_datasets"], + "credential_headers": ["x-demo-token"], } ], } diff --git a/packages/mcp-runtime/tests/test_server.py b/packages/mcp-runtime/tests/test_server.py index c1a1cd0..1711dc8 100644 --- a/packages/mcp-runtime/tests/test_server.py +++ b/packages/mcp-runtime/tests/test_server.py @@ -1,9 +1,13 @@ +import sys +import types + import pytest from pydantic import ValidationError from mcp_runtime.server import ( RuntimeSettings, build_server, + load_credential_headers, load_tools, toolset_module_name, ) @@ -52,6 +56,19 @@ def test_load_tools_missing_export(): load_tools("mcp_runtime.server") +def test_load_credential_headers(): + assert load_credential_headers("credential_demo.tools") == ["x-demo-token"] + assert load_credential_headers("dataset_search.tools") == [] + + +def test_load_credential_headers_rejects_bad_export(monkeypatch): + module = types.ModuleType("bad_toolset_tools") + setattr(module, "CREDENTIAL_HEADERS", "x-demo-token") # noqa: B010 - not a list + monkeypatch.setitem(sys.modules, "bad_toolset_tools", module) + with pytest.raises(RuntimeError, match="CREDENTIAL_HEADERS"): + load_credential_headers("bad_toolset_tools") + + async def test_build_server_exposes_tools(): server = build_server("dataset-search") tools = await server.list_tools() diff --git a/pyproject.toml b/pyproject.toml index 341fde0..80550ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "dataset-search", "aoi-generator", "cds", + "credential-demo", ] [dependency-groups] @@ -34,6 +35,7 @@ mcp-agent = { workspace = true } dataset-search = { workspace = true } aoi-generator = { workspace = true } cds = { workspace = true } +credential-demo = { workspace = true } [tool.ruff] line-length = 88 diff --git a/toolsets/credential-demo/pyproject.toml b/toolsets/credential-demo/pyproject.toml new file mode 100644 index 0000000..53666f8 --- /dev/null +++ b/toolsets/credential-demo/pyproject.toml @@ -0,0 +1,16 @@ +[project] +name = "credential-demo" +version = "0.1.0" +description = "Example toolset demonstrating per-user credentials via HTTP headers." +requires-python = ">=3.12,<3.14" +dependencies = [ + "langchain-core<2.0.0,>=1.4.6", + "mcp-runtime", +] + +[tool.uv.sources] +mcp-runtime = { workspace = true } + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" diff --git a/toolsets/credential-demo/src/credential_demo/__init__.py b/toolsets/credential-demo/src/credential_demo/__init__.py new file mode 100644 index 0000000..b9f7a56 --- /dev/null +++ b/toolsets/credential-demo/src/credential_demo/__init__.py @@ -0,0 +1 @@ +"""credential-demo toolset.""" diff --git a/toolsets/credential-demo/src/credential_demo/tools.py b/toolsets/credential-demo/src/credential_demo/tools.py new file mode 100644 index 0000000..75568e9 --- /dev/null +++ b/toolsets/credential-demo/src/credential_demo/tools.py @@ -0,0 +1,49 @@ +"""LangChain tools demonstrating per-user credentials. + +Example of a credential-using toolset: the user's token arrives as an HTTP +header on each MCP call (never as a tool argument, so the model and the chat +history never see it) and is read with +``mcp_runtime.credentials.credential_from_header``. The "account lookup" is +stubbed — a real implementation would pass the token to an upstream API. +""" + +import hashlib +import logging +from typing import Any + +from langchain_core.tools import tool + +from mcp_runtime.credentials import MissingCredentialError, credential_from_header + +logger = logging.getLogger(__name__) + +DEMO_TOKEN_HEADER = "x-demo-token" # noqa: S105 - the header's name, not a secret + + +@tool +def whoami() -> dict[str, Any]: + """Report which account the calling user's credential belongs to. + + Requires the caller's token in the `x-demo-token` HTTP header of the MCP + request — ask the user to configure it if missing; it cannot be passed + as an argument. + """ + try: + token = credential_from_header(DEMO_TOKEN_HEADER) + except MissingCredentialError: + logger.info("whoami: no %s credential on this request", DEMO_TOKEN_HEADER) + raise + # Log presence only — never the value. + logger.info( + "whoami: %s credential present (%d chars)", DEMO_TOKEN_HEADER, len(token) + ) + # Stub: derive a stable account id instead of calling an upstream API. + account = hashlib.sha256(token.encode()).hexdigest()[:8] + return {"account": f"user-{account}", "status": "ok"} + + +TOOLS = [whoami] + +# Advertised via /health and the index: clients send this credential to this +# toolset's connection only, never to unrelated toolsets. +CREDENTIAL_HEADERS = [DEMO_TOKEN_HEADER] diff --git a/toolsets/credential-demo/tests/test_credential_demo.py b/toolsets/credential-demo/tests/test_credential_demo.py new file mode 100644 index 0000000..dcb0a44 --- /dev/null +++ b/toolsets/credential-demo/tests/test_credential_demo.py @@ -0,0 +1,31 @@ +import pytest +from mcp_runtime.credentials import MissingCredentialError, header_context + +from credential_demo.tools import DEMO_TOKEN_HEADER, whoami + + +def test_account_reported_with_credential(): + with header_context({DEMO_TOKEN_HEADER: "secret"}): + result = whoami.invoke({}) + assert result["status"] == "ok" + assert result["account"].startswith("user-") + + +def test_account_stable_per_user(): + with header_context({DEMO_TOKEN_HEADER: "secret"}): + first, second = (whoami.invoke({}) for _ in range(2)) + with header_context({DEMO_TOKEN_HEADER: "other"}): + other_user = whoami.invoke({}) + assert first == second + assert first["account"] != other_user["account"] + + +def test_token_never_appears_in_the_result(): + with header_context({DEMO_TOKEN_HEADER: "secret"}): + result = whoami.invoke({}) + assert "secret" not in str(result) + + +def test_missing_credential_names_the_header(): + with pytest.raises(MissingCredentialError, match=DEMO_TOKEN_HEADER): + whoami.invoke({}) diff --git a/toolsets/credential-demo/toolset.yaml b/toolsets/credential-demo/toolset.yaml new file mode 100644 index 0000000..a4fd46a --- /dev/null +++ b/toolsets/credential-demo/toolset.yaml @@ -0,0 +1,2 @@ +# Optional Helm values overrides for this toolset (see charts/mcp-toolset/values.yaml). +{} diff --git a/uv.lock b/uv.lock index b1f9b40..e9591f0 100644 --- a/uv.lock +++ b/uv.lock @@ -10,6 +10,7 @@ resolution-markers = [ members = [ "aoi-generator", "cds", + "credential-demo", "dataset-search", "mcp-agent", "mcp-cli", @@ -396,6 +397,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "credential-demo" +version = "0.1.0" +source = { editable = "toolsets/credential-demo" } +dependencies = [ + { name = "langchain-core" }, + { name = "mcp-runtime" }, +] + +[package.metadata] +requires-dist = [ + { name = "langchain-core", specifier = ">=1.4.6,<2.0.0" }, + { name = "mcp-runtime", editable = "packages/mcp-runtime" }, +] + [[package]] name = "cryptography" version = "48.0.1" @@ -1113,6 +1129,7 @@ dependencies = [ { name = "langchain" }, { name = "langchain-mcp-adapters" }, { name = "langchain-mistralai" }, + { name = "mcp" }, { name = "pydantic-settings" }, { name = "rich" }, { name = "typer" }, @@ -1125,6 +1142,7 @@ requires-dist = [ { name = "langchain", specifier = ">=1.3.7,<2.0.0" }, { name = "langchain-mcp-adapters", specifier = ">=0.3.0,<0.4.0" }, { name = "langchain-mistralai", specifier = ">=1.1.5,<2.0.0" }, + { name = "mcp", specifier = ">=1.27.2,<2.0.0" }, { name = "pydantic-settings", specifier = ">=2.12.0,<3.0.0" }, { name = "rich", specifier = ">=15.0.0,<16.0.0" }, { name = "typer", specifier = ">=0.26.7,<0.27.0" }, @@ -1175,6 +1193,7 @@ source = { virtual = "." } dependencies = [ { name = "aoi-generator" }, { name = "cds" }, + { name = "credential-demo" }, { name = "dataset-search" }, { name = "mcp-agent" }, { name = "mcp-cli" }, @@ -1193,6 +1212,7 @@ dev = [ requires-dist = [ { name = "aoi-generator", editable = "toolsets/aoi-generator" }, { name = "cds", editable = "toolsets/cds" }, + { name = "credential-demo", editable = "toolsets/credential-demo" }, { name = "dataset-search", editable = "toolsets/dataset-search" }, { name = "mcp-agent", editable = "packages/mcp-agent" }, { name = "mcp-cli", editable = "packages/mcp-cli" }, From e260c3ba24a6dfa1b6f8ccf48832e50f9b6adcf9 Mon Sep 17 00:00:00 2001 From: Ciaran Sweet Date: Fri, 12 Jun 2026 12:11:26 +0100 Subject: [PATCH 2/2] fix!: refactor cds tools to use per-user credentials not long lived secret --- toolsets/cds/src/cds/client.py | 14 ++++++- toolsets/cds/src/cds/settings.py | 6 +-- toolsets/cds/src/cds/tools/__init__.py | 11 +++++ toolsets/cds/src/cds/tools/_client.py | 12 ------ .../cds/src/cds/tools/apply_constraints.py | 15 +++---- .../cds/src/cds/tools/check_credentials.py | 14 ++++--- .../cds/src/cds/tools/get_dataset_schema.py | 9 ++-- toolsets/cds/src/cds/tools/get_job_status.py | 9 ++-- toolsets/cds/src/cds/tools/get_results.py | 37 ++++++++--------- toolsets/cds/src/cds/tools/list_jobs.py | 9 ++-- toolsets/cds/src/cds/tools/submit_request.py | 15 +++---- toolsets/cds/tests/test_cds.py | 41 ++++++++++++++++--- toolsets/cds/toolset.yaml | 8 ++-- 13 files changed, 122 insertions(+), 78 deletions(-) delete mode 100644 toolsets/cds/src/cds/tools/_client.py diff --git a/toolsets/cds/src/cds/client.py b/toolsets/cds/src/cds/client.py index 2c33019..4306d19 100644 --- a/toolsets/cds/src/cds/client.py +++ b/toolsets/cds/src/cds/client.py @@ -1,13 +1,25 @@ import httpx +from mcp_runtime.credentials import credential_from_header + from .settings import settings +CDS_TOKEN_HEADER = "x-cds-token" # noqa: S105 - the header's name, not a secret + def make_client() -> httpx.AsyncClient: + """A retrieve-API client authenticated as the calling user. + + Reads the user's CDS key from the `x-cds-token` MCP request header, + raising MissingCredentialError (which tells the caller how to supply it) + when absent. Created per tool call and closed with it (`async with`), so + one user's auth — or any other client state — can never leak into + another user's requests. + """ return httpx.AsyncClient( base_url=settings.cds_api_url, headers={ - "PRIVATE-TOKEN": settings.cds_api_key, + "PRIVATE-TOKEN": credential_from_header(CDS_TOKEN_HEADER), "Content-Type": "application/json", }, timeout=httpx.Timeout(connect=10.0, read=120.0, write=10.0, pool=10.0), diff --git a/toolsets/cds/src/cds/settings.py b/toolsets/cds/src/cds/settings.py index 840f117..b7ca2f6 100644 --- a/toolsets/cds/src/cds/settings.py +++ b/toolsets/cds/src/cds/settings.py @@ -4,16 +4,14 @@ class Settings(BaseSettings): """CDS API configuration, read from the environment (or a .env in the cwd). - `cds_api_key` defaults to empty so the toolset imports and serves without - credentials; calls then return structured auth errors until CDS_API_KEY - is provided (in k8s, via the Secret listed in toolset.yaml). + There is no API-key setting: the CDS key is per calling user, sent as the + `x-cds-token` HTTP header on each MCP request (see `tools._client`). """ model_config = SettingsConfigDict( env_file=".env", env_file_encoding="utf-8", extra="ignore" ) - cds_api_key: str = "" cds_api_url: str = "https://cds.climate.copernicus.eu/api/retrieve/v1" cds_catalogue_url: str = "https://cds.climate.copernicus.eu/api/catalogue/v1" diff --git a/toolsets/cds/src/cds/tools/__init__.py b/toolsets/cds/src/cds/tools/__init__.py index 76ef4dd..74d1e88 100644 --- a/toolsets/cds/src/cds/tools/__init__.py +++ b/toolsets/cds/src/cds/tools/__init__.py @@ -3,8 +3,14 @@ Covers the full retrieval workflow: search the catalogue, inspect a dataset's schema and constraints, submit requests, then poll jobs and fetch download links. + +Retrieve-API tools act as the calling user: their CDS key arrives as the +`x-cds-token` HTTP header on each MCP request (advertised below via +``CREDENTIAL_HEADERS``). Only `search_datasets` (public catalogue) works +without it. """ +from ..client import CDS_TOKEN_HEADER from .apply_constraints import apply_constraints from .check_credentials import check_credentials from .get_dataset_schema import get_dataset_schema @@ -15,6 +21,7 @@ from .submit_request import submit_request __all__ = [ + "CREDENTIAL_HEADERS", "TOOLS", "apply_constraints", "check_credentials", @@ -36,3 +43,7 @@ list_jobs, check_credentials, ] + +# Advertised via /health and the index: clients send this credential to this +# toolset's connection only, never to unrelated toolsets. +CREDENTIAL_HEADERS = [CDS_TOKEN_HEADER] diff --git a/toolsets/cds/src/cds/tools/_client.py b/toolsets/cds/src/cds/tools/_client.py deleted file mode 100644 index 3307917..0000000 --- a/toolsets/cds/src/cds/tools/_client.py +++ /dev/null @@ -1,12 +0,0 @@ -import httpx - -from ..client import make_client - -_client: httpx.AsyncClient | None = None - - -def get_client() -> httpx.AsyncClient: - global _client - if _client is None: - _client = make_client() - return _client diff --git a/toolsets/cds/src/cds/tools/apply_constraints.py b/toolsets/cds/src/cds/tools/apply_constraints.py index d85761d..729a0be 100644 --- a/toolsets/cds/src/cds/tools/apply_constraints.py +++ b/toolsets/cds/src/cds/tools/apply_constraints.py @@ -2,19 +2,20 @@ from langchain_core.tools import tool -from ._client import get_client +from ..client import make_client from ._errors import classify_http_error, not_found_error, transient_error from ._retry import TRANSIENT_EXC, with_retry @with_retry async def _call(dataset: str, partial_request: dict[str, Any]) -> dict[str, Any]: - resp = await get_client().post( - f"/processes/{dataset}/constraints", - json={"inputs": partial_request}, - ) - if resp.status_code >= 500: - resp.raise_for_status() + async with make_client() as client: + resp = await client.post( + f"/processes/{dataset}/constraints", + json={"inputs": partial_request}, + ) + if resp.status_code >= 500: + resp.raise_for_status() if resp.status_code == 404: return not_found_error(f"Dataset {dataset!r} not found.") if resp.status_code != 200: diff --git a/toolsets/cds/src/cds/tools/check_credentials.py b/toolsets/cds/src/cds/tools/check_credentials.py index 0636d9c..b6f9808 100644 --- a/toolsets/cds/src/cds/tools/check_credentials.py +++ b/toolsets/cds/src/cds/tools/check_credentials.py @@ -2,16 +2,17 @@ from langchain_core.tools import tool -from ._client import get_client +from ..client import make_client from ._errors import classify_http_error, transient_error from ._retry import TRANSIENT_EXC, with_retry @with_retry async def _call() -> dict[str, Any]: - resp = await get_client().get("/jobs", params={"limit": 1}) - if resp.status_code >= 500: - resp.raise_for_status() + async with make_client() as client: + resp = await client.get("/jobs", params={"limit": 1}) + if resp.status_code >= 500: + resp.raise_for_status() if resp.status_code == 200: return {"ok": True} return classify_http_error(resp) @@ -19,7 +20,10 @@ async def _call() -> dict[str, Any]: @tool async def check_credentials() -> dict[str, Any]: - """Validate the CDS API key. Returns {"ok": True} on success or a structured error dict.""" + """Validate the calling user's CDS API key (the `x-cds-token` MCP header). + + Returns {"ok": True} on success or a structured error dict. + """ try: return await _call() except TRANSIENT_EXC as exc: diff --git a/toolsets/cds/src/cds/tools/get_dataset_schema.py b/toolsets/cds/src/cds/tools/get_dataset_schema.py index 0521501..25d62ff 100644 --- a/toolsets/cds/src/cds/tools/get_dataset_schema.py +++ b/toolsets/cds/src/cds/tools/get_dataset_schema.py @@ -2,7 +2,7 @@ from langchain_core.tools import tool -from ._client import get_client +from ..client import make_client from ._errors import classify_http_error, not_found_error, transient_error from ._retry import TRANSIENT_EXC, with_retry @@ -84,9 +84,10 @@ def _parse_parameter(spec: dict[str, Any]) -> dict[str, Any]: @with_retry async def _call(dataset: str) -> dict[str, Any]: - resp = await get_client().get(f"/processes/{dataset}") - if resp.status_code >= 500: - resp.raise_for_status() + async with make_client() as client: + resp = await client.get(f"/processes/{dataset}") + if resp.status_code >= 500: + resp.raise_for_status() if resp.status_code == 404: return not_found_error(f"Dataset {dataset!r} not found.") if resp.status_code != 200: diff --git a/toolsets/cds/src/cds/tools/get_job_status.py b/toolsets/cds/src/cds/tools/get_job_status.py index fd9f510..804ce61 100644 --- a/toolsets/cds/src/cds/tools/get_job_status.py +++ b/toolsets/cds/src/cds/tools/get_job_status.py @@ -2,16 +2,17 @@ from langchain_core.tools import tool -from ._client import get_client +from ..client import make_client from ._errors import classify_http_error, transient_error from ._retry import TRANSIENT_EXC, with_retry @with_retry async def _call(job_id: str) -> dict[str, Any]: - resp = await get_client().get(f"/jobs/{job_id}") - if resp.status_code >= 500: - resp.raise_for_status() + async with make_client() as client: + resp = await client.get(f"/jobs/{job_id}") + if resp.status_code >= 500: + resp.raise_for_status() if resp.status_code != 200: return classify_http_error(resp) diff --git a/toolsets/cds/src/cds/tools/get_results.py b/toolsets/cds/src/cds/tools/get_results.py index 7ce5c8c..c3e73bb 100644 --- a/toolsets/cds/src/cds/tools/get_results.py +++ b/toolsets/cds/src/cds/tools/get_results.py @@ -2,31 +2,30 @@ from langchain_core.tools import tool -from ._client import get_client +from ..client import make_client from ._errors import classify_http_error, not_ready_error, transient_error from ._retry import TRANSIENT_EXC, with_retry @with_retry async def _call(job_id: str) -> dict[str, Any]: - client = get_client() - - # Confirm the job is successful before fetching assets. - status_resp = await client.get(f"/jobs/{job_id}") - if status_resp.status_code >= 500: - status_resp.raise_for_status() - if status_resp.status_code != 200: - return classify_http_error(status_resp) - - status: str = status_resp.json().get("status", "unknown") - if status != "successful": - return not_ready_error(job_id, status) - - results_resp = await client.get(f"/jobs/{job_id}/results") - if results_resp.status_code >= 500: - results_resp.raise_for_status() - if results_resp.status_code != 200: - return classify_http_error(results_resp) + async with make_client() as client: + # Confirm the job is successful before fetching assets. + status_resp = await client.get(f"/jobs/{job_id}") + if status_resp.status_code >= 500: + status_resp.raise_for_status() + if status_resp.status_code != 200: + return classify_http_error(status_resp) + + status: str = status_resp.json().get("status", "unknown") + if status != "successful": + return not_ready_error(job_id, status) + + results_resp = await client.get(f"/jobs/{job_id}/results") + if results_resp.status_code >= 500: + results_resp.raise_for_status() + if results_resp.status_code != 200: + return classify_http_error(results_resp) try: val = results_resp.json()["asset"]["value"] diff --git a/toolsets/cds/src/cds/tools/list_jobs.py b/toolsets/cds/src/cds/tools/list_jobs.py index 25c0c1b..8fe10ad 100644 --- a/toolsets/cds/src/cds/tools/list_jobs.py +++ b/toolsets/cds/src/cds/tools/list_jobs.py @@ -2,7 +2,7 @@ from langchain_core.tools import tool -from ._client import get_client +from ..client import make_client from ._errors import classify_http_error, transient_error from ._retry import TRANSIENT_EXC, with_retry @@ -13,9 +13,10 @@ async def _call(status: list[str] | None, limit: int) -> dict[str, Any]: if status: params["status"] = status # httpx repeats the key for each value - resp = await get_client().get("/jobs", params=params) - if resp.status_code >= 500: - resp.raise_for_status() + async with make_client() as client: + resp = await client.get("/jobs", params=params) + if resp.status_code >= 500: + resp.raise_for_status() if resp.status_code != 200: return classify_http_error(resp) diff --git a/toolsets/cds/src/cds/tools/submit_request.py b/toolsets/cds/src/cds/tools/submit_request.py index 045c576..064842c 100644 --- a/toolsets/cds/src/cds/tools/submit_request.py +++ b/toolsets/cds/src/cds/tools/submit_request.py @@ -2,7 +2,7 @@ from langchain_core.tools import tool -from ._client import get_client +from ..client import make_client from ._errors import classify_http_error, classify_submit_failure, transient_error from ._retry import TRANSIENT_EXC, with_retry @@ -20,12 +20,13 @@ def _failure_message(data: dict[str, Any]) -> str: @with_retry async def _call(dataset: str, request: dict[str, Any]) -> dict[str, Any]: - resp = await get_client().post( - f"/processes/{dataset}/execution", - json={"inputs": request}, - ) - if resp.status_code >= 500: - resp.raise_for_status() + async with make_client() as client: + resp = await client.post( + f"/processes/{dataset}/execution", + json={"inputs": request}, + ) + if resp.status_code >= 500: + resp.raise_for_status() if resp.status_code not in (200, 201): return classify_http_error(resp, dataset) diff --git a/toolsets/cds/tests/test_cds.py b/toolsets/cds/tests/test_cds.py index e4caa9e..9d44b07 100644 --- a/toolsets/cds/tests/test_cds.py +++ b/toolsets/cds/tests/test_cds.py @@ -3,8 +3,11 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx +import pytest +from mcp_runtime.credentials import MissingCredentialError, header_context -from cds.tools import TOOLS +from cds.client import CDS_TOKEN_HEADER, make_client +from cds.tools import CREDENTIAL_HEADERS, TOOLS from cds.tools.apply_constraints import _call as _apply_call from cds.tools.get_dataset_schema import _call as _schema_call from cds.tools.get_dataset_schema import _parse_parameter @@ -58,8 +61,11 @@ def _mock_client(method: str, response: httpx.Response) -> MagicMock: + """A stand-in for make_client(): async-with-able, one canned response.""" client = MagicMock() setattr(client, method, AsyncMock(return_value=response)) + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) return client @@ -133,7 +139,7 @@ def test_parse_optional_field_with_default_and_enum() -> None: async def test_get_dataset_schema_parses_all_fields() -> None: client = _mock_client("get", httpx.Response(200, json=ERA5_PROCESS_RESPONSE)) - with patch("cds.tools.get_dataset_schema.get_client", return_value=client): + with patch("cds.tools.get_dataset_schema.make_client", return_value=client): result = await _schema_call("reanalysis-era5-land") assert result["dataset"] == "reanalysis-era5-land" @@ -158,7 +164,7 @@ async def test_get_dataset_schema_parses_all_fields() -> None: async def test_get_dataset_schema_not_found() -> None: client = _mock_client("get", httpx.Response(404, json={"detail": "not found"})) - with patch("cds.tools.get_dataset_schema.get_client", return_value=client): + with patch("cds.tools.get_dataset_schema.make_client", return_value=client): result = await _schema_call("bad-dataset") assert result["error"] == "not_found" @@ -173,7 +179,7 @@ async def test_get_dataset_schema_not_found() -> None: async def test_apply_constraints_filters_empty_lists() -> None: client = _mock_client("post", httpx.Response(200, json=CONSTRAINTS_RESPONSE)) - with patch("cds.tools.apply_constraints.get_client", return_value=client): + with patch("cds.tools.apply_constraints.make_client", return_value=client): result = await _apply_call( "reanalysis-era5-land", {"year": "2024", "month": "01"} ) @@ -186,7 +192,7 @@ async def test_apply_constraints_filters_empty_lists() -> None: async def test_apply_constraints_sends_inputs_wrapper() -> None: client = _mock_client("post", httpx.Response(200, json=CONSTRAINTS_RESPONSE)) - with patch("cds.tools.apply_constraints.get_client", return_value=client): + with patch("cds.tools.apply_constraints.make_client", return_value=client): await _apply_call("reanalysis-era5-land", {"year": "2024", "month": "01"}) _, call_kwargs = client.post.call_args @@ -196,7 +202,7 @@ async def test_apply_constraints_sends_inputs_wrapper() -> None: async def test_apply_constraints_not_found() -> None: client = _mock_client("post", httpx.Response(404, json={"detail": "not found"})) - with patch("cds.tools.apply_constraints.get_client", return_value=client): + with patch("cds.tools.apply_constraints.make_client", return_value=client): result = await _apply_call("bad-dataset", {}) assert result["error"] == "not_found" @@ -324,3 +330,26 @@ def test_parse_no_schema_key() -> None: assert result["type"] == "unknown" assert result["required"] is True assert result["title"] == "Mystery field" + + +# --------------------------------------------------------------------------- +# Per-user credentials +# --------------------------------------------------------------------------- + + +def test_credential_headers_exported() -> None: + assert CREDENTIAL_HEADERS == [CDS_TOKEN_HEADER] == ["x-cds-token"] + + +async def test_client_authenticates_as_the_calling_user() -> None: + with header_context({CDS_TOKEN_HEADER: "user-a-key"}): + async with make_client() as client: + assert client.headers["PRIVATE-TOKEN"] == "user-a-key" + with header_context({CDS_TOKEN_HEADER: "user-b-key"}): + async with make_client() as client: + assert client.headers["PRIVATE-TOKEN"] == "user-b-key" + + +def test_client_requires_the_credential() -> None: + with pytest.raises(MissingCredentialError, match=CDS_TOKEN_HEADER): + make_client() diff --git a/toolsets/cds/toolset.yaml b/toolsets/cds/toolset.yaml index 7edd03d..4b07eea 100644 --- a/toolsets/cds/toolset.yaml +++ b/toolsets/cds/toolset.yaml @@ -1,6 +1,4 @@ # Optional Helm values overrides for this toolset (see charts/mcp-toolset/values.yaml). -# The cds-credentials Secret must exist in the namespace and provide CDS_API_KEY -# (create it out-of-band: kubectl -n mcp-toolsets create secret generic -# cds-credentials --from-literal=CDS_API_KEY=...). -secrets: - - cds-credentials +# No secrets: the CDS API key is per calling user, sent as the x-cds-token +# HTTP header on each MCP request (see the README's "Per-user credentials"). +{}