diff --git a/.github/workflows/deploy-docs.yml b/.github/workflows/deploy-docs.yml index eaf1c167f5..6f84154747 100644 --- a/.github/workflows/deploy-docs.yml +++ b/.github/workflows/deploy-docs.yml @@ -17,6 +17,7 @@ on: - mkdocs.yml - src/mcp/** - src/mcp-types/** + - src/mcp-client/** - scripts/build-docs.sh - scripts/docs/** - pyproject.toml diff --git a/.github/workflows/publish-pypi.yml b/.github/workflows/publish-pypi.yml index 73b8f87adc..df7a4bdc5d 100644 --- a/.github/workflows/publish-pypi.yml +++ b/.github/workflows/publish-pypi.yml @@ -29,6 +29,7 @@ jobs: - name: Build run: | uv build --package mcp + uv build --package mcp-client uv build --package mcp-types - name: Upload artifacts diff --git a/.github/workflows/shared.yml b/.github/workflows/shared.yml index ef6eb25a33..e778394426 100644 --- a/.github/workflows/shared.yml +++ b/.github/workflows/shared.yml @@ -48,6 +48,38 @@ jobs: uv run --isolated --no-project --with ./src/mcp-types python -c \ "import mcp_types, mcp_types.jsonrpc, mcp_types.methods, mcp_types.version, mcp_types._v2025_11_25, mcp_types._v2026_07_28" + packages: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + with: + version: 0.9.5 + - name: Build all distributions + run: | + uv build --package mcp-types + uv build --package mcp-client + uv build --package mcp + - name: Exercise the client wheel and sdist without the server SDK + run: | + for package in dist/mcp_client-*.whl dist/mcp_client-*.tar.gz; do + uv run --isolated --no-project --find-links dist --with "$package" \ + python scripts/check_client_package.py + done + - name: Type-check the standalone client package + run: | + uv export --frozen --package mcp-client --no-default-groups --no-dev --no-emit-workspace \ + --output-file "$RUNNER_TEMP/client-requirements.txt" + uv run --isolated --no-project --find-links dist --with dist/mcp_client-*.whl \ + --with-requirements "$RUNNER_TEMP/client-requirements.txt" --with pyright==1.1.405 \ + python scripts/check_client_types.py + - name: Import the full SDK with the client package first + run: | + uv run --isolated --no-project --find-links dist --with dist/mcp-*.whl python -c \ + 'import mcp_client, mcp; from typing import get_type_hints; assert mcp.Client is mcp_client.Client; get_type_hints(mcp.Client)' + test: name: test (${{ matrix.python-version }}, ${{ matrix.dep-resolution.name }}, ${{ matrix.os }}) runs-on: ${{ matrix.os }} diff --git a/DEPENDENCY_POLICY.md b/DEPENDENCY_POLICY.md index 961157b378..cc08dcd4de 100644 --- a/DEPENDENCY_POLICY.md +++ b/DEPENDENCY_POLICY.md @@ -4,7 +4,7 @@ ## How requirements are declared -Every runtime dependency is a `>=` floor set to the oldest version that provides what the SDK uses, with no upper bound unless a dependency's next major is known to break the SDK. The one exception is `mcp-types`, the wire-types package released in lockstep with `mcp`: each `mcp` release requires exactly its own version of it, so it is the other half of the SDK rather than an independent constraint. +Every runtime dependency is a `>=` floor set to the oldest version that provides what the SDK uses, with no upper bound unless a dependency's next major is known to break the SDK. The exceptions are `mcp-client` and `mcp-types`, which release in lockstep with `mcp`: each `mcp` release requires exactly its own version of both, and `mcp-client` requires the matching `mcp-types`. They are parts of the SDK rather than independent constraints. ## When a floor moves diff --git a/README.md b/README.md index cc067aca25..9f715814b7 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,10 @@ uv add "mcp[cli]" # or: pip install "mcp[cli]" The `cli` extra adds the `mcp` command-line tool (`mcp dev`, `mcp run`, `mcp install`) on top of the SDK; install plain `mcp` if you don't need it. For one-off commands, `uv run --with "mcp[cli]" mcp ...` works without a project. +For a client-only project, use `uv add mcp-client` and `from mcp_client import Client`. +It includes the client transports and OAuth support without the HTTP server dependencies. +See [client-only installation](https://py.sdk.modelcontextprotocol.io/get-started/installation/#client-only-installation). + ## A server in 15 lines Create a `server.py`: diff --git a/RELEASE.md b/RELEASE.md index cada37480c..dae27a0cf1 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -7,7 +7,8 @@ move; this is the mechanics. 1. Change the dependency version in `pyproject.toml`. The root `mcp` project's runtime dependencies are dynamic and live under - `[tool.hatch.metadata.hooks.uv-dynamic-versioning].dependencies`. + `[tool.hatch.metadata.hooks.uv-dynamic-versioning].dependencies`, as do + `mcp-client`'s dependencies in `src/mcp-client/pyproject.toml`. 2. Regenerate the lock with `uv lock` (or `uv lock --upgrade-package ` to move just that package's locked version). The committed `uv.lock` is a normal (default-strategy) resolution; the `lowest-direct` resolution that @@ -19,8 +20,8 @@ move; this is the mechanics. Two branches ship, and the package version comes from the git tag (`uv-dynamic-versioning`). Publishing a GitHub release runs `publish-pypi.yml` **from the tagged commit**, so the workflow that fires is the tagged branch's -own: a `main` tag builds and publishes two distributions (`mcp` and -`mcp-types`, lock-stepped via `Requires-Dist: mcp-types=={{ version }}`), and a +own: a `main` tag builds and publishes three distributions (`mcp`, +`mcp-client`, and `mcp-types`, with exact matching-version dependencies), and a `v1.x` tag builds and publishes `mcp` only. | Line | Branch | Tag | GitHub release flags | @@ -29,11 +30,13 @@ own: a `main` tag builds and publishes two distributions (`mcp` and | Maintenance (previous major) | `v1.x` | `v1.X.Y` | not a pre-release; **not** Latest | | Pre-releases | `main` | `v2.X.YaN` / `bN` / `rcN` | **Pre-release** ticked, never Latest | -The `Development Status` classifier in both `pyproject.toml` files is +The `Development Status` classifier in all three `pyproject.toml` files is permanently `5 - Production/Stable`; it is not bumped as part of any release. The `mcp-types` PyPI project carries the same trusted publisher as `mcp` (this -repository, workflow `publish-pypi.yml`, environment `release`). For a release -cut from `main`, if only some of the four files upload, fix the cause and +repository, workflow `publish-pypi.yml`, environment `release`). Before the +first `mcp-client` release, verify ownership of the existing PyPI project and +configure that same trusted publisher for it too. For a release cut from `main`, if only some of the six files upload, +correct the cause and re-run the publish job — its `skip-existing` setting makes it skip whatever already landed (the `v1.x` workflow publishes a single distribution and has no such setting). @@ -77,7 +80,7 @@ before the tag. URLs (relative links don't resolve in GitHub release bodies). 5. If a stable release turns out to be broken, yank it on PyPI and release the fix as the next patch version. Never delete a release from PyPI — version - numbers cannot be reused. Yank `mcp` and `mcp-types` together (they are one + numbers cannot be reused. Yank `mcp`, `mcp-client`, and `mcp-types` together (they are one release), and set the yank reason and the GitHub release notes to point at the replacement version, since yanking doesn't stop `==` pins from installing the broken version. @@ -134,6 +137,6 @@ specifier that names a pre-release version, or `--pre`. 4. Curate the release notes: what changed since the previous pre-release, what is known-incomplete, the install line (`pip install mcp==2.X.YbN`), and a link to the migration guide, with absolute URLs. -5. If a pre-release turns out to be broken, yank both `mcp` and `mcp-types` on PyPI +5. If a pre-release turns out to be broken, yank `mcp`, `mcp-client`, and `mcp-types` on PyPI and cut the next one, pointing the yank reason and the GitHub release notes at the replacement version. diff --git a/VERSIONING.md b/VERSIONING.md index bc80d51aef..475aeff9d5 100644 --- a/VERSIONING.md +++ b/VERSIONING.md @@ -4,11 +4,11 @@ What a version number of `mcp` promises: which changes can arrive in a minor rel ## The version number -[Semantic Versioning](https://semver.org/) semantics in [PEP 440](https://peps.python.org/pep-0440/) syntax, taken from the git tag: in `2.X.Y`, **X** (minor) carries new functionality and every non-breaking change, **Y** (patch) carries bug fixes only, and a breaking change to the public API lands only in a new **major**. Pre-releases are cut from `main` as `aN`/`bN`/`rcN`; installers prefer final releases by default, so an unpinned `pip install mcp` stays on a stable release whenever one satisfies your requirement. `mcp` and its wire-types package `mcp-types` release in lockstep, each `mcp` requiring exactly the matching `mcp-types`. +[Semantic Versioning](https://semver.org/) semantics in [PEP 440](https://peps.python.org/pep-0440/) syntax, taken from the git tag: in `2.X.Y`, **X** (minor) carries new functionality and every non-breaking change, **Y** (patch) carries bug fixes only, and a breaking change to the public API lands only in a new **major**. Pre-releases are cut from `main` as `aN`/`bN`/`rcN`; installers prefer final releases by default, so an unpinned `pip install mcp` stays on a stable release whenever one satisfies your requirement. `mcp`, `mcp-client`, and `mcp-types` release in lockstep. Each `mcp` requires exactly the matching `mcp-client` and `mcp-types`; `mcp-client` also requires exactly the matching `mcp-types`. ## The public API -The promise covers every name exported by `mcp` and `mcp_types` (their `__all__`), the import paths, signatures, and behavior documented on the [documentation site](https://py.sdk.modelcontextprotocol.io/) and in its [API Reference](https://py.sdk.modelcontextprotocol.io/api/mcp/). It does not cover underscore-prefixed names, undocumented modules, or the wording of log lines, warnings, and exception messages (their types and documented raise conditions are covered). APIs labelled **provisional** (for example the middleware chain) may still change in a minor release; **experimental** APIs are opt-in previews. +The promise covers every name exported by `mcp`, `mcp_client`, and `mcp_types` (their `__all__`), the import paths, signatures, and behavior documented on the [documentation site](https://py.sdk.modelcontextprotocol.io/) and in its [API Reference](https://py.sdk.modelcontextprotocol.io/api/mcp/). It does not cover underscore-prefixed names, undocumented modules, or the wording of log lines, warnings, and exception messages (their types and documented raise conditions are covered). APIs labelled **provisional** (for example the middleware chain) may still change in a minor release; **experimental** APIs are opt-in previews. ## Breaking and non-breaking changes diff --git a/docs/get-started/installation.md b/docs/get-started/installation.md index 728a86cfd0..22b49973cd 100644 --- a/docs/get-started/installation.md +++ b/docs/get-started/installation.md @@ -21,10 +21,45 @@ These docs describe **v2**, the current stable release line: covers every one. If your *package* depends on `mcp` and isn't ready to migrate, keep a `<2` upper bound (for example `mcp>=1.28,<2`) so an unpinned resolve stays on the 1.x line. +## Client-only installation + +```bash +uv add mcp-client +``` + +```python +import anyio + +from mcp_client import Client + + +async def main() -> None: + async with Client("http://localhost:8000/mcp") as client: + tools = await client.list_tools() + for tool in tools.tools: + print(tool.name) + + +anyio.run(main) +``` + +Run this example against an MCP server listening at `http://localhost:8000/mcp`. + +Use `mcp-client` when you only connect to servers. It includes the client transports, +OAuth support, and shared protocol machinery without installing Starlette, Uvicorn, +`sse-starlette`, or `python-multipart`. Import client APIs from `mcp_client`, OAuth +support from `mcp_client.client.auth`, and protocol types from `mcp_types`. + +Install `mcp` if you also build servers, use the CLI, or pass a server instance to +`Client(server)` for in-process testing. Existing `mcp` imports keep working and +refer to the same client implementation. All three distributions release together; +`mcp` requires its exact `mcp-client` version, which requires its exact `mcp-types` version. + ## What gets installed You don't need to know any of this to use the SDK, but if you're wondering what each dependency is for: +* `mcp-client`: the client API, transports, OAuth support, and shared protocol machinery, versioned in lockstep with the SDK. * `mcp-types`: every protocol type (requests, results, content blocks) as its own package, versioned in lockstep with the SDK. Code that depends on `mcp` imports it through the `mcp.types` alias (every `from mcp.types import ...` in these docs); import `mcp_types` directly only in a project that installs `mcp-types` without the SDK. * [`anyio`](https://anyio.readthedocs.io/): the async runtime. The whole SDK is written against anyio, so it runs on either `asyncio` or `trio`. * [`pydantic`](https://docs.pydantic.dev/): what every `mcp.types` model is built on, plus all schema generation and validation. diff --git a/mkdocs.yml b/mkdocs.yml index a75053326f..272c5c84d8 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -192,7 +192,7 @@ plugins: - mkdocstrings: handlers: python: - paths: [src, src/mcp-types] + paths: [src, src/mcp-client, src/mcp-types] # Zensical renders pages in undefined (filesystem-dependent) order # against one shared griffe collection, so a cross-package re-export # (`mcp` -> `mcp_types`) resolves only if its target package happens diff --git a/pyproject.toml b/pyproject.toml index b2f26da55f..27dcf6baea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -132,6 +132,7 @@ dependencies = [ "anyio>=4.10; python_version >= '3.14'", "anyio>=4.9; python_version < '3.14'", "httpx2>=2.5.0", + "mcp-client=={{ version }}", "mcp-types=={{ version }}", "pydantic>=2.12.0", "starlette>=0.48.0; python_version >= '3.14'", @@ -140,7 +141,6 @@ dependencies = [ "sse-starlette>=3.0.0", "uvicorn>=0.31.1; sys_platform != 'emscripten'", "jsonschema>=4.20.0", - "pywin32>=311; sys_platform == 'win32'", "pyjwt[crypto]>=2.10.1", "typing-extensions>=4.13.0", "typing-inspection>=0.4.1", @@ -160,6 +160,7 @@ packages = ["src/mcp"] typeCheckingMode = "strict" include = [ "src/mcp", + "src/mcp-client/mcp_client", "src/mcp-types/mcp_types", "tests", "docs_src", @@ -167,6 +168,8 @@ include = [ "examples/servers", "examples/snippets", "examples/clients", + "scripts/check_client_package.py", + "scripts/check_client_types.py", "scripts/docs/build_config.py", "scripts/docs/translations.py", ] @@ -198,6 +201,10 @@ executionEnvironments = [ # docs_src/ holds the complete, runnable code examples included into docs/*.md. # Decorated (@mcp.tool/...) module-level functions are never called by name. { root = "docs_src", reportUnusedFunction = false }, + # Compatibility modules re-export the original module namespaces, including private names. + { root = "src/mcp/client", reportPrivateUsage = false }, + { root = "src/mcp/shared", reportPrivateUsage = false }, + { root = "src/mcp/os", reportPrivateUsage = false }, ] [tool.ruff] @@ -245,10 +252,11 @@ max-returns = 13 # Default is 6 max-statements = 102 # Default is 50 [tool.uv.workspace] -members = ["src/mcp-types", "examples", "examples/clients/*", "examples/servers/*", "examples/snippets"] +members = ["src/mcp-client", "src/mcp-types", "examples", "examples/clients/*", "examples/servers/*", "examples/snippets"] [tool.uv.sources] mcp = { workspace = true } +mcp-client = { workspace = true } mcp-example-stories = { workspace = true } mcp-types = { workspace = true } strict-no-cover = { git = "https://github.com/pydantic/strict-no-cover" } @@ -302,12 +310,13 @@ MD059 = false # descriptive-link-text branch = true patch = ["subprocess"] concurrency = ["multiprocessing", "thread"] -source = ["src", "src/mcp-types/mcp_types", "tests"] +source = ["src", "src/mcp-client/mcp_client", "src/mcp-types/mcp_types", "tests"] omit = [ "src/mcp/client/__main__.py", + "src/mcp-client/mcp_client/client/__main__.py", "src/mcp/server/__main__.py", - "src/mcp/os/posix/utilities.py", - "src/mcp/os/win32/utilities.py", + "src/mcp-client/mcp_client/os/posix/utilities.py", + "src/mcp-client/mcp_client/os/win32/utilities.py", ] # https://coverage.readthedocs.io/en/latest/config.html#report diff --git a/scripts/check_client_package.py b/scripts/check_client_package.py new file mode 100644 index 0000000000..d65673a4a1 --- /dev/null +++ b/scripts/check_client_package.py @@ -0,0 +1,71 @@ +"""Exercise an installed client distribution without the SDK's server dependencies. + +The peer uses raw messages because the server package must be absent from this environment. +""" + +import importlib +import importlib.util +import pkgutil +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import get_type_hints + +import anyio +import mcp_client +from mcp_client.shared import exceptions +from mcp_client.shared.memory import MessageStream, create_client_server_memory_streams +from mcp_client.shared.message import SessionMessage +from mcp_types import JSONRPCRequest, JSONRPCResponse, ListToolsResult, Tool + +get_type_hints(mcp_client.Client.__init__) + +for error_type in ( + exceptions.MCPError, + exceptions.MCPDeprecationWarning, + exceptions.NoBackChannelError, + exceptions.UrlElicitationRequiredError, +): + assert error_type.__module__ == "mcp_client.shared.exceptions" + +for name in ("mcp", "starlette", "uvicorn", "sse_starlette", "multipart"): + assert importlib.util.find_spec(name) is None, name + +for info in pkgutil.walk_packages(mcp_client.__path__, prefix="mcp_client."): + if not any(part.startswith("_") for part in info.name.split(".")): + importlib.import_module(info.name) + +RESULT = ListToolsResult(tools=[Tool(name="example", input_schema={"type": "object"})]) + + +@asynccontextmanager +async def transport() -> AsyncIterator[MessageStream]: + async with create_client_server_memory_streams() as (client_streams, server_streams): + read, write = server_streams + + async def respond() -> None: + received = await read.receive() + assert isinstance(received, SessionMessage) + message = received.message + assert isinstance(message, JSONRPCRequest) + assert message.method == "tools/list" + await write.send( + SessionMessage( + JSONRPCResponse( + jsonrpc="2.0", id=message.id, result=RESULT.model_dump(by_alias=True, exclude_none=True) + ) + ) + ) + + async with anyio.create_task_group() as group: + group.start_soon(respond) + yield client_streams + + +async def main() -> None: + """Verify a client request and response through the installed public API.""" + with anyio.fail_after(5): + async with mcp_client.Client(transport(), mode="2026-07-28") as client: + assert await client.list_tools() == RESULT + + +anyio.run(main) diff --git a/scripts/check_client_types.py b/scripts/check_client_types.py new file mode 100644 index 0000000000..bdc62340ba --- /dev/null +++ b/scripts/check_client_types.py @@ -0,0 +1,48 @@ +"""Type-check the installed client package and reject an invalid constructor argument.""" + +import importlib.util +import json +import subprocess +import sys +import tempfile +from pathlib import Path + +import mcp_client + +assert importlib.util.find_spec("mcp") is None + +with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + consumer = root / "consumer.py" + consumer.write_text( + "from mcp_client import Client, StdioServerParameters\n" + 'Client("https://example.com/mcp")\n' + 'Client(StdioServerParameters(command="python"))\n' + "Client(42)\n", + encoding="utf-8", + ) + config = root / "pyrightconfig.json" + config.write_text( + json.dumps( + { + "pythonPath": sys.executable, + "typeCheckingMode": "strict", + } + ), + encoding="utf-8", + ) + result = subprocess.run( + ["pyright", "--project", str(config), "--outputjson", *mcp_client.__path__, str(consumer)], + cwd=root, + capture_output=True, + text=True, + encoding="utf-8", + check=False, + ) + diagnostics = json.loads(result.stdout)["generalDiagnostics"] + errors = [diagnostic for diagnostic in diagnostics if diagnostic["severity"] == "error"] + assert result.returncode == 1, result.stdout + result.stderr + assert len(errors) == 1, result.stdout + assert errors[0]["file"] == str(consumer), result.stdout + assert errors[0]["range"]["start"]["line"] == 3, result.stdout + assert errors[0]["rule"] == "reportArgumentType", result.stdout diff --git a/scripts/docs/gen_ref_pages.py b/scripts/docs/gen_ref_pages.py index c368ff89f7..0314d4710a 100644 --- a/scripts/docs/gen_ref_pages.py +++ b/scripts/docs/gen_ref_pages.py @@ -29,14 +29,18 @@ # `src/mcp-types` is a distribution directory, not an import package, so each # package's dotted module path is taken relative to its own parent: deriving # it from `src/` would emit the unimportable `mcp-types.mcp_types.*`. -PACKAGES = (ROOT / "src" / "mcp", ROOT / "src" / "mcp-types" / "mcp_types") +PACKAGES = ( + ROOT / "src" / "mcp", + ROOT / "src" / "mcp-client" / "mcp_client", + ROOT / "src" / "mcp-types" / "mcp_types", +) # Module paths that get no page, and neither does anything under them: alias -# packages that mirror another package's namespaces (`mcp.types` mirrors -# `mcp_types`), whose canonical rendering is the mirrored package's pages; and +# packages that mirror the extracted `mcp_client` and `mcp_types` namespaces, +# whose canonical rendering is the extracted package's pages; and # removed v1 import paths (`mcp.server.fastmcp`) that only raise a pointer to # the migration guide and carry no API. -EXCLUDED = frozenset({"mcp.types", "mcp.server.fastmcp"}) +EXCLUDED = frozenset({"mcp.types", "mcp.client", "mcp.shared", "mcp.os", "mcp.server.fastmcp"}) _KIND_SECTIONS = { griffe.Kind.MODULE: "Modules", diff --git a/scripts/docs/llms_txt.py b/scripts/docs/llms_txt.py index 0ac09399a8..b3bb561572 100644 --- a/scripts/docs/llms_txt.py +++ b/scripts/docs/llms_txt.py @@ -36,6 +36,11 @@ # Pages with no markdown source, linked as HTML under "## Optional". _OPTIONAL_PAGES = [ ("api/mcp/index.md", "mcp API reference", "Auto-generated API reference for the mcp package (rendered HTML)"), + ( + "api/mcp_client/index.md", + "mcp-client API reference", + "Auto-generated API reference for the mcp-client package (rendered HTML)", + ), ( "api/mcp_types/index.md", "mcp-types API reference", diff --git a/src/mcp-client/README.md b/src/mcp-client/README.md new file mode 100644 index 0000000000..195722a28d --- /dev/null +++ b/src/mcp-client/README.md @@ -0,0 +1,41 @@ +# MCP Client + +```bash +uv add mcp-client +``` + +```python +import anyio + +from mcp_client import Client + + +async def main() -> None: + async with Client("http://localhost:8000/mcp") as client: + tools = await client.list_tools() + for tool in tools.tools: + print(tool.name) + + +anyio.run(main) +``` + +Run this example against an MCP server listening at `http://localhost:8000/mcp`. + +`mcp-client` provides the official Python SDK's clients, transports, OAuth support, +and shared protocol machinery without installing Starlette, Uvicorn, +`sse-starlette`, or `python-multipart`. It depends on the matching `mcp-types` +release for protocol models. + +`mcp_client` exports the high-level client API. Import OAuth support from +`mcp_client.client.auth` and protocol models from `mcp_types`. +The full `mcp` distribution preserves the existing `mcp.client`, `mcp.shared`, +and `mcp.os` import paths as aliases to the same implementation. + +Install `mcp` instead if you also build servers, use the CLI, or pass a server +instance to `Client(server)` for in-process testing. It includes the matching +`mcp-client` release and preserves existing imports such as `from mcp import Client`. + +See the [client documentation](https://py.sdk.modelcontextprotocol.io/client/) +for usage and the [repository](https://github.com/modelcontextprotocol/python-sdk) +for development instructions. diff --git a/src/mcp-client/mcp_client/__init__.py b/src/mcp-client/mcp_client/__init__.py new file mode 100644 index 0000000000..324e5c6264 --- /dev/null +++ b/src/mcp-client/mcp_client/__init__.py @@ -0,0 +1,53 @@ +from mcp_client.client import ( + CacheConfig, + CacheEntry, + CacheKey, + CacheMode, + ClaimContext, + Client, + ClientExtension, + ClientRequestContext, + ClientSession, + IncomingMessage, + InMemoryResponseCacheStore, + InputRequiredRoundsExceededError, + NotificationBinding, + ResponseCacheStore, + ResultClaim, + Transport, + UnexpectedClaimedResult, + advertise, +) +from mcp_client.client.session_group import ClientSessionGroup +from mcp_client.client.stdio import StdioServerParameters, stdio_client +from mcp_client.shared.exceptions import MCPDeprecationWarning, MCPError, UrlElicitationRequiredError +from mcp_client.shared.uri_template import InvalidUriTemplate, UriTemplate + +__all__ = [ + "CacheConfig", + "CacheEntry", + "CacheKey", + "CacheMode", + "ClaimContext", + "Client", + "ClientExtension", + "ClientRequestContext", + "ClientSession", + "ClientSessionGroup", + "IncomingMessage", + "InMemoryResponseCacheStore", + "InputRequiredRoundsExceededError", + "InvalidUriTemplate", + "MCPDeprecationWarning", + "MCPError", + "NotificationBinding", + "ResponseCacheStore", + "ResultClaim", + "StdioServerParameters", + "Transport", + "UnexpectedClaimedResult", + "UriTemplate", + "UrlElicitationRequiredError", + "advertise", + "stdio_client", +] diff --git a/src/mcp-client/mcp_client/client/__init__.py b/src/mcp-client/mcp_client/client/__init__.py new file mode 100644 index 0000000000..cd453f83fc --- /dev/null +++ b/src/mcp-client/mcp_client/client/__init__.py @@ -0,0 +1,44 @@ +"""MCP Client module.""" + +from mcp_client.client._input_required import InputRequiredRoundsExceededError +from mcp_client.client._transport import Transport +from mcp_client.client.caching import ( + CacheConfig, + CacheEntry, + CacheKey, + CacheMode, + InMemoryResponseCacheStore, + ResponseCacheStore, +) +from mcp_client.client.client import Client +from mcp_client.client.context import ClientRequestContext +from mcp_client.client.extension import ( + ClaimContext, + ClientExtension, + NotificationBinding, + ResultClaim, + UnexpectedClaimedResult, + advertise, +) +from mcp_client.client.session import ClientSession, IncomingMessage + +__all__ = [ + "CacheConfig", + "CacheEntry", + "CacheKey", + "CacheMode", + "ClaimContext", + "Client", + "ClientExtension", + "ClientRequestContext", + "ClientSession", + "IncomingMessage", + "InMemoryResponseCacheStore", + "InputRequiredRoundsExceededError", + "NotificationBinding", + "ResponseCacheStore", + "ResultClaim", + "Transport", + "UnexpectedClaimedResult", + "advertise", +] diff --git a/src/mcp-client/mcp_client/client/__main__.py b/src/mcp-client/mcp_client/client/__main__.py new file mode 100644 index 0000000000..65a98319ed --- /dev/null +++ b/src/mcp-client/mcp_client/client/__main__.py @@ -0,0 +1,81 @@ +import argparse +import logging +import sys +import warnings +from functools import partial +from urllib.parse import urlparse + +import anyio +import mcp_types as types + +from mcp_client.client._transport import ReadStream, WriteStream +from mcp_client.client.session import ClientSession, IncomingMessage +from mcp_client.client.sse import sse_client +from mcp_client.client.stdio import StdioServerParameters, stdio_client +from mcp_client.shared.message import SessionMessage + +if not sys.warnoptions: + warnings.simplefilter("ignore") + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("client") + + +async def message_handler(message: IncomingMessage) -> None: + if isinstance(message, Exception): + logger.error("Error: %s", message) + return + + logger.info("Received message from server: %s", message) + + +async def run_session( + read_stream: ReadStream[SessionMessage | Exception], + write_stream: WriteStream[SessionMessage], + client_info: types.Implementation | None = None, +): + async with ClientSession( + read_stream, + write_stream, + message_handler=message_handler, + client_info=client_info, + ) as session: + logger.info("Initializing session") + await session.initialize() + logger.info("Initialized") + + +async def main(command_or_url: str, args: list[str], env: list[tuple[str, str]]): + env_dict = dict(env) + + if urlparse(command_or_url).scheme in ("http", "https"): + # Use SSE client for HTTP(S) URLs + async with sse_client(command_or_url) as streams: + await run_session(*streams) + else: + # Use stdio client for commands + server_parameters = StdioServerParameters(command=command_or_url, args=args, env=env_dict) + async with stdio_client(server_parameters) as streams: + await run_session(*streams) + + +def cli(): + parser = argparse.ArgumentParser() + parser.add_argument("command_or_url", help="Command or URL to connect to") + parser.add_argument("args", nargs="*", help="Additional arguments") + parser.add_argument( + "-e", + "--env", + nargs=2, + action="append", + metavar=("KEY", "VALUE"), + help="Environment variables to set. Can be used multiple times.", + default=[], + ) + + args = parser.parse_args() + anyio.run(partial(main, args.command_or_url, args.args, args.env), backend="trio") + + +if __name__ == "__main__": + cli() diff --git a/src/mcp-client/mcp_client/client/_input_required.py b/src/mcp-client/mcp_client/client/_input_required.py new file mode 100644 index 0000000000..6f4302dc34 --- /dev/null +++ b/src/mcp-client/mcp_client/client/_input_required.py @@ -0,0 +1,127 @@ +"""SEP-2322 client-side multi-round-trip driver. + +When a server returns `InputRequiredResult` instead of the normal result of a +`tools/call` / `prompts/get` / `resources/read`, the client fulfils the +embedded `input_requests` (sampling, elicitation, roots) and retries the +original request carrying the responses and the echoed opaque `request_state`. +This module implements that retry loop as a pure function so it can drive any +of the three methods identically; `Client` builds the `dispatch` and `retry` +closures, `ClientSession` stays mechanics-only. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import TypeVar + +import anyio +import anyio.abc +from mcp_types import ErrorData, InputRequest, InputRequiredResult, InputResponse, InputResponses + +from mcp_client.shared.exceptions import MCPError + +DEFAULT_INPUT_REQUIRED_MAX_ROUNDS = 10 +"""Default cap on `InputRequiredResult` retry rounds before the driver gives up. + +Matches the typescript-sdk default; csharp-sdk and go-sdk use the same value +as a hard constant. +""" + +_STATE_ONLY_BACKOFF_INITIAL_SECONDS = 0.05 +"""First sleep when an `InputRequiredResult` carries only `request_state` (no input requests).""" + +_STATE_ONLY_BACKOFF_CAP_SECONDS = 0.25 +"""Upper bound on the state-only backoff sleep; reached after three consecutive state-only legs.""" + + +ResultT = TypeVar("ResultT") + + +class InputRequiredRoundsExceededError(RuntimeError): + """The server kept returning `InputRequiredResult` past the configured `max_rounds`.""" + + def __init__(self, max_rounds: int) -> None: + super().__init__( + f"Server returned InputRequiredResult for more than {max_rounds} rounds; " + "raise input_required_max_rounds on the Client, or use " + "client.session.(..., allow_input_required=True) to drive the loop manually." + ) + self.max_rounds = max_rounds + + +async def run_input_required_driver( + first: InputRequiredResult, + *, + dispatch: Callable[[str, InputRequest], Awaitable[InputResponse | ErrorData]], + retry: Callable[[InputResponses | None, str | None], Awaitable[ResultT | InputRequiredResult]], + max_rounds: int = DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, +) -> ResultT: + """Resolve an `InputRequiredResult` to its terminal result. + + Loops until `retry` returns a non-`InputRequiredResult`, or `max_rounds` is + exhausted. Each round either dispatches all `input_requests` concurrently + and retries with the collected responses, or — when the server sent only + `request_state` — sleeps with exponential backoff (50ms doubling to a 250ms + cap, reset by any leg that carries input requests) and retries empty. + `request_state` is passed through byte-exact and never inspected. + + Args: + first: The `InputRequiredResult` the original call returned. + dispatch: Runs one embedded `InputRequest` through the client's + sampling / elicitation / roots callbacks. Called concurrently per + request key. An `ErrorData` return aborts the loop as an `MCPError`. + retry: Re-issues the original request with the collected responses and + the latest `request_state`. Each call mints a fresh JSON-RPC id. + max_rounds: Cap on retry rounds. + + Raises: + InputRequiredRoundsExceededError: `max_rounds` exhausted. + MCPError: A `dispatch` call returned `ErrorData`. + """ + rounds = 0 + state_only_delay = _STATE_ONLY_BACKOFF_INITIAL_SECONDS + current: ResultT | InputRequiredResult = first + while isinstance(current, InputRequiredResult): + rounds += 1 + if rounds > max_rounds: + raise InputRequiredRoundsExceededError(max_rounds) + if current.input_requests: + state_only_delay = _STATE_ONLY_BACKOFF_INITIAL_SECONDS + responses: InputResponses | None = await _dispatch_all(current.input_requests, dispatch) + else: + await anyio.sleep(state_only_delay) + state_only_delay = min(state_only_delay * 2, _STATE_ONLY_BACKOFF_CAP_SECONDS) + responses = None + current = await retry(responses, current.request_state) + return current + + +async def _dispatch_all( + requests: dict[str, InputRequest], + dispatch: Callable[[str, InputRequest], Awaitable[InputResponse | ErrorData]], +) -> InputResponses: + """Run `dispatch` concurrently for every key, raising `MCPError` on the first `ErrorData`. + + The first task to return `ErrorData` cancels its siblings via the task + group's cancel scope, so a refused input does not wait on a slow peer. + A callback that *raises* propagates as an `ExceptionGroup` like any other + task-group failure. + """ + responses: InputResponses = {} + refused: ErrorData | None = None + + async def run_one(tg: anyio.abc.TaskGroup, key: str, req: InputRequest) -> None: + nonlocal refused + result = await dispatch(key, req) + if isinstance(result, ErrorData): + refused = result + tg.cancel_scope.cancel() + else: + responses[key] = result + + async with anyio.create_task_group() as tg: + for key, req in requests.items(): + tg.start_soon(run_one, tg, key, req) + if refused is not None: + raise MCPError.from_error_data(refused) + return responses diff --git a/src/mcp-client/mcp_client/client/_probe.py b/src/mcp-client/mcp_client/client/_probe.py new file mode 100644 index 0000000000..5b08dc3b04 --- /dev/null +++ b/src/mcp-client/mcp_client/client/_probe.py @@ -0,0 +1,114 @@ +"""Connect-time era negotiation for ``mode='auto'``. + +The ``server/discover`` probe is sent at the newest modern version. Anything +that is not positive evidence the peer is a modern MCP server falls back to +the legacy ``initialize`` handshake — a *denylist* (only the disjoint-modern +case raises) rather than an allowlist of fallback codes. + +Every ``MCPError`` falls back except ``-32022`` with a disjoint modern-only +``supported`` list. The streamable-HTTP transport already maps HTTP-layer +4xx rejections (no JSON-RPC body) into ``MCPError`` codes, so those reach +the same path. Any non-``MCPError`` exception (network/connection errors, +anyio cancellation) propagates to the caller; an outage or in-process bug +is never an era verdict. + +A successful ``DiscoverResult`` whose ``supportedVersions`` shares no modern +version with this client is treated the same way: the server speaks discover +but advertises only handshake-era versions, which is a legacy advertisement, +not an incompatibility. + +The fallback handshake itself can be answered with ``-32022`` — e.g. a probe +that timed out client-side but succeeded on a slow-starting server locked the +connection modern before the pipelined ``initialize`` arrived. That code is +itself positive modern evidence (it names the server's versions), so it +triggers one re-probe at a mutual version instead of failing the connect. +""" + +from __future__ import annotations + +from typing import Any + +import mcp_types as types +from mcp_types import UNSUPPORTED_PROTOCOL_VERSION +from mcp_types.version import ( + HANDSHAKE_PROTOCOL_VERSIONS, + LATEST_MODERN_VERSION, + MODERN_PROTOCOL_VERSIONS, +) +from pydantic import ValidationError + +from mcp_client.client.session import ClientSession +from mcp_client.shared.exceptions import MCPError + + +def _parse_supported(data: Any) -> list[str] | None: + """Pull ``data.supported`` off a -32022 error, or ``None`` if not actionable.""" + try: + return types.UnsupportedProtocolVersionErrorData.model_validate(data).supported + except ValidationError: + return None + + +async def negotiate_auto(session: ClientSession) -> None: + """Drive the ``mode='auto'`` connect-time policy on ``session``. + + Probes ``server/discover`` once (twice if the server names a mutual + modern version via -32022), then either ``adopt()``s the result or falls + back to ``initialize()``. Idempotent only in the sense that one of + ``session.discover_result`` / ``session.initialize_result`` is set on + return. + + Raises: + MCPError: The server is modern-only and shares no version with this + client (-32022 with a disjoint ``supported`` list), or the + fallback handshake failed and one corrective re-probe did too. + Exception: Any transport/network error from the probe propagates as-is. + """ + version = LATEST_MODERN_VERSION + for attempt in range(2): + try: + raw = await session.send_discover(version) + except MCPError as e: + if e.code == UNSUPPORTED_PROTOCOL_VERSION: + supported = _parse_supported(e.error.data) + mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in (supported or ())] + if mutual and attempt == 0: + version = mutual[-1] + continue + if supported is not None and not any(v in HANDSHAKE_PROTOCOL_VERSIONS for v in supported): + raise # server is modern-only and disjoint — real incompatibility + try: + await session.initialize() # every other rpc-error → legacy (the denylist) + except MCPError as handshake_exc: + if handshake_exc.code != UNSUPPORTED_PROTOCOL_VERSION or attempt != 0: + raise + # -32022 from the handshake is itself modern evidence: a probe + # that timed out client-side but succeeded on the server locked + # the connection modern before this initialize arrived. Re-probe + # once at a version the server names; the era is already + # settled, so the second probe answers without the slow start. + supported = _parse_supported(handshake_exc.error.data) + mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in (supported or ())] + if not mutual: + raise + version = mutual[-1] + continue + return + # any other exception (httpx2.TransportError, ConnectionError, + # anyio errors) → propagate + try: + result = types.DiscoverResult.model_validate(raw) + except ValidationError: + await session.initialize() # unparseable result → not modern evidence + return + if not any(v in result.supported_versions for v in MODERN_PROTOCOL_VERSIONS): + # A discover-answering server that advertises no modern version + # (go-sdk's stateful streamable default does this) is an explicit + # legacy advertisement: fall back like the -32022 branch above + # instead of letting `adopt()` raise. The ts and go clients fall + # back here too. + await session.initialize() + return + session.adopt(result) + return + raise AssertionError("unreachable") # pragma: no cover — loop body always returns or raises diff --git a/src/mcp-client/mcp_client/client/_transport.py b/src/mcp-client/mcp_client/client/_transport.py new file mode 100644 index 0000000000..96865a656b --- /dev/null +++ b/src/mcp-client/mcp_client/client/_transport.py @@ -0,0 +1,21 @@ +"""Transport protocol for MCP clients.""" + +from __future__ import annotations + +from contextlib import AbstractAsyncContextManager +from typing import Protocol + +from mcp_client.shared._stream_protocols import ReadStream, WriteStream +from mcp_client.shared.message import SessionMessage + +__all__ = ["ReadStream", "WriteStream", "Transport", "TransportStreams"] + +TransportStreams = tuple[ReadStream[SessionMessage | Exception], WriteStream[SessionMessage]] + + +class Transport(AbstractAsyncContextManager[TransportStreams], Protocol): + """Protocol for MCP transports. + + A transport is an async context manager that yields read and write streams + for bidirectional communication with an MCP server. + """ diff --git a/src/mcp-client/mcp_client/client/auth/__init__.py b/src/mcp-client/mcp_client/client/auth/__init__.py new file mode 100644 index 0000000000..39a82f55af --- /dev/null +++ b/src/mcp-client/mcp_client/client/auth/__init__.py @@ -0,0 +1,22 @@ +"""OAuth2 Authentication implementation for httpx2. + +Implements authorization code flow with PKCE and automatic token refresh. +""" + +from mcp_client.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError +from mcp_client.client.auth.oauth2 import ( + OAuthClientProvider, + PKCEParameters, + TokenStorage, +) +from mcp_client.shared.auth import AuthorizationCodeResult + +__all__ = [ + "AuthorizationCodeResult", + "OAuthClientProvider", + "OAuthFlowError", + "OAuthRegistrationError", + "OAuthTokenError", + "PKCEParameters", + "TokenStorage", +] diff --git a/src/mcp-client/mcp_client/client/auth/exceptions.py b/src/mcp-client/mcp_client/client/auth/exceptions.py new file mode 100644 index 0000000000..5ce8777b86 --- /dev/null +++ b/src/mcp-client/mcp_client/client/auth/exceptions.py @@ -0,0 +1,10 @@ +class OAuthFlowError(Exception): + """Base exception for OAuth flow errors.""" + + +class OAuthTokenError(OAuthFlowError): + """Raised when token operations fail.""" + + +class OAuthRegistrationError(OAuthFlowError): + """Raised when client registration fails.""" diff --git a/src/mcp-client/mcp_client/client/auth/extensions/__init__.py b/src/mcp-client/mcp_client/client/auth/extensions/__init__.py new file mode 100644 index 0000000000..a9a2c5b3bb --- /dev/null +++ b/src/mcp-client/mcp_client/client/auth/extensions/__init__.py @@ -0,0 +1 @@ +__all__ = [] diff --git a/src/mcp-client/mcp_client/client/auth/extensions/client_credentials.py b/src/mcp-client/mcp_client/client/auth/extensions/client_credentials.py new file mode 100644 index 0000000000..f9131d4006 --- /dev/null +++ b/src/mcp-client/mcp_client/client/auth/extensions/client_credentials.py @@ -0,0 +1,416 @@ +"""OAuth client credential extensions for MCP. + +Provides OAuth providers for machine-to-machine authentication flows: +- ClientCredentialsOAuthProvider: For client_credentials with client_id + client_secret +- PrivateKeyJWTOAuthProvider: For client_credentials with private_key_jwt authentication + (typically using a pre-built JWT from workload identity federation) +""" + +import time +import warnings +from collections.abc import Awaitable, Callable +from typing import Any, Literal +from urllib.parse import urlparse +from uuid import uuid4 + +import httpx2 +import jwt +from pydantic import BaseModel, Field + +from mcp_client.client.auth import OAuthClientProvider, OAuthFlowError, TokenStorage +from mcp_client.client.auth.oauth2 import OAuthContext +from mcp_client.client.auth.utils import issuers_match +from mcp_client.shared.auth import OAuthClientInformationFull, OAuthClientMetadata +from mcp_client.shared.exceptions import MCPDeprecationWarning + + +def _checked_issuer(issuer: str | None) -> str | None: + if issuer is None: + warnings.warn( + "Omitting `issuer` is deprecated and it will be required in 3.0. Without it, the MCP server " + "decides which authorization server receives this client's credentials; pass " + "issuer= so they are only ever sent there.", + MCPDeprecationWarning, + stacklevel=3, + ) + return None + if urlparse(issuer).scheme not in ("http", "https"): + raise ValueError(f"issuer must be the authorization server's http(s) issuer URL, got {issuer!r}") + return issuer + + +def _preferred_authorization_server(advertised: list[str], issuer: str | None) -> str: + """The advertised server matching the configured issuer if there is one, else the first.""" + return next( + (server for server in advertised if issuer is not None and issuers_match(server, issuer)), advertised[0] + ) + + +def _require_metadata_for_configured_issuer(context: OAuthContext, issuer: str | None) -> None: + """With an issuer configured, a token request is only built from metadata discovered for that issuer. + + Anything else held is dropped along with the tokens, so the next request starts discovery afresh + rather than refreshing against it. + """ + if issuer is None: + return + metadata = context.oauth_metadata + if metadata is not None and issuers_match(str(metadata.issuer), issuer): + return + context.oauth_metadata = None + context.clear_tokens() + if metadata is None: + raise OAuthFlowError(f"No authorization server metadata discovered for configured issuer {issuer}") + raise OAuthFlowError(f"Authorization server metadata issuer mismatch: {metadata.issuer} != {issuer}") + + +class ClientCredentialsOAuthProvider(OAuthClientProvider): + """OAuth provider for client_credentials grant with client_id + client_secret. + + This provider sets client_info directly, bypassing dynamic client registration. + Use this when you already have client credentials (client_id and client_secret). + Pass `issuer` to name the authorization server those credentials belong to: token + requests are then only built from authorization server metadata for that issuer, and + the flow stops if the MCP server leads anywhere else. + + Example: + ```python + provider = ClientCredentialsOAuthProvider( + server_url="https://api.example.com", + storage=my_token_storage, + client_id="my-client-id", + client_secret="my-client-secret", + issuer="https://auth.example.com", + ) + ``` + """ + + def __init__( + self, + server_url: str, + storage: TokenStorage, + client_id: str, + client_secret: str, + token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_basic", + scope: str | None = None, + issuer: str | None = None, + ) -> None: + """Initialize client_credentials OAuth provider. + + Args: + server_url: The MCP server URL. + storage: Token storage implementation. + client_id: The OAuth client ID. + client_secret: The OAuth client secret. + token_endpoint_auth_method: Authentication method for token endpoint. + Either "client_secret_basic" (default) or "client_secret_post". + scope: Optional space-separated list of scopes to request. + issuer: The issuer identifier of the authorization server that issued + `client_id` and `client_secret`. When set, token requests are only built from + discovered authorization server metadata whose `issuer` is exactly this string; + otherwise the flow stops with `OAuthFlowError`. Omitting it is deprecated + (`MCPDeprecationWarning`) and it will be required in 3.0; until then, whichever + authorization server discovery yields is used. + """ + # Build minimal client_metadata for the base class + client_metadata = OAuthClientMetadata( + redirect_uris=None, + grant_types=["client_credentials"], + token_endpoint_auth_method=token_endpoint_auth_method, + scope=scope, + ) + super().__init__(server_url, client_metadata, storage, None, None) + self._issuer = _checked_issuer(issuer) + # Store client_info to be set during _initialize - no dynamic registration needed + self._fixed_client_info = OAuthClientInformationFull( + redirect_uris=None, + client_id=client_id, + client_secret=client_secret, + grant_types=["client_credentials"], + token_endpoint_auth_method=token_endpoint_auth_method, + scope=scope, + ) + + async def _initialize(self) -> None: + """Load stored tokens and set pre-configured client_info.""" + self.context.current_tokens = await self.context.storage.get_tokens() + self.context.client_info = self._fixed_client_info + self._initialized = True + + def _select_authorization_server(self, advertised: list[str]) -> str: + return _preferred_authorization_server(advertised, self._issuer) + + async def _perform_authorization(self) -> httpx2.Request: + """Perform client_credentials authorization.""" + return await self._exchange_token_client_credentials() + + async def _exchange_token_client_credentials(self) -> httpx2.Request: + """Build token exchange request for client_credentials grant.""" + _require_metadata_for_configured_issuer(self.context, self._issuer) + + token_data: dict[str, Any] = { + "grant_type": "client_credentials", + } + + headers: dict[str, str] = {"Content-Type": "application/x-www-form-urlencoded"} + + # Use standard auth methods (client_secret_basic, client_secret_post, none) + token_data, headers = self.context.prepare_token_auth(token_data, headers) + + if self.context.should_include_resource_param(self.context.protocol_version): + token_data["resource"] = self.context.get_resource_url() + + if self.context.client_metadata.scope: + token_data["scope"] = self.context.client_metadata.scope + + token_url = self._get_token_endpoint() + return httpx2.Request("POST", token_url, data=token_data, headers=headers) + + +def static_assertion_provider(token: str) -> Callable[[str], Awaitable[str]]: + """Create an assertion provider that returns a static JWT token. + + Use this when you have a pre-built JWT (e.g., from workload identity federation) + that doesn't need the audience parameter. + + Example: + ```python + provider = PrivateKeyJWTOAuthProvider( + server_url="https://api.example.com", + storage=my_token_storage, + client_id="my-client-id", + assertion_provider=static_assertion_provider(my_prebuilt_jwt), + issuer="https://auth.example.com", + ) + ``` + + Args: + token: The pre-built JWT assertion string. + + Returns: + An async callback suitable for use as an assertion_provider. + """ + + async def provider(audience: str) -> str: + return token + + return provider + + +class SignedJWTParameters(BaseModel): + """Parameters for creating SDK-signed JWT assertions. + + Use `create_assertion_provider()` to create an assertion provider callback + for use with `PrivateKeyJWTOAuthProvider`. + + Example: + ```python + jwt_params = SignedJWTParameters( + issuer="my-client-id", + subject="my-client-id", + signing_key=private_key_pem, + ) + provider = PrivateKeyJWTOAuthProvider( + server_url="https://api.example.com", + storage=my_token_storage, + client_id="my-client-id", + assertion_provider=jwt_params.create_assertion_provider(), + issuer="https://auth.example.com", + ) + ``` + """ + + issuer: str = Field(description="Issuer for JWT assertions (typically client_id).") + subject: str = Field(description="Subject identifier for JWT assertions (typically client_id).") + signing_key: str = Field(description="Private key for JWT signing (PEM format).") + signing_algorithm: str = Field(default="RS256", description="Algorithm for signing JWT assertions.") + lifetime_seconds: int = Field(default=300, description="Lifetime of generated JWT in seconds.") + additional_claims: dict[str, Any] | None = Field(default=None, description="Additional claims.") + + def create_assertion_provider(self) -> Callable[[str], Awaitable[str]]: + """Create an assertion provider callback for use with PrivateKeyJWTOAuthProvider. + + Returns: + An async callback that takes the audience (authorization server issuer URL) + and returns a signed JWT assertion. + """ + + async def provider(audience: str) -> str: + now = int(time.time()) + claims: dict[str, Any] = { + "iss": self.issuer, + "sub": self.subject, + "aud": audience, + "exp": now + self.lifetime_seconds, + "iat": now, + "jti": str(uuid4()), + } + if self.additional_claims: + claims.update(self.additional_claims) + + return jwt.encode(claims, self.signing_key, algorithm=self.signing_algorithm) + + return provider + + +class PrivateKeyJWTOAuthProvider(OAuthClientProvider): + """OAuth provider for client_credentials grant with private_key_jwt authentication. + + Uses RFC 7523 Section 2.2 for client authentication via JWT assertion. + + The JWT assertion's audience MUST be the authorization server's issuer identifier + (per RFC 7523bis security updates). The `assertion_provider` callback receives + this audience value and must return a JWT with that audience. Pass `issuer` to name + the authorization server this client is registered with: an assertion is then only + minted once metadata for that issuer has been discovered, and token requests are only + built from that metadata. + + **Option 1: Pre-built JWT via Workload Identity Federation** + + In production scenarios, the JWT assertion is typically obtained from a workload + identity provider (e.g., GCP, AWS IAM, Azure AD): + + ```python + async def get_workload_identity_token(audience: str) -> str: + # Fetch JWT from your identity provider + # The JWT's audience must match the provided audience parameter + return await fetch_token_from_identity_provider(audience=audience) + + provider = PrivateKeyJWTOAuthProvider( + server_url="https://api.example.com", + storage=my_token_storage, + client_id="my-client-id", + assertion_provider=get_workload_identity_token, + issuer="https://auth.example.com", + ) + ``` + + **Option 2: Static pre-built JWT** + + If you have a static JWT that doesn't need the audience parameter: + + ```python + provider = PrivateKeyJWTOAuthProvider( + server_url="https://api.example.com", + storage=my_token_storage, + client_id="my-client-id", + assertion_provider=static_assertion_provider(my_prebuilt_jwt), + issuer="https://auth.example.com", + ) + ``` + + **Option 3: SDK-signed JWT (for testing/simple setups)** + + For testing or simple deployments, use `SignedJWTParameters.create_assertion_provider()`: + + ```python + jwt_params = SignedJWTParameters( + issuer="my-client-id", + subject="my-client-id", + signing_key=private_key_pem, + ) + provider = PrivateKeyJWTOAuthProvider( + server_url="https://api.example.com", + storage=my_token_storage, + client_id="my-client-id", + assertion_provider=jwt_params.create_assertion_provider(), + issuer="https://auth.example.com", + ) + ``` + """ + + def __init__( + self, + server_url: str, + storage: TokenStorage, + client_id: str, + assertion_provider: Callable[[str], Awaitable[str]], + scope: str | None = None, + issuer: str | None = None, + ) -> None: + """Initialize private_key_jwt OAuth provider. + + Args: + server_url: The MCP server URL. + storage: Token storage implementation. + client_id: The OAuth client ID. + assertion_provider: Async callback that takes the audience (authorization + server's issuer identifier) and returns a JWT assertion. Use + `SignedJWTParameters.create_assertion_provider()` for SDK-signed JWTs, + `static_assertion_provider()` for pre-built JWTs, or provide your own + callback for workload identity federation. + scope: Optional space-separated list of scopes to request. + issuer: The issuer identifier of the authorization server `client_id` is + registered with. When set, an assertion is only minted, and token requests + are only built, once authorization server metadata whose `issuer` is exactly this + string has been discovered; otherwise the flow stops with `OAuthFlowError`. + Omitting it is deprecated (`MCPDeprecationWarning`) and it will be required in + 3.0; until then, whichever authorization server discovery yields is used. + """ + # Build minimal client_metadata for the base class + client_metadata = OAuthClientMetadata( + redirect_uris=None, + grant_types=["client_credentials"], + token_endpoint_auth_method="private_key_jwt", + scope=scope, + ) + super().__init__(server_url, client_metadata, storage, None, None) + self._assertion_provider = assertion_provider + self._issuer = _checked_issuer(issuer) + # Store client_info to be set during _initialize - no dynamic registration needed + self._fixed_client_info = OAuthClientInformationFull( + redirect_uris=None, + client_id=client_id, + grant_types=["client_credentials"], + token_endpoint_auth_method="private_key_jwt", + scope=scope, + ) + + async def _initialize(self) -> None: + """Load stored tokens and set pre-configured client_info.""" + self.context.current_tokens = await self.context.storage.get_tokens() + self.context.client_info = self._fixed_client_info + self._initialized = True + + def _select_authorization_server(self, advertised: list[str]) -> str: + return _preferred_authorization_server(advertised, self._issuer) + + async def _perform_authorization(self) -> httpx2.Request: + """Perform client_credentials authorization with private_key_jwt.""" + return await self._exchange_token_client_credentials() + + async def _add_client_authentication_jwt(self, *, token_data: dict[str, Any]) -> None: + """Add JWT assertion for client authentication to token endpoint parameters.""" + if not self.context.oauth_metadata: + raise OAuthFlowError("Missing OAuth metadata for private_key_jwt flow") # pragma: no cover + + # Audience MUST be the issuer identifier of the authorization server + # https://datatracker.ietf.org/doc/html/draft-ietf-oauth-rfc7523bis-01 + audience = str(self.context.oauth_metadata.issuer) + assertion = await self._assertion_provider(audience) + + # RFC 7523 Section 2.2: client authentication via JWT + token_data["client_assertion"] = assertion + token_data["client_assertion_type"] = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer" + + async def _exchange_token_client_credentials(self) -> httpx2.Request: + """Build token exchange request for client_credentials grant with private_key_jwt.""" + _require_metadata_for_configured_issuer(self.context, self._issuer) + + token_data: dict[str, Any] = { + "grant_type": "client_credentials", + } + + headers: dict[str, str] = {"Content-Type": "application/x-www-form-urlencoded"} + + # Add JWT client authentication (RFC 7523 Section 2.2) + await self._add_client_authentication_jwt(token_data=token_data) + + if self.context.should_include_resource_param(self.context.protocol_version): + token_data["resource"] = self.context.get_resource_url() + + if self.context.client_metadata.scope: + token_data["scope"] = self.context.client_metadata.scope + + token_url = self._get_token_endpoint() + return httpx2.Request("POST", token_url, data=token_data, headers=headers) diff --git a/src/mcp-client/mcp_client/client/auth/extensions/identity_assertion.py b/src/mcp-client/mcp_client/client/auth/extensions/identity_assertion.py new file mode 100644 index 0000000000..6756ffd713 --- /dev/null +++ b/src/mcp-client/mcp_client/client/auth/extensions/identity_assertion.py @@ -0,0 +1,216 @@ +"""SEP-990 Identity Assertion Authorization Grant (RFC 7523 jwt-bearer) client provider. + +`IdentityAssertionOAuthProvider` is the client side of SEP-990 leg 2: it presents an Identity +Assertion Authorization Grant (ID-JAG) - a signed JWT issued by the enterprise identity provider - +to the MCP authorization server's token endpoint using the RFC 7523 jwt-bearer grant +(`grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer`, ID-JAG as `assertion`), and receives an +MCP access token. + +The authorization server is configuration, not discovery. SEP-990's trust model is the inverse of +the default OAuth client's: the AS issuer is supplied at construction, authorization-server metadata +is fetched from that issuer's own RFC 8414 well-known, and the resource server is never asked which +AS to use - so it cannot redirect the ID-JAG or client secret elsewhere. There is no protected +resource metadata fetch, no dynamic client registration, and no server-driven scope selection. + +Obtaining the ID-JAG (logging into the IdP and the leg-1 token exchange against it) is +deployment-specific and out of scope for the SDK. The caller supplies it through the +`assertion_provider` callback, which receives the configured issuer (the `aud` the ID-JAG must +carry) and the MCP server's resource identifier (the `resource` claim it must carry, per ext-auth +section 4.3), and returns the ID-JAG. +""" + +import base64 +import time +from collections.abc import AsyncGenerator, Awaitable, Callable +from typing import Literal +from urllib.parse import quote, urlsplit + +import anyio +import httpx2 + +from mcp_client.client.auth import OAuthFlowError, OAuthTokenError, TokenStorage +from mcp_client.client.auth.utils import ( + build_oauth_authorization_server_metadata_discovery_urls, + create_oauth_metadata_request, + extract_field_from_www_auth, + extract_scope_from_www_auth, + handle_auth_metadata_response, + handle_token_response_scopes, + union_scopes, + validate_metadata_issuer, +) +from mcp_client.shared._httpx_utils import RedirectAwareAuth, redirect_note +from mcp_client.shared.auth import JWT_BEARER_GRANT_TYPE, OAuthClientInformationFull, OAuthToken +from mcp_client.shared.auth_utils import calculate_token_expiry, resource_url_from_server_url + +_DEFAULT_PORTS = {"https": 443, "http": 80} + + +def _origin(url: str) -> tuple[str, str, int | None]: + """Return the (scheme, host, port) origin of a URL for same-origin comparison. + + The port is normalized to the scheme's default so an explicit `:443`/`:80` compares equal to the + same origin written without a port. + """ + parsed = urlsplit(url) + port = parsed.port if parsed.port is not None else _DEFAULT_PORTS.get(parsed.scheme) + return (parsed.scheme, parsed.hostname or "", port) + + +class IdentityAssertionOAuthProvider(RedirectAwareAuth): + """`httpx2.Auth` for the SEP-990 ID-JAG flow (RFC 7523 jwt-bearer grant) against a configured AS. + + The authorization server `issuer` is fixed at construction; metadata is fetched from its + RFC 8414 well-known and the ID-JAG and client secret are sent only to that issuer's token + endpoint. The resource server is never consulted for AS selection. The ID-JAG is fetched lazily + from `assertion_provider` so a fresh assertion is used on each exchange. + + Example: + ```python + async def fetch_id_jag(audience: str, resource: str) -> str: + # `audience` is the configured issuer (the ID-JAG `aud`); `resource` is the MCP + # server's identifier (the ID-JAG `resource` claim). Obtaining the ID-JAG from the + # enterprise IdP is deployment-specific and not handled by the SDK. + return await my_idp.issue_id_jag(audience=audience, resource=resource) + + + provider = IdentityAssertionOAuthProvider( + server_url="https://mcp.example.com/mcp", + storage=my_token_storage, + client_id="my-client-id", + client_secret="my-client-secret", + issuer="https://auth.example.com", + assertion_provider=fetch_id_jag, + ) + ``` + """ + + requires_response_body = True + + def __init__( + self, + server_url: str, + storage: TokenStorage, + client_id: str, + client_secret: str, + issuer: str, + assertion_provider: Callable[[str, str], Awaitable[str]], + scope: str | None = None, + token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_post", + ) -> None: + """Initialize the identity-assertion OAuth provider. + + Args: + server_url: The MCP server URL. + storage: Token storage implementation. + client_id: The OAuth client ID registered with the MCP authorization server. + client_secret: The client secret. SEP-990 section 5.1 requires a confidential client. + issuer: The issuer identifier of the MCP authorization server this client is provisioned + for. Authorization-server metadata is fetched from this issuer's well-known and the + ID-JAG and secret are sent only to its token endpoint. + assertion_provider: Async callback taking `(audience, resource)` - the configured issuer + and the MCP server's resource identifier - and returning the ID-JAG. + scope: Optional space-separated list of scopes to request. + token_endpoint_auth_method: Confidential-client auth method, either `client_secret_post` + (default) or `client_secret_basic`. + """ + if not client_secret: + raise ValueError("client_secret is required: SEP-990 mandates a confidential client") + if not issuer: + raise ValueError("issuer is required: the authorization server is configuration, not discovery") + self._resource = resource_url_from_server_url(server_url) + self._storage = storage + self._issuer = issuer + self._assertion_provider = assertion_provider + self._scope = scope + self._client = OAuthClientInformationFull( + client_id=client_id, + client_secret=client_secret, + redirect_uris=None, + grant_types=[JWT_BEARER_GRANT_TYPE], + token_endpoint_auth_method=token_endpoint_auth_method, + issuer=issuer, + ) + self._token_endpoint: str | None = None + self._tokens: OAuthToken | None = None + self._expiry: float | None = None + self._lock = anyio.Lock() + self._initialized = False + + def _build_token_request(self, scope: str | None, assertion: str) -> httpx2.Request: + """Build the RFC 7523 jwt-bearer token request, applying confidential-client auth.""" + assert self._token_endpoint is not None + assert self._client.client_id is not None and self._client.client_secret is not None + data: dict[str, str] = { + "grant_type": JWT_BEARER_GRANT_TYPE, + "assertion": assertion, + "client_id": self._client.client_id, + "resource": self._resource, + } + if scope: + data["scope"] = scope + headers = {"Content-Type": "application/x-www-form-urlencoded"} + if self._client.token_endpoint_auth_method == "client_secret_basic": + # RFC 6749 section 2.3.1: URL-encode each part, then base64 the colon-joined pair. + encoded_id = quote(self._client.client_id, safe="") + encoded_secret = quote(self._client.client_secret, safe="") + credentials = base64.b64encode(f"{encoded_id}:{encoded_secret}".encode()).decode() + headers["Authorization"] = f"Basic {credentials}" + else: + data["client_secret"] = self._client.client_secret + return httpx2.Request("POST", self._token_endpoint, data=data, headers=headers) + + async def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + async with self._lock: + if not self._initialized: + self._tokens = await self._storage.get_tokens() + self._expiry = calculate_token_expiry(self._tokens.expires_in) if self._tokens else None + self._initialized = True + + if self._tokens and (self._expiry is None or time.time() <= self._expiry): + request.headers["Authorization"] = f"Bearer {self._tokens.access_token}" + response = yield request + + if response.status_code == 401: + scope_to_request = self._scope + elif response.status_code == 403 and extract_field_from_www_auth(response, "error") == "insufficient_scope": + scope_to_request = union_scopes(self._scope, extract_scope_from_www_auth(response)) + else: + return + + # Discover ASM from the configured issuer's well-known. The RS is not consulted: both + # arguments are the issuer, so even the helper's legacy fallback resolves there. + if self._token_endpoint is None: + for url in build_oauth_authorization_server_metadata_discovery_urls(self._issuer, self._issuer): + asm_response = yield create_oauth_metadata_request(url) + ok, asm = await handle_auth_metadata_response(asm_response) + if not ok: + break + if asm is not None: + validate_metadata_issuer(asm, self._issuer) + token_endpoint = str(asm.token_endpoint) + if _origin(token_endpoint) != _origin(self._issuer): + raise OAuthFlowError( + f"Token endpoint {token_endpoint} is not on the configured issuer origin {self._issuer}" + ) + self._token_endpoint = token_endpoint + break + if self._token_endpoint is None: + raise OAuthFlowError(f"No authorization server metadata at configured issuer {self._issuer}") + + assertion = await self._assertion_provider(self._issuer, self._resource) + token_response = yield self._build_token_request(scope_to_request, assertion) + if token_response.status_code != 200: + body = (await token_response.aread()).decode(errors="replace") + raise OAuthTokenError( + f"Token exchange failed ({token_response.status_code}){redirect_note(token_response)}: {body}" + ) + tokens = await handle_token_response_scopes(token_response) + if tokens.scope is None: + tokens.scope = scope_to_request + self._tokens = tokens + self._expiry = calculate_token_expiry(tokens.expires_in) + await self._storage.set_tokens(tokens) + + request.headers["Authorization"] = f"Bearer {tokens.access_token}" + yield request diff --git a/src/mcp-client/mcp_client/client/auth/oauth2.py b/src/mcp-client/mcp_client/client/auth/oauth2.py new file mode 100644 index 0000000000..3152e5203d --- /dev/null +++ b/src/mcp-client/mcp_client/client/auth/oauth2.py @@ -0,0 +1,793 @@ +"""OAuth2 Authentication implementation for httpx2. + +Implements authorization code flow with PKCE and automatic token refresh. +""" + +import base64 +import hashlib +import logging +import secrets +import string +import time +from collections.abc import AsyncGenerator, Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any, Protocol, get_args +from urllib.parse import quote, urlencode, urljoin, urlparse + +import anyio +import httpx2 +from mcp_types.version import is_version_at_least +from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +from mcp_client.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError +from mcp_client.client.auth.utils import ( + build_oauth_authorization_server_metadata_discovery_urls, + build_protected_resource_metadata_discovery_urls, + create_client_info_from_metadata_url, + create_client_registration_request, + create_oauth_metadata_request, + credentials_match_issuer, + extract_field_from_www_auth, + extract_resource_metadata_from_www_auth, + extract_scope_from_www_auth, + get_client_metadata_scopes, + handle_auth_metadata_response, + handle_protected_resource_response, + handle_registration_response, + handle_token_response_scopes, + is_valid_client_metadata_url, + issuers_match, + should_use_client_metadata_url, + union_scopes, + validate_authorization_response_iss, + validate_metadata_issuer, +) +from mcp_client.shared._httpx_utils import RedirectAwareAuth, redirect_note +from mcp_client.shared.auth import ( + AuthorizationCodeResult, + OAuthClientInformationFull, + OAuthClientMetadata, + OAuthMetadata, + OAuthToken, + ProtectedResourceMetadata, + TokenEndpointAuthMethod, +) +from mcp_client.shared.auth_utils import ( + calculate_token_expiry, + check_resource_allowed, + resource_url_from_server_url, +) +from mcp_client.shared.inbound import MCP_PROTOCOL_VERSION_HEADER + +logger = logging.getLogger("mcp.client.auth.oauth2") + +# Methods a registered client's record may carry without a token request being an error, +# derived from the set the SDK is willing to request so the two cannot drift. `None`/"none" +# send no client secret. `private_key_jwt` sends none from here either: only +# `PrivateKeyJWTOAuthProvider` signs the assertion, and only in its client-credentials +# exchange, so its inherited refresh path must pass through here without raising - a refresh +# the server then rejects falls back to a fresh client-credentials exchange, which signs. +# Anything else is a method no client here can apply. +_KNOWN_TOKEN_ENDPOINT_AUTH_METHODS: tuple[str | None, ...] = (None, *get_args(TokenEndpointAuthMethod)) + +# Methods that authenticate the token request with the minted `client_secret`; a +# registration assigning one is only usable if the server issued that secret. +_SECRET_TOKEN_ENDPOINT_AUTH_METHODS = ("client_secret_post", "client_secret_basic") + +# Methods a registration completed by the authorization-code flow can act on. That flow +# authenticates the token request with the minted client secret (or nothing); it holds no key +# to sign a `private_key_jwt` assertion, so a server assigning that method has registered a +# client this flow cannot use. `PrivateKeyJWTOAuthProvider` never registers dynamically. +_REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS: tuple[str | None, ...] = tuple( + method for method in _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS if method != "private_key_jwt" +) + + +def check_registration_usable(client_info: OAuthClientInformationFull) -> None: + """Confirm a registration this flow completed is one it can act on. + + RFC 7591 §3.2.1 lets the authorization server replace requested metadata and leaves it to + the client to "check the values in the response to determine if the registration is + sufficient for use". Two substitutions make the minted credentials unusable, and both are + judged here - before the record is persisted or any interactive authorization begins - + rather than surfacing later as an opaque failure at the token endpoint: a token-endpoint + auth method the authorization-code flow cannot apply (one it does not implement, or + `private_key_jwt`, whose assertion this flow has no key to sign), and a secret-based + method the flow could apply but for which the server issued no `client_secret`. + + Raises: + OAuthRegistrationError: The server registered the client with a + `token_endpoint_auth_method` this flow cannot apply, or with a secret-based + method but no `client_secret`. + """ + method = client_info.token_endpoint_auth_method + if method not in _REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS: + raise OAuthRegistrationError( + f"Authorization server registered the client with unsupported token_endpoint_auth_method {method!r}" + ) + if method in _SECRET_TOKEN_ENDPOINT_AUTH_METHODS and client_info.client_secret is None: + raise OAuthRegistrationError( + f"Authorization server registered the client for {method!r} but issued no client_secret" + ) + + +class PKCEParameters(BaseModel): + """PKCE (Proof Key for Code Exchange) parameters.""" + + code_verifier: str = Field(..., min_length=43, max_length=128) + code_challenge: str = Field(..., min_length=43, max_length=128) + + @classmethod + def generate(cls) -> "PKCEParameters": + """Generate new PKCE parameters.""" + code_verifier = "".join(secrets.choice(string.ascii_letters + string.digits + "-._~") for _ in range(128)) + digest = hashlib.sha256(code_verifier.encode()).digest() + code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=") + return cls(code_verifier=code_verifier, code_challenge=code_challenge) + + +class TokenStorage(Protocol): + """Protocol for token storage implementations.""" + + async def get_tokens(self) -> OAuthToken | None: + """Get stored tokens.""" + ... + + async def set_tokens(self, tokens: OAuthToken) -> None: + """Store tokens.""" + ... + + async def get_client_info(self) -> OAuthClientInformationFull | None: + """Get stored client information.""" + ... + + async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + """Store client information.""" + ... + + +@dataclass +class OAuthContext: + """OAuth flow context.""" + + server_url: str + client_metadata: OAuthClientMetadata + storage: TokenStorage + redirect_handler: Callable[[str], Awaitable[None]] | None + callback_handler: Callable[[], Awaitable[AuthorizationCodeResult]] | None + client_metadata_url: str | None = None + + # Discovered metadata + protected_resource_metadata: ProtectedResourceMetadata | None = None + oauth_metadata: OAuthMetadata | None = None + auth_server_url: str | None = None + protocol_version: str | None = None + + # Client registration + client_info: OAuthClientInformationFull | None = None + + # Token management + current_tokens: OAuthToken | None = None + token_expiry_time: float | None = None + + # State + lock: anyio.Lock = field(default_factory=anyio.Lock) + + def get_authorization_base_url(self, server_url: str) -> str: + """Extract base URL by removing path component.""" + parsed = urlparse(server_url) + return f"{parsed.scheme}://{parsed.netloc}" + + def update_token_expiry(self, token: OAuthToken) -> None: + """Update token expiry time using shared util function.""" + self.token_expiry_time = calculate_token_expiry(token.expires_in) + + def is_token_valid(self) -> bool: + """Check if current token is valid.""" + return bool( + self.current_tokens + and self.current_tokens.access_token + and (not self.token_expiry_time or time.time() <= self.token_expiry_time) + ) + + def can_refresh_token(self) -> bool: + """Check if token can be refreshed.""" + return bool(self.current_tokens and self.current_tokens.refresh_token and self.client_info) + + def clear_tokens(self) -> None: + """Clear current tokens.""" + self.current_tokens = None + self.token_expiry_time = None + + def get_resource_url(self) -> str: + """Get resource URL for RFC 8707. + + Uses PRM resource if it's a valid parent, otherwise uses canonical server URL. + """ + resource = resource_url_from_server_url(self.server_url) + + # If PRM provides a resource that's a valid parent, use it + if self.protected_resource_metadata and self.protected_resource_metadata.resource: + prm_resource = str(self.protected_resource_metadata.resource) + if check_resource_allowed(requested_resource=resource, configured_resource=prm_resource): + resource = prm_resource + + return resource + + def should_include_resource_param(self, protocol_version: str | None = None) -> bool: + """Determine if the resource parameter should be included in OAuth requests. + + Returns True if: + - Protected resource metadata is available, OR + - MCP-Protocol-Version header is 2025-06-18 or later + """ + # If we have protected resource metadata, include the resource param + if self.protected_resource_metadata is not None: + return True + + # If no protocol version provided, don't include resource param + if not protocol_version: + return False + + return is_version_at_least(protocol_version, "2025-06-18") + + def prepare_token_auth( + self, data: dict[str, str], headers: dict[str, str] | None = None + ) -> tuple[dict[str, str], dict[str, str]]: + """Prepare authentication for token requests. + + Args: + data: The form data to send + headers: Optional headers dict to update + + Returns: + Tuple of (updated_data, updated_headers) + + Raises: + OAuthTokenError: The client record carries a `token_endpoint_auth_method` this + client does not know. A dynamic registration assigning an unusable method is + rejected earlier, by `check_registration_usable`; this fires for a stored or + pre-registered record that reaches a token request with such a method. + """ + if headers is None: + headers = {} # pragma: no cover + + if not self.client_info: + return data, headers + + auth_method = self.client_info.token_endpoint_auth_method + + if auth_method == "client_secret_basic" and self.client_info.client_secret: + # URL-encode client ID and secret per RFC 6749 Section 2.3.1 + encoded_id = quote(self.client_info.client_id, safe="") + encoded_secret = quote(self.client_info.client_secret, safe="") + credentials = f"{encoded_id}:{encoded_secret}" + encoded_credentials = base64.b64encode(credentials.encode()).decode() + headers["Authorization"] = f"Basic {encoded_credentials}" + # Don't include client_secret in body for basic auth + data = {k: v for k, v in data.items() if k != "client_secret"} + elif auth_method == "client_secret_post" and self.client_info.client_secret: + # Include client_id and client_secret in request body (RFC 6749 §2.3.1) + data["client_id"] = self.client_info.client_id + data["client_secret"] = self.client_info.client_secret + elif auth_method not in _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS: + raise OAuthTokenError(f"Registered client uses unsupported token_endpoint_auth_method {auth_method!r}") + # For "none" (or absent), don't add any client_secret; "private_key_jwt" adds its + # assertion in the provider that implements it, not here. + + return data, headers + + +_ORIGIN_URL = TypeAdapter(AnyHttpUrl, config=ConfigDict(url_preserve_empty_path=True)) + + +def _origin_issuer(server_url: str) -> str: + """The resource server's origin as an issuer identifier: `scheme://authority`, rendered the way + `OAuthMetadata.issuer` renders URLs (host case, default ports) so the two compare as strings.""" + parsed = urlparse(server_url) + return str(_ORIGIN_URL.validate_python(f"{parsed.scheme}://{parsed.netloc}")) + + +class OAuthClientProvider(RedirectAwareAuth): + """OAuth2 authentication for httpx2. + + Handles OAuth flow with automatic client registration and token storage. + """ + + requires_response_body = True + + def __init__( + self, + server_url: str, + client_metadata: OAuthClientMetadata, + storage: TokenStorage, + redirect_handler: Callable[[str], Awaitable[None]] | None = None, + callback_handler: Callable[[], Awaitable[AuthorizationCodeResult]] | None = None, + client_metadata_url: str | None = None, + validate_resource_url: Callable[[str, str | None], Awaitable[None]] | None = None, + ): + """Initialize OAuth2 authentication. + + Args: + server_url: The MCP server URL. + client_metadata: OAuth client metadata for registration. + storage: Token storage implementation. + redirect_handler: Handler for authorization redirects. + callback_handler: Handler for authorization callbacks. + client_metadata_url: URL-based client ID. When provided and the server + advertises client_id_metadata_document_supported=True, this URL will be + used as the client_id instead of performing dynamic client registration. + Must be a valid HTTPS URL with a non-root pathname. + validate_resource_url: Optional callback to override resource URL validation. + Called with (server_url, prm_resource) where prm_resource is the resource + from Protected Resource Metadata (or None if not present). If not provided, + default validation rejects mismatched resources per RFC 8707. + + Raises: + ValueError: If client_metadata_url is provided but not a valid HTTPS URL + with a non-root pathname. + """ + # Validate client_metadata_url if provided + if client_metadata_url is not None and not is_valid_client_metadata_url(client_metadata_url): + raise ValueError( + f"client_metadata_url must be a valid HTTPS URL with a non-root pathname, got: {client_metadata_url}" + ) + + self.context = OAuthContext( + server_url=server_url, + client_metadata=client_metadata, + storage=storage, + redirect_handler=redirect_handler, + callback_handler=callback_handler, + client_metadata_url=client_metadata_url, + ) + self._validate_resource_url_callback = validate_resource_url + self._initialized = False + + async def _handle_protected_resource_response(self, response: httpx2.Response) -> bool: + """Handle protected resource metadata discovery response. + + Per SEP-985, supports fallback when discovery fails at one URL. + + Returns: + True if metadata was successfully discovered, False if we should try next URL + """ + if response.status_code == 200: + try: + content = await response.aread() + metadata = ProtectedResourceMetadata.model_validate_json(content) + self.context.protected_resource_metadata = metadata + if metadata.authorization_servers: # pragma: no branch + self.context.auth_server_url = str(metadata.authorization_servers[0]) + return True + + except ValidationError: # pragma: no cover + # Invalid metadata - try next URL + logger.warning(f"Invalid protected resource metadata at {response.request.url}") + return False + elif response.status_code == 404: # pragma: no cover + # Not found - try next URL in fallback chain + logger.debug(f"Protected resource metadata not found at {response.request.url}, trying next URL") + return False + else: + # Other error - fail immediately + raise OAuthFlowError( + f"Protected Resource Metadata request failed: {response.status_code}" + ) # pragma: no cover + + async def _perform_authorization(self) -> httpx2.Request: + """Perform the authorization flow.""" + auth_code, code_verifier = await self._perform_authorization_code_grant() + token_request = await self._exchange_token_authorization_code(auth_code, code_verifier) + return token_request + + async def _perform_authorization_code_grant(self) -> tuple[str, str]: + """Perform the authorization redirect and get auth code.""" + if self.context.client_metadata.redirect_uris is None: + raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover + if not self.context.redirect_handler: + raise OAuthFlowError("No redirect handler provided for authorization code grant") # pragma: no cover + if not self.context.callback_handler: + raise OAuthFlowError("No callback handler provided for authorization code grant") # pragma: no cover + + if self.context.oauth_metadata and self.context.oauth_metadata.authorization_endpoint: + auth_endpoint = str(self.context.oauth_metadata.authorization_endpoint) + else: + auth_base_url = self.context.get_authorization_base_url(self.context.server_url) + auth_endpoint = urljoin(auth_base_url, "/authorize") + + if not self.context.client_info: + raise OAuthFlowError("No client info available for authorization") # pragma: no cover + + # Generate PKCE parameters + pkce_params = PKCEParameters.generate() + state = secrets.token_urlsafe(32) + + auth_params = { + "response_type": "code", + "client_id": self.context.client_info.client_id, + "redirect_uri": str(self.context.client_metadata.redirect_uris[0]), + "state": state, + "code_challenge": pkce_params.code_challenge, + "code_challenge_method": "S256", + } + + # Only include resource param if conditions are met + if self.context.should_include_resource_param(self.context.protocol_version): + auth_params["resource"] = self.context.get_resource_url() # RFC 8707 + + if self.context.client_metadata.scope: # pragma: no branch + auth_params["scope"] = self.context.client_metadata.scope + + # OIDC requires prompt=consent when offline_access is requested + # https://openid.net/specs/openid-connect-core-1_0.html#OfflineAccess + if "offline_access" in self.context.client_metadata.scope.split(): + auth_params["prompt"] = "consent" + + authorization_url = f"{auth_endpoint}?{urlencode(auth_params)}" + await self.context.redirect_handler(authorization_url) + + # Wait for callback + result = await self.context.callback_handler() + + if result.state is None or not secrets.compare_digest(result.state, state): + raise OAuthFlowError(f"State parameter mismatch: {result.state} != {state}") + + # RFC 9207: validate the authorization-response issuer + validate_authorization_response_iss(result.iss, self.context.oauth_metadata) + + if not result.code: + raise OAuthFlowError("No authorization code received") + + # Return auth code and code verifier for token exchange + return result.code, pkce_params.code_verifier + + def _get_token_endpoint(self) -> str: + if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint: + token_url = str(self.context.oauth_metadata.token_endpoint) + else: + auth_base_url = self.context.get_authorization_base_url(self.context.server_url) + token_url = urljoin(auth_base_url, "/token") + return token_url + + async def _exchange_token_authorization_code(self, auth_code: str, code_verifier: str) -> httpx2.Request: + """Build token exchange request for authorization_code flow.""" + if self.context.client_metadata.redirect_uris is None: + raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover + if not self.context.client_info: + raise OAuthFlowError("Missing client info") # pragma: no cover + + token_url = self._get_token_endpoint() + token_data: dict[str, Any] = { + "grant_type": "authorization_code", + "code": auth_code, + "redirect_uri": str(self.context.client_metadata.redirect_uris[0]), + "client_id": self.context.client_info.client_id, + "code_verifier": code_verifier, + } + + # Only include resource param if conditions are met + if self.context.should_include_resource_param(self.context.protocol_version): + token_data["resource"] = self.context.get_resource_url() # RFC 8707 + + # Prepare authentication based on preferred method + headers = {"Content-Type": "application/x-www-form-urlencoded"} + token_data, headers = self.context.prepare_token_auth(token_data, headers) + + return httpx2.Request("POST", token_url, data=token_data, headers=headers) + + async def _handle_token_response(self, response: httpx2.Response) -> None: + """Handle token exchange response.""" + if response.status_code not in {200, 201}: + body = await response.aread() + body_text = body.decode("utf-8") + raise OAuthTokenError( + f"Token exchange failed ({response.status_code}){redirect_note(response)}: {body_text}" + ) + + # Parse and validate response with scope validation + token_response = await handle_token_response_scopes(response) + + # RFC 6749 §5.1: an omitted scope means the granted scope equals the requested + # scope. Record it explicitly so the persisted token is self-describing — the + # SEP-2350 step-up union reads it after a restart, when client_metadata.scope + # has reverted to its constructor value. + if token_response.scope is None: + token_response.scope = self.context.client_metadata.scope + + # Store tokens in context + self.context.current_tokens = token_response + self.context.update_token_expiry(token_response) + await self.context.storage.set_tokens(token_response) + + async def _refresh_token(self) -> httpx2.Request: + """Build token refresh request.""" + if not self.context.current_tokens or not self.context.current_tokens.refresh_token: + raise OAuthTokenError("No refresh token available") # pragma: no cover + + if not self.context.client_info or not self.context.client_info.client_id: + raise OAuthTokenError("No client info available") # pragma: no cover + + if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint: + token_url = str(self.context.oauth_metadata.token_endpoint) + else: + auth_base_url = self.context.get_authorization_base_url(self.context.server_url) + token_url = urljoin(auth_base_url, "/token") + + refresh_data: dict[str, str] = { + "grant_type": "refresh_token", + "refresh_token": self.context.current_tokens.refresh_token, + "client_id": self.context.client_info.client_id, + } + + # Only include resource param if conditions are met + if self.context.should_include_resource_param(self.context.protocol_version): + refresh_data["resource"] = self.context.get_resource_url() # RFC 8707 + + # Prepare authentication based on preferred method + headers = {"Content-Type": "application/x-www-form-urlencoded"} + refresh_data, headers = self.context.prepare_token_auth(refresh_data, headers) + + return httpx2.Request("POST", token_url, data=refresh_data, headers=headers) + + async def _handle_refresh_response(self, response: httpx2.Response) -> bool: + """Handle token refresh response. Returns True if successful.""" + if response.status_code != 200: + logger.warning(f"Token refresh failed: {response.status_code}{redirect_note(response)}") + self.context.clear_tokens() + return False + + try: + content = await response.aread() + token_response = OAuthToken.model_validate_json(content) + + # RFC 6749 §6: a refresh response may omit scope (unchanged) and refresh_token + # (the AS does not rotate). Carry both forward so the persisted token stays + # self-describing for the SEP-2350 step-up union and the next expiry can + # still refresh instead of forcing a full re-authorization. + prior = self.context.current_tokens + if token_response.scope is None and prior is not None: + token_response.scope = prior.scope + if token_response.refresh_token is None and prior is not None: + token_response.refresh_token = prior.refresh_token + + self.context.current_tokens = token_response + self.context.update_token_expiry(token_response) + await self.context.storage.set_tokens(token_response) + + return True + except ValidationError: # pragma: no cover + logger.exception("Invalid refresh response") + self.context.clear_tokens() + return False + + async def _initialize(self) -> None: + """Load stored tokens and client info.""" + self.context.current_tokens = await self.context.storage.get_tokens() + self.context.client_info = await self.context.storage.get_client_info() + self._initialized = True + + def _add_auth_header(self, request: httpx2.Request) -> None: + """Add authorization header to request if we have valid tokens.""" + if self.context.current_tokens and self.context.current_tokens.access_token: # pragma: no branch + request.headers["Authorization"] = f"Bearer {self.context.current_tokens.access_token}" + + async def _handle_oauth_metadata_response(self, response: httpx2.Response) -> None: + content = await response.aread() + metadata = OAuthMetadata.model_validate_json(content) + self.context.oauth_metadata = metadata + + async def _validate_resource_match(self, prm: ProtectedResourceMetadata) -> None: + """Validate that PRM resource matches the server URL per RFC 8707.""" + prm_resource = str(prm.resource) if prm.resource else None + + if self._validate_resource_url_callback is not None: + await self._validate_resource_url_callback(self.context.server_url, prm_resource) + return + + if not prm_resource: + return # pragma: no cover + default_resource = resource_url_from_server_url(self.context.server_url) + if not check_resource_allowed(requested_resource=default_resource, configured_resource=prm_resource): + raise OAuthFlowError(f"Protected resource {prm_resource} does not match expected {default_resource}") + + def _select_authorization_server(self, advertised: list[str]) -> str: + """Which of the servers listed in protected resource metadata to use: the first (the list is never empty).""" + return advertised[0] + + def _expected_issuer(self) -> str: + """The issuer that authorization server metadata and client credentials must belong to: the + PRM-advertised server, or on the legacy no-PRM path the resource server's origin, which is what + the 2025-03-26 well-known URL is built from (RFC 8414 §3.3).""" + return self.context.auth_server_url or _origin_issuer(self.context.server_url) + + async def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + """The OAuth flow proper; `async_auth_flow` drives it (see `RedirectAwareAuth`).""" + async with self.context.lock: + if not self._initialized: + await self._initialize() + + # Capture protocol version from request headers + self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER) + + if not self.context.is_token_valid() and self.context.can_refresh_token(): + # Try to refresh token + refresh_request = await self._refresh_token() + refresh_response = yield refresh_request + + if not await self._handle_refresh_response(refresh_response): + # Refresh failed, need full re-authentication + self._initialized = False + + if self.context.is_token_valid(): + self._add_auth_header(request) + + response = yield request + + step_up = ( + response.status_code == 403 and extract_field_from_www_auth(response, "error") == "insufficient_scope" + ) + + if response.status_code == 401 or step_up: + # Perform full OAuth flow + try: + # Read before discovery, which may clear the tokens: on a restart the stored + # token's scope is the only record of what was granted (see Step 3). + granted_scope = self.context.current_tokens.scope if self.context.current_tokens else None + + # OAuth flow must be inline due to generator constraints. + # Steps 1-2 run on every 401. A scope step-up reuses the metadata discovered earlier + # in this process, and discovers it first when none is held yet (for example when + # tokens were loaded from storage), so re-authorization targets the right server. + if response.status_code == 401 or self.context.oauth_metadata is None: + www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response) + + # Step 1: Discover protected resource metadata (SEP-985 with fallback support) + prm_discovery_urls = build_protected_resource_metadata_discovery_urls( + www_auth_resource_metadata_url, self.context.server_url + ) + + prm_request_failed: int | None = None + for url in prm_discovery_urls: + discovery_request = create_oauth_metadata_request(url) + + discovery_response = yield discovery_request # sending request + + if discovery_response.status_code >= 500 or discovery_response.status_code == 429: + prm_request_failed = discovery_response.status_code + prm = await handle_protected_resource_response(discovery_response) + if prm: + # Validate PRM resource matches server URL (RFC 8707) + await self._validate_resource_match(prm) + self.context.protected_resource_metadata = prm + + self.context.auth_server_url = self._select_authorization_server( + [str(url) for url in prm.authorization_servers] + ) + break + else: + logger.debug(f"Protected resource metadata discovery failed: {url}") + else: + if prm_request_failed is not None: + # A server error says nothing about whether the resource publishes + # metadata, so it must not send the flow down the legacy path. + raise OAuthFlowError( + f"Protected resource metadata request failed: HTTP {prm_request_failed}" + ) + + expected_issuer = self._expected_issuer() + + # SEP-2352: stored credentials are bound to the issuer that registered them. + # Decided before any metadata is fetched: if the expected issuer is a different + # server, drop them (and the old tokens) so the flow re-registers instead of + # presenting another server's credentials. + if self.context.client_info is not None and not credentials_match_issuer( + self.context.client_info, expected_issuer, self.context.client_metadata_url + ): + logger.debug( + "Authorization server changed; discarding bound credentials and re-registering" + ) + self.context.client_info = None + self.context.clear_tokens() + # Any cached AS metadata is for the old server; drop it so a failed + # rediscovery cannot leak the old registration/token endpoints into Step 4. + self.context.oauth_metadata = None + + asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls( + self.context.auth_server_url, self.context.server_url + ) + + # Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers) + for url in asm_discovery_urls: # pragma: no branch + oauth_metadata_request = create_oauth_metadata_request(url) + oauth_metadata_response = yield oauth_metadata_request + + ok, asm = await handle_auth_metadata_response(oauth_metadata_response) + if not ok: + break + if ok and asm: + # SEP-2468 / RFC 8414 §3.3: the metadata must name the expected issuer. + # On the legacy path a root issuer rendered with its trailing slash + # names the same origin. + if self.context.auth_server_url is None and issuers_match( + str(asm.issuer), expected_issuer + ): + expected_issuer = str(asm.issuer) + validate_metadata_issuer(asm, expected_issuer) + self.context.oauth_metadata = asm + break + else: + logger.debug(f"OAuth metadata discovery failed: {url}") + + # Step 3: Apply scope selection strategy + challenged_scope = get_client_metadata_scopes( + extract_scope_from_www_auth(response), + self.context.protected_resource_metadata, + self.context.oauth_metadata, + self.context.client_metadata.grant_types, + ) + if step_up: + # SEP-2350: union previously requested scopes with the newly challenged ones so + # escalating one operation keeps the others' grants, folding in the granted + # scope read above since client_metadata.scope is not reloaded on a restart. + prior_scope = union_scopes(self.context.client_metadata.scope, granted_scope) + self.context.client_metadata.scope = union_scopes(prior_scope, challenged_scope) + else: + self.context.client_metadata.scope = challenged_scope + + # Step 4: Register client or use URL-based client ID (CIMD) + if not self.context.client_info: + # SEP-2352: the issuer to bind these credentials to, once metadata for it + # was actually found. + discovered_issuer = self._expected_issuer() if self.context.oauth_metadata is not None else None + + if should_use_client_metadata_url( + self.context.oauth_metadata, self.context.client_metadata_url + ): + # Use URL-based client ID (CIMD). CIMD records are portable across + # authorization servers, so the issuer stamp is informational. + logger.debug(f"Using URL-based client ID (CIMD): {self.context.client_metadata_url}") + client_information = create_client_info_from_metadata_url( + self.context.client_metadata_url, # type: ignore[arg-type] + redirect_uris=self.context.client_metadata.redirect_uris, + ) + client_information.issuer = discovered_issuer + self.context.client_info = client_information + await self.context.storage.set_client_info(client_information) + else: + # Fallback to Dynamic Client Registration + fallback_base = self.context.get_authorization_base_url(self.context.server_url) + registration_request = create_client_registration_request( + self.context.oauth_metadata, self.context.client_metadata, fallback_base + ) + registration_response = yield registration_request + client_information = await handle_registration_response(registration_response) + check_registration_usable(client_information) + # Only record the issuer when the registration above actually targeted + # the discovered AS — either via its published registration_endpoint, + # or because the resource-origin /register fallback is on the issuer's + # own host (legacy same-origin embedded AS). Otherwise the fallback hit + # a different server and recording a binding to the PRM-advertised AS + # would persist a binding that was never established. + if ( + self.context.oauth_metadata is not None + and discovered_issuer is not None + and ( + self.context.oauth_metadata.registration_endpoint is not None + or self.context.get_authorization_base_url(discovered_issuer) == fallback_base + ) + ): + client_information.issuer = discovered_issuer + self.context.client_info = client_information + await self.context.storage.set_client_info(client_information) + + # Step 5: Perform authorization and complete token exchange + token_response = yield await self._perform_authorization() + await self._handle_token_response(token_response) + except Exception: + logger.exception("OAuth flow error") + raise + + # Retry with new tokens + self._add_auth_header(request) + yield request diff --git a/src/mcp-client/mcp_client/client/auth/utils.py b/src/mcp-client/mcp_client/client/auth/utils.py new file mode 100644 index 0000000000..fd64cc46b7 --- /dev/null +++ b/src/mcp-client/mcp_client/client/auth/utils.py @@ -0,0 +1,442 @@ +import re +from typing import Any, cast +from urllib.parse import urljoin, urlparse + +from httpx2 import Request, Response +from mcp_types import LATEST_PROTOCOL_VERSION +from pydantic import AnyUrl, ValidationError +from pydantic_core import from_json + +from mcp_client.client.auth import OAuthFlowError, OAuthRegistrationError, OAuthTokenError +from mcp_client.shared._httpx_utils import redirect_note +from mcp_client.shared.auth import ( + OAuthClientInformationFull, + OAuthClientMetadata, + OAuthMetadata, + OAuthToken, + ProtectedResourceMetadata, +) +from mcp_client.shared.inbound import MCP_PROTOCOL_VERSION_HEADER + + +def extract_field_from_www_auth(response: Response, field_name: str) -> str | None: + """Extract field from WWW-Authenticate header. + + Returns: + Field value if found in WWW-Authenticate header, None otherwise + """ + www_auth_header = response.headers.get("WWW-Authenticate") + if not www_auth_header: + return None + + # Pattern matches: field_name="value" or field_name=value (unquoted) + pattern = rf'{field_name}=(?:"([^"]+)"|([^\s,]+))' + match = re.search(pattern, www_auth_header) + + if match: + # Return quoted value if present, otherwise unquoted value + return match.group(1) or match.group(2) + + return None + + +def extract_scope_from_www_auth(response: Response) -> str | None: + """Extract scope parameter from WWW-Authenticate header as per RFC 6750. + + Returns: + Scope string if found in WWW-Authenticate header, None otherwise + """ + return extract_field_from_www_auth(response, "scope") + + +def extract_resource_metadata_from_www_auth(response: Response) -> str | None: + """Extract protected resource metadata URL from WWW-Authenticate header as per RFC 9728. + + Returns: + Resource metadata URL if found in WWW-Authenticate header, None otherwise + """ + if not response or response.status_code not in (401, 403): + return None # pragma: no cover + + return extract_field_from_www_auth(response, "resource_metadata") + + +def build_protected_resource_metadata_discovery_urls(www_auth_url: str | None, server_url: str) -> list[str]: + """Build ordered list of URLs to try for protected resource metadata discovery. + + Per SEP-985, the client MUST: + 1. Try resource_metadata from WWW-Authenticate header (if present) + 2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path} + 3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource + + Args: + www_auth_url: Optional resource_metadata URL extracted from the WWW-Authenticate header + server_url: Server URL + + Returns: + Ordered list of URLs to try for discovery + """ + urls: list[str] = [] + + # Priority 1: WWW-Authenticate header with resource_metadata parameter + if www_auth_url: + urls.append(www_auth_url) + + # Priority 2-3: Well-known URIs (RFC 9728) + parsed = urlparse(server_url) + base_url = f"{parsed.scheme}://{parsed.netloc}" + + # Priority 2: Path-based well-known URI (if server has a path component) + if parsed.path and parsed.path != "/": + path_based_url = urljoin(base_url, f"/.well-known/oauth-protected-resource{parsed.path}") + urls.append(path_based_url) + + # Priority 3: Root-based well-known URI + root_based_url = urljoin(base_url, "/.well-known/oauth-protected-resource") + urls.append(root_based_url) + + return urls + + +def get_client_metadata_scopes( + www_authenticate_scope: str | None, + protected_resource_metadata: ProtectedResourceMetadata | None, + authorization_server_metadata: OAuthMetadata | None = None, + client_grant_types: list[str] | None = None, +) -> str | None: + """Select effective scopes and augment for refresh token support.""" + selected_scope: str | None = None + + # MCP spec scope selection priority: + # 1. WWW-Authenticate header scope + # 2. PRM scopes_supported + # 3. AS scopes_supported (SDK fallback) + # 4. Omit scope parameter + if www_authenticate_scope is not None: + selected_scope = www_authenticate_scope + elif protected_resource_metadata is not None and protected_resource_metadata.scopes_supported is not None: + selected_scope = " ".join(protected_resource_metadata.scopes_supported) + elif authorization_server_metadata is not None and authorization_server_metadata.scopes_supported is not None: + selected_scope = " ".join(authorization_server_metadata.scopes_supported) + + # SEP-2207: append offline_access when the AS supports it and the client can use refresh tokens + if ( + selected_scope is not None + and authorization_server_metadata is not None + and authorization_server_metadata.scopes_supported is not None + and "offline_access" in authorization_server_metadata.scopes_supported + and client_grant_types is not None + and "refresh_token" in client_grant_types + and "offline_access" not in selected_scope.split() + ): + selected_scope = f"{selected_scope} offline_access" + + return selected_scope + + +def union_scopes(previous_scope: str | None, new_scope: str | None) -> str | None: + """Merge two space-delimited scope strings, preserving order and dropping duplicates. + + SEP-2350: on step-up re-authorization the client requests the union of previously requested + scopes and the newly challenged scopes, so escalating one operation does not drop the + permissions granted for another. Previously requested scopes come first; new scopes are + appended in order. + """ + if not previous_scope: + return new_scope + if not new_scope: + return previous_scope + + merged = previous_scope.split() + seen = set(merged) + for scope in new_scope.split(): + if scope not in seen: + merged.append(scope) + seen.add(scope) + return " ".join(merged) + + +def build_oauth_authorization_server_metadata_discovery_urls(auth_server_url: str | None, server_url: str) -> list[str]: + """Generate an ordered list of URLs for authorization server metadata discovery. + + Args: + auth_server_url: OAuth Authorization Server Metadata URL if found, otherwise None + server_url: URL for the MCP server, used as a fallback if auth_server_url is None + """ + + if not auth_server_url: + # Legacy path using the 2025-03-26 spec: + # link: https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization + parsed = urlparse(server_url) + return [f"{parsed.scheme}://{parsed.netloc}/.well-known/oauth-authorization-server"] + + urls: list[str] = [] + parsed = urlparse(auth_server_url) + base_url = f"{parsed.scheme}://{parsed.netloc}" + + # RFC 8414: Path-aware OAuth discovery + if parsed.path and parsed.path != "/": + oauth_path = f"/.well-known/oauth-authorization-server{parsed.path.rstrip('/')}" + urls.append(urljoin(base_url, oauth_path)) + + # RFC 8414 section 5: Path-aware OIDC discovery + # See https://www.rfc-editor.org/rfc/rfc8414.html#section-5 + oidc_path = f"/.well-known/openid-configuration{parsed.path.rstrip('/')}" + urls.append(urljoin(base_url, oidc_path)) + + # https://openid.net/specs/openid-connect-discovery-1_0.html + oidc_path = f"{parsed.path.rstrip('/')}/.well-known/openid-configuration" + urls.append(urljoin(base_url, oidc_path)) + return urls + + # OAuth root + urls.append(urljoin(base_url, "/.well-known/oauth-authorization-server")) + + # OIDC 1.0 fallback (appends to full URL per OIDC spec) + # https://openid.net/specs/openid-connect-discovery-1_0.html + urls.append(urljoin(base_url, "/.well-known/openid-configuration")) + + return urls + + +async def handle_protected_resource_response( + response: Response, +) -> ProtectedResourceMetadata | None: + """Handle protected resource metadata discovery response. + + Per SEP-985, supports fallback when discovery fails at one URL. + + Returns: + ProtectedResourceMetadata if successfully discovered, None if we should try next URL + """ + if response.status_code == 200: + try: + content = await response.aread() + metadata = ProtectedResourceMetadata.model_validate_json(content) + return metadata + + except ValidationError: # pragma: no cover + # Invalid metadata - try next URL + return None + else: + # Not found - try next URL in fallback chain + return None + + +async def handle_auth_metadata_response(response: Response) -> tuple[bool, OAuthMetadata | None]: + if response.status_code == 200: + try: + content = await response.aread() + asm = OAuthMetadata.model_validate_json(content) + return True, asm + except ValidationError: # pragma: no cover + return True, None + elif 300 <= response.status_code < 500: + return True, None # Not served at this URL (redirects are not followed) - try the next candidate + return False, None # Server error or unexpected status, stop trying + + +def validate_authorization_response_iss(iss: str | None, oauth_metadata: OAuthMetadata | None) -> None: + """Validate the RFC 9207 `iss` authorization-response parameter. + + Per RFC 9207 section 2.4, the client compares `iss` against the issuer of the + authorization server the request was sent to, using simple string comparison + (RFC 3986 section 6.2.1, i.e. without URL normalization), and rejects on mismatch. + A response that omits `iss` is rejected only when the server advertised support via + `authorization_response_iss_parameter_supported`. + + Raises: + OAuthFlowError: If `iss` is present and does not match, or is absent when the + authorization server advertised support. + """ + expected = str(oauth_metadata.issuer) if oauth_metadata else None + + if iss is not None: + if iss != expected: + raise OAuthFlowError(f"Authorization response iss mismatch: {iss} != {expected}") + return + + if oauth_metadata is not None and oauth_metadata.authorization_response_iss_parameter_supported: + raise OAuthFlowError("Authorization response missing iss parameter advertised by the authorization server") + + +def validate_metadata_issuer(oauth_metadata: OAuthMetadata, expected_issuer: str) -> None: + """Validate that authorization server metadata `issuer` matches the discovery issuer. + + Per RFC 8414 section 3.3 / SEP-2468, the `issuer` in the metadata must match the issuer + used to construct the well-known URL, compared as a simple string (RFC 3986 section 6.2.1). + + Raises: + OAuthFlowError: If the metadata issuer does not match `expected_issuer`. + """ + if str(oauth_metadata.issuer) != expected_issuer: + raise OAuthFlowError( + f"Authorization server metadata issuer mismatch: {oauth_metadata.issuer} != {expected_issuer}" + ) + + +def create_oauth_metadata_request(url: str) -> Request: + return Request("GET", url, headers={MCP_PROTOCOL_VERSION_HEADER: LATEST_PROTOCOL_VERSION}) + + +def create_client_registration_request( + auth_server_metadata: OAuthMetadata | None, client_metadata: OAuthClientMetadata, auth_base_url: str +) -> Request: + """Build a client registration request.""" + + if auth_server_metadata and auth_server_metadata.registration_endpoint: + registration_url = str(auth_server_metadata.registration_endpoint) + else: + registration_url = urljoin(auth_base_url, "/register") + + registration_data = client_metadata.model_dump(by_alias=True, mode="json", exclude_none=True) + + return Request("POST", registration_url, json=registration_data, headers={"Content-Type": "application/json"}) + + +async def handle_registration_response(response: Response) -> OAuthClientInformationFull: + """Handle registration response.""" + if response.status_code not in (200, 201): + await response.aread() + raise OAuthRegistrationError( + f"Registration failed: {response.status_code}{redirect_note(response)} {response.text}" + ) + + try: + content = await response.aread() + body = from_json(content) + # `issuer` is the SDK's own binding of these credentials to the server they were + # registered with (SEP-2352), stamped by the auth flow - never sourced from the + # wire, so it is dropped before the body is parsed rather than trusted or cleared. + if isinstance(body, dict): + cast(dict[str, Any], body).pop("issuer", None) + return OAuthClientInformationFull.model_validate(body) + except ValueError as e: + # `from_json` reports malformed bytes/JSON as ValueError, and pydantic's + # ValidationError is itself a ValueError, so both parse layers surface here. + raise OAuthRegistrationError(f"Invalid registration response: {e}") from e + + +def is_valid_client_metadata_url(url: str | None) -> bool: + """Validate that a URL is suitable for use as a client_id (CIMD). + + The URL must be HTTPS with a non-root pathname. + + Args: + url: The URL to validate + + Returns: + True if the URL is a valid HTTPS URL with a non-root pathname + """ + if not url: + return False + try: + parsed = urlparse(url) + return parsed.scheme == "https" and parsed.path not in ("", "/") + except Exception: + return False + + +def credentials_match_issuer( + client_info: OAuthClientInformationFull, issuer: str, client_metadata_url: str | None +) -> bool: + """Whether stored client credentials may be reused against `issuer` (SEP-2352). + + A URL-based client ID (CIMD) is portable across authorization servers — the same self-hosted + document is resolved by whichever server is in use — so it always matches; CIMD is identified + by the client ID being the configured `client_metadata_url`, not by URL shape (a registration + server may also issue URL-shaped IDs that are bound to it). Credentials with a recorded issuer + match only when it equals `issuer` (simple string comparison; a root issuer with and without + its trailing slash count as equal). Credentials with no recorded + issuer (pre-registered, or stored before issuer binding existed) carry no binding to enforce + and are left as-is. + """ + if client_metadata_url is not None and client_info.client_id == client_metadata_url: + return True + if client_info.issuer is None: + return True + return issuers_match(client_info.issuer, issuer) + + +def issuers_match(a: str, b: str) -> bool: + """Simple string comparison of two issuer identifiers (RFC 8414 section 3.3), except that a root + issuer with and without its trailing slash (`scheme://authority` and `scheme://authority/`) name + the same server.""" + if a == b: + return True + shorter, longer = sorted((a, b), key=len) + parsed = urlparse(shorter) + return longer == f"{shorter}/" and shorter == f"{parsed.scheme}://{parsed.netloc}" + + +def should_use_client_metadata_url( + oauth_metadata: OAuthMetadata | None, + client_metadata_url: str | None, +) -> bool: + """Determine if URL-based client ID (CIMD) should be used instead of DCR. + + URL-based client IDs should be used when: + 1. The server advertises client_id_metadata_document_supported=True + 2. The client has a valid client_metadata_url configured + + Args: + oauth_metadata: OAuth authorization server metadata + client_metadata_url: URL-based client ID (already validated) + + Returns: + True if CIMD should be used, False if DCR should be used + """ + if not client_metadata_url: + return False + + if not oauth_metadata: + return False + + return oauth_metadata.client_id_metadata_document_supported is True + + +def create_client_info_from_metadata_url( + client_metadata_url: str, redirect_uris: list[AnyUrl] | None = None +) -> OAuthClientInformationFull: + """Create client information using a URL-based client ID (CIMD). + + When using URL-based client IDs, the URL itself becomes the client_id + and no client_secret is used (token_endpoint_auth_method="none"). + + Args: + client_metadata_url: The URL to use as the client_id + redirect_uris: The redirect URIs from the client metadata, recorded on the client + information alongside the client_id + + Returns: + OAuthClientInformationFull with the URL as client_id + """ + return OAuthClientInformationFull( + client_id=client_metadata_url, + token_endpoint_auth_method="none", + redirect_uris=redirect_uris, + ) + + +async def handle_token_response_scopes( + response: Response, +) -> OAuthToken: + """Parse and validate a token response. + + Parses token response JSON. Callers should check response.status_code before calling. + + Args: + response: HTTP response from token endpoint (status already checked by caller) + + Returns: + Validated OAuthToken model + + Raises: + OAuthTokenError: If response JSON is invalid + """ + try: + content = await response.aread() + token_response = OAuthToken.model_validate_json(content) + return token_response + except ValidationError as e: # pragma: no cover + raise OAuthTokenError(f"Invalid token response: {e}") diff --git a/src/mcp-client/mcp_client/client/caching.py b/src/mcp-client/mcp_client/client/caching.py new file mode 100644 index 0000000000..444bd43b98 --- /dev/null +++ b/src/mcp-client/mcp_client/client/caching.py @@ -0,0 +1,387 @@ +"""Client-side response caching primitives (SEP-2549, protocol revision 2026-07-28).""" + +from __future__ import annotations + +import json +import logging +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Final, Literal, Protocol + +import anyio +import anyio.lowlevel +from mcp_types import ( + CacheableResult, + PromptListChangedNotification, + ResourceListChangedNotification, + ResourceUpdatedNotification, + ServerNotification, + ToolListChangedNotification, +) +from mcp_types.version import MODERN_PROTOCOL_VERSIONS + +__all__ = [ + "MAX_TTL_MS", + "CacheConfig", + "CacheEntry", + "CacheKey", + "CacheMode", + "InMemoryResponseCacheStore", + "ResponseCacheStore", +] + +logger = logging.getLogger("mcp.client.caching") + +CacheMode = Literal["use", "refresh", "bypass"] +"""Per-call cache behavior: `"use"` serves and stores, `"refresh"` stores +without serving, `"bypass"` skips the cache entirely.""" + +MAX_TTL_MS: Final[int] = 24 * 60 * 60 * 1000 +"""Cap on any entry's time-to-live (24 hours, in milliseconds); larger `ttlMs` values are clamped down.""" + + +@dataclass(frozen=True, slots=True) +class CacheKey: + """Identity of one cached response; compare as the field tuple, never a flattened string (collision hazard).""" + + method: str + + params_key: str = "" + """Result-affecting params discriminator: the uri for `resources/read`, `""` for the list methods.""" + + partition: str = "" + """Coordinator-computed arm identifier; opaque to stores.""" + + +@dataclass(frozen=True, slots=True) +class CacheEntry: + """One cached response with its freshness and sharing metadata.""" + + value: Any + """The cached result; the SDK deep-copies on write and on serve, so a store may hold it as-is.""" + + scope: Literal["public", "private"] + """Server-asserted `cacheScope`: only `"public"` entries may be shared across authorization contexts.""" + + expires_at: float | None + """Epoch seconds after which the entry is stale; `None` is never fresh.""" + + +class ResponseCacheStore(Protocol): + """Storage contract for the client response cache. + + Each `Client` calls its store from a single event loop; per-operation + atomicity is the implementation's responsibility. Operations may raise - + the SDK degrades to a miss rather than failing the call. A serializing + store must round-trip `value` back to the result model object (a + wrong-shape entry is a miss, never an error). A lookup may issue two + sequential `get` calls (private arm, then public). + """ + + async def get(self, key: CacheKey) -> CacheEntry | None: ... + + async def set(self, key: CacheKey, entry: CacheEntry) -> None: ... + + async def delete(self, key: CacheKey) -> None: ... + + async def clear(self) -> None: ... + + +@dataclass(frozen=True, slots=True) +class CacheConfig: + """Configuration for a `Client`'s response cache. + + Raises: + ValueError: On a custom `store` without `partition`, an empty `target_id`, or a negative `default_ttl_ms`. + """ + + store: ResponseCacheStore | None = None + """Backing store; `None` means a per-client `InMemoryResponseCacheStore`. + A custom store requires an explicit `partition`.""" + + partition: str = "" + """Authorization-context identifier isolating `"private"`-scoped entries + within a shared store. Derive it from a verified credential - never from + request-supplied data or the server URL. Fixed for the `Client`'s + lifetime: construct a new `Client` when the principal changes.""" + + target_id: str | None = None + """Server-identity override for custom transports and proxies where the + SDK cannot derive one from a URL; must be non-empty when provided.""" + + default_ttl_ms: int = 0 + """TTL in milliseconds for results carrying no `ttlMs` hint; the default `0` leaves them uncached.""" + + clock: Callable[[], float] = time.time + """Wall-clock source returning epoch seconds; injectable for expiry tests.""" + + share_public: bool = False + """Serve server-marked `"public"` entries across every partition in the store. + + WARNING: this trusts the server's `"public"` classification for every + principal sharing the store - a mislabeled response leaks across tenants. + Constructor-level only: the per-call `cache_mode` can never widen sharing.""" + + def __post_init__(self) -> None: + if self.store is not None and not self.partition: + raise ValueError("a custom store requires an explicit partition") + if self.target_id == "": + raise ValueError("target_id must be a non-empty string or omitted") + if self.default_ttl_ms < 0: + raise ValueError(f"default_ttl_ms must be >= 0, got {self.default_ttl_ms}") + + +class InMemoryResponseCacheStore: + """Default in-process `ResponseCacheStore`. + + Method bodies are synchronous, so concurrent tasks never observe a torn + write. `max_entries` caps the whole store, evicting least-recently-used + at the cap (`0` disables it); `get` and `set` both refresh recency, so a + hot entry survives churn from other keys. + + Raises: + ValueError: If `max_entries` is negative. + """ + + def __init__(self, *, max_entries: int = 1024) -> None: + if max_entries < 0: + raise ValueError(f"max_entries must be >= 0, got {max_entries}") + self._max_entries = max_entries + self._entries: dict[CacheKey, CacheEntry] = {} + + async def get(self, key: CacheKey) -> CacheEntry | None: + entry = self._entries.get(key) + if entry is not None: + # Pop-and-reinsert moves the key to the back: the dict's insertion order is the LRU ledger. + self._entries[key] = self._entries.pop(key) + return entry + + async def set(self, key: CacheKey, entry: CacheEntry) -> None: + self._entries.pop(key, None) + self._entries[key] = entry + if self._max_entries and len(self._entries) > self._max_entries: + del self._entries[next(iter(self._entries))] + + async def delete(self, key: CacheKey) -> None: + self._entries.pop(key, None) + + async def clear(self) -> None: + self._entries.clear() + + +_GENERATION_MAP_CAP: Final[int] = 4096 +"""Cap on the generation map; at the cap the oldest key's eviction-race guard is dropped (FIFO).""" + +_STORE_CLEANUP_TIMEOUT: Final[float] = 5 +"""Bound for must-complete store cleanup deletes (mirrors the dispatcher's final-write bound); +a wedged store delete must not hold client teardown uncancellably.""" + + +class ClientResponseCache: + """Coordinates the `Client` caching verbs with a `ResponseCacheStore`: keys, era gate, TTL/scope, eviction.""" + + def __init__( + self, + *, + store: ResponseCacheStore, + partition: str, + arm_id: str, + default_ttl_ms: int, + clock: Callable[[], float], + share_public: bool, + negotiated_version: Callable[[], str | None], + generation_map_cap: int = _GENERATION_MAP_CAP, + store_cleanup_timeout: float = _STORE_CLEANUP_TIMEOUT, + ) -> None: + self._store = store + self._partition = partition + self._arm_id = arm_id + self._share_public = share_public + self._default_ttl_ms = default_ttl_ms + self._clock = clock + self._negotiated_version = negotiated_version + # A key is eviction-race-guarded iff registered here. + self._generations: dict[tuple[str, str], int] = {} + self._generation_map_cap = generation_map_cap + self._store_cleanup_timeout = store_cleanup_timeout + self._warned_store_ops: set[str] = set() + + def _arm(self, scope: Literal["public", "private"]) -> str: + # JSON arrays so crafted arm_id/partition values cannot collide across field boundaries. + # The negotiated version era-scopes every arm: a session never serves an entry written + # under a different protocol era (its content differs - sieve-stripped fields, header + # filtering). Every caller runs post-connect; were that ever untrue, the supplier's + # None still partitions harmlessly. + fields: list[str | None] = [scope, self._negotiated_version(), self._arm_id] + if scope == "private" or not self._share_public: + fields.append(self._partition) + return json.dumps(fields) + + async def read(self, method: str, params_key: str) -> CacheableResult | None: + """Serve a fresh entry for the key, or `None`; the served result is a deep copy.""" + # A hit completes without any other yielding await, so checkpoint here: a poll + # loop over a fresh entry must not starve spawned tasks (eviction dispatch). + await anyio.lowlevel.checkpoint() + # A wrong-shape entry raises as late as the copy, so the boundary wraps the whole read path. + try: + entry = await self._get_fresh(CacheKey(method, params_key, self._arm("private"))) + if entry is None: + # After a scope flip, a stale private entry must not shadow a fresh public one. + entry = await self._get_fresh(CacheKey(method, params_key, self._arm("public"))) + if entry is not None and entry.scope != "public": + # Never serve an entry the server scoped "private" out of the shared arm. + entry = None + copied: CacheableResult | None = None if entry is None else entry.value.model_copy(deep=True) + except Exception: # boundary around user store code: any read-path failure is a miss, never a failed call + self._warn_store_failure("get") + return None + self._warned_store_ops.discard("get") + return copied + + async def _get_fresh(self, key: CacheKey) -> CacheEntry | None: + entry = await self._store.get(key) + if entry is None or entry.expires_at is None or entry.expires_at <= self._clock(): + return None + return entry + + def capture(self, method: str, params_key: str) -> int: + """Register the key for eviction-race detection before the fetch; `write` takes the returned generation.""" + gen_key = (method, params_key) + if gen_key not in self._generations: + if len(self._generations) >= self._generation_map_cap: + # FIFO overflow: the dropped key's race guard degrades to the accepted co-tenant class. + del self._generations[next(iter(self._generations))] + self._generations[gen_key] = 0 + return self._generations[gen_key] + + async def write( + self, + method: str, + params_key: str, + result: CacheableResult, + gen_at_capture: int, + mode: Literal["use", "refresh"], + ) -> None: + """Store a fetched result under the arm its resolved scope selects.""" + gen_key = (method, params_key) + if self._generation_moved(gen_key, gen_at_capture): + return # the key was evicted while the fetch was in flight + ttl_ms, scope = self._resolve(result) + private_key = CacheKey(method, params_key, self._arm("private")) + public_key = CacheKey(method, params_key, self._arm("public")) + if ttl_ms <= 0: + if mode == "refresh": + # The refetch superseded the warm entry, which a cancellation must not leave serving. + await self._cleanup_delete(private_key, public_key) + return + own, opposite = (public_key, private_key) if scope == "public" else (private_key, public_key) + # Opposite arm first: a failed delete aborts before the set - never two arms answering for one key. + if not await self._delete(opposite): + # The own arm's entry is superseded too: best-effort delete, degrading to a full miss. + await self._cleanup_delete(own) + return + entry = CacheEntry(value=result.model_copy(deep=True), scope=scope, expires_at=self._clock() + ttl_ms / 1000) + try: + if not await self._set(own, entry): + # The fetch superseded any pre-existing own-arm entry, and the failed set + # left it in place: purge it (mirrors the opposite-arm-failure path). + await self._cleanup_delete(own) + finally: + # An eviction can land while the set commits - even when the await + # is cancelled - so re-check on every exit; the delete must complete + # so the pending cancellation cannot resurrect the evicted entry. + if self._generation_moved(gen_key, gen_at_capture): + await self._cleanup_delete(own) + + async def evict_method(self, method: str) -> None: + """Evict the method's cursor-less entry.""" + await self.evict_key(method, "") + + async def evict_key(self, method: str, params_key: str) -> None: + """Evict one key from both arms. + + Only the current era's arms are touched; other-era entries in a persistent store age out by TTL. + """ + gen_key = (method, params_key) + # Bump first so an in-flight fetch cannot write the evicted entry back. + # Unregistered keys skip the bump (uris must not grow the map) but not + # the deletes - a persistent store may hold uncaptured entries. + if gen_key in self._generations: + self._generations[gen_key] += 1 + # Must complete: a cancellation between the deletes would leave one arm serving the evicted entry. + await self._cleanup_delete( + CacheKey(method, params_key, self._arm("private")), + CacheKey(method, params_key, self._arm("public")), + ) + + async def evict_for_notification(self, notification: ServerNotification) -> None: + """Map a server notification to the entries it makes stale. + + Eviction is eventual (spawned-task dispatch): the generation bump closes + the write-back race; a racing read may briefly serve the old entry. + """ + match notification: + case ToolListChangedNotification(): + await self.evict_method("tools/list") + case PromptListChangedNotification(): + await self.evict_method("prompts/list") + case ResourceListChangedNotification(): + # Templates enumerate the same changed resource space. + await self.evict_method("resources/list") + await self.evict_method("resources/templates/list") + case ResourceUpdatedNotification(): + await self.evict_key("resources/read", notification.params.uri) + case _: + pass + + def _resolve(self, result: CacheableResult) -> tuple[int, Literal["public", "private"]]: + # A legacy peer can also put `ttlMs`/`cacheScope` keys on the wire, so + # wire presence is not a peer-era signal - hints count only when modern. + modern = self._negotiated_version() in MODERN_PROTOCOL_VERSIONS + if modern and "ttl_ms" in result.model_fields_set: + # An explicit `ttlMs: 0` stays 0, and negatives are unconstructible + # upstream (model ge=0, parse-seam floor) - only the cap applies. + ttl_ms = result.ttl_ms + else: + ttl_ms = self._default_ttl_ms + scope: Literal["public", "private"] = "public" if modern and result.cache_scope == "public" else "private" + return min(ttl_ms, MAX_TTL_MS), scope + + def _generation_moved(self, gen_key: tuple[str, str], gen_at_capture: int) -> bool: + # A FIFO-dropped key fails open (the accepted co-tenant race) rather than discarding the fetch. + return self._generations.get(gen_key, gen_at_capture) != gen_at_capture + + async def _set(self, key: CacheKey, entry: CacheEntry) -> bool: + try: + await self._store.set(key, entry) + except Exception: # boundary around user store code: nothing cached, the fetch already succeeded + self._warn_store_failure("set") + return False + self._warned_store_ops.discard("set") + return True + + async def _cleanup_delete(self, *keys: CacheKey) -> None: + # Must-complete cleanup: shielded so a pending cancellation cannot skip the deletes, + # bounded so a wedged store delete cannot hold client teardown uncancellably. + with anyio.move_on_after(self._store_cleanup_timeout, shield=True) as scope: + for key in keys: + await self._delete(key) + if scope.cancelled_caught: + logger.warning("Response cache store delete timed out; the entry will age out by TTL") + + async def _delete(self, key: CacheKey) -> bool: + try: + await self._store.delete(key) + except Exception: # boundary around user store code: callers decide whether a failed delete aborts + self._warn_store_failure("delete") + return False + self._warned_store_ops.discard("delete") + return True + + def _warn_store_failure(self, kind: Literal["get", "set", "delete"]) -> None: + # One warning per failure burst, per op kind; re-armed only when that + # same kind succeeds, so a healthy delete cannot re-arm a broken set. + if kind not in self._warned_store_ops: + self._warned_store_ops.add(kind) + logger.warning("Response cache store operation failed; continuing without the cache", exc_info=True) diff --git a/src/mcp-client/mcp_client/client/client.py b/src/mcp-client/mcp_client/client/client.py new file mode 100644 index 0000000000..f486f30f24 --- /dev/null +++ b/src/mcp-client/mcp_client/client/client.py @@ -0,0 +1,921 @@ +"""Unified MCP Client that wraps ClientSession with transport management.""" + +from __future__ import annotations + +import hashlib +import logging +import uuid +from collections.abc import Awaitable, Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, AsyncExitStack +from dataclasses import KW_ONLY, dataclass, field +from typing import Any, Literal, TypeAlias, TypeVar, cast + +import anyio +import anyio.lowlevel +import mcp_types as types +from mcp_types import ( + INVALID_PARAMS, + CacheableResult, + CallToolResult, + CompleteResult, + EmptyResult, + ErrorData, + GetPromptResult, + Implementation, + InputRequest, + InputRequiredResult, + InputResponse, + InputResponses, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + ListToolsResult, + LoggingLevel, + PaginatedRequestParams, + PromptReference, + ReadResourceResult, + RequestParamsMeta, + ResourceTemplateReference, + Result, + ServerCapabilities, +) +from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS +from typing_extensions import Protocol, deprecated, runtime_checkable + +from mcp_client.client._input_required import DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, run_input_required_driver +from mcp_client.client._probe import negotiate_auto +from mcp_client.client._transport import Transport +from mcp_client.client.caching import CacheConfig, CacheMode, ClientResponseCache, InMemoryResponseCacheStore +from mcp_client.client.extension import ClaimContext, ClientExtension, NotificationBinding, ResultClaim +from mcp_client.client.session import ( + ClientRequestContext, + ClientSession, + ElicitationFnT, + IncomingMessage, + ListRootsFnT, + LoggingFnT, + MessageHandlerFnT, + SamplingFnT, +) +from mcp_client.client.stdio import StdioServerParameters, stdio_client +from mcp_client.client.streamable_http import streamable_http_client +from mcp_client.client.subscriptions import ServerEvent, Subscription +from mcp_client.client.subscriptions import listen as _listen +from mcp_client.shared.dispatcher import Dispatcher, ProgressFnT +from mcp_client.shared.exceptions import MCPDeprecationWarning, MCPError +from mcp_client.shared.extension import validate_extension_identifier +from mcp_client.shared.jsonrpc_dispatcher import JSONRPCDispatcher +from mcp_client.shared.subscriptions import event_to_notification + +logger = logging.getLogger("mcp.client.client") + +ConnectMode = Literal["legacy", "auto"] | str +"""``mode=`` value: ``"legacy"`` (initialize handshake), ``"auto"`` (discover, fall back to +initialize), or a modern protocol-version string (adopt directly). The ``str`` arm is for +forward-compat; ``Client.__post_init__`` rejects anything outside that set at construction.""" + +_T = TypeVar("_T") +_ResultT = TypeVar("_ResultT") +_CacheableT = TypeVar("_CacheableT", bound=CacheableResult) + +_Connector = Callable[[AsyncExitStack, ConnectMode, bool], Awaitable["Dispatcher[Any]"]] +"""Resolved at ``__post_init__`` from the shape of ``server`` alone: enter whatever resources +are needed onto the exit stack and hand back the ``Dispatcher`` ``ClientSession`` will drive. +``mode`` and ``raise_exceptions`` are passed at call time so they're read at the same moment +``__aenter__`` reads them for the handshake step.""" + + +def _connect_transport(transport: Transport) -> _Connector: + """Connector for the stream-backed paths (URL, user-supplied ``Transport``).""" + + async def connect(exit_stack: AsyncExitStack, _mode: ConnectMode, _raise_exceptions: bool) -> Dispatcher[Any]: + read_stream, write_stream = await exit_stack.enter_async_context(transport) + return JSONRPCDispatcher(read_stream, write_stream) + + return connect + + +@runtime_checkable +class _ServerConnector(Protocol): + async def __mcp_client_connect__( + self, exit_stack: AsyncExitStack, mode: str, raise_exceptions: bool + ) -> Dispatcher[Any]: ... + + +# The full SDK rebinds this annotation alias, not the runtime-checkable protocol. +_InProcessServer: TypeAlias = _ServerConnector + + +def _connected(value: _T | None) -> _T: + """Narrow a post-handshake session attribute from ``T | None`` to ``T``. + + ``Client.__aenter__`` only assigns ``_session`` after the handshake succeeds, so inside + ``async with Client(...)`` these attributes are always populated; the ``.session`` gate + raises before this is reached otherwise. The guard exists for pyright, not runtime. + """ + if value is None: # pragma: no cover + raise RuntimeError("Client must be used within an async context manager") + return value + + +def _strip_userinfo(url: str) -> str: + """Drop any userinfo from the URL's authority component; byte-exact otherwise. + + Credentials must not enter cache-key material; any further normalization could merge distinct servers. + """ + # Pure text, no urlsplit: it strips embedded tab/CR/LF before parsing, which would misalign slices. + sep = url.find("//") + if sep == -1: + return url + start = sep + 2 + end = len(url) + for delimiter in "/?#": + if (found := url.find(delimiter, start)) != -1: + end = min(end, found) + authority = url[start:end] + if "@" not in authority: + return url + return url[:start] + authority.rpartition("@")[2] + url[end:] + + +def _evicting_message_handler(cache: ClientResponseCache, user_handler: MessageHandlerFnT | None) -> MessageHandlerFnT: + """Wrap the session message handler with cache eviction on server notifications.""" + + async def handler(message: IncomingMessage) -> None: + if isinstance(message, types.ServerNotification): + try: + await cache.evict_for_notification(message) + except Exception: # boundary: eviction reaches user store code; a cache fault must not block delivery + logger.exception("Response cache eviction failed; the notification is still delivered") + if user_handler is not None: + await user_handler(message) + else: + # Mirrors ClientSession's default handler (session._default_message_handler). + await anyio.lowlevel.checkpoint() + + return handler + + +def _synthesize_discover(protocol_version: str) -> types.DiscoverResult: + return types.DiscoverResult( + supported_versions=[protocol_version], + capabilities=types.ServerCapabilities(), + result_type="complete", + ttl_ms=0, + cache_scope="public", + ) + + +@dataclass(frozen=True) +class _FoldedExtensions: + """`Client.extensions` instances folded into the shapes `ClientSession` consumes.""" + + ad: dict[str, dict[str, Any]] | None + claims: dict[str, tuple[ResultClaim[Any], ...]] | None + bindings: tuple[NotificationBinding[Any], ...] | None + by_model: Mapping[type[Result], ResultClaim[Any]] + + +def _fold_extensions(extensions: Sequence[ClientExtension] | None) -> _FoldedExtensions: + """Fold extension contributions at construction, naming both owners on duplicate tags or methods.""" + if isinstance(extensions, Mapping): + raise TypeError( + "extensions= takes a sequence of ClientExtension instances. The mapping form was " + "replaced: use advertise(identifier, settings) for advertise-only entries" + ) + if not extensions: + return _FoldedExtensions(ad=None, claims=None, bindings=None, by_model={}) + ad: dict[str, dict[str, Any]] = {} + claims: dict[str, tuple[ResultClaim[Any], ...]] = {} + bindings: list[NotificationBinding[Any]] = [] + by_model: dict[type[Result], ResultClaim[Any]] = {} + claim_owners: dict[str, str] = {} + binding_owners: dict[str, str] = {} + for extension in extensions: + identifier = getattr(extension, "identifier", None) + if identifier is None: + raise ValueError( + f"{type(extension).__name__} has no `identifier`; a ClientExtension must set the " + "`identifier` class attribute (or assign one in `__init__`) before it can be used" + ) + validate_extension_identifier(identifier, owner=type(extension).__name__) + if identifier in ad: + raise ValueError(f"extension identifier {identifier!r} is passed more than once") + ad[identifier] = extension.settings() + extension_claims = tuple(extension.claims()) + for claim in extension_claims: + tag = claim.result_type + if tag in claim_owners: + owner = claim_owners[tag] + both = ( + f"extension {identifier!r} claims" + if owner == identifier + else (f"extensions {owner!r} and {identifier!r} both claim") + ) + raise ValueError(f"{both} resultType {tag!r}; a wire tag can have only one resolver") + claim_owners[tag] = identifier + # Each model pins its result_type Literal to one tag, so this index cannot collide. + by_model[claim.model] = claim + if extension_claims: + claims[identifier] = extension_claims + for binding in extension.notifications(): + if binding.method in binding_owners: + owner = binding_owners[binding.method] + both = ( + f"extension {identifier!r} binds" + if owner == identifier + else (f"extensions {owner!r} and {identifier!r} both bind") + ) + raise ValueError(f"{both} notification method {binding.method!r}; a method can have only one observer") + binding_owners[binding.method] = identifier + bindings.append(binding) + return _FoldedExtensions(ad=ad, claims=claims or None, bindings=tuple(bindings) or None, by_model=by_model) + + +@dataclass +class Client: + """A high-level MCP client for connecting to MCP servers. + + Pass a URL string (Streamable HTTP), a `StdioServerParameters` (launch the command as a + subprocess and talk over its stdin/stdout), any `Transport`, or - in tests - a `Server` or + `MCPServer` instance to connect to it in-process. + + Example: + ```python + import asyncio + + from mcp_client import Client + + async def main(): + async with Client("http://localhost:8000/mcp") as client: + result = await client.call_tool("add", {"a": 1, "b": 2}) + + asyncio.run(main()) + ``` + """ + + server: _InProcessServer | Transport | StdioServerParameters | str + """The MCP server to connect to. + + If the server is a URL string, it will be used as the URL for a `streamable_http_client` transport. + If the server is a `StdioServerParameters`, the command is launched with `stdio_client`. + If the server is a `Transport` instance, it will be used directly. + If the server is a `Server` or `MCPServer` instance, it will be connected in-process. + """ + + _: KW_ONLY + + # TODO(Marcelo): When do `raise_exceptions=True` actually raises? + raise_exceptions: bool = False + """Whether to raise exceptions from the server.""" + + read_timeout_seconds: float | None = None + """Timeout for read operations.""" + + sampling_callback: SamplingFnT | None = None + """Callback for handling sampling requests.""" + + sampling_capabilities: types.SamplingCapability | None = None + """Sampling sub-capabilities (e.g. tools) declared alongside `sampling_callback`; no effect without it.""" + + list_roots_callback: ListRootsFnT | None = None + """Callback for handling list roots requests.""" + + logging_callback: LoggingFnT | None = None + """Callback for handling logging notifications.""" + + log_level: LoggingLevel | None = None + """The log level to opt in to on 2026-07-28+ connections (deprecated logging feature, SEP-2577). + + Modern (2026-07-28+) servers send `notifications/message` only for requests that opt in by + carrying `io.modelcontextprotocol/logLevel` in `_meta`, and only at or above that level. Setting + this stamps that opt-in on every request; `None` (the default) means no opt-in, so no log + messages arrive - a `logging_callback` alone is not an opt-in. No effect on handshake-era + connections, where the deprecated `logging/setLevel` request governs delivery instead. A + per-request `_meta` entry with the same key overrides this default.""" + + # TODO(Marcelo): Why do we have both "callback" and "handler"? + message_handler: MessageHandlerFnT | None = None + """Callback for handling raw messages.""" + + client_info: Implementation | None = None + """Client implementation info to send to server.""" + + mode: ConnectMode = "auto" + """How to negotiate the protocol version. + + 'auto' (the default) probes `server/discover` and falls back to the initialize handshake on legacy servers; + for an in-process `Server`/`MCPServer` it dispatches directly without JSON-RPC framing. 'legacy' forces the + initialize handshake (byte-identical pre-2026 behavior). A modern protocol-version string (e.g. '2026-07-28') + adopts that version directly without a probe — supply `prior_discover` to reuse a known DiscoverResult, or + omit it to synthesize a minimal one.""" + + prior_discover: types.DiscoverResult | None = None + """A previously-obtained DiscoverResult to install via .adopt() when mode is a version pin. + Ignored when mode='legacy'.""" + + elicitation_callback: ElicitationFnT | None = None + """Callback for handling elicitation requests.""" + + input_required_max_rounds: int = DEFAULT_INPUT_REQUIRED_MAX_ROUNDS + """Cap on `InputRequiredResult` retry rounds before `call_tool` / `get_prompt` / + `read_resource` give up. Use `client.session.(..., allow_input_required=True)` + to drive the loop manually instead.""" + + extensions: Sequence[ClientExtension] | None = None + """Opt-in client extensions (SEP-2133). + + Each instance contributes its capability ad, its result claims (resolved + transparently by `call_tool`), and its notification bindings. For an + ad-only entry use `mcp_client.client.advertise(identifier, settings)`.""" + + cache: CacheConfig | None = field(default_factory=CacheConfig) + """Client-side response caching for the SEP-2549 cacheable methods (2026-07-28). + + The default `CacheConfig()` honors server `ttlMs`/`cacheScope` hints with a + per-client in-memory store; pass a customized `CacheConfig`, or `None` to + disable. The cacheable verbs take a per-call `cache_mode` (see `CacheMode`); + calls carrying `meta` always reach the server. A `CacheConfig` with a custom + `store` requires `target_id` when the server is not a URL (no identity can be + derived).""" + + _entered: bool = field(init=False, default=False) + _session: ClientSession | None = field(init=False, default=None) + _exit_stack: AsyncExitStack | None = field(init=False, default=None) + _connect: _Connector = field(init=False, repr=False, compare=False) + _response_cache: ClientResponseCache | None = field(init=False, default=None, repr=False, compare=False) + _folded_extensions: _FoldedExtensions = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + if self.mode not in ("legacy", "auto") and self.mode not in MODERN_PROTOCOL_VERSIONS: + hint = ( + f" ({self.mode!r} is a handshake-era version; use mode='legacy')" + if self.mode in HANDSHAKE_PROTOCOL_VERSIONS + else "" + ) + raise ValueError( + f"mode must be 'legacy', 'auto', or one of {list(MODERN_PROTOCOL_VERSIONS)}; got {self.mode!r}{hint}" + ) + + self._folded_extensions = _fold_extensions(self.extensions) + + srv = self.server + if isinstance(srv, _ServerConnector): + self._connect = srv.__mcp_client_connect__ + elif isinstance(srv, str): + self._connect = _connect_transport(streamable_http_client(srv)) + elif isinstance(srv, StdioServerParameters): + self._connect = _connect_transport(stdio_client(srv)) + else: + self._connect = _connect_transport(srv) + + if self.cache is not None: + config = self.cache + # Only the hash below leaves this scope - the raw identity may carry credentials; never log or store it. + target_id = config.target_id + if target_id is None and isinstance(self.server, str): + target_id = _strip_userinfo(self.server) + if target_id is None: + if config.store is not None: + raise ValueError( + "a custom cache store requires CacheConfig.target_id when the server is not a URL: " + "in-process servers and Transport instances get a random per-client identity, so " + "their entries in a shared store could never be served to another client" + ) + target_id = uuid.uuid4().hex + self._response_cache = ClientResponseCache( + store=config.store if config.store is not None else InMemoryResponseCacheStore(), + partition=config.partition, + arm_id=hashlib.sha256(target_id.encode()).hexdigest(), + default_ttl_ms=config.default_ttl_ms, + clock=config.clock, + share_public=config.share_public, + # Lazy: the negotiated version is unknown until __aenter__'s handshake. + negotiated_version=lambda: self._session.protocol_version if self._session is not None else None, + ) + + async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession: + """Enter the resolved connector and return an un-entered ClientSession.""" + dispatcher = await self._connect(exit_stack, self.mode, self.raise_exceptions) + message_handler = self.message_handler + if self._response_cache is not None: + message_handler = _evicting_message_handler(self._response_cache, self.message_handler) + return ClientSession( + dispatcher=dispatcher, + read_timeout_seconds=self.read_timeout_seconds, + sampling_callback=self.sampling_callback, + sampling_capabilities=self.sampling_capabilities, + list_roots_callback=self.list_roots_callback, + logging_callback=self.logging_callback, + log_level=self.log_level, + message_handler=message_handler, + client_info=self.client_info, + elicitation_callback=self.elicitation_callback, + extensions=self._folded_extensions.ad, + result_claims=self._folded_extensions.claims, + notification_bindings=self._folded_extensions.bindings, + ) + + async def __aenter__(self) -> Client: + """Enter the async context manager.""" + if self._entered: + raise RuntimeError("Client is already entered; cannot reenter") + self._entered = True + + async with AsyncExitStack() as exit_stack: + session = await self._build_session(exit_stack) + session = await exit_stack.enter_async_context(session) + + if self.mode == "legacy": + await session.initialize() + elif self.mode == "auto": + await negotiate_auto(session) + else: + session.adopt(self.prior_discover or _synthesize_discover(self.mode)) + + # Only publish the session after the handshake succeeds, so `_session is not None` + # implies the protocol_version/server_capabilities are populated (server_info + # stays optional: 2026-era servers may not identify themselves). If the + # handshake raised above, the local exit_stack unwinds the transport for us. + self._session = session + self._exit_stack = exit_stack.pop_all() + return self + + async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: + """Exit the async context manager.""" + if self._exit_stack: # pragma: no branch + await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb) + self._session = None + + @property + def session(self) -> ClientSession: + """Get the underlying ClientSession. + + This provides access to the full ClientSession API for advanced use cases. + + Raises: + RuntimeError: If accessed before entering the context manager. + """ + if self._session is None: + raise RuntimeError("Client must be used within an async context manager") + return self._session + + # TODO(maxisbey): the by-construction shape is for __aenter__ to return a connected-view + # type whose protocol_version/server_capabilities are non-Optional fields, + # eliminating these guards (and the one in .session). Same family as resolving the + # transport/connector at __post_init__ so the Optional internal fields disappear. + # (server_info stays Optional even connected: the 2026-era stamp is optional.) + @property + def protocol_version(self) -> str: + """Negotiated protocol version (set by initialize/discover/adopt during ``__aenter__``).""" + return _connected(self.session.protocol_version) + + @property + def server_info(self) -> Implementation | None: + """Server name/version, or `None` when the server did not identify itself. + + Legacy connections always carry it (`InitializeResult.serverInfo` is + required); on 2026-era connections the `_meta` `serverInfo` stamp is + optional, so an anonymous server reads as `None`. + """ + return self.session.server_info + + @property + def server_capabilities(self) -> ServerCapabilities: + """Server capabilities (set by initialize/discover/adopt during ``__aenter__``).""" + return _connected(self.session.server_capabilities) + + @property + def instructions(self) -> str | None: + """Server-provided instructions text, if any.""" + return self.session.instructions + + @deprecated( + "ping is removed as of 2026-07-28; the method only works under mode='legacy'.", + category=MCPDeprecationWarning, + ) + async def send_ping(self, *, meta: RequestParamsMeta | None = None) -> EmptyResult: + """Send a ping request to the server.""" + return await self.session.send_ping(meta=meta) + + @deprecated( + "Client-to-server progress is deprecated as of 2026-07-28; progress is server-to-client only.", + category=MCPDeprecationWarning, + ) + async def send_progress_notification( + self, + progress_token: str | int, + progress: float, + total: float | None = None, + message: str | None = None, + ) -> None: + """Send a progress notification to the server.""" + await self.session.send_progress_notification( # pyright: ignore[reportDeprecated] + progress_token=progress_token, + progress=progress, + total=total, + message=message, + ) + + @deprecated("The logging capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def set_logging_level(self, level: LoggingLevel, *, meta: RequestParamsMeta | None = None) -> EmptyResult: + """Set the logging level on the server.""" + return await self.session.set_logging_level(level=level, meta=meta) # pyright: ignore[reportDeprecated] + + async def _cached_fetch( + self, + method: str, + *, + cursor: str | None, + meta: RequestParamsMeta | None, + cache_mode: CacheMode, + send: Callable[[], Awaitable[_CacheableT]], + absorb: Callable[[_CacheableT], _CacheableT] | None = None, + ) -> _CacheableT: + """Serve one of the four list verbs through the response cache. + + `absorb` (tools/list only) re-applies session-side derived state to a served cache hit. + """ + cache = self._response_cache + if cache is None or cache_mode == "bypass": + return await send() + # A closed (or never-entered) client must raise, never serve cached entries. + _ = self.session + if meta is not None and cache_mode == "use": + # meta (a progress token, tracing fields) expects a wire request; fetch and replace the entry. + cache_mode = "refresh" + if cursor is not None: + # Continuation pages skip the cache, but an expired cursor means the listing changed (spec SHOULD evict). + try: + return await send() + except MCPError as e: + if e.code == INVALID_PARAMS: + await cache.evict_method(method) + raise + if cache_mode == "use" and (hit := await cache.read(method, "")) is not None: + # The hit is a private deep copy, so absorption may mutate it freely. + served = cast(_CacheableT, hit) + return served if absorb is None else absorb(served) + gen = cache.capture(method, "") + result = await send() + await cache.write(method, "", result, gen, cache_mode) + return result + + async def list_resources( + self, + *, + cursor: str | None = None, + meta: RequestParamsMeta | None = None, + cache_mode: CacheMode = "use", + ) -> ListResourcesResult: + """List available resources from the server.""" + return await self._cached_fetch( + "resources/list", + cursor=cursor, + meta=meta, + cache_mode=cache_mode, + send=lambda: self.session.list_resources(params=PaginatedRequestParams(cursor=cursor, _meta=meta)), + ) + + async def list_resource_templates( + self, + *, + cursor: str | None = None, + meta: RequestParamsMeta | None = None, + cache_mode: CacheMode = "use", + ) -> ListResourceTemplatesResult: + """List available resource templates from the server.""" + return await self._cached_fetch( + "resources/templates/list", + cursor=cursor, + meta=meta, + cache_mode=cache_mode, + send=lambda: self.session.list_resource_templates(params=PaginatedRequestParams(cursor=cursor, _meta=meta)), + ) + + async def read_resource( + self, + uri: str, + *, + input_responses: InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + cache_mode: CacheMode = "use", + ) -> ReadResourceResult: + """Read a resource from the server. + + If the server returns an `InputRequiredResult`, the embedded input + requests are dispatched to this client's sampling / elicitation / roots + callbacks and the read is retried automatically (up to + `input_required_max_rounds`). + + Args: + uri: The URI of the resource to read. + input_responses: Responses to seed the first call with (e.g. when + resuming from a persisted `InputRequiredResult`). + request_state: Opaque state to seed the first call with. + meta: Additional metadata for the request. + cache_mode: Cache behavior for this call (see `CacheMode`); seeded + calls (`input_responses` or `request_state` set) ignore it. + + Returns: + The resource content. + + Raises: + InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted. + MCPError: A callback returned `ErrorData` for an embedded input request. + pydantic.ValidationError: The server returned a result that does not + conform to the negotiated protocol version. + """ + + async def retry(r: InputResponses | None, s: str | None) -> ReadResourceResult | InputRequiredResult: + return await self.session.read_resource( + uri, input_responses=r, request_state=s, meta=meta, allow_input_required=True + ) + + # Seeded calls resume a specific exchange and must never be cached (spec MUST). + seeded = input_responses is not None or request_state is not None + cache = None if seeded else self._response_cache + if cache is None or cache_mode == "bypass": + return await self._drive_input_required(await retry(input_responses, request_state), retry) + # A closed (or never-entered) client must raise, never serve cached entries. + _ = self.session + if meta is not None and cache_mode == "use": + # Calls carrying meta always reach the server (mirrors `_cached_fetch`). + cache_mode = "refresh" + if cache_mode == "use" and (hit := await cache.read("resources/read", uri)) is not None: + # Only terminal first-round results are stored, so a hit legitimately skips the driver. + return cast(ReadResourceResult, hit) + gen = cache.capture("resources/read", uri) + first = await retry(None, None) + if not isinstance(first, InputRequiredResult): + await cache.write("resources/read", uri, first, gen, cache_mode) + elif cache_mode == "refresh": + # The refresh superseded whatever was cached, but an input_required resolution + # cannot be stored: purge the warm entry so it cannot be served again. + await cache.evict_key("resources/read", uri) + # Driver rounds carry inputResponses, so a terminal result reached through them is never cached (spec MUST). + return await self._drive_input_required(first, retry) + + def listen( + self, + *, + tools_list_changed: bool = False, + prompts_list_changed: bool = False, + resources_list_changed: bool = False, + resource_subscriptions: Sequence[str] = (), + ) -> AbstractAsyncContextManager[Subscription]: + """Open a `subscriptions/listen` stream of typed change events (2026-07-28 only). + + Keyword args mirror the wire `SubscriptionFilter`; entering waits for the ack (honored subset: `sub.honored`): + + async with client.listen(tools_list_changed=True) as sub: + async for event in sub: + tools = await client.list_tools() # refetch on change + + A graceful close ends the loop; an abrupt drop raises `SubscriptionLost`. No replay: re-listen and refetch. + + Raises: + ListenNotSupportedError: The negotiated protocol version predates 2026-07-28. + MCPError: The server rejected the request or the connection failed first. + SubscriptionLost: The stream ended before it was acknowledged. + TimeoutError: The read timeout elapsed before the acknowledgment. + """ + return _listen( + self.session, + tools_list_changed=tools_list_changed, + prompts_list_changed=prompts_list_changed, + resources_list_changed=resources_list_changed, + resource_subscriptions=resource_subscriptions, + on_event=self._evict_for_listen_event if self._response_cache is not None else None, + ) + + async def _evict_for_listen_event(self, event: ServerEvent) -> None: + """Finish response-cache eviction before a listen consumer can refetch. + + Without it the iterator wakes first and refetches a still-warm entry, with no + corrective wake (events are deduplicated level triggers). The tee path repeats + the eviction; deliberate: idempotent, and it covers non-iterating consumers. + """ + cache = self._response_cache + assert cache is not None # installed as the event barrier only when a cache exists + try: + await cache.evict_for_notification(event_to_notification(event, {})) + except Exception: # boundary: eviction reaches user store code; a cache fault must not block delivery + logger.exception("Response cache eviction failed; the event is still delivered") + + @deprecated( + "resources/subscribe is removed as of 2026-07-28; use Client.listen() instead.", + category=MCPDeprecationWarning, + ) + async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> EmptyResult: + """Subscribe to resource updates (2025-era servers only).""" + return await self.session.subscribe_resource(uri, meta=meta) # pyright: ignore[reportDeprecated] + + @deprecated( + "resources/unsubscribe is removed as of 2026-07-28; use Client.listen() instead.", + category=MCPDeprecationWarning, + ) + async def unsubscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> EmptyResult: + """Unsubscribe from resource updates (2025-era servers only).""" + return await self.session.unsubscribe_resource(uri, meta=meta) # pyright: ignore[reportDeprecated] + + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + ) -> CallToolResult: + """Call a tool on the server. + + If the server returns an `InputRequiredResult`, the embedded input + requests are dispatched to this client's sampling / elicitation / roots + callbacks and the call is retried automatically (up to + `input_required_max_rounds`). To drive the loop yourself — e.g. to + persist `request_state` across process restarts — use + `client.session.call_tool(..., allow_input_required=True)`. Persisted + state is still subject to the server's TTL, request binding, and key + lifetime; a server on the default process-local key rejects it after a restart. + + Result shapes claimed by this client's `extensions` are finished by the + owning claim's resolver, whose `CallToolResult` is returned; resolver + exceptions propagate as-is. To receive the claimed shape yourself, use + `client.session.call_tool(..., allow_claimed=True)`. + + Args: + name: The name of the tool to call. + arguments: Arguments to pass to the tool. + read_timeout_seconds: Timeout for each underlying `tools/call` round. + progress_callback: Callback for progress updates. + input_responses: Responses to seed the first call with (e.g. when + resuming from a persisted `InputRequiredResult`). + request_state: Opaque state to seed the first call with. + meta: Additional metadata for the request. + + Returns: + The tool result. + + Raises: + InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted. + MCPError: A callback returned `ErrorData` for an embedded input request. + pydantic.ValidationError: The server returned a result that does not + conform to the negotiated protocol version. + """ + + async def retry(r: InputResponses | None, s: str | None) -> CallToolResult | InputRequiredResult | Result: + return await self.session.call_tool( + name, + arguments, + read_timeout_seconds=read_timeout_seconds, + progress_callback=progress_callback, + input_responses=r, + request_state=s, + meta=meta, + allow_input_required=True, + # Input rounds resolve before a claimed result, so a claim may end any round. + allow_claimed=True, + ) + + result = await self._drive_input_required(await retry(input_responses, request_state), retry) + if isinstance(result, CallToolResult): + return result + # Only claimed shapes reach this point, so the lookup is total. + claim = self._folded_extensions.by_model[type(result)] + final = await claim.resolve( + result, + ClaimContext(session=self.session, tool_name=name, read_timeout_seconds=read_timeout_seconds), + ) + if not final.is_error: + # Match the direct path: revalidate the output schema, but never for isError results. + await self.session.validate_tool_result(name, final) + return final + + async def list_prompts( + self, + *, + cursor: str | None = None, + meta: RequestParamsMeta | None = None, + cache_mode: CacheMode = "use", + ) -> ListPromptsResult: + """List available prompts from the server.""" + return await self._cached_fetch( + "prompts/list", + cursor=cursor, + meta=meta, + cache_mode=cache_mode, + send=lambda: self.session.list_prompts(params=PaginatedRequestParams(cursor=cursor, _meta=meta)), + ) + + async def get_prompt( + self, + name: str, + arguments: dict[str, str] | None = None, + *, + input_responses: InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + ) -> GetPromptResult: + """Get a prompt from the server. + + If the server returns an `InputRequiredResult`, the embedded input + requests are dispatched to this client's sampling / elicitation / roots + callbacks and the get is retried automatically (up to + `input_required_max_rounds`). + + Args: + name: The name of the prompt. + arguments: Arguments to pass to the prompt. + input_responses: Responses to seed the first call with (e.g. when + resuming from a persisted `InputRequiredResult`). + request_state: Opaque state to seed the first call with. + meta: Additional metadata for the request. + + Returns: + The prompt content. + + Raises: + InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted. + MCPError: A callback returned `ErrorData` for an embedded input request. + pydantic.ValidationError: The server returned a result that does not + conform to the negotiated protocol version. + """ + + async def retry(r: InputResponses | None, s: str | None) -> GetPromptResult | InputRequiredResult: + return await self.session.get_prompt( + name, arguments, input_responses=r, request_state=s, meta=meta, allow_input_required=True + ) + + return await self._drive_input_required(await retry(input_responses, request_state), retry) + + async def _drive_input_required( + self, + first: _ResultT | InputRequiredResult, + retry: Callable[[InputResponses | None, str | None], Awaitable[_ResultT | InputRequiredResult]], + ) -> _ResultT: + """Hand an `InputRequiredResult` to the SEP-2322 driver, or pass a terminal result through. + + `dispatch` routes each embedded request through the same callback table + that serves legacy server→client RPCs, so the two paths stay + behaviourally identical by construction. + """ + if not isinstance(first, InputRequiredResult): + return first + session = self.session + + async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData: + ctx = ClientRequestContext(session=session, request_id=key, meta=req.params.meta if req.params else None) + return await session.dispatch_input_request(ctx, req) + + return await run_input_required_driver( + first, dispatch=dispatch, retry=retry, max_rounds=self.input_required_max_rounds + ) + + async def complete( + self, + ref: ResourceTemplateReference | PromptReference, + argument: dict[str, str], + context_arguments: dict[str, str] | None = None, + ) -> CompleteResult: + """Get completions for a prompt or resource template argument. + + Args: + ref: Reference to the prompt or resource template + argument: The argument to complete + context_arguments: Additional context arguments + + Returns: + Completion suggestions. + """ + return await self.session.complete(ref=ref, argument=argument, context_arguments=context_arguments) + + async def list_tools( + self, + *, + cursor: str | None = None, + meta: RequestParamsMeta | None = None, + cache_mode: CacheMode = "use", + ) -> ListToolsResult: + """List available tools from the server.""" + return await self._cached_fetch( + "tools/list", + cursor=cursor, + meta=meta, + cache_mode=cache_mode, + send=lambda: self.session.list_tools(params=PaginatedRequestParams(cursor=cursor, _meta=meta)), + # A cache hit skips session.list_tools, so the session re-absorbs the served + # listing to rebuild its derived per-tool state. Hits are cursorless, but a + # cached page 1 can carry next_cursor - never prune on a partial listing. + absorb=lambda hit: self.session._absorb_tool_listing( # pyright: ignore[reportPrivateUsage] + hit, complete=hit.next_cursor is None + ), + ) + + @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def send_roots_list_changed(self) -> None: + """Send a notification that the roots list has changed.""" + await self.session.send_roots_list_changed() # pyright: ignore[reportDeprecated] diff --git a/src/mcp-client/mcp_client/client/context.py b/src/mcp-client/mcp_client/client/context.py new file mode 100644 index 0000000000..cc4d718d87 --- /dev/null +++ b/src/mcp-client/mcp_client/client/context.py @@ -0,0 +1,5 @@ +"""Request context for MCP client handlers.""" + +from mcp_client.client.session import ClientRequestContext + +__all__ = ["ClientRequestContext"] diff --git a/src/mcp-client/mcp_client/client/extension.py b/src/mcp-client/mcp_client/client/extension.py new file mode 100644 index 0000000000..08c03ef72b --- /dev/null +++ b/src/mcp-client/mcp_client/client/extension.py @@ -0,0 +1,196 @@ +"""Opt-in extension interface for MCP clients. + +Subclass `ClientExtension`, set `identifier`, override the hooks you need, and +pass instances to `Client(extensions=[...])`. For an identifier-only +capability ad, use `advertise()`. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeVar, get_args + +from mcp_types import CORE_RESULT_TYPES, CallToolResult, InputRequiredResult, Result +from mcp_types.version import MODERN_PROTOCOL_VERSIONS +from pydantic import AliasChoices, AliasPath, BaseModel +from pydantic.fields import FieldInfo + +from mcp_client.shared.extension import validate_extension_identifier + +if TYPE_CHECKING: + from mcp_client.client.session import ClientSession + +__all__ = [ + "ClaimContext", + "ClientExtension", + "NotificationBinding", + "ResultClaim", + "UnexpectedClaimedResult", + "advertise", +] + +_CLAIM_METHODS: Final[frozenset[str]] = frozenset({"tools/call"}) +"""The closed set of verbs a claim may attach to; widen together with the `method` Literal.""" + +_RESERVED_WIRE_ALIASES: Final[frozenset[str]] = frozenset({"requestState", "inputRequests"}) +"""Typed optional fields of the core result surface that pre-validates every inbound result.""" + + +def _wire_keys(name: str, field: FieldInfo) -> frozenset[str]: + """Every top-level wire key this field can read from or write to.""" + keys = {field.alias or name} + if field.serialization_alias: + keys.add(field.serialization_alias) + validation_alias = field.validation_alias + choices = validation_alias.choices if isinstance(validation_alias, AliasChoices) else [validation_alias] + for choice in choices: + if isinstance(choice, AliasPath): + choice = choice.path[0] + if isinstance(choice, str): + keys.add(choice) + return frozenset(keys) + + +ClaimedT = TypeVar("ClaimedT", bound=Result) +NotifyParamsT = TypeVar("NotifyParamsT", bound=BaseModel) + + +@dataclass(frozen=True, kw_only=True) +class ClaimContext: + """Host-injected context for one `ResultClaim.resolve` call.""" + + session: ClientSession + tool_name: str + read_timeout_seconds: float | None + + +@dataclass(frozen=True, kw_only=True) +class ResultClaim(Generic[ClaimedT]): + """One extra result shape on one spec verb, keyed by the wire `resultType`. + + Active only while the declaring extension is constructed into the client and + the negotiated protocol version admits it. `resolve` finishes a claimed + result, may send follow-ups through `ctx.session`, and must return the + verb's ordinary result. All field constraints are enforced at construction. + """ + + result_type: str + model: type[ClaimedT] + resolve: Callable[[ClaimedT, ClaimContext], Awaitable[CallToolResult]] + method: Literal["tools/call"] = "tools/call" + protocol_versions: frozenset[str] | None = None + + def __post_init__(self) -> None: + if self.method not in _CLAIM_METHODS: + raise ValueError(f"claims attach to {sorted(_CLAIM_METHODS)} only; got method {self.method!r}") + if self.result_type in CORE_RESULT_TYPES: + raise ValueError(f"resultType {self.result_type!r} is core protocol vocabulary") + if Result not in self.model.__mro__: # runtime guard; the ClaimedT bound only constrains checked callers + raise ValueError(f"{self.model.__name__} must subclass mcp_types.Result") + if issubclass(self.model, CallToolResult | InputRequiredResult): + raise ValueError("claim models must not subclass core result types") + for name, model_field in self.model.model_fields.items(): + for clash in sorted(_wire_keys(name, model_field) & _RESERVED_WIRE_ALIASES): + raise ValueError( + f"{self.model.__name__}.{name} aliases {clash!r}, a typed field of the core " + "result surface; a colliding value would fail core validation before the " + "claim adapter runs" + ) + field = self.model.model_fields.get("result_type") + if field is None or get_args(field.annotation) != (self.result_type,): + raise ValueError(f"{self.model.__name__}.result_type must be Literal[{self.result_type!r}]") + if self.protocol_versions is not None and not self.protocol_versions: + raise ValueError("empty protocol_versions could never activate; use None for all") + if self.protocol_versions is not None and not self.protocol_versions.issubset(MODERN_PROTOCOL_VERSIONS): + unrecognized = sorted(self.protocol_versions.difference(MODERN_PROTOCOL_VERSIONS)) + raise ValueError( + f"protocol_versions {unrecognized} are not modern protocol revisions; claimed shapes " + "cannot be delivered on a legacy wire (None means every modern version)" + ) + + +class UnexpectedClaimedResult(RuntimeError): + """A claimed (extension) result arrived on a `call_tool` that did not opt in. + + The parsed value is carried as `result`; the server may already hold state it + references. Opt in via `Client(extensions=[...])` or `allow_claimed=True`. + """ + + def __init__(self, result: Result) -> None: + super().__init__( + f"Server returned a claimed result ({type(result).__name__}); pass the owning extension to " + "Client(extensions=[...]) for transparent resolution, or call with allow_claimed=True " + "and handle the shape. The carried result may reference server-side state needing cleanup." + ) + self.result = result + + +@dataclass(frozen=True, kw_only=True) +class NotificationBinding(Generic[NotifyParamsT]): + """Deliver server notifications for `method` (the bare wire name) to `handler`. + + Observation-only: validated params arrive one at a time per binding, in + dispatch order, through a bounded queue that drops the oldest with a warning + on overflow. Stream transports dispatch each notification independently, so + near-simultaneous notifications may be dispatched out of wire order. Methods + the negotiated version's core tables handle are never delivered to bindings. + """ + + method: str + params_type: type[NotifyParamsT] + handler: Callable[[NotifyParamsT], Awaitable[None]] + + +class ClientExtension: + """Base class for an opt-in client extension; override only what you need. + + The surface is declarative, fixed at construction, and never receives the client. + """ + + #: Reverse-DNS extension identifier, advertised under `ClientCapabilities.extensions`. + identifier: str + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + # Per-instance identifiers (assigned in __init__) are validated at consumption instead. + if (identifier := cls.__dict__.get("identifier")) is not None: + validate_extension_identifier(identifier, owner=cls.__name__) + + def settings(self) -> dict[str, Any]: + """Per-extension settings advertised at `ClientCapabilities.extensions[identifier]`. + + Read once at `Client` construction. A claim-bearing extension is + advertised only at protocol versions where at least one of its claims + is active. + """ + return {} + + def claims(self) -> Sequence[ResultClaim[Any]]: + """Extra result shapes this extension claims, with their resolvers.""" + return () + + def notifications(self) -> Sequence[NotificationBinding[Any]]: + """Server notifications this extension observes.""" + return () + + +class _AdvertiseOnly(ClientExtension): + """Ad-only extension returned by `advertise()`.""" + + def __init__(self, identifier: str, settings: dict[str, Any]) -> None: + self.identifier = identifier + self._settings = settings + + def settings(self) -> dict[str, Any]: + return self._settings + + +def advertise(identifier: str, settings: dict[str, Any] | None = None) -> ClientExtension: + """Advertise an extension identifier (with optional settings) and nothing else. + + Advertising an extension you do not implement asserts wire support you do + not have; for behavioral extensions construct the real extension instead. + """ + validate_extension_identifier(identifier, owner="advertise") + return _AdvertiseOnly(identifier, {} if settings is None else settings) diff --git a/src/mcp-client/mcp_client/client/session.py b/src/mcp-client/mcp_client/client/session.py new file mode 100644 index 0000000000..692dccaafc --- /dev/null +++ b/src/mcp-client/mcp_client/client/session.py @@ -0,0 +1,1506 @@ +from __future__ import annotations + +import json +import logging +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from functools import cache, reduce +from operator import or_ +from types import TracebackType, UnionType +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, cast, get_args, overload + +import anyio +import anyio.abc +import anyio.lowlevel +import mcp_types as types +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream +from mcp_types import ( + CLIENT_CAPABILITIES_META_KEY, + CLIENT_INFO_META_KEY, + CONNECTION_CLOSED, + INTERNAL_ERROR, + LOG_LEVEL_META_KEY, + METHOD_NOT_FOUND, + PROTOCOL_VERSION_META_KEY, + SERVER_INFO_META_KEY, + UNSUPPORTED_PROTOCOL_VERSION, + RequestId, + RequestParamsMeta, +) +from mcp_types import methods as _methods +from mcp_types.version import ( + HANDSHAKE_PROTOCOL_VERSIONS, + KNOWN_PROTOCOL_VERSIONS, + LATEST_HANDSHAKE_VERSION, + LATEST_MODERN_VERSION, + MODERN_PROTOCOL_VERSIONS, +) +from pydantic import BaseModel, Discriminator, Tag, TypeAdapter, ValidationError +from typing_extensions import Self, TypeVar, deprecated + +from mcp_client.client._transport import ReadStream, WriteStream +from mcp_client.client.extension import NotificationBinding, ResultClaim, UnexpectedClaimedResult +from mcp_client.client.subscriptions import ListenRoute +from mcp_client.shared._compat import resync_tracer +from mcp_client.shared.dispatcher import CallOptions, DispatchContext, Dispatcher, ProgressFnT, as_request_id +from mcp_client.shared.exceptions import MCPDeprecationWarning, MCPError +from mcp_client.shared.inbound import ( + MCP_METHOD_HEADER, + MCP_NAME_HEADER, + MCP_PROTOCOL_VERSION_HEADER, + NAME_BEARING_METHODS, + encode_header_value, + find_invalid_x_mcp_header, + mcp_param_headers, + x_mcp_header_map, +) +from mcp_client.shared.jsonrpc_dispatcher import JSONRPCDispatcher, cancelled_request_id_from_params +from mcp_client.shared.message import ClientMessageMetadata, SessionMessage +from mcp_client.shared.subscriptions import SUBSCRIPTION_ID_META_KEY, event_from_wire +from mcp_client.shared.transport_context import TransportContext + +if TYPE_CHECKING: + # `jsonschema` is imported lazily inside `validate_tool_result`: pulling it (and its + # `attrs`/`referencing` tree) in at module scope costs every client that never validates. + from jsonschema.protocols import Validator + +DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0") +DISCOVER_TIMEOUT_SECONDS = 10.0 +_NOTIFICATION_QUEUE_SIZE: Final = 256 + +logger = logging.getLogger("client") + + +def _clamp_inbound_ttl(raw: dict[str, Any]) -> None: + """Floor a negative inbound `ttlMs` to 0 before `ge=0` validation fails the call (2026-07-28 caching SHOULD).""" + ttl = raw.get("ttlMs") + if isinstance(ttl, int | float) and not isinstance(ttl, bool) and ttl < 0: + raw["ttlMs"] = 0 + + +@cache +def _wire_fields(target: type[BaseModel] | UnionType) -> frozenset[str]: + """Top-level wire keys `target` declares (its members', for a union).""" + members: tuple[Any, ...] = get_args(target) if isinstance(target, UnionType) else (target,) + models = [m for m in members if isinstance(m, type) and issubclass(m, BaseModel)] + fields: set[str] = set() + for model in models: + fields.update(field.alias or name for name, field in model.model_fields.items()) + return frozenset(fields) + + +@cache +def _later_revision_fields(method: str, version: str) -> frozenset[str]: + """Result keys a revision newer than `version` declares for `method` but `version` doesn't. + + The version-free result types carry every revision's fields, so such a key + (e.g. 2026-07-28 `ttlMs`/`cacheScope` on a pre-2026 session) is outside the + negotiated contract yet would still parse into the model and trip that later + revision's constraints. Empty at the newest known revision. + """ + current = _methods.SERVER_RESULTS.get((method, version)) + if current is None or version not in KNOWN_PROTOCOL_VERSIONS: + return frozenset() + newer = KNOWN_PROTOCOL_VERSIONS[KNOWN_PROTOCOL_VERSIONS.index(version) + 1 :] + later: set[str] = set() + for revision in newer: + row = _methods.SERVER_RESULTS.get((method, revision)) + if row is not None: + later |= _wire_fields(row) + return frozenset(later) - _wire_fields(current) + + +def _same_schema(a: dict[str, Any] | None, b: dict[str, Any] | None) -> bool: + """JSON equality for two output schemas. + + Python `==` is not JSON equality: it conflates `True`/`1` and `False`/`0`, which JSON + Schema keeps distinct (`const: true` vs `const: 1`). Canonical serialization compares as + JSON does; where it is stricter (`1` vs `1.0`), erring toward "changed" only costs a + recompile, never a stale validator. + """ + return json.dumps(a, sort_keys=True) == json.dumps(b, sort_keys=True) + + +def _preconnect_stamp(data: dict[str, Any], opts: CallOptions) -> None: + # initialize/discover forbid cancellation; other pre-handshake requests (lowlevel + # ClientSession callers may skip the handshake entirely) keep the courtesy cancel. + if data["method"] in ("initialize", "server/discover"): + opts["cancel_on_abandon"] = False + + +def _parse_server_info_stamp(result: types.DiscoverResult) -> types.Implementation | None: + """The typed identity from a discover result's `_meta` serverInfo stamp. + + The stamp is display-only per the spec, so absent and malformed both read + as `None` rather than failing the connection. + """ + raw = (result.meta or {}).get(SERVER_INFO_META_KEY) + if raw is None: + return None + try: + return types.Implementation.model_validate(raw) + except ValidationError: + return None + + +def _make_handshake_stamp(protocol_version: str) -> Callable[[dict[str, Any], CallOptions], None]: + def stamp(data: dict[str, Any], opts: CallOptions) -> None: + opts.setdefault("headers", {})[MCP_PROTOCOL_VERSION_HEADER] = protocol_version + + return stamp + + +def _make_modern_stamp( + protocol_version: str, + client_info: dict[str, Any], + capabilities: dict[str, Any], + resolve_param_headers: Callable[[str, Mapping[str, Any]], dict[str, str]], + *, + log_level: types.LoggingLevel | None = None, +) -> Callable[[dict[str, Any], CallOptions], None]: + def stamp(data: dict[str, Any], opts: CallOptions) -> None: + params = data.setdefault("params", {}) + meta = params.setdefault("_meta", {}) + meta[PROTOCOL_VERSION_META_KEY] = protocol_version + meta[CLIENT_INFO_META_KEY] = client_info + meta[CLIENT_CAPABILITIES_META_KEY] = capabilities + # The per-request log-delivery opt-in (2026 logging is opt-in per + # request). A default the caller can override on any single call by + # supplying the key in that request's `_meta`, hence setdefault. + if log_level is not None: + meta.setdefault(LOG_LEVEL_META_KEY, log_level) + # `cancel_on_abandon` stays at the dispatcher default (True): the + # courtesy `notifications/cancelled` is the abandon signal. On the + # stream transports it is the 2026 wire's cancellation spelling; the + # streamable-HTTP transport translates it into aborting the request's + # own POST instead of writing it (the 2026 HTTP wire has no + # client-to-server notifications - closing the stream is the signal). + # The negotiation methods still opt out, mirroring `_preconnect_stamp`: + # the spec forbids cancelling them. + if data["method"] in ("initialize", "server/discover"): + opts["cancel_on_abandon"] = False + headers = opts.setdefault("headers", {}) + headers[MCP_PROTOCOL_VERSION_HEADER] = protocol_version + headers[MCP_METHOD_HEADER] = data["method"] + name_key = NAME_BEARING_METHODS.get(data["method"]) + if name_key is not None and isinstance(name := params.get(name_key), str): + headers[MCP_NAME_HEADER] = encode_header_value(name) + if data["method"] == "tools/call" and isinstance(name := params.get("name"), str): + headers.update(resolve_param_headers(name, params.get("arguments") or {})) + + return stamp + + +ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel) + + +@dataclass(kw_only=True) +class ClientRequestContext: + """Context for a server-initiated request, passed to the sampling/elicitation/list-roots callbacks.""" + + session: ClientSession + request_id: RequestId + meta: RequestParamsMeta | None = None + + +class SamplingFnT(Protocol): + async def __call__( + self, + context: ClientRequestContext, + params: types.CreateMessageRequestParams, + ) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: ... # pragma: no branch + + +class ElicitationFnT(Protocol): + async def __call__( + self, + context: ClientRequestContext, + params: types.ElicitRequestParams, + ) -> types.ElicitResult | types.ErrorData: ... # pragma: no branch + + +class ListRootsFnT(Protocol): + async def __call__( + self, context: ClientRequestContext + ) -> types.ListRootsResult | types.ErrorData: ... # pragma: no branch + + +class LoggingFnT(Protocol): + async def __call__(self, params: types.LoggingMessageNotificationParams) -> None: ... # pragma: no branch + + +IncomingMessage: TypeAlias = types.ServerNotification | Exception +"""What `message_handler` receives: the server notifications the session surfaces, plus transport-level exceptions. + +`notifications/cancelled` is applied by the dispatcher and never surfaced, and a +`notifications/subscriptions/acknowledged` for a live `listen()` stream is consumed by that +stream, so neither reaches the handler. +""" + + +class MessageHandlerFnT(Protocol): + async def __call__(self, message: IncomingMessage) -> None: ... # pragma: no branch + + +async def _default_message_handler(message: IncomingMessage) -> None: + await anyio.lowlevel.checkpoint() + + +async def _default_sampling_callback( + context: ClientRequestContext, + params: types.CreateMessageRequestParams, +) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: + return types.ErrorData( + code=types.INVALID_REQUEST, + message="Sampling not supported", + ) + + +async def _default_elicitation_callback( + context: ClientRequestContext, + params: types.ElicitRequestParams, +) -> types.ElicitResult | types.ErrorData: + return types.ErrorData( + code=types.INVALID_REQUEST, + message="Elicitation not supported", + ) + + +async def _default_list_roots_callback( + context: ClientRequestContext, +) -> types.ListRootsResult | types.ErrorData: + return types.ErrorData( + code=types.INVALID_REQUEST, + message="List roots not supported", + ) + + +async def _default_logging_callback( + params: types.LoggingMessageNotificationParams, +) -> None: + pass + + +ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(types.ClientResult | types.ErrorData) + +# Typed against the wide parse union so adopt-built claim adapters share this attribute type. +_CallToolResultAdapter: TypeAdapter[types.CallToolResult | types.InputRequiredResult | types.Result] = TypeAdapter( + types.CallToolResult | types.InputRequiredResult +) +_GetPromptResultAdapter: TypeAdapter[types.GetPromptResult | types.InputRequiredResult] = TypeAdapter( + types.GetPromptResult | types.InputRequiredResult +) +_ReadResourceResultAdapter: TypeAdapter[types.ReadResourceResult | types.InputRequiredResult] = TypeAdapter( + types.ReadResourceResult | types.InputRequiredResult +) + + +def _claim_active(claim: ResultClaim[Any], version: str) -> bool: + """A claim is active at modern versions only, narrowed by its optional version subset.""" + return version in MODERN_PROTOCOL_VERSIONS and ( + claim.protocol_versions is None or version in claim.protocol_versions + ) + + +def _active_claims_at( + claims_by_extension: Mapping[str, tuple[ResultClaim[Any], ...]], version: str +) -> dict[str, ResultClaim[Any]]: + """Claims active at `version`, keyed by wire tag; empty at any legacy version.""" + return { + claim.result_type: claim + for claims in claims_by_extension.values() + for claim in claims + if _claim_active(claim, version) + } + + +def _build_call_tool_adapter( + active: Mapping[str, ResultClaim[Any]], +) -> TypeAdapter[types.CallToolResult | types.InputRequiredResult | types.Result]: + """Build a discriminated tools/call adapter: a core arm plus one arm per active claim.""" + if not active: + return _CallToolResultAdapter + tags = frozenset(active) + core_arm = "core" + while core_arm in tags: # the routing sentinel must never collide with a claimed tag + core_arm += "-" + + def _route(value: Any) -> str: + # pydantic hands the discriminator either the raw dict or an already-built model. + # Unknown or non-string tags route to the core arm and fail core validation there. + if isinstance(value, dict): + tag = cast("dict[str, Any]", value).get("resultType") + else: + tag = getattr(value, "result_type", None) + return tag if isinstance(tag, str) and tag in tags else core_arm + + arms: list[Any] = [Annotated[types.CallToolResult | types.InputRequiredResult, Tag(core_arm)]] + arms += [Annotated[claim.model, Tag(tag)] for tag, claim in active.items()] + # reduce(or_) rather than Union star-unpack, which needs py3.11+. + return TypeAdapter(Annotated[reduce(or_, arms), Discriminator(_route)]) + + +def _index_claims( + result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None, + extensions: dict[str, dict[str, Any]] | None, +) -> dict[str, tuple[ResultClaim[Any], ...]]: + """Validate and copy the claims-by-extension mapping.""" + indexed: dict[str, tuple[ResultClaim[Any], ...]] = {} + seen: set[str] = set() + for identifier, claims in (result_claims or {}).items(): + if extensions is None or identifier not in extensions: + raise ValueError( + f"result_claims key {identifier!r} has no extensions entry; a claim is only " + "advertised through its extension's capability ad" + ) + if not claims: + raise ValueError( + f"result_claims[{identifier!r}] is empty and would drop the extension from " + "the capability ad at every version. Omit the key instead" + ) + for claim in claims: + if claim.result_type in seen: + raise ValueError(f"duplicate result claim for resultType {claim.result_type!r}") + seen.add(claim.result_type) + indexed[identifier] = tuple(claims) + return indexed + + +def _index_bindings( + notification_bindings: Sequence[NotificationBinding[Any]] | None, +) -> dict[str, NotificationBinding[Any]]: + """Index bindings by wire method, rejecting duplicates.""" + indexed: dict[str, NotificationBinding[Any]] = {} + for binding in notification_bindings or (): + if binding.method in indexed: + raise ValueError(f"duplicate notification binding for method {binding.method!r}") + indexed[binding.method] = binding + return indexed + + +def _input_required_unexpected(method: str) -> RuntimeError: + return RuntimeError( + "Server returned InputRequiredResult; pass allow_input_required=True to receive it " + f"and retry {method}(..., input_responses=..., request_state=result.request_state)." + ) + + +class ClientSession: + """Client half of an MCP connection, running on a `Dispatcher`. + + Construct it over a transport's stream pair (or pass a pre-built + `dispatcher=`), enter as an async context manager, then call + `initialize()`. The dispatcher owns the receive loop and request + correlation; this class owns the typed MCP layer and the constructor + callbacks. Transport `Exception` items reach `message_handler` on any + stream-backed dispatcher (`JSONRPCDispatcher`), whether built here from a + stream pair or supplied without a stream-exception hook of its own; an + in-process `DirectDispatcher` carries none. + + Extension `result_claims` fold into tools/call parsing at `adopt()`; + `notification_bindings` observe vendor notifications via bounded FIFOs. + """ + + def __init__( + self, + read_stream: ReadStream[SessionMessage | Exception] | None = None, + write_stream: WriteStream[SessionMessage] | None = None, + read_timeout_seconds: float | None = None, + sampling_callback: SamplingFnT | None = None, + elicitation_callback: ElicitationFnT | None = None, + list_roots_callback: ListRootsFnT | None = None, + logging_callback: LoggingFnT | None = None, + message_handler: MessageHandlerFnT | None = None, + client_info: types.Implementation | None = None, + *, + log_level: types.LoggingLevel | None = None, + sampling_capabilities: types.SamplingCapability | None = None, + extensions: dict[str, dict[str, Any]] | None = None, + result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None = None, + notification_bindings: Sequence[NotificationBinding[Any]] | None = None, + dispatcher: Dispatcher[Any] | None = None, + ) -> None: + self._session_read_timeout_seconds = read_timeout_seconds + self._client_info = client_info or DEFAULT_CLIENT_INFO + self._sampling_callback = sampling_callback or _default_sampling_callback + self._sampling_capabilities = sampling_capabilities + self._extensions = dict(extensions) if extensions is not None else None + self._result_claims = _index_claims(result_claims, extensions) + self._notification_bindings = _index_bindings(notification_bindings) + self._active_claims: dict[str, ResultClaim[Any]] = {} + self._call_tool_adapter = _CallToolResultAdapter + self._binding_queues: dict[ + str, tuple[MemoryObjectSendStream[BaseModel], MemoryObjectReceiveStream[BaseModel]] + ] = {} + self._elicitation_callback = elicitation_callback or _default_elicitation_callback + self._list_roots_callback = list_roots_callback or _default_list_roots_callback + self._logging_callback = logging_callback or _default_logging_callback + self._log_level: types.LoggingLevel | None = log_level + self._message_handler = message_handler or _default_message_handler + self._tool_output_schemas: dict[str, dict[str, Any] | None] = {} + # Compiled output-schema validators, derived from `_tool_output_schemas` and owned by + # `_absorb_tool_listing`, which evicts a tool's entry whenever its schema changes. + self._tool_output_validators: dict[str, Validator] = {} + self._x_mcp_header_maps: dict[str, dict[tuple[str, ...], str]] = {} + self._initialize_result: types.InitializeResult | None = None + self._discover_result: types.DiscoverResult | None = None + self._discover_server_info: types.Implementation | None = None + self._negotiated_version: str | None = None + self._stamp: Callable[[dict[str, Any], CallOptions], None] = _preconnect_stamp + self._task_group: anyio.abc.TaskGroup | None = None + # subscriptions/listen demux routes; membership decides ack consumption (raw listens are never registered) + self._listen_routes: dict[RequestId, ListenRoute] = {} + if dispatcher is not None: + if read_stream is not None or write_stream is not None: + raise ValueError("pass read_stream/write_stream or dispatcher, not both") + self._dispatcher: Dispatcher[Any] = dispatcher + if isinstance(dispatcher, JSONRPCDispatcher) and dispatcher.on_stream_exception is None: + # Route transport-level Exception items into message_handler — only + # stream-backed dispatchers carry these; DirectDispatcher has none. + # Don't clobber a caller-supplied hook. + # TODO(L78): this leaves a bound-method ref on the dispatcher after the + # session exits (memory pin) and a second wrap of the same dispatcher would + # skip install. The Transport-as-Dispatcher rework (L77) removes this seam. + dispatcher.on_stream_exception = self._on_stream_exception + else: + if read_stream is None or write_stream is None: + raise ValueError("read_stream and write_stream are required when no dispatcher is given") + # Built eagerly so notifications can be sent before entering the context manager. + self._dispatcher = JSONRPCDispatcher( + read_stream, write_stream, on_stream_exception=self._on_stream_exception + ) + + async def __aenter__(self) -> Self: + self._task_group = anyio.create_task_group() + await self._task_group.__aenter__() + try: + # Queues must exist before the dispatcher starts: _on_notify enqueues into this dict. + for binding in self._notification_bindings.values(): + send, receive = anyio.create_memory_object_stream[BaseModel](_NOTIFICATION_QUEUE_SIZE) + self._binding_queues[binding.method] = (send, receive) + await self._task_group.start( + self._dispatcher.run, self._on_request, self._on_notify, self._intercept_notification + ) + for binding in self._notification_bindings.values(): + _, receive = self._binding_queues[binding.method] + self._task_group.start_soon(self._deliver_bound_notifications, binding, receive) + except BaseException: + # Unwind the entered task group before propagating: a cancellation + # landing here (e.g. `move_on_after` around connect) would abandon + # it and anyio would later raise "exited non-innermost cancel scope". + task_group = self._task_group + self._task_group = None + task_group.cancel_scope.cancel() + # Shield the group's own scope (a new one would break LIFO exit) + # so a pending outer cancellation cannot re-fire inside __aexit__. + task_group.cancel_scope.shield = True + try: + await task_group.__aexit__(None, None, None) + finally: + self._close_binding_queues() + raise + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> bool | None: + # Exit must not block: cancel the dispatcher, binding consumers, and in-flight callbacks. + assert self._task_group is not None + self._task_group.cancel_scope.cancel() + try: + result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb) + finally: + self._close_binding_queues() + self._settle_listen_routes_closed() + await resync_tracer() + return result + + def _close_binding_queues(self) -> None: + # Unclosed memory object streams warn at garbage collection; close is idempotent. + for send, receive in self._binding_queues.values(): + send.close() + receive.close() + self._binding_queues.clear() + + async def _deliver_bound_notifications( + self, binding: NotificationBinding[Any], receive: MemoryObjectReceiveStream[BaseModel] + ) -> None: + """Consume one binding's FIFO, decoupled from the dispatcher so handlers can do session I/O.""" + while True: + params = await receive.receive() + try: + await binding.handler(params) + except Exception: + # A raising handler costs only that delivery, as in _on_notify. + logger.exception("notification binding handler for %r raised", binding.method) + + async def send_request( + self, + request: types.ClientRequest | types.Request[Any, Any], + result_type: type[ReceiveResultT] | TypeAdapter[ReceiveResultT], + request_read_timeout_seconds: float | None = None, + metadata: ClientMessageMetadata | None = None, + progress_callback: ProgressFnT | None = None, + ) -> ReceiveResultT: + """Send a request and wait for its typed result. + + Args: + metadata: Streamable HTTP resumption hints. + + Raises: + MCPError: Error response, read timeout, or connection closed. + RuntimeError: Called before entering the context manager. + ValueError: The request declares `name_param` but its params carry no string name. + pydantic.ValidationError: The server returned a result that does not + conform to the negotiated protocol version. + """ + data = request.model_dump(by_alias=True, mode="json", exclude_none=True) + method: str = data["method"] + opts: CallOptions = {} + self._stamp(data, opts) + # The stamp runs first, so its NAME_BEARING_METHODS rows win; a missing name fails loud. + headers = opts.setdefault("headers", {}) + if (key := type(request).name_param) is not None and MCP_NAME_HEADER not in headers: + params_data: dict[str, Any] = data.get("params") or {} + name = params_data.get(key) + if not isinstance(name, str): + raise ValueError(f"{method} requires params[{key!r}] for Mcp-Name") + headers[MCP_NAME_HEADER] = encode_header_value(name) + timeout = ( + request_read_timeout_seconds + if request_read_timeout_seconds is not None + else self._session_read_timeout_seconds + ) + if timeout is not None: + opts["timeout"] = timeout + if progress_callback is not None: + opts["on_progress"] = progress_callback + if metadata is not None: + if metadata.resumption_token is not None: + opts["resumption_token"] = metadata.resumption_token + if metadata.on_resumption_token_update is not None: + opts["on_resumption_token"] = metadata.on_resumption_token_update + raw = await self._dispatcher.send_raw_request(method, data.get("params"), opts) + _clamp_inbound_ttl(raw) + # Literal fallback covers pre-handshake and stateless; matches runner.py. + version = self._negotiated_version or "2025-11-25" + try: + _methods.validate_server_result(method, version, raw) + except KeyError: + pass + # Drop a later revision's fields (e.g. 2026-07-28 cache hints on a pre-2026 + # session): they are outside the negotiated contract, and the version-free + # result type would otherwise apply that revision's constraints to them. + if not (foreign := _later_revision_fields(method, version)).isdisjoint(raw): + raw = {key: value for key, value in raw.items() if key not in foreign} + if isinstance(result_type, TypeAdapter): + return result_type.validate_python(raw, by_name=False) + return result_type.model_validate(raw, by_name=False) + + async def send_notification(self, notification: types.ClientNotification) -> None: + """Send a one-way notification. Usable before entering the context manager. + + Fire-and-forget: after the connection has closed, the notification is + dropped with a debug log instead of raising. + """ + data = notification.model_dump(by_alias=True, mode="json", exclude_none=True) + opts: CallOptions = {} + self._stamp(data, opts) + await self._dispatcher.notify(data["method"], data.get("params"), opts) + + def _build_capabilities(self, version: str) -> types.ClientCapabilities: + """Build the capability ad for a wire speaking `version`. + + Claim-bearing identifiers whose claims are all inactive at `version` drop, so + the client never advertises result shapes it would reject; claim-less + identifiers always advertise. + """ + extensions = self._extensions + if extensions is not None and self._result_claims: + extensions = { + identifier: settings + for identifier, settings in extensions.items() + if identifier not in self._result_claims + or any(_claim_active(claim, version) for claim in self._result_claims[identifier]) + } or None + sampling = ( + (self._sampling_capabilities or types.SamplingCapability()) + if self._sampling_callback is not _default_sampling_callback + else None + ) + elicitation = ( + types.ElicitationCapability(form=types.FormElicitationCapability(), url=types.UrlElicitationCapability()) + if self._elicitation_callback is not _default_elicitation_callback + else None + ) + roots = ( + # TODO: Should this be based on whether we + # _will_ send notifications, or only whether + # they're supported? + types.RootsCapability(list_changed=True) + if self._list_roots_callback is not _default_list_roots_callback + else None + ) + return types.ClientCapabilities( + sampling=sampling, elicitation=elicitation, experimental=None, extensions=extensions, roots=roots + ) + + async def initialize(self) -> types.InitializeResult: + if self._initialize_result is not None: + return self._initialize_result + result = await self.send_request( + types.InitializeRequest( + params=types.InitializeRequestParams( + protocol_version=LATEST_HANDSHAKE_VERSION, + # The handshake negotiates only legacy versions, where no claim is active. + capabilities=self._build_capabilities(LATEST_HANDSHAKE_VERSION), + client_info=self._client_info, + ), + ), + types.InitializeResult, + ) + + if result.protocol_version not in HANDSHAKE_PROTOCOL_VERSIONS: + raise RuntimeError(f"Unsupported protocol version from the server: {result.protocol_version}") + + self.adopt(result) + + await self.send_notification(types.InitializedNotification()) + + return result + + def adopt(self, result: types.InitializeResult | types.DiscoverResult) -> None: + """Install negotiated state from a result the caller already holds (no wire traffic). + + Clears the opposite slot, so at most one of `initialize_result` / + `discover_result` is ever non-None. + + Raises: + RuntimeError: `result` is a `DiscoverResult` whose `supported_versions` + shares nothing with this client's `MODERN_PROTOCOL_VERSIONS`. + """ + if isinstance(result, types.DiscoverResult): + # ordered oldest→newest via MODERN_PROTOCOL_VERSIONS + mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in result.supported_versions] + if not mutual: + raise RuntimeError( + f"No mutually supported modern protocol version " + f"(server: {result.supported_versions}, client: {list(MODERN_PROTOCOL_VERSIONS)})" + ) + version = mutual[-1] + client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True) + capabilities = self._build_capabilities(version).model_dump(by_alias=True, mode="json", exclude_none=True) + self._stamp = _make_modern_stamp( + version, client_info, capabilities, self._resolve_param_headers, log_level=self._log_level + ) + self._discover_result = result + self._discover_server_info = _parse_server_info_stamp(result) + self._initialize_result = None + else: + version = result.protocol_version + self._stamp = _make_handshake_stamp(version) + self._initialize_result = result + self._discover_result = None + self._discover_server_info = None + self._negotiated_version = version + # Both arms reach here, so re-adoption resets cleanly; legacy versions activate no claims. + # Core-vocabulary tags are unconstructible (ResultClaim.__post_init__), so no exclusion needed. + self._active_claims = _active_claims_at(self._result_claims, version) + self._call_tool_adapter = _build_call_tool_adapter(self._active_claims) + for method in self._notification_bindings: + # Bindings are consulted only for methods core does not know, so this one can never fire. + if (method, version) in _methods.SERVER_NOTIFICATIONS: + logger.warning( + "notification binding for %r will never fire at %s: the core protocol defines this method", + method, + version, + ) + + async def send_discover(self, version: str) -> dict[str, Any]: + """Send a single ``server/discover`` at ``version`` and return the raw result dict. + + No retry, no ``adopt()``. The ``_meta`` envelope and the + ``Mcp-Protocol-Version`` header are stamped at ``version`` so the + server-side era router sees a coherent probe. Used by ``discover()`` and + the connect-time auto-negotiation policy. + + Raises: + MCPError: The server returned a JSON-RPC error, or the transport + bounced the request at its own layer (a bare HTTP 4xx is + synthesized into a JSON-RPC error by the transport). + """ + client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True) + capabilities = self._build_capabilities(version).model_dump(by_alias=True, mode="json", exclude_none=True) + request = types.DiscoverRequest( + params=types.RequestParams( + _meta={ + PROTOCOL_VERSION_META_KEY: version, + CLIENT_INFO_META_KEY: client_info, + CLIENT_CAPABILITIES_META_KEY: capabilities, + } + ) + ) + data = request.model_dump(by_alias=True, mode="json", exclude_none=True) + opts: CallOptions = { + "timeout": DISCOVER_TIMEOUT_SECONDS, + "cancel_on_abandon": False, + "headers": {MCP_PROTOCOL_VERSION_HEADER: version, MCP_METHOD_HEADER: data["method"]}, + } + raw = await self._dispatcher.send_raw_request(data["method"], data.get("params"), opts) + # Un-floored, a negative ttl fails the mode='auto' probe's validation and silently downgrades the handshake. + _clamp_inbound_ttl(raw) + return raw + + async def discover(self) -> types.DiscoverResult: + """Probe `server/discover` and adopt the result. + + Sends a single `server/discover` proposing the newest modern protocol + version. On `UNSUPPORTED_PROTOCOL_VERSION` (-32022) the server's + `supported` list is intersected with `MODERN_PROTOCOL_VERSIONS` and the + probe is retried once at the highest mutual version. Any other error — + including `METHOD_NOT_FOUND` (-32601) and `REQUEST_TIMEOUT` (-32001) — + propagates; the legacy `initialize()` fallback is the caller's policy. + + Raises: + MCPError: The server rejected `server/discover`, the probe timed + out, or the -32022 retry found no mutual version / failed again. + RuntimeError: `adopt()` found no mutual version in the returned + `supported_versions`. + """ + if self._discover_result is not None: + return self._discover_result + + try: + raw = await self.send_discover(LATEST_MODERN_VERSION) + except MCPError as e: + if e.code != UNSUPPORTED_PROTOCOL_VERSION: + raise + try: + data = types.UnsupportedProtocolVersionErrorData.model_validate(e.error.data) + except ValidationError: + raise e from None + # ordered oldest→newest via MODERN_PROTOCOL_VERSIONS + mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in data.supported] + if not mutual: + raise + raw = await self.send_discover(mutual[-1]) + + result = types.DiscoverResult.model_validate(raw) + self.adopt(result) + return result + + @property + def initialize_result(self) -> types.InitializeResult | None: + """The server's InitializeResult. None unless `initialize()` ran (or was adopted).""" + return self._initialize_result + + @property + def discover_result(self) -> types.DiscoverResult | None: + """The server's DiscoverResult. None unless `discover()` ran (or was adopted). + + Retained intact (supported_versions, ttl_ms, cache_scope) so callers + can round-trip it as ``prior_discover=``. + """ + return self._discover_result + + @property + def protocol_version(self) -> str | None: + """Negotiated protocol version. None until `initialize()`, `discover()`, or `adopt()`.""" + return self._negotiated_version + + @property + def server_info(self) -> types.Implementation | None: + """Server name/version. None until `initialize()`, `discover()`, or `adopt()`. + + On 2026-era connections this is the discover result's optional `_meta` + `serverInfo` stamp, parsed once at adopt time; `None` when the server + did not identify itself. The stamp is display-only per the spec, so a + malformed value reads as absent rather than failing the connection. + """ + if self._discover_result is not None: + return self._discover_server_info + if self._initialize_result is not None: + return self._initialize_result.server_info + return None + + @property + def server_capabilities(self) -> types.ServerCapabilities | None: + """Server capabilities. None until `initialize()`, `discover()`, or `adopt()`.""" + if self._discover_result is not None: + return self._discover_result.capabilities + if self._initialize_result is not None: + return self._initialize_result.capabilities + return None + + @property + def instructions(self) -> str | None: + """Server-provided instructions text, if any.""" + if self._discover_result is not None: + return self._discover_result.instructions + if self._initialize_result is not None: + return self._initialize_result.instructions + return None + + async def send_ping(self, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: + """Send a ping request.""" + return await self.send_request(types.PingRequest(params=types.RequestParams(_meta=meta)), types.EmptyResult) + + @deprecated( + "Client-to-server progress is deprecated as of 2026-07-28; progress is server-to-client only.", + category=MCPDeprecationWarning, + ) + async def send_progress_notification( + self, + progress_token: str | int, + progress: float, + total: float | None = None, + message: str | None = None, + *, + meta: RequestParamsMeta | None = None, + ) -> None: + """Send a progress notification.""" + await self.send_notification( + types.ProgressNotification( + params=types.ProgressNotificationParams( + progress_token=progress_token, + progress=progress, + total=total, + message=message, + _meta=meta, + ), + ) + ) + + @deprecated("The logging capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def set_logging_level( + self, + level: types.LoggingLevel, + *, + meta: RequestParamsMeta | None = None, + ) -> types.EmptyResult: + """Send a logging/setLevel request.""" + return await self.send_request( + types.SetLevelRequest(params=types.SetLevelRequestParams(level=level, _meta=meta)), + types.EmptyResult, + ) + + async def list_resources(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListResourcesResult: + """Send a resources/list request. + + Args: + params: Full pagination parameters including cursor and any future fields + """ + return await self.send_request(types.ListResourcesRequest(params=params), types.ListResourcesResult) + + async def list_resource_templates( + self, *, params: types.PaginatedRequestParams | None = None + ) -> types.ListResourceTemplatesResult: + """Send a resources/templates/list request. + + Args: + params: Full pagination parameters including cursor and any future fields + """ + return await self.send_request( + types.ListResourceTemplatesRequest(params=params), + types.ListResourceTemplatesResult, + ) + + @overload + async def read_resource( + self, + uri: str, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: Literal[False] = False, + ) -> types.ReadResourceResult: ... + + @overload + async def read_resource( + self, + uri: str, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: bool, + ) -> types.ReadResourceResult | types.InputRequiredResult: ... + + async def read_resource( + self, + uri: str, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: bool = False, + ) -> types.ReadResourceResult | types.InputRequiredResult: + """Send a resources/read request. + + Args: + input_responses: Responses to a prior `InputRequiredResult.input_requests`. + request_state: Opaque state echoed from a prior `InputRequiredResult`. + allow_input_required: When `False` (default), an `InputRequiredResult` + from the server raises `RuntimeError`; when `True`, it is returned + so the caller can resolve the requests and retry. + + Raises: + RuntimeError: If the server returns an `InputRequiredResult` and + `allow_input_required` is `False`. + """ + result = await self.send_request( + types.ReadResourceRequest( + params=types.ReadResourceRequestParams( + uri=uri, + input_responses=input_responses, + request_state=request_state, + _meta=meta, + ), + ), + _ReadResourceResultAdapter, + ) + if isinstance(result, types.InputRequiredResult) and not allow_input_required: + raise _input_required_unexpected("read_resource") + return result + + @deprecated( + "resources/subscribe is removed as of 2026-07-28; use Client.listen() instead.", + category=MCPDeprecationWarning, + ) + async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: + """Send a resources/subscribe request (2025-era servers only).""" + return await self.send_request( + types.SubscribeRequest(params=types.SubscribeRequestParams(uri=uri, _meta=meta)), + types.EmptyResult, + ) + + @deprecated( + "resources/unsubscribe is removed as of 2026-07-28; use Client.listen() instead.", + category=MCPDeprecationWarning, + ) + async def unsubscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: + """Send a resources/unsubscribe request (2025-era servers only).""" + return await self.send_request( + types.UnsubscribeRequest(params=types.UnsubscribeRequestParams(uri=uri, _meta=meta)), + types.EmptyResult, + ) + + @overload + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: Literal[False] = False, + allow_claimed: Literal[False] = False, + ) -> types.CallToolResult: ... + + @overload + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: bool, + allow_claimed: Literal[False] = False, + ) -> types.CallToolResult | types.InputRequiredResult: ... + + @overload + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: Literal[False] = False, + allow_claimed: bool, + ) -> types.CallToolResult | types.Result: ... + + @overload + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: bool, + allow_claimed: bool, + ) -> types.CallToolResult | types.InputRequiredResult | types.Result: ... + + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: bool = False, + allow_claimed: bool = False, + ) -> types.CallToolResult | types.InputRequiredResult | types.Result: + """Send a tools/call request with optional progress callback support. + + On a modern (2026-07-28) connection, arguments annotated with `x-mcp-header` + in the tool's input schema are mirrored into `Mcp-Param-*` request headers. + The annotations are read from the tool's last `list_tools` entry, so list + the tool before calling it to enable header emission. + + Args: + input_responses: Responses to a prior `InputRequiredResult.input_requests`. + request_state: Opaque state echoed from a prior `InputRequiredResult`. + allow_input_required: When ``False`` (default), an `InputRequiredResult` + from the server raises `RuntimeError`; when ``True``, it is returned + so the caller can resolve the requests and retry. + allow_claimed: When `False` (default), a claimed extension result raises + `UnexpectedClaimedResult`; when `True`, the parsed claim model is returned. + + Raises: + RuntimeError: If the server returns an `InputRequiredResult` and + ``allow_input_required`` is ``False``. + UnexpectedClaimedResult: Claimed result with `allow_claimed` False; carries the parsed value. + """ + result = await self.send_request( + types.CallToolRequest( + params=types.CallToolRequestParams( + name=name, + arguments=arguments, + input_responses=input_responses, + request_state=request_state, + _meta=meta, + ), + ), + self._call_tool_adapter, + request_read_timeout_seconds=read_timeout_seconds, + progress_callback=progress_callback, + ) + + if isinstance(result, types.CallToolResult) and not result.is_error: + await self.validate_tool_result(name, result) + + # The input_required arm stays first; a claimed shape is terminal for the multi-round-trip driver. + if isinstance(result, types.InputRequiredResult) and not allow_input_required: + raise _input_required_unexpected("call_tool") + if not isinstance(result, types.CallToolResult | types.InputRequiredResult) and not allow_claimed: + raise UnexpectedClaimedResult(result) + return result + + def _resolve_param_headers(self, name: str, arguments: Mapping[str, Any]) -> dict[str, str]: + """`Mcp-Param-*` headers for a `tools/call`, or empty when the tool was never listed.""" + header_map = self._x_mcp_header_maps.get(name) + if header_map is None: + return {} + return mcp_param_headers(header_map, arguments) + + async def validate_tool_result(self, name: str, result: types.CallToolResult) -> None: + """Revalidate a `CallToolResult` against the tool's declared output schema. + + Raises: + RuntimeError: Structured content is missing or does not conform to the schema, or the + schema is invalid or has a `$ref` that does not resolve within the schema document. + """ + if name not in self._tool_output_schemas: + # refresh output schema cache + await self.list_tools() + + output_schema = None + if name in self._tool_output_schemas: + output_schema = self._tool_output_schemas.get(name) + else: + logger.warning(f"Tool {name} not listed by server, cannot validate any structured content") + + if output_schema is not None: + from jsonschema import exceptions as jsonschema_exceptions + from referencing.exceptions import Unresolvable + + if result.structured_content is None: + raise RuntimeError(f"Tool {name} has an output schema but did not return structured content") + validator = self._output_schema_validator(name, output_schema) + # `best_match` picks the same error the previous `jsonschema.validate()` call raised, + # so the message a caller sees is unchanged. It is untyped upstream. + errors = validator.iter_errors(result.structured_content) + try: + error = cast( + "Exception | None", + jsonschema_exceptions.best_match(errors), # pyright: ignore[reportUnknownMemberType] + ) + except Unresolvable as e: + # A `$ref` did not resolve within the schema document. + raise RuntimeError(f"Invalid schema for tool {name}: {e}") from e + if error is not None: + raise RuntimeError(f"Invalid structured content returned by tool {name}: {error}") from error + + def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) -> Validator: + """Compiled validator for the tool's cached output schema, built once per schema value. + + Compiling is ~60x the cost of validating, so a one-shot `jsonschema.validate()` per + result dominates `call_tool`; the compiled validator is cached instead. It stays valid + because `_absorb_tool_listing` evicts a tool's validator whenever it absorbs a different + schema for that tool, so a cached entry always matches `output_schema`. + + Raises: + RuntimeError: The schema is not a valid JSON Schema. Raised on every call, since a + failed compile is never cached. + """ + from jsonschema import SchemaError + from jsonschema.validators import validator_for + from referencing import Registry + + if (validator := self._tool_output_validators.get(name)) is not None: + return validator + + validator_cls = validator_for(output_schema) + try: + validator_cls.check_schema(output_schema) + except SchemaError as e: + raise RuntimeError(f"Invalid schema for tool {name}: {e}") + # An explicit empty registry: `$ref`s resolve within the schema document and the bundled metaschemas. + validator = validator_cls(output_schema, registry=Registry()) + self._tool_output_validators[name] = validator + return validator + + async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult: + """Send a prompts/list request. + + Args: + params: Full pagination parameters including cursor and any future fields + """ + return await self.send_request(types.ListPromptsRequest(params=params), types.ListPromptsResult) + + @overload + async def get_prompt( + self, + name: str, + arguments: dict[str, str] | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: Literal[False] = False, + ) -> types.GetPromptResult: ... + + @overload + async def get_prompt( + self, + name: str, + arguments: dict[str, str] | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: bool, + ) -> types.GetPromptResult | types.InputRequiredResult: ... + + async def get_prompt( + self, + name: str, + arguments: dict[str, str] | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: RequestParamsMeta | None = None, + allow_input_required: bool = False, + ) -> types.GetPromptResult | types.InputRequiredResult: + """Send a prompts/get request. + + Args: + input_responses: Responses to a prior `InputRequiredResult.input_requests`. + request_state: Opaque state echoed from a prior `InputRequiredResult`. + allow_input_required: When `False` (default), an `InputRequiredResult` + from the server raises `RuntimeError`; when `True`, it is returned + so the caller can resolve the requests and retry. + + Raises: + RuntimeError: If the server returns an `InputRequiredResult` and + `allow_input_required` is `False`. + """ + result = await self.send_request( + types.GetPromptRequest( + params=types.GetPromptRequestParams( + name=name, + arguments=arguments, + input_responses=input_responses, + request_state=request_state, + _meta=meta, + ), + ), + _GetPromptResultAdapter, + ) + if isinstance(result, types.InputRequiredResult) and not allow_input_required: + raise _input_required_unexpected("get_prompt") + return result + + async def complete( + self, + ref: types.ResourceTemplateReference | types.PromptReference, + argument: dict[str, str], + context_arguments: dict[str, str] | None = None, + ) -> types.CompleteResult: + """Send a completion/complete request.""" + context = None + if context_arguments is not None: + context = types.CompletionContext(arguments=context_arguments) + + return await self.send_request( + types.CompleteRequest( + params=types.CompleteRequestParams( + ref=ref, + argument=types.CompletionArgument(**argument), + context=context, + ), + ), + types.CompleteResult, + ) + + async def list_tools(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListToolsResult: + """Send a tools/list request. + + Args: + params: Full pagination parameters including cursor and any future fields + """ + result = await self.send_request( + types.ListToolsRequest(params=params), + types.ListToolsResult, + ) + complete = (params is None or params.cursor is None) and result.next_cursor is None + return self._absorb_tool_listing(result, complete=complete) + + def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool) -> types.ListToolsResult: + """Filter the listing per the 2026 x-mcp-header MUST and rebuild derived per-tool state, in place. + + Idempotent: cached values are already post-filter, so the response cache can re-absorb a served listing. + `complete` (an uncursored single-page listing) prunes per-tool state down to the listing's tools. + """ + if self._negotiated_version in MODERN_PROTOCOL_VERSIONS: + # 2026-07-28: clients MUST drop tools whose x-mcp-header annotations are invalid. + kept: list[types.Tool] = [] + for tool in result.tools: + if (reason := find_invalid_x_mcp_header(tool.input_schema)) is not None: + logger.warning("dropping tool %r: invalid x-mcp-header (%s)", tool.name, reason) + # Evict any map cached from a prior valid listing so a stale entry can't + # mirror headers for a tool this listing dropped. + self._x_mcp_header_maps.pop(tool.name, None) + continue + # Cache the arg→header map so a later tools/call mirrors it into Mcp-Param-* headers. + self._x_mcp_header_maps[tool.name] = x_mcp_header_map(tool.input_schema) + kept.append(tool) + result.tools = kept + + # Cache tool output schemas for future validation; cursor pages only ever add. A + # changed schema evicts its compiled validator; an unchanged one (a re-listing, or the + # response cache re-absorbing a served hit) keeps it. Only validated tools pay the check. + for tool in result.tools: + if tool.name in self._tool_output_validators and not _same_schema( + self._tool_output_schemas.get(tool.name), tool.output_schema + ): + del self._tool_output_validators[tool.name] + self._tool_output_schemas[tool.name] = tool.output_schema + + if complete: + # The listing is the full tool universe, so state for unlisted tools is stale + # (the server dropped them, or a shared-cache writer's filter did). + names = {tool.name for tool in result.tools} + self._x_mcp_header_maps = {k: v for k, v in self._x_mcp_header_maps.items() if k in names} + self._tool_output_schemas = {k: v for k, v in self._tool_output_schemas.items() if k in names} + self._tool_output_validators = {k: v for k, v in self._tool_output_validators.items() if k in names} + + return result + + @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def send_roots_list_changed(self) -> None: + """Send a roots/list_changed notification.""" + await self.send_notification(types.RootsListChangedNotification()) + + async def _on_request( + self, dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None + ) -> dict[str, Any]: + """Answer a server-initiated request via the registered callbacks.""" + # Literal, not LATEST_PROTOCOL_VERSION: the fallback covers the initialize + # handshake (which only exists at <=2025) and stateless until the header + # is plumbed; its meaning is fixed regardless of LATEST bumps. + version = self._negotiated_version or "2025-11-25" + try: + request = cast(types.ServerRequest, _methods.parse_server_request(method, version, params)) + except KeyError: + raise MCPError(code=METHOD_NOT_FOUND, message="Method not found", data=method) from None + + response: types.ClientResult | types.ErrorData + if isinstance(request, types.PingRequest): + # Answered without a context: ping has no callback that would need one. + response = types.EmptyResult() + else: + assert dctx.request_id is not None # the callback-driving dispatchers always assign ids + ctx = ClientRequestContext( + session=self, request_id=dctx.request_id, meta=request.params.meta if request.params else None + ) + response = await self.dispatch_input_request(ctx, request) + client_response = ClientResponse.validate_python(response) + if isinstance(client_response, types.ErrorData): + raise MCPError.from_error_data(client_response) + dumped = client_response.model_dump(by_alias=True, mode="json", exclude_none=True) + try: + _methods.validate_client_result(method, version, dumped) + except ValidationError: + logger.exception("client callback for %r returned an invalid result", method) + raise MCPError(code=INTERNAL_ERROR, message="Client callback returned an invalid result") from None + return dumped + + async def dispatch_input_request( + self, ctx: ClientRequestContext, request: types.InputRequest + ) -> types.InputResponse | types.ErrorData: + """Route an input request through the client's callback table. + + Shared by the legacy server→client RPC path (`_on_request`) and the + 2026-07-28 multi-round-trip driver, which dispatches the embedded + `InputRequiredResult.input_requests` through the same callbacks. + + Returns the callback's `InputResponse`, or `ErrorData` when the callback declines. + """ + match request: + case types.CreateMessageRequest(params=p): + return await self._sampling_callback(ctx, p) + case types.ElicitRequest(params=p): + return await self._elicitation_callback(ctx, p) + case types.ListRootsRequest(): # pragma: no branch + return await self._list_roots_callback(ctx) + + def _register_listen_route(self, request_id: RequestId) -> ListenRoute: + """Create the demux route for a listen request id; the caller registers BEFORE sending.""" + route = ListenRoute() + self._listen_routes[request_id] = route + return route + + def _unregister_listen_route(self, request_id: RequestId) -> None: + """Drop a listen route; the handle owns membership, so a missing key is a no-op.""" + self._listen_routes.pop(request_id, None) + + def _settle_listen_routes_closed(self) -> None: + """Settle all open listen routes as lost on session exit; cancelled driver tasks cannot.""" + closed = MCPError(code=CONNECTION_CLOSED, message="Connection closed") + for route in self._listen_routes.values(): + route.settle("lost", error=closed) + self._listen_routes.clear() + + def _intercept_notification(self, method: str, params: Mapping[str, Any] | None) -> bool: + """Wire-order listen demux, run synchronously on the dispatcher's receive path. + + Bookkeeping must advance in receive order with the listen result (resolved on + this same path); the spawned `_on_notify` path would race it and drop events. + Returns True to consume the frame: a live route's ack is driver state, never surfaced. + """ + if not self._listen_routes: + return False + if method == "notifications/cancelled": + request_id = cancelled_request_id_from_params(params) + if request_id is not None and (listen_route := self._listen_routes.get(request_id)) is not None: + # a server-sent cancel naming a listen request is that stream's teardown signal + listen_route.settle("lost") + return False # _on_notify swallows every cancelled either way (v1 parity) + if params is None: + return False + meta = params.get("_meta") + if not isinstance(meta, Mapping): + return False + # as_request_id is not a tripwire: raw wire _meta can carry a non-id (even unhashable) value + subscription_id = as_request_id(cast("Mapping[str, Any]", meta).get(SUBSCRIPTION_ID_META_KEY)) + if subscription_id is None or (listen_route := self._listen_routes.get(subscription_id)) is None: + return False + if method == "notifications/subscriptions/acknowledged": + raw_filter = params.get("notifications") + if raw_filter is None: + # malformed, not an empty filter: leave it to the spawned path's validation warning + return False + try: + honored = types.SubscriptionFilter.model_validate(raw_filter) + except ValidationError: + return False + listen_route.set_acked(honored) + return True + if (event := event_from_wire(method, params)) is not None: + listen_route.deliver(event) + return False # events (and any other stamped frame) still tee as usual + + async def _on_notify( + self, dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None + ) -> None: + """Route a server notification: validate, run the typed callback, tee to message_handler.""" + # Same fallback as `_on_request`: covers pre-handshake and stateless. + version = self._negotiated_version or "2025-11-25" + try: + notification = cast(types.ServerNotification, _methods.parse_server_notification(method, version, params)) + except KeyError: + # Only methods unknown to the negotiated version's core tables reach the bindings. + binding = self._notification_bindings.get(method) + if binding is None: + logger.debug("dropped %r: not defined at %s", method, version) + return + try: + bound_params = binding.params_type.model_validate(params or {}) + except ValidationError: + logger.warning("Failed to validate notification: %s", method, exc_info=True) + return + send, receive = self._binding_queues[method] + try: + # Must not await: DirectDispatcher calls _on_notify inline; blocking deadlocks in-process servers. + send.send_nowait(bound_params) + except anyio.WouldBlock: + # Evict the oldest event; no checkpoint since the failed send, + # so the buffer is still full and the retry cannot block. + receive.receive_nowait() + logger.warning("notification queue for %r is full; dropped the oldest event", method) + send.send_nowait(bound_params) + return + except ValidationError: + logger.warning("Failed to validate notification: %s", method, exc_info=True) + return + if isinstance(notification, types.CancelledNotification): + # Never surfaced (v1 parity): the dispatcher already applied it; listen cancels settled by the intercept. + return + try: + if isinstance(notification, types.LoggingMessageNotification): + await self._logging_callback(notification.params) + await self._message_handler(notification) + except Exception: + # Contain here, not in the dispatcher: DirectDispatcher awaits this + # handler inline in the peer's notify() call, so a raising callback + # would otherwise fail the peer's send. A raising logging_callback + # skips the message_handler tee for that notification (v1 parity). + logger.exception("notification callback for %r raised", method) + + async def _on_stream_exception(self, exc: Exception) -> None: + """Deliver a transport-level fault to message_handler via a spawned task. + + Running the handler inline would park the dispatcher's read loop and + deadlock handlers that await session I/O. + """ + assert self._task_group is not None + self._task_group.start_soon(self._deliver_stream_exception, exc) + + async def _deliver_stream_exception(self, exc: Exception) -> None: + try: + await self._message_handler(exc) + except Exception: + logger.exception("message_handler raised on transport exception") diff --git a/src/mcp-client/mcp_client/client/session_group.py b/src/mcp-client/mcp_client/client/session_group.py new file mode 100644 index 0000000000..74e71740e9 --- /dev/null +++ b/src/mcp-client/mcp_client/client/session_group.py @@ -0,0 +1,453 @@ +"""SessionGroup concurrently manages multiple MCP session connections. + +Tools, resources, and prompts are aggregated across servers. Servers may +be connected to or disconnected from at any point after initialization. + +This abstraction can handle naming collisions using a custom user-provided hook. +""" + +import contextlib +import logging +from collections.abc import Callable +from dataclasses import dataclass +from types import TracebackType +from typing import Any, Literal, TypeAlias, overload + +import anyio +import httpx2 +import mcp_types as types +from pydantic import BaseModel, Field +from typing_extensions import Self + +from mcp_client.client.session import ( + ClientSession, + ElicitationFnT, + ListRootsFnT, + LoggingFnT, + MessageHandlerFnT, + SamplingFnT, +) +from mcp_client.client.sse import sse_client +from mcp_client.client.stdio import StdioServerParameters, stdio_client +from mcp_client.client.streamable_http import streamable_http_client +from mcp_client.shared._httpx_utils import create_mcp_http_client +from mcp_client.shared.dispatcher import ProgressFnT +from mcp_client.shared.exceptions import MCPError + + +class SseServerParameters(BaseModel): + """Parameters for initializing an sse_client.""" + + # The endpoint URL. + url: str + + # Optional headers to include in requests. + headers: dict[str, Any] | None = None + + # HTTP timeout for regular operations (in seconds). + timeout: float = 5.0 + + # Timeout for SSE read operations (in seconds). + sse_read_timeout: float = 300.0 + + +class StreamableHttpParameters(BaseModel): + """Parameters for initializing a streamable_http_client.""" + + # The endpoint URL. + url: str + + # Optional headers to include in requests. + headers: dict[str, Any] | None = None + + # HTTP timeout for regular operations (in seconds). + timeout: float = 30.0 + + # Timeout for SSE read operations (in seconds). + sse_read_timeout: float = 300.0 + + # Close the client session when the transport closes. + terminate_on_close: bool = True + + +ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters + + +# Use dataclass instead of Pydantic BaseModel +# because Pydantic BaseModel cannot handle Protocol fields. +@dataclass +class ClientSessionParameters: + """Parameters for establishing a client session to an MCP server.""" + + read_timeout_seconds: float | None = None + sampling_callback: SamplingFnT | None = None + elicitation_callback: ElicitationFnT | None = None + list_roots_callback: ListRootsFnT | None = None + logging_callback: LoggingFnT | None = None + message_handler: MessageHandlerFnT | None = None + client_info: types.Implementation | None = None + + +class ClientSessionGroup: + """Client for managing connections to multiple MCP servers. + + This class is responsible for encapsulating management of server connections. + It aggregates tools, resources, and prompts from all connected servers. + + For auxiliary handlers, such as resource subscription, this is delegated to + the client and can be accessed via the session. + + Example: + ```python + name_fn = lambda name, server_info: f"{(server_info.name)}_{name}" + async with ClientSessionGroup(component_name_hook=name_fn) as group: + for server_param in server_params: + await group.connect_to_server(server_param) + ... + ``` + """ + + class _ComponentNames(BaseModel): + """Used for reverse index to find components.""" + + prompts: set[str] = Field(default_factory=set) + resources: set[str] = Field(default_factory=set) + tools: set[str] = Field(default_factory=set) + + # Standard MCP components. + _prompts: dict[str, types.Prompt] + _resources: dict[str, types.Resource] + _tools: dict[str, types.Tool] + + # Client-server connection management. + _sessions: dict[ClientSession, _ComponentNames] + _tool_to_session: dict[str, ClientSession] + _exit_stack: contextlib.AsyncExitStack + _session_exit_stacks: dict[ClientSession, contextlib.AsyncExitStack] + + # Optional fn consuming (component_name, server_info) for custom names. + # This is to provide a means to mitigate naming conflicts across servers. + # Example: (tool_name, server_info) => "{result.server_info.name}.{tool_name}" + _ComponentNameHook: TypeAlias = Callable[[str, types.Implementation], str] + _component_name_hook: _ComponentNameHook | None + + def __init__( + self, + exit_stack: contextlib.AsyncExitStack | None = None, + component_name_hook: _ComponentNameHook | None = None, + ) -> None: + """Initializes the MCP client.""" + + self._tools = {} + self._resources = {} + self._prompts = {} + + self._sessions = {} + self._tool_to_session = {} + if exit_stack is None: + self._exit_stack = contextlib.AsyncExitStack() + self._owns_exit_stack = True + else: + self._exit_stack = exit_stack + self._owns_exit_stack = False + self._session_exit_stacks = {} + self._component_name_hook = component_name_hook + + async def __aenter__(self) -> Self: # pragma: no cover + # Enter the exit stack only if we created it ourselves + if self._owns_exit_stack: + await self._exit_stack.__aenter__() + return self + + async def __aexit__( + self, + _exc_type: type[BaseException] | None, + _exc_val: BaseException | None, + _exc_tb: TracebackType | None, + ) -> bool | None: # pragma: no cover + """Closes session exit stacks and main exit stack upon completion.""" + + # Only close the main exit stack if we created it + if self._owns_exit_stack: + await self._exit_stack.aclose() + + # Concurrently close session stacks. + async with anyio.create_task_group() as tg: + for exit_stack in self._session_exit_stacks.values(): + tg.start_soon(exit_stack.aclose) + + @property + def sessions(self) -> list[ClientSession]: + """Returns the list of sessions being managed.""" + return list(self._sessions.keys()) # pragma: no cover + + @property + def prompts(self) -> dict[str, types.Prompt]: + """Returns the prompts as a dictionary of names to prompts.""" + return self._prompts + + @property + def resources(self) -> dict[str, types.Resource]: + """Returns the resources as a dictionary of names to resources.""" + return self._resources + + @property + def tools(self) -> dict[str, types.Tool]: + """Returns the tools as a dictionary of names to tools.""" + return self._tools + + @overload + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: Literal[False] = False, + ) -> types.CallToolResult: ... + + @overload + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: bool, + ) -> types.CallToolResult | types.InputRequiredResult: ... + + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + read_timeout_seconds: float | None = None, + progress_callback: ProgressFnT | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: bool = False, + ) -> types.CallToolResult | types.InputRequiredResult: + """Executes a tool given its name and arguments. + + Raises: + RuntimeError: If the server returns an `InputRequiredResult` and + ``allow_input_required`` is ``False``. + """ + session = self._tool_to_session[name] + session_tool_name = self.tools[name].name + return await session.call_tool( + session_tool_name, + arguments=arguments, + read_timeout_seconds=read_timeout_seconds, + progress_callback=progress_callback, + input_responses=input_responses, + request_state=request_state, + meta=meta, + allow_input_required=allow_input_required, + ) + + async def disconnect_from_server(self, session: ClientSession) -> None: + """Disconnects from a single MCP server.""" + + session_known_for_components = session in self._sessions + session_known_for_stack = session in self._session_exit_stacks + + if not session_known_for_components and not session_known_for_stack: + raise MCPError( + code=types.INVALID_PARAMS, + message="Provided session is not managed or already disconnected.", + ) + + if session_known_for_components: # pragma: no branch + component_names = self._sessions.pop(session) # Pop from _sessions tracking + + # Remove prompts associated with the session. + for name in component_names.prompts: + if name in self._prompts: # pragma: no branch + del self._prompts[name] + # Remove resources associated with the session. + for name in component_names.resources: + if name in self._resources: # pragma: no branch + del self._resources[name] + # Remove tools associated with the session. + for name in component_names.tools: + if name in self._tools: # pragma: no branch + del self._tools[name] + if name in self._tool_to_session: # pragma: no branch + del self._tool_to_session[name] + + # Clean up the session's resources via its dedicated exit stack + if session_known_for_stack: + session_stack_to_close = self._session_exit_stacks.pop(session) # pragma: no cover + await session_stack_to_close.aclose() # pragma: no cover + + async def connect_with_session(self, server_info: types.Implementation, session: ClientSession) -> ClientSession: + """Connects to a single MCP server.""" + await self._aggregate_components(server_info, session) + return session + + async def connect_to_server( + self, + server_params: ServerParameters, + session_params: ClientSessionParameters | None = None, + ) -> ClientSession: + """Connects to a single MCP server.""" + server_info, session = await self._establish_session(server_params, session_params or ClientSessionParameters()) + return await self.connect_with_session(server_info, session) + + async def _establish_session( + self, + server_params: ServerParameters, + session_params: ClientSessionParameters, + ) -> tuple[types.Implementation, ClientSession]: + """Establish a client session to an MCP server.""" + + session_stack = contextlib.AsyncExitStack() + try: + # Create read and write streams that facilitate io with the server. + if isinstance(server_params, StdioServerParameters): + client = stdio_client(server_params) + read, write = await session_stack.enter_async_context(client) + elif isinstance(server_params, SseServerParameters): + client = sse_client( + url=server_params.url, + headers=server_params.headers, + timeout=server_params.timeout, + sse_read_timeout=server_params.sse_read_timeout, + ) + read, write = await session_stack.enter_async_context(client) + else: + httpx_client = create_mcp_http_client( + headers=server_params.headers, + timeout=httpx2.Timeout( + server_params.timeout, + read=server_params.sse_read_timeout, + ), + ) + await session_stack.enter_async_context(httpx_client) + + client = streamable_http_client( + url=server_params.url, + http_client=httpx_client, + terminate_on_close=server_params.terminate_on_close, + ) + read, write = await session_stack.enter_async_context(client) + + session = await session_stack.enter_async_context( + ClientSession( + read, + write, + read_timeout_seconds=session_params.read_timeout_seconds, + sampling_callback=session_params.sampling_callback, + elicitation_callback=session_params.elicitation_callback, + list_roots_callback=session_params.list_roots_callback, + logging_callback=session_params.logging_callback, + message_handler=session_params.message_handler, + client_info=session_params.client_info, + ) + ) + + result = await session.initialize() + + # Session successfully initialized. + # Store its stack and register the stack with the main group stack. + self._session_exit_stacks[session] = session_stack + # session_stack itself becomes a resource managed by the + # main _exit_stack. + await self._exit_stack.enter_async_context(session_stack) + + return result.server_info, session + except Exception: # pragma: no cover + # If anything during this setup fails, ensure the session-specific + # stack is closed. + await session_stack.aclose() + raise + + async def _aggregate_components(self, server_info: types.Implementation, session: ClientSession) -> None: + """Aggregates prompts, resources, and tools from a given session.""" + + # Create a reverse index so we can find all prompts, resources, and + # tools belonging to this session. Used for removing components from + # the session group via self.disconnect_from_server. + component_names = self._ComponentNames() + + # Temporary components dicts. We do not want to modify the aggregate + # lists in case of an intermediate failure. + prompts_temp: dict[str, types.Prompt] = {} + resources_temp: dict[str, types.Resource] = {} + tools_temp: dict[str, types.Tool] = {} + tool_to_session_temp: dict[str, ClientSession] = {} + + # Query the server for its prompts and aggregate to list. + try: + prompts = (await session.list_prompts()).prompts + for prompt in prompts: + name = self._component_name(prompt.name, server_info) + prompts_temp[name] = prompt + component_names.prompts.add(name) + except MCPError as err: # pragma: no cover + logging.warning(f"Could not fetch prompts: {err}") + + # Query the server for its resources and aggregate to list. + try: + resources = (await session.list_resources()).resources + for resource in resources: + name = self._component_name(resource.name, server_info) + resources_temp[name] = resource + component_names.resources.add(name) + except MCPError as err: # pragma: no cover + logging.warning(f"Could not fetch resources: {err}") + + # Query the server for its tools and aggregate to list. + try: + tools = (await session.list_tools()).tools + for tool in tools: + name = self._component_name(tool.name, server_info) + tools_temp[name] = tool + tool_to_session_temp[name] = session + component_names.tools.add(name) + except MCPError as err: # pragma: no cover + logging.warning(f"Could not fetch tools: {err}") + + # Clean up exit stack for session if we couldn't retrieve anything + # from the server. + if not any((prompts_temp, resources_temp, tools_temp)): + del self._session_exit_stacks[session] # pragma: no cover + + # Check for duplicates. + matching_prompts = prompts_temp.keys() & self._prompts.keys() + if matching_prompts: + raise MCPError( # pragma: no cover + code=types.INVALID_PARAMS, + message=f"{matching_prompts} already exist in group prompts.", + ) + matching_resources = resources_temp.keys() & self._resources.keys() + if matching_resources: + raise MCPError( # pragma: no cover + code=types.INVALID_PARAMS, + message=f"{matching_resources} already exist in group resources.", + ) + matching_tools = tools_temp.keys() & self._tools.keys() + if matching_tools: + raise MCPError(code=types.INVALID_PARAMS, message=f"{matching_tools} already exist in group tools.") + + # Aggregate components. + self._sessions[session] = component_names + self._prompts.update(prompts_temp) + self._resources.update(resources_temp) + self._tools.update(tools_temp) + self._tool_to_session.update(tool_to_session_temp) + + def _component_name(self, name: str, server_info: types.Implementation) -> str: + if self._component_name_hook: + return self._component_name_hook(name, server_info) + return name diff --git a/src/mcp-client/mcp_client/client/sse.py b/src/mcp-client/mcp_client/client/sse.py new file mode 100644 index 0000000000..57f492eec3 --- /dev/null +++ b/src/mcp-client/mcp_client/client/sse.py @@ -0,0 +1,171 @@ +import logging +from collections.abc import Callable +from contextlib import asynccontextmanager +from typing import Any +from urllib.parse import parse_qs, urljoin, urlparse + +import anyio +import httpx2 +import mcp_types as types +from anyio.abc import TaskStatus +from httpx2 import SSEError + +from mcp_client.shared._compat import resync_tracer +from mcp_client.shared._context_streams import create_context_streams +from mcp_client.shared._httpx_utils import ( + McpHttpClientFactory, + create_mcp_http_client, + request_within_origin, + sse_within_origin, +) +from mcp_client.shared.message import SessionMessage + +logger = logging.getLogger("mcp.client.sse") + + +def remove_request_params(url: str) -> str: + return urljoin(url, urlparse(url).path) + + +def _extract_session_id_from_endpoint(endpoint_url: str) -> str | None: + query_params = parse_qs(urlparse(endpoint_url).query) + return query_params.get("sessionId", [None])[0] or query_params.get("session_id", [None])[0] + + +@asynccontextmanager +async def sse_client( + url: str, + headers: dict[str, Any] | None = None, + timeout: float = 5.0, + sse_read_timeout: float = 300.0, + httpx_client_factory: McpHttpClientFactory = create_mcp_http_client, + auth: httpx2.Auth | None = None, + on_session_created: Callable[[str], None] | None = None, +): + """Client transport for SSE. + + `sse_read_timeout` determines how long (in seconds) the client will wait for a new + event before disconnecting. All other HTTP operations are controlled by `timeout`. + + Args: + url: The SSE endpoint URL. + headers: Optional headers to include in requests. + timeout: HTTP timeout for regular operations (in seconds). + sse_read_timeout: Timeout for SSE read operations (in seconds). + httpx_client_factory: Factory function for creating the httpx2 client. Whichever client it + returns, MCP requests follow a redirect only when it stays on the endpoint's origin + (same scheme, host and port, or http to https on the same host with default ports) and + keeps the request method (any status for the SSE GET, 307/308 for a message POST); any + other redirect is not followed, so connecting fails with + `httpx2.HTTPStatusError` for the redirect response. The client's `follow_redirects` + setting is not consulted; the SDK's OAuth providers apply the same rule to the requests + they make. + auth: Optional httpx2 authentication handler. + on_session_created: Optional callback invoked with the session ID when received. + """ + logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}") + async with httpx_client_factory( + headers=headers, auth=auth, timeout=httpx2.Timeout(timeout, read=sse_read_timeout) + ) as client: + async with sse_within_origin(client, url) as event_source: + event_source.response.raise_for_status() + logger.debug("SSE connection established") + + read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0) + write_stream, write_stream_reader = create_context_streams[SessionMessage](0) + + async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED): + try: + async for sse in event_source: # pragma: no branch + logger.debug(f"Received SSE event: {sse.event}") + match sse.event: + case "endpoint": + endpoint_url = urljoin(url, sse.data) + logger.debug(f"Received endpoint URL: {endpoint_url}") + + url_parsed = urlparse(url) + endpoint_parsed = urlparse(endpoint_url) + if ( # pragma: no cover + url_parsed.netloc != endpoint_parsed.netloc + or url_parsed.scheme != endpoint_parsed.scheme + ): + error_msg = ( # pragma: no cover + f"Endpoint origin does not match connection origin: {endpoint_url}" + ) + logger.error(error_msg) # pragma: no cover + raise ValueError(error_msg) # pragma: no cover + + if on_session_created: + session_id = _extract_session_id_from_endpoint(endpoint_url) + if session_id: + on_session_created(session_id) + + task_status.started(endpoint_url) + + case "message": + # Skip empty data (keep-alive pings) + if not sse.data: + continue + try: + message = types.jsonrpc_message_adapter.validate_json(sse.data, by_name=False) + logger.debug(f"Received server message: {message}") + except Exception as exc: # pragma: no cover + logger.exception("Error parsing server message") # pragma: no cover + await read_stream_writer.send(exc) # pragma: no cover + continue # pragma: no cover + + session_message = SessionMessage(message) + await read_stream_writer.send(session_message) + case _: # pragma: no cover + logger.warning(f"Unknown SSE event: {sse.event}") # pragma: no cover + except SSEError as sse_exc: # pragma: lax no cover + logger.exception("Encountered SSE exception") + raise sse_exc + except Exception as exc: # pragma: lax no cover + logger.exception("Error in sse_reader") + await read_stream_writer.send(exc) + finally: + await read_stream_writer.aclose() + + async def post_writer(endpoint_url: str): + try: + async with write_stream_reader, write_stream: + + async def _send_message(session_message: SessionMessage) -> None: + logger.debug(f"Sending client message: {session_message}") + response = await request_within_origin( + client, + "POST", + endpoint_url, + json=session_message.message.model_dump(by_alias=True, mode="json", exclude_unset=True), + ) + response.raise_for_status() + logger.debug(f"Client message sent successfully: {response.status_code}") + + async for session_message in write_stream_reader: + sender_ctx = write_stream_reader.last_context + if sender_ctx is not None: + async with anyio.create_task_group() as tg: + sender_ctx.run(tg.start_soon, _send_message, session_message) + else: + await _send_message(session_message) # pragma: no cover + except Exception: # pragma: lax no cover + logger.exception("Error in post_writer") + + # On Python 3.14, coverage.py reports a phantom branch arc on this + # line (->yield) when nested two async-with levels deep. The branch + # is the unreachable "did __aexit__ suppress?" arm for memory streams. + async with ( # pragma: no branch + read_stream_writer, + read_stream, + write_stream, + write_stream_reader, + anyio.create_task_group() as tg, + ): + endpoint_url = await tg.start(sse_reader) + logger.debug(f"Starting post writer with endpoint URL: {endpoint_url}") + tg.start_soon(post_writer, endpoint_url) + + yield read_stream, write_stream + tg.cancel_scope.cancel() + await resync_tracer() diff --git a/src/mcp-client/mcp_client/client/stdio.py b/src/mcp-client/mcp_client/client/stdio.py new file mode 100644 index 0000000000..755885118a --- /dev/null +++ b/src/mcp-client/mcp_client/client/stdio.py @@ -0,0 +1,355 @@ +"""stdio client transport. + +Runs an MCP server as a subprocess and exchanges newline-delimited JSON-RPC +messages with it over stdin/stdout. Two pipe tasks bridge the server's pipes +to the session's in-memory streams; shutdown follows the MCP spec sequence +(close stdin, wait, then kill the process tree) inside a cancellation shield +with every wait bounded, so a cancelled caller can neither leak a live server +process nor hang on one. +""" + +import logging +import os +import sys +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager, suppress +from pathlib import Path +from typing import Literal, TextIO + +import anyio +import anyio.lowlevel +import anyio.to_thread +import mcp_types as types +from anyio.abc import AsyncResource, Process +from anyio.streams.text import TextReceiveStream +from pydantic import BaseModel, Field + +from mcp_client.client._transport import TransportStreams +from mcp_client.os.posix.utilities import terminate_posix_process_tree +from mcp_client.os.win32.utilities import ( + ServerProcess, + close_process_job, + create_windows_process, + get_windows_executable_command, + terminate_windows_process_tree, +) +from mcp_client.shared.message import SessionMessage + +logger = logging.getLogger("mcp.client.stdio") + +# Environment variables to inherit by default +DEFAULT_INHERITED_ENV_VARS = ( + [ + "APPDATA", + "HOMEDRIVE", + "HOMEPATH", + "LOCALAPPDATA", + "PATH", + "PATHEXT", + "PROCESSOR_ARCHITECTURE", + "SYSTEMDRIVE", + "SYSTEMROOT", + "TEMP", + "USERNAME", + "USERPROFILE", + ] + if sys.platform == "win32" + else ["HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER"] +) + +# Grace period for the server to exit on its own after its stdin closes. +PROCESS_TERMINATION_TIMEOUT = 2.0 + +# Extra time after SIGTERM before SIGKILL; POSIX only (Windows kills hard). +FORCE_KILL_TIMEOUT = 2.0 + +# Time for the event loop to observe a kill; only an unkillable process runs this out. +_KILL_REAP_TIMEOUT = 2.0 + +# Time for the writer to flush accepted messages before stdin closes. +_WRITER_FLUSH_TIMEOUT = 0.5 + +# How often to poll returncode while waiting for the process to die. +_EXIT_POLL_INTERVAL = 0.01 + + +def get_default_environment() -> dict[str, str]: + """Returns only the environment variables that are safe to inherit.""" + env: dict[str, str] = {} + + for key in DEFAULT_INHERITED_ENV_VARS: + value = os.environ.get(key) + if value is None: # pragma: lax no cover + continue + + if value.startswith("()"): # pragma: no cover + # Skip functions, which are a security risk + continue # pragma: no cover + + env[key] = value + + return env + + +class StdioServerParameters(BaseModel): + command: str + """The executable to run to start the server.""" + + args: list[str] = Field(default_factory=list) + """Command line arguments to pass to the executable.""" + + env: dict[str, str] | None = None + """Extra environment variables, merged over get_default_environment().""" + + cwd: str | Path | None = None + """The working directory to use when spawning the process.""" + + encoding: str = "utf-8" + """Text encoding for messages to and from the server.""" + + encoding_error_handler: Literal["strict", "ignore", "replace"] = "strict" + """Encoding error handler; see https://docs.python.org/3/library/codecs.html#error-handlers.""" + + +@asynccontextmanager +async def stdio_client( + server: StdioServerParameters, errlog: TextIO = sys.stderr +) -> AsyncGenerator[TransportStreams, None]: + """Spawns an MCP server subprocess and connects to it over stdin/stdout. + + Raises: + OSError: If the server process cannot be spawned. + ValueError: If the spawn parameters are invalid (embedded NUL bytes). + """ + command = await _get_executable_command(server.command) + + process = await _create_platform_compatible_process( + command=command, + args=server.args, + env=get_default_environment() | (server.env or {}), + errlog=errlog, + cwd=server.cwd, + ) + + # The spawn succeeded; no awaits until the task group is entered, or a + # cancellation delivered in the gap would leak the live process. + read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0) + write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0) + + shutting_down = False + writer_done = anyio.Event() + + async def stdout_reader() -> None: + assert process.stdout, "Opened process is missing stdout" + + stdout = TextReceiveStream(process.stdout, encoding=server.encoding, errors=server.encoding_error_handler) + try: + async with read_stream_writer: + try: + # One line at a time; no read-ahead while a delivery is blocked. + buffer = "" + async for chunk in stdout: + lines = (buffer + chunk).split("\n") + buffer = lines.pop() + for line in lines: + try: + await read_stream_writer.send(_parse_line(line)) + except (anyio.ClosedResourceError, anyio.BrokenResourceError): + return # the session is gone; only the drain below remains + finally: + await _drain_stdout(process) + except anyio.ClosedResourceError: + pass # our own shutdown closed the stdout stream under the read + except (anyio.BrokenResourceError, ConnectionError): + # Teardown noise during shutdown, a real failure otherwise; either way + # the session sees clean closure when the read stream closes. + if not shutting_down: + logger.exception("Reading from the MCP server's stdout failed mid-session") + + async def stdin_writer() -> None: + assert process.stdin, "Opened process is missing stdin" + + try: + async with write_stream_reader: + async for session_message in write_stream_reader: + json = session_message.message.model_dump_json(by_alias=True, exclude_unset=True) + data = (json + "\n").encode(encoding=server.encoding, errors=server.encoding_error_handler) + await process.stdin.send(data) + except (anyio.ClosedResourceError, anyio.BrokenResourceError, OSError): + # The server may still be alive: close the read stream so the session + # sees the connection end instead of a request hanging forever. + await read_stream_writer.aclose() + finally: + writer_done.set() + + async def shutdown() -> None: + """Winds the transport down: stop traffic, flush, stop the server, release the streams.""" + # Unblock the reader into its drain: a server stuck writing stdout cannot + # read its stdin, so draining is what lets the flush below complete. + read_stream.close() + # Bounded window for the writer to flush already-accepted messages. + write_stream.close() + with anyio.move_on_after(_WRITER_FLUSH_TIMEOUT) as flush_scope: + await writer_done.wait() + if flush_scope.cancelled_caught: + await anyio.lowlevel.cancel_shielded_checkpoint() # resync coverage on 3.11 (gh-106749) + await _stop_server_process(process) + await _aclose_all(read_stream, write_stream, read_stream_writer, write_stream_reader) + # One pass so unblocked tasks exit via their except paths before the cancel. + await anyio.lowlevel.checkpoint() + + async with anyio.create_task_group() as tg: + tg.start_soon(stdout_reader) + tg.start_soon(stdin_writer) + try: + yield read_stream, write_stream + finally: + shutting_down = True + # Shutdown must finish even under caller cancellation, or the server + # process would leak; every wait inside is bounded. (Native + # task.cancel() and the fallback's worker threads can still defeat it.) + with anyio.CancelScope(shield=True): + await shutdown() + # Unstick pipe tasks a kill survivor's open pipe end could still block. + tg.cancel_scope.cancel() + # The cancel lands via throw(); one yield resyncs 3.11 coverage (gh-106749). + await anyio.lowlevel.cancel_shielded_checkpoint() + + +def _parse_line(line: str) -> SessionMessage | Exception: + """Parses one stdout line, returning parse errors as values for the session to surface.""" + try: + message = types.jsonrpc_message_adapter.validate_json(line, by_name=False) + except ValueError as exc: + logger.exception("Failed to parse JSONRPC message from server") + return exc + return SessionMessage(message) + + +async def _drain_stdout(process: ServerProcess) -> None: + """Consumes and discards the server's remaining stdout. + + Keeps a server flushing buffered output from blocking on a full pipe and + missing its chance to exit; shielded, raw bytes, ends when shutdown closes + the pipe. + """ + assert process.stdout + with anyio.CancelScope(shield=True): + with suppress( + anyio.EndOfStream, + anyio.ClosedResourceError, + anyio.BrokenResourceError, + ConnectionError, + OSError, + ): + while True: + await process.stdout.receive() + + +async def _stop_server_process(process: ServerProcess) -> None: + """Closes stdin, waits out the grace period, then kills the whole tree. + + The escalation order is spec text; timeouts and tree-wide scope are SDK policy: + https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle#shutdown + """ + assert process.stdin and process.stdout, "server process is spawned with pipes" + + await _close_pipe(process.stdin) + if not await _wait_for_process_exit(process, PROCESS_TERMINATION_TIMEOUT): + await _terminate_process_tree(process) + # Until the event loop observes the death, the transport cannot close. + if not await _wait_for_process_exit(process, _KILL_REAP_TIMEOUT): + logger.warning("MCP server process %d is still alive after the kill escalation; abandoning it", process.pid) + + # Reaps surviving Windows job members now, not at GC; no-op on POSIX. + close_process_job(process) + # A kill survivor can hold the stdout pipe open; poison the reader anyway. + await _close_pipe(process.stdout) + _close_subprocess_transport(process) + + +async def _close_pipe(stream: AsyncResource) -> None: + """Closes a pipe stream, tolerating one already closed, broken, or contended.""" + with suppress(OSError, anyio.BrokenResourceError, anyio.ClosedResourceError): + await stream.aclose() + + +async def _wait_for_process_exit(process: ServerProcess, timeout: float) -> bool: + """Returns whether the process died within the timeout, by polling returncode. + + Not process.wait(): on asyncio 3.11+ it also waits for pipe EOF, and a + child that inherited the pipes makes an exited server look hung. + """ + deadline = anyio.current_time() + timeout + while process.returncode is None: + if anyio.current_time() >= deadline: + return False + await anyio.sleep(_EXIT_POLL_INTERVAL) + return True + + +async def _terminate_process_tree(process: ServerProcess) -> None: + """Kills the process and all its descendants. + + POSIX: SIGTERM to the process group, SIGKILL after FORCE_KILL_TIMEOUT. + Windows: immediate Job Object termination (already a hard kill). + """ + if sys.platform == "win32": # pragma: no cover + await terminate_windows_process_tree(process) + else: # pragma: lax no cover + # The Windows-only FallbackProcess never reaches the POSIX path. + assert isinstance(process, Process) + await terminate_posix_process_tree(process, FORCE_KILL_TIMEOUT) + + +def _close_subprocess_transport(process: ServerProcess) -> None: + """Closes the asyncio subprocess transport, if there is one. + + The transport otherwise stays open (and warns at GC) while a surviving + descendant holds a pipe end; nothing public exposes it, hence the attribute + walk. No-op on trio and the Windows fallback. + """ + transport = getattr(getattr(process, "_process", None), "_transport", None) + # Duck-typed: uvloop's UVProcessTransport is not an asyncio.SubprocessTransport. + close = getattr(transport, "close", None) + if callable(close): + # close() on <=3.12 can raise PermissionError re-killing a setuid child. + with suppress(PermissionError): + close() + + +async def _get_executable_command(command: str) -> str: + """Normalizes the command for the current platform.""" + if sys.platform == "win32": + return await anyio.to_thread.run_sync(get_windows_executable_command, command, abandon_on_cancel=True) + else: # pragma: lax no cover + return command + + +async def _create_platform_compatible_process( + command: str, + args: list[str], + env: dict[str, str] | None = None, + errlog: TextIO = sys.stderr, + cwd: Path | str | None = None, +) -> ServerProcess: + """Spawns the server in its own kill scope. + + A new session/process group on POSIX, a Job Object on Windows. + """ + if sys.platform == "win32": # pragma: no cover + return await create_windows_process(command, args, env, errlog, cwd) + else: # pragma: lax no cover + return await anyio.open_process( + [command, *args], + env=env, + stderr=errlog, + cwd=cwd, + start_new_session=True, + ) + + +async def _aclose_all(*streams: AsyncResource) -> None: + """Closes every given stream.""" + for stream in streams: + await stream.aclose() diff --git a/src/mcp-client/mcp_client/client/streamable_http.py b/src/mcp-client/mcp_client/client/streamable_http.py new file mode 100644 index 0000000000..362de5802e --- /dev/null +++ b/src/mcp-client/mcp_client/client/streamable_http.py @@ -0,0 +1,758 @@ +"""Implements StreamableHTTP transport for MCP clients.""" + +from __future__ import annotations as _annotations + +import contextlib +import logging +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import asynccontextmanager +from dataclasses import dataclass + +import anyio +import httpx2 +from anyio.abc import TaskGroup +from httpx2 import EventSource, ServerSentEvent +from mcp_types import ( + CONNECTION_CLOSED, + INTERNAL_ERROR, + INVALID_REQUEST, + METHOD_NOT_FOUND, + PARSE_ERROR, + ErrorData, + JSONRPCError, + JSONRPCMessage, + JSONRPCNotification, + JSONRPCRequest, + JSONRPCResponse, + RequestId, + jsonrpc_message_adapter, +) +from mcp_types.version import MODERN_PROTOCOL_VERSIONS +from pydantic import ValidationError + +from mcp_client.client._transport import TransportStreams +from mcp_client.shared._compat import resync_tracer +from mcp_client.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams +from mcp_client.shared._httpx_utils import ( + create_mcp_http_client, + redirect_location, + request_within_origin, + sse_within_origin, + stream_within_origin, +) +from mcp_client.shared.inbound import MCP_PROTOCOL_VERSION_HEADER +from mcp_client.shared.jsonrpc_dispatcher import cancelled_request_id_from_params +from mcp_client.shared.message import ClientMessageMetadata, SessionMessage + +logger = logging.getLogger("mcp.client.streamable_http") + + +# TODO(Marcelo): Put the TransportStreams in a module under shared, so we can import here. +SessionMessageOrError = SessionMessage | Exception +StreamWriter = ContextSendStream[SessionMessageOrError] +StreamReader = ContextReceiveStream[SessionMessage] + +MCP_SESSION_ID = "mcp-session-id" +LAST_EVENT_ID = "last-event-id" + +# Reconnection defaults +DEFAULT_RECONNECTION_DELAY_MS = 1000 # 1 second fallback when server doesn't provide retry +MAX_RECONNECTION_ATTEMPTS = 2 # Max retry attempts before giving up + + +class StreamableHTTPError(Exception): + """Base exception for StreamableHTTP transport errors.""" + + +class ResumptionError(StreamableHTTPError): + """Raised when resumption request is invalid.""" + + +def _unfollowed_redirect(response: httpx2.Response) -> str | None: + """Describe a redirect `stream_within_origin` left unfollowed, or None if `response` is not one.""" + location = redirect_location(response) + if location is None: + return None + if response.request.url.scheme == "https" and location.scheme == "http": + return ( + f"Redirect to {location} not followed: it would downgrade this HTTPS endpoint to plain HTTP.\n" + "The server is likely behind a TLS-terminating proxy whose forwarded headers it does not trust,\n" + f"often combined with a trailing-slash difference. Try {location.copy_with(scheme='https')} instead, " + "or fix the proxy settings." + ) + return f"Redirect to {location} not followed; use that URL as the endpoint if it is the intended server" + + +@dataclass +class RequestContext: + """Context for a request operation.""" + + client: httpx2.AsyncClient + session_id: str | None + session_message: SessionMessage + metadata: ClientMessageMetadata | None + read_stream_writer: StreamWriter + + +@dataclass(slots=True) +class _InFlightPost: + """A request POST in flight: its abort scope and the era it was sent under. + + `modern` is the negotiated-version cache as of this request's dequeue, so a + later cancel frame is interpreted under the era the request actually ran + with, not whatever the cache says by then. + """ + + scope: anyio.CancelScope + modern: bool + + +class StreamableHTTPTransport: + """StreamableHTTP client transport implementation.""" + + def __init__(self, url: str) -> None: + """Initialize the StreamableHTTP transport. + + Args: + url: The endpoint URL. + """ + self.url = url + self.session_id: str | None = None + # Captured from each stamped message's metadata, synchronously in the + # post_writer loop so the cache always reflects wire order (a POST task's + # scheduling is arbitrary). Reused on outbound HTTP that carries no + # per-message header (transport-internal GET/DELETE, and dispatcher-written + # response/error POSTs that bypass the session's stamp), and consulted by + # `_consume_modern_cancellation`. Cleared when an `initialize` message is + # dequeued so a probe-stamped value cannot leak onto the handshake. + self._protocol_version_header: str | None = None + # Every request's POST runs inside one of these so an outbound + # `notifications/cancelled` at 2026 can abort it; see + # `_consume_modern_cancellation`. Keys are verbatim-typed ("1" is not 1). + self._in_flight_posts: dict[RequestId, _InFlightPost] = {} + + def _prepare_headers(self) -> dict[str, str]: + """Build MCP-specific request headers for any outbound HTTP request. + + These are merged with the ``httpx2.AsyncClient`` defaults (these take + precedence). The cached ``MCP-Protocol-Version`` is included whenever + present so messages that don't pass through the session's stamp — + response/error POSTs, legacy cancel frames, transport-internal + GET/DELETE — still carry the negotiated version. Per-message headers + are layered on top by the caller. + """ + headers: dict[str, str] = { + "accept": "application/json, text/event-stream", + "content-type": "application/json", + } + if self.session_id: + headers[MCP_SESSION_ID] = self.session_id + if self._protocol_version_header: + headers[MCP_PROTOCOL_VERSION_HEADER] = self._protocol_version_header + return headers + + def _is_initialization_request(self, message: JSONRPCMessage) -> bool: + """Check if the message is an initialization request.""" + return isinstance(message, JSONRPCRequest) and message.method == "initialize" + + def _is_initialized_notification(self, message: JSONRPCMessage) -> bool: + """Check if the message is an initialized notification.""" + return isinstance(message, JSONRPCNotification) and message.method == "notifications/initialized" + + def _maybe_extract_session_id_from_response(self, response: httpx2.Response) -> None: + """Extract and store session ID from response headers.""" + new_session_id = response.headers.get(MCP_SESSION_ID) + if new_session_id: + self.session_id = new_session_id + logger.info(f"Received session ID: {self.session_id}") + + async def _handle_sse_event( + self, + sse: ServerSentEvent, + read_stream_writer: StreamWriter, + original_request_id: RequestId | None = None, + resumption_callback: Callable[[str], Awaitable[None]] | None = None, + ) -> bool: + """Handle an SSE event, returning True if the response is complete.""" + if sse.event == "message": + # Handle priming events (empty data with ID) for resumability + if not sse.data: + # Call resumption callback for priming events that have an ID + if sse.id and resumption_callback: + await resumption_callback(sse.id) + return False + try: + message = jsonrpc_message_adapter.validate_json(sse.data, by_name=False) + logger.debug(f"SSE message: {message}") + + # If this is a response and we have original_request_id, replace it + if original_request_id is not None and isinstance(message, JSONRPCResponse | JSONRPCError): + message.id = original_request_id + + session_message = SessionMessage(message) + await read_stream_writer.send(session_message) + + # Call resumption token callback if we have an ID + if sse.id and resumption_callback: + await resumption_callback(sse.id) + + # If this is a response or error return True indicating completion + # Otherwise, return False to continue listening + return isinstance(message, JSONRPCResponse | JSONRPCError) + + # Forwarding to a closed read stream lands here when the caller cancels mid-SSE + # (BrokenResourceError, not a parse failure); coverage is timing-dependent in the + # streaming story's modern HTTP cancellation leg. + except Exception as exc: # pragma: lax no cover + logger.exception("Error parsing SSE message") + if original_request_id is not None: + error_data = ErrorData(code=PARSE_ERROR, message=f"Failed to parse SSE message: {exc}") + error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=original_request_id, error=error_data)) + await read_stream_writer.send(error_msg) + return True + await read_stream_writer.send(exc) + return False + else: # pragma: no cover + logger.warning(f"Unknown SSE event: {sse.event}") + return False + + async def handle_get_stream(self, client: httpx2.AsyncClient, read_stream_writer: StreamWriter) -> None: + """Handle GET stream for server-initiated messages with auto-reconnect.""" + last_event_id: str | None = None + retry_interval_ms: int | None = None + attempt: int = 0 + + while attempt < MAX_RECONNECTION_ATTEMPTS: # pragma: no branch + try: + if not self.session_id: + return + + headers = self._prepare_headers() + if last_event_id: + headers[LAST_EVENT_ID] = last_event_id + + async with sse_within_origin(client, self.url, headers=headers) as event_source: + if (redirect := _unfollowed_redirect(event_source.response)) is not None: + # The same GET would be redirected again, so retrying cannot help. + logger.warning(f"GET stream not opened: {redirect}") + return + event_source.response.raise_for_status() + logger.debug("GET SSE connection established") + + async for sse in event_source: + # Track last event ID for reconnection + if sse.id: + last_event_id = sse.id + # Track retry interval from server + if sse.retry is not None: + retry_interval_ms = sse.retry + + await self._handle_sse_event(sse, read_stream_writer) + + # Stream ended normally (server closed) - reset attempt counter + attempt = 0 + + except Exception: + logger.debug("GET stream error", exc_info=True) + attempt += 1 + + if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover + logger.debug(f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded") + return + + # Wait before reconnecting + delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS + logger.info(f"GET stream disconnected, reconnecting in {delay_ms}ms...") + await anyio.sleep(delay_ms / 1000.0) + + async def _handle_resumption_request(self, ctx: RequestContext) -> None: + """Handle a resumption request using GET with SSE.""" + headers = self._prepare_headers() + if ctx.metadata and ctx.metadata.resumption_token: + headers[LAST_EVENT_ID] = ctx.metadata.resumption_token + else: + raise ResumptionError("Resumption request requires a resumption token") # pragma: no cover + + # Extract original request ID to map responses + original_request_id = None + if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch + original_request_id = ctx.session_message.message.id + + async with sse_within_origin(ctx.client, self.url, headers=headers) as event_source: + if (redirect := _unfollowed_redirect(event_source.response)) is not None: + logger.warning(redirect) + assert original_request_id is not None + await self._resolve_abandoned_request( + ctx.read_stream_writer, original_request_id, redirect, code=INVALID_REQUEST + ) + return + event_source.response.raise_for_status() + logger.debug("Resumption GET SSE connection established") + + async for sse in event_source: # pragma: no branch + is_complete = await self._handle_sse_event( + sse, + ctx.read_stream_writer, + original_request_id, + ctx.metadata.on_resumption_token_update if ctx.metadata else None, + ) + if is_complete: + await event_source.response.aclose() + break + + def _consume_modern_cancellation(self, session_message: SessionMessage) -> bool: + """Translate an outbound `notifications/cancelled` at 2026; True means "do not POST". + + The 2026 wire defines no client-to-server notifications over streamable + HTTP: closing a request's response stream IS its cancellation signal. + The dispatcher still emits the courtesy frame as its abandon signal + (every outbound cancel names one of our own request ids - the spec + forbids cancelling a request the sender did not issue), so this + transport translates it: when the named request's POST is in flight, + that POST's own recorded era decides - abort-and-swallow at 2026, POST + the frame below it (where the frame is the signal and a disconnect + explicitly is not). With no POST to consult, the cached negotiated + version decides; at 2026 the frame is swallowed even unmatched, so a + late cancel racing the response cannot leak onto the wire. + """ + message = session_message.message + if not (isinstance(message, JSONRPCNotification) and message.method == "notifications/cancelled"): + return False + request_id = cancelled_request_id_from_params(message.params) + post = self._in_flight_posts.get(request_id) if request_id is not None else None + if post is not None: + if not post.modern: + return False + logger.debug("aborting in-flight POST for cancelled request %r", request_id) + post.scope.cancel() + return True + return self._protocol_version_header in MODERN_PROTOCOL_VERSIONS + + async def _run_request_post( + self, + post_fn: Callable[[], Awaitable[None]], + post: _InFlightPost, + request_id: RequestId, + ) -> None: + """Run one request's POST inside its abort scope (see `_consume_modern_cancellation`).""" + try: + with post.scope: + await post_fn() + finally: + # Identity-guarded: a reused id may already have a successor + # registered while this task unwinds - popping by key alone would + # evict the live entry and leave the new POST unabortable. + if self._in_flight_posts.get(request_id) is post: + del self._in_flight_posts[request_id] + + async def _handle_post_request(self, ctx: RequestContext) -> None: + """Handle a POST request with response processing.""" + message = ctx.session_message.message + headers = self._prepare_headers() + if ctx.metadata is not None and ctx.metadata.headers is not None: + headers.update(ctx.metadata.headers) + + async with stream_within_origin( + ctx.client, + "POST", + self.url, + json=message.model_dump(by_alias=True, mode="json", exclude_unset=True), + headers=headers, + ) as response: + if response.status_code == 202: + logger.debug("Received 202 Accepted") + if isinstance(message, JSONRPCRequest): + # A request's response arrives on this POST's body; 202 says + # none will follow. Resolve rather than park the caller forever. + await self._resolve_abandoned_request( + ctx.read_stream_writer, + message.id, + "server answered a request with 202 Accepted", + code=INVALID_REQUEST, + ) + return + + if (redirect := _unfollowed_redirect(response)) is not None: + logger.warning(redirect) + if isinstance(message, JSONRPCRequest): + await self._resolve_abandoned_request( + ctx.read_stream_writer, message.id, redirect, code=INVALID_REQUEST + ) + return + + if response.status_code >= 400: + if isinstance(message, JSONRPCRequest): + # A spec-correct server may return the JSON-RPC error in the + # body at a non-2xx status (e.g. 400 for INVALID_PARAMS, 404 + # for METHOD_NOT_FOUND). Surface that error rather than the + # status-derived stand-in below. + if response.headers.get("content-type", "").lower().startswith("application/json"): + try: + body = await response.aread() + parsed = jsonrpc_message_adapter.validate_json(body, by_name=False) + if isinstance(parsed, JSONRPCError): + # The server may have set `id: null` (request rejected before its + # id was parsed); use this request's id so correlation works. + reply = JSONRPCError(jsonrpc="2.0", id=message.id, error=parsed.error) + await ctx.read_stream_writer.send(SessionMessage(reply)) + return + except (httpx2.StreamError, ValidationError): + pass + logger.debug("Non-2xx body was not a JSON-RPC error; using fallback") + if response.status_code == 404: + if self.session_id is None: + # No session yet → 404 is the HTTP-level spelling of + # METHOD_NOT_FOUND (gateway / legacy server doesn't know + # this method); "Session terminated" would be a lie here. + error_data = ErrorData(code=METHOD_NOT_FOUND, message="Not Found") + else: + error_data = ErrorData(code=INVALID_REQUEST, message="Session terminated") + else: + error_data = ErrorData(code=INTERNAL_ERROR, message="Server returned an error response") + session_message = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data)) + await ctx.read_stream_writer.send(session_message) + return + + if self._is_initialization_request(message): + self._maybe_extract_session_id_from_response(response) + + # Per https://modelcontextprotocol.io/specification/2025-06-18/basic#notifications: + # The server MUST NOT send a response to notifications. + if isinstance(message, JSONRPCRequest): + content_type = response.headers.get("content-type", "").lower() + if content_type.startswith("application/json"): + await self._handle_json_response(response, ctx.read_stream_writer, request_id=message.id) + elif content_type.startswith("text/event-stream"): + await self._handle_sse_response(response, ctx) + else: + logger.error(f"Unexpected content type: {content_type}") + error_data = ErrorData(code=INVALID_REQUEST, message=f"Unexpected content type: {content_type}") + error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data)) + await ctx.read_stream_writer.send(error_msg) + + async def _handle_json_response( + self, + response: httpx2.Response, + read_stream_writer: StreamWriter, + *, + request_id: RequestId, + ) -> None: + """Handle JSON response from the server.""" + try: + content = await response.aread() + message = jsonrpc_message_adapter.validate_json(content, by_name=False) + session_message = SessionMessage(message) + await read_stream_writer.send(session_message) + except (httpx2.StreamError, ValidationError) as exc: + logger.exception("Error parsing JSON response") + error_data = ErrorData(code=PARSE_ERROR, message=f"Failed to parse JSON response: {exc}") + error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=request_id, error=error_data)) + await read_stream_writer.send(error_msg) + + async def _handle_sse_response( + self, + response: httpx2.Response, + ctx: RequestContext, + ) -> None: + """Handle SSE response from the server.""" + last_event_id: str | None = None + retry_interval_ms: int | None = None + + # The caller (_handle_post_request) only reaches here inside + # isinstance(message, JSONRPCRequest), so this is always a JSONRPCRequest. + assert isinstance(ctx.session_message.message, JSONRPCRequest) + original_request_id = ctx.session_message.message.id + + try: + event_source = EventSource(response) + async for sse in event_source: # pragma: no branch + # Track last event ID for potential reconnection + if sse.id: + last_event_id = sse.id + + # Track retry interval from server + if sse.retry is not None: + retry_interval_ms = sse.retry + + is_complete = await self._handle_sse_event( + sse, + ctx.read_stream_writer, + original_request_id=original_request_id, + resumption_callback=(ctx.metadata.on_resumption_token_update if ctx.metadata else None), + ) + # If the SSE event indicates completion, like returning response/error + # break the loop + if is_complete: + await response.aclose() + return # Normal completion, no reconnect needed + except Exception: + logger.debug("SSE stream ended", exc_info=True) # pragma: lax no cover + + # Stream ended without response - reconnect if we received an event with ID + if last_event_id is not None: + logger.info("SSE stream disconnected, reconnecting...") + await self._handle_reconnection(ctx, last_event_id, retry_interval_ms) + else: + # Not resumable: resolve the waiter, else a listen stream's consumer + # would hang forever instead of learning the subscription is lost. + await self._resolve_abandoned_request( + ctx.read_stream_writer, original_request_id, "SSE stream ended without a response" + ) + + async def _resolve_abandoned_request( + self, read_stream_writer: StreamWriter, request_id: RequestId, message: str, *, code: int = CONNECTION_CLOSED + ) -> None: + """Resolve a request whose response can never arrive with a synthesized error. + + Best-effort: a closed read stream means the session is tearing down. + """ + error_data = ErrorData(code=code, message=message) + error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=request_id, error=error_data)) + try: + await read_stream_writer.send(error_msg) + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + logger.debug("read stream closed before request %r could be resolved", request_id) + + async def _handle_reconnection( + self, + ctx: RequestContext, + last_event_id: str, + retry_interval_ms: int | None = None, + attempt: int = 0, + ) -> None: + """Reconnect with Last-Event-ID to resume stream after server disconnect.""" + # Only requests reconnect: every caller arrives from a request's response stream. + assert isinstance(ctx.session_message.message, JSONRPCRequest) + original_request_id = ctx.session_message.message.id + + if attempt >= MAX_RECONNECTION_ATTEMPTS: + # Resolve on give-up: a request with no read timeout (a listen + # stream) would otherwise hang its caller forever. + logger.debug(f"Max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded") + await self._resolve_abandoned_request( + ctx.read_stream_writer, original_request_id, "SSE stream ended and reconnection attempts were exhausted" + ) + return + + # Always wait - use server value or default + delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS + await anyio.sleep(delay_ms / 1000.0) + + headers = self._prepare_headers() + headers[LAST_EVENT_ID] = last_event_id + + try: + async with sse_within_origin(ctx.client, self.url, headers=headers) as event_source: + event_source.response.raise_for_status() + logger.info("Reconnected to SSE stream") + + # Track for potential further reconnection + reconnect_last_event_id: str = last_event_id + reconnect_retry_ms = retry_interval_ms + + async for sse in event_source: + if sse.id: # pragma: no branch + reconnect_last_event_id = sse.id + if sse.retry is not None: + reconnect_retry_ms = sse.retry + + is_complete = await self._handle_sse_event( + sse, + ctx.read_stream_writer, + original_request_id, + ctx.metadata.on_resumption_token_update if ctx.metadata else None, + ) + if is_complete: + await event_source.response.aclose() + return + + # Stream ended again without response - reconnect again (reset attempt counter) + logger.info("SSE stream disconnected, reconnecting...") + await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, 0) + except Exception as e: # pragma: no cover + logger.debug(f"Reconnection failed: {e}") + # Try to reconnect again if we still have an event ID + await self._handle_reconnection(ctx, last_event_id, retry_interval_ms, attempt + 1) + + async def post_writer( + self, + client: httpx2.AsyncClient, + write_stream_reader: StreamReader, + read_stream_writer: StreamWriter, + write_stream: ContextSendStream[SessionMessage], + start_get_stream: Callable[[], None], + tg: TaskGroup, + ) -> None: + """Handle writing requests to the server.""" + try: + async with write_stream_reader, read_stream_writer, write_stream: + + async def _handle_message(session_message: SessionMessage) -> None: + message = session_message.message + if self._consume_modern_cancellation(session_message): + return + metadata = ( + session_message.metadata + if isinstance(session_message.metadata, ClientMessageMetadata) + else None + ) + + # Check if this is a resumption request + is_resumption = bool(metadata and metadata.resumption_token) + + logger.debug(f"Sending client message: {message}") + + # Handle initialized notification + if self._is_initialized_notification(message): + start_get_stream() + + if self._is_initialization_request(message): + # `initialize` is the negotiation, not a "subsequent request" — discard any + # probe-stamped value so the discover→fallback path can't leak it onto the handshake. + self._protocol_version_header = None + elif metadata is not None and metadata.headers is not None: + stamped_version = metadata.headers.get(MCP_PROTOCOL_VERSION_HEADER) + if stamped_version is not None: + self._protocol_version_header = stamped_version + + ctx = RequestContext( + client=client, + session_id=self.session_id, + session_message=session_message, + metadata=metadata, + read_stream_writer=read_stream_writer, + ) + + async def handle_request_async(): + if is_resumption: + await self._handle_resumption_request(ctx) + else: + await self._handle_post_request(ctx) + + # If this is a request, start a new task to handle it + if isinstance(message, JSONRPCRequest): + # Register the abort scope before the spawn: the next + # message through this loop can already be the abandon + # signal for this id, ahead of the task ever running. + post = _InFlightPost( + scope=anyio.CancelScope(), + modern=self._protocol_version_header in MODERN_PROTOCOL_VERSIONS, + ) + superseded = self._in_flight_posts.get(message.id) + if superseded is not None: + # A reused id means the waiter belongs to this attempt now: + # sever the old POST so its zombie stream cannot answer, + # fail, or resolve the successor's request. + superseded.scope.cancel() + self._in_flight_posts[message.id] = post + tg.start_soon(self._run_request_post, handle_request_async, post, message.id) + else: + await handle_request_async() + + async for session_message in write_stream_reader: + sender_ctx = write_stream_reader.last_context + if sender_ctx is not None: + async with anyio.create_task_group() as tg_local: + sender_ctx.run(tg_local.start_soon, _handle_message, session_message) + else: + await _handle_message(session_message) # pragma: no cover + + except Exception: # pragma: lax no cover + logger.exception("Error in post_writer") + + async def terminate_session(self, client: httpx2.AsyncClient) -> None: + """Terminate the session by sending a DELETE request.""" + if not self.session_id: + return # pragma: no cover + + try: + headers = self._prepare_headers() + response = await request_within_origin(client, "DELETE", self.url, headers=headers) + + if response.status_code == 405: + logger.debug("Server does not allow session termination") + elif response.status_code not in (200, 204): + logger.warning(f"Session termination failed: {response.status_code}") # pragma: no cover + except Exception as exc: # pragma: no cover + logger.warning(f"Session termination failed: {exc}") + + +@asynccontextmanager +async def streamable_http_client( + url: str, + *, + http_client: httpx2.AsyncClient | None = None, + terminate_on_close: bool = True, +) -> AsyncGenerator[TransportStreams, None]: + """Client transport for StreamableHTTP. + + Args: + url: The MCP server endpoint URL. + http_client: Optional pre-configured httpx2.AsyncClient. If None, a default + client with recommended MCP timeouts will be created. To configure headers, + authentication, or other HTTP settings, create an httpx2.AsyncClient and pass it here. + Whichever client is used, MCP requests follow a redirect only when it stays on the + endpoint's origin (same scheme, host and port, or http to https on the same host with + default ports) and keeps the request method (307/308 for a POST; any status for the GET + stream); any other redirect is not followed and the message it answered fails with an + error naming the location. The + client's `follow_redirects` setting is not consulted; the SDK's OAuth providers apply the + same rule to the requests they make. + terminate_on_close: If True, send a DELETE request to terminate the session when the context exits. + + Yields: + Tuple containing: + - read_stream: Stream for reading messages from the server + - write_stream: Stream for sending messages to the server + + Example: + See examples/snippets/clients/ for usage patterns. + """ + # Determine if we need to create and manage the client + client_provided = http_client is not None + client = http_client + + if client is None: + # Create default client with recommended MCP timeouts + client = create_mcp_http_client() + + transport = StreamableHTTPTransport(url) + + logger.debug(f"Connecting to StreamableHTTP endpoint: {url}") + + async with contextlib.AsyncExitStack() as stack: + # Only manage client lifecycle if we created it + if not client_provided: + await stack.enter_async_context(client) + + read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0) + write_stream, write_stream_reader = create_context_streams[SessionMessage](0) + + async with ( + read_stream_writer, + read_stream, + write_stream, + write_stream_reader, + anyio.create_task_group() as tg, + ): + + def start_get_stream() -> None: + tg.start_soon(transport.handle_get_stream, client, read_stream_writer) + + tg.start_soon( + transport.post_writer, + client, + write_stream_reader, + read_stream_writer, + write_stream, + start_get_stream, + tg, + ) + + try: + yield read_stream, write_stream + finally: + if transport.session_id and terminate_on_close: + await transport.terminate_session(client) + tg.cancel_scope.cancel() + await resync_tracer() diff --git a/src/mcp-client/mcp_client/client/subscriptions.py b/src/mcp-client/mcp_client/client/subscriptions.py new file mode 100644 index 0000000000..8794c88fb2 --- /dev/null +++ b/src/mcp-client/mcp_client/client/subscriptions.py @@ -0,0 +1,282 @@ +"""Client-side `subscriptions/listen` driver (2026-07-28, SEP-2575). + +`listen()` opens the stream as an async context manager: entering waits for +the server's acknowledgment, iteration yields typed change events, a graceful +server close ends the loop, and an abrupt drop raises `SubscriptionLost`. +There is no replay and no automatic re-listen: a client that re-opens a +subscription refetches what it depends on. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from contextlib import asynccontextmanager +from itertools import count +from typing import TYPE_CHECKING, Literal + +import anyio +import mcp_types as types +from mcp_types.version import MODERN_PROTOCOL_VERSIONS + +from mcp_client.shared.dispatcher import CallOptions +from mcp_client.shared.exceptions import MCPError +from mcp_client.shared.subscriptions import ( + PromptsListChanged, + ResourcesListChanged, + ResourceUpdated, + ServerEvent, + ToolsListChanged, + event_matches, +) + +if TYPE_CHECKING: + from mcp_client.client.session import ClientSession + +__all__ = [ + "ListenNotSupportedError", + "OnEvent", + "PromptsListChanged", + "ResourceUpdated", + "ResourcesListChanged", + "ServerEvent", + "Subscription", + "SubscriptionLost", + "ToolsListChanged", + "listen", +] + +_listen_ids = count(1) +"""Process-wide `listen-N` sequence: string ids can never collide with a dispatcher's minted ints.""" + +_MAX_PENDING_EVENTS = 1024 +"""Backlog backstop: the spec allows sub-resource URIs, so distinct pending +`ResourceUpdated` events are unbounded; overflowing this cap settles the +subscription lost rather than growing client memory.""" + +_SubscriptionEnd = Literal["graceful", "lost", "local"] + + +class ListenNotSupportedError(RuntimeError): + """`subscriptions/listen` requires a 2026-07-28 connection.""" + + def __init__(self, negotiated_version: str | None) -> None: + self.negotiated_version = negotiated_version + super().__init__( + f"subscriptions/listen is not available at protocol version {negotiated_version!r}; it requires " + "2026-07-28. On earlier versions use subscribe_resource() and the change notifications delivered " + "through message_handler." + ) + + +class SubscriptionLost(RuntimeError): + """The stream ended without the server's graceful close; re-listen and refetch.""" + + +class ListenRoute: + """Package-internal demux state for one listen stream, fed synchronously in receive order by the session.""" + + def __init__(self) -> None: + self.honored: types.SubscriptionFilter | None = None + self.acked = anyio.Event() + self.error: MCPError | None = None + self.end: _SubscriptionEnd | None = None + self._honored_uris: frozenset[str] = frozenset() + self._pending: dict[ServerEvent, None] = {} + self._wake = anyio.Event() + + def set_acked(self, honored: types.SubscriptionFilter) -> None: + """Record the acknowledged filter; the first ack wins.""" + if not self.acked.is_set(): + self.honored = honored + self._honored_uris = frozenset(honored.resource_subscriptions or ()) + self.acked.set() + + def deliver(self, event: ServerEvent) -> None: + """Queue an event within the honored filter, deduplicated against the backlog. + + Any `ResourceUpdated` is admitted once URI subscriptions were honored at + all: the spec allows the stamped URI to be a sub-resource of a subscribed one. + """ + if self.end is not None or self.honored is None: + return + if isinstance(event, ResourceUpdated): + admitted = bool(self._honored_uris) + else: + admitted = event_matches(self.honored, self._honored_uris, event) + if not admitted or event in self._pending: + return + if len(self._pending) >= _MAX_PENDING_EVENTS: + self.settle( + "lost", + error=MCPError( + types.INTERNAL_ERROR, + f"subscription backlog exceeded {_MAX_PENDING_EVENTS} unconsumed events; re-listen and refetch", + ), + ) + return + self._pending[event] = None + self._wake.set() + + def settle(self, end: _SubscriptionEnd, error: MCPError | None = None) -> None: + """Record the stream's end; the first reason wins and wakes both waiters.""" + if self.end is None: + self.end = end + self.error = error + self.acked.set() + self._wake.set() + + async def next_event(self) -> ServerEvent | _SubscriptionEnd: + """Peek the next pending event, or the stream's end once the backlog drains. + + A "local" end short-circuits the backlog; the other endings drain it first, + so a graceful close never swallows events that preceded it. + """ + while True: + # Snapshot the wake event before checking state so a deliver landing after the checks cannot be missed. + wake = self._wake + if self.end == "local": + return self.end + if self._pending: + return next(iter(self._pending)) + if self.end is not None: + return self.end + await wake.wait() + self._wake = anyio.Event() + + def consume(self, event: ServerEvent) -> None: + """Remove a peeked event from the backlog.""" + self._pending.pop(event, None) + + +OnEvent = Callable[[ServerEvent], Awaitable[None]] +"""Per-event barrier awaited before a `Subscription` returns each event to its consumer.""" + + +class Subscription: + """One open `subscriptions/listen` stream: an async iterator of typed events. + + Produced by `listen()` / `Client.listen()`, not constructed directly. + """ + + def __init__( + self, + route: ListenRoute, + subscription_id: types.RequestId, + honored: types.SubscriptionFilter, + on_event: OnEvent | None = None, + ): + self._route = route + self._on_event = on_event + self.subscription_id = subscription_id + """The listen request's JSON-RPC id, stamped into every frame's `_meta`.""" + self.honored = honored + """The subset of the requested filter the server agreed to deliver.""" + + def __aiter__(self) -> Subscription: + return self + + async def __anext__(self) -> ServerEvent: + """Yield the next change event; the loop ends when the stream does. + + Raises: + SubscriptionLost: the stream dropped without the server's graceful close. + """ + outcome = await self._route.next_event() + if isinstance(outcome, str): + if outcome == "lost": + raise SubscriptionLost( + f"subscription {self.subscription_id!r} ended without the server's graceful close;" + " re-listen and refetch" + ) from self._route.error + raise StopAsyncIteration + if self._on_event is not None: + # The event stays pending while the barrier runs: a cancellation or a + # raising barrier leaves it for the next anext instead of dropping it. + await self._on_event(outcome) + self._route.consume(outcome) + return outcome + + +@asynccontextmanager +async def listen( + session: ClientSession, + *, + tools_list_changed: bool = False, + prompts_list_changed: bool = False, + resources_list_changed: bool = False, + resource_subscriptions: Sequence[str] = (), + on_event: OnEvent | None = None, +) -> AsyncIterator[Subscription]: + """Open one `subscriptions/listen` stream on `session` (2026-07-28 only). + + Entering sends the request and returns once the server's acknowledgment + arrives; exiting ends the subscription. `on_event` is awaited before each + event is returned - the seam `Client.listen` uses to finish cache eviction + before the consumer can refetch. + + Raises: + ListenNotSupportedError: negotiated version predates 2026-07-28. + MCPError: the server rejected the request, or the connection failed pre-ack. + SubscriptionLost: the stream ended before it was acknowledged. + TimeoutError: the session's read timeout elapsed before the acknowledgment. + """ + if session.protocol_version not in MODERN_PROTOCOL_VERSIONS: + raise ListenNotSupportedError(session.protocol_version) + if isinstance(resource_subscriptions, str): + raise TypeError("resource_subscriptions takes a sequence of URIs, not a bare string") + request = types.SubscriptionsListenRequest( + params=types.SubscriptionsListenRequestParams( + notifications=types.SubscriptionFilter( + tools_list_changed=tools_list_changed or None, + prompts_list_changed=prompts_list_changed or None, + resources_list_changed=resources_list_changed or None, + resource_subscriptions=list(resource_subscriptions) or None, + ) + ) + ) + task_group = session._task_group # pyright: ignore[reportPrivateUsage] + if task_group is None: + raise RuntimeError("listen() requires an entered session") + request_id: types.RequestId = f"listen-{next(_listen_ids)}" + data = request.model_dump(by_alias=True, mode="json", exclude_none=True) + opts: CallOptions = {"request_id": request_id} + session._stamp(data, opts) # pyright: ignore[reportPrivateUsage] + driver_scope = anyio.CancelScope() + + async def drive() -> None: + # Deliberately no result timeout: the response arrives when the stream ends. + with driver_scope: + try: + await session._dispatcher.send_raw_request( # pyright: ignore[reportPrivateUsage] + data["method"], data.get("params"), opts + ) + except MCPError as error: + route.settle("lost", error=error) + return + except ValueError as error: + # A raw request id collided with our minted listen id: fail this subscription + # and release the route in this same slice, so it cannot consume the raw caller's ack. + session._unregister_listen_route(request_id) # pyright: ignore[reportPrivateUsage] + route.settle("lost", error=MCPError(types.INTERNAL_ERROR, str(error))) + return + # A result, whatever its body, is the spec's graceful close; with no prior ack + # it opens the subscription already closed. + route.set_acked(types.SubscriptionFilter()) + route.settle("graceful") + + # Register the demux route before the request is written so the ack cannot race it. + route = session._register_listen_route(request_id) # pyright: ignore[reportPrivateUsage] + try: + task_group.start_soon(drive) + with anyio.fail_after(session._session_read_timeout_seconds): # pyright: ignore[reportPrivateUsage] + await route.acked.wait() + if route.honored is None: + # Only reachable on failure paths: a graceful no-ack result acked an empty filter in drive(). + if route.error is not None: + raise route.error + raise SubscriptionLost(f"subscription {request_id!r} ended before it was acknowledged") + yield Subscription(route, request_id, route.honored, on_event) + finally: + route.settle("local") + driver_scope.cancel() + session._unregister_listen_route(request_id) # pyright: ignore[reportPrivateUsage] diff --git a/src/mcp-client/mcp_client/os/__init__.py b/src/mcp-client/mcp_client/os/__init__.py new file mode 100644 index 0000000000..a9a2c5b3bb --- /dev/null +++ b/src/mcp-client/mcp_client/os/__init__.py @@ -0,0 +1 @@ +__all__ = [] diff --git a/src/mcp-client/mcp_client/os/posix/__init__.py b/src/mcp-client/mcp_client/os/posix/__init__.py new file mode 100644 index 0000000000..a9a2c5b3bb --- /dev/null +++ b/src/mcp-client/mcp_client/os/posix/__init__.py @@ -0,0 +1 @@ +__all__ = [] diff --git a/src/mcp-client/mcp_client/os/posix/utilities.py b/src/mcp-client/mcp_client/os/posix/utilities.py new file mode 100644 index 0000000000..311f2a7051 --- /dev/null +++ b/src/mcp-client/mcp_client/os/posix/utilities.py @@ -0,0 +1,63 @@ +"""POSIX-specific functionality for stdio client operations.""" + +import logging +import os +import signal +from contextlib import suppress + +import anyio +from anyio.abc import Process + +logger = logging.getLogger("mcp.os.posix.utilities") + +# How often to probe for surviving group members between SIGTERM and SIGKILL. +_GROUP_POLL_INTERVAL = 0.01 + + +async def terminate_posix_process_tree(process: Process, timeout_seconds: float = 2.0) -> None: + """Terminates a process and all its descendants on POSIX. + + SIGTERMs the process group, waits up to timeout_seconds for it to + disappear, then SIGKILLs whatever remains. killpg reaches every descendant + atomically, even ones whose parent already exited; daemonizers that left + the group escape by design. A group only disappears once every member is + dead and reaped, so a client running as PID 1 should reap orphans (e.g. + docker run --init) or the wait below runs its full timeout. + """ + # The leader's pid is the pgid (start_new_session). Never use getpgid(): + # it fails once the leader is reaped, even with live members left. + pgid = process.pid + + try: + os.killpg(pgid, signal.SIGTERM) + except ProcessLookupError: + return # the whole group is already gone + except PermissionError: + # EPERM never proves the group is gone (macOS raises it for zombie or + # foreign-euid members), so keep waiting and escalating. + logger.warning( + "No permission to signal some of process group %d; waiting for it to exit anyway", pgid, exc_info=True + ) + + with anyio.move_on_after(timeout_seconds): + while _group_alive(pgid): + # Reading returncode reaps the leader on trio; a zombie leader would + # otherwise keep the group alive for the full timeout. + _ = process.returncode + await anyio.sleep(_GROUP_POLL_INTERVAL) + return + + # ESRCH: died since the last probe. EPERM: we killed what we were allowed to. + with suppress(ProcessLookupError, PermissionError): + os.killpg(pgid, signal.SIGKILL) + + +def _group_alive(pgid: int) -> bool: + """Probes the group with signal 0; only ESRCH proves it is gone.""" + try: + os.killpg(pgid, 0) + except ProcessLookupError: + return False + except PermissionError: + pass # unsignalable survivors or unreaped zombies; EPERM is ambiguous + return True diff --git a/src/mcp-client/mcp_client/os/win32/__init__.py b/src/mcp-client/mcp_client/os/win32/__init__.py new file mode 100644 index 0000000000..a9a2c5b3bb --- /dev/null +++ b/src/mcp-client/mcp_client/os/win32/__init__.py @@ -0,0 +1 @@ +__all__ = [] diff --git a/src/mcp-client/mcp_client/os/win32/utilities.py b/src/mcp-client/mcp_client/os/win32/utilities.py new file mode 100644 index 0000000000..c096f770ee --- /dev/null +++ b/src/mcp-client/mcp_client/os/win32/utilities.py @@ -0,0 +1,292 @@ +"""Windows-specific functionality for stdio transport operations.""" + +import logging +import shutil +import subprocess +import sys +import weakref +from contextlib import suppress +from pathlib import Path +from typing import BinaryIO, TextIO, TypeAlias, cast + +import anyio +from anyio.abc import Process +from anyio.streams.file import FileReadStream, FileWriteStream + +logger = logging.getLogger("mcp.os.win32.utilities") + +# Windows-specific imports for Job Objects +if sys.platform == "win32": + import msvcrt + + import pywintypes + import win32api + import win32con + import win32job +else: + # Type stubs for non-Windows platforms + win32api = None + win32con = None + msvcrt = None + win32job = None + pywintypes = None + + +def rebind_std_handle_to_fd(fd: int) -> None: + """Points the Win32 standard-handle slot for fd 0, 1, or 2 at fd's current OS handle. + + os.dup2 updates only the CRT descriptor table; subprocess handle inheritance + reads the Win32 slot, so it must be repointed too. + + Raises: + OSError: The slot could not be set. + """ + if sys.platform != "win32" or not win32api or not msvcrt or not pywintypes: + return + std_ids = {0: win32api.STD_INPUT_HANDLE, 1: win32api.STD_OUTPUT_HANDLE, 2: win32api.STD_ERROR_HANDLE} + try: + win32api.SetStdHandle(std_ids[fd], msvcrt.get_osfhandle(fd)) + except pywintypes.error as exc: + # Normalized so callers' OSError-based best-effort handling covers it. + raise OSError(f"SetStdHandle failed for fd {fd}") from exc + + +# How often FallbackProcess polls the underlying Popen for exit. +_EXIT_POLL_INTERVAL = 0.01 + +# Job Object handle per spawned process, for tree termination at shutdown. +# Values stay pywin32 PyHANDLEs: if no pop site ever runs, the dying weak entry +# drops the last reference and the PyHANDLE destructor closes the handle, which +# is what makes KILL_ON_JOB_CLOSE reap an abandoned tree. +_process_jobs: "weakref.WeakKeyDictionary[Process | FallbackProcess, object]" = weakref.WeakKeyDictionary() + + +def get_windows_executable_command(command: str) -> str: + """Resolves the command to a Windows executable path. + + Tries the bare name first, then the common script extensions (.cmd, .bat, + .exe, .ps1). + """ + try: + if command_path := shutil.which(command): + return command_path + + for ext in [".cmd", ".bat", ".exe", ".ps1"]: + ext_version = f"{command}{ext}" + if ext_path := shutil.which(ext_version): + return ext_path + + return command + except OSError: + return command # path probing failed (permissions, broken symlinks) + + +class FallbackProcess: + """Async wrapper around subprocess.Popen for SelectorEventLoop. + + Windows event loops without async subprocess support get this Popen-backed + fallback, with anyio file streams wrapping the pipes. + """ + + def __init__(self, popen_obj: subprocess.Popen[bytes]) -> None: + self.popen: subprocess.Popen[bytes] = popen_obj + stdin = popen_obj.stdin + stdout = popen_obj.stdout + + self.stdin = FileWriteStream(cast(BinaryIO, stdin)) if stdin else None + self.stdout = FileReadStream(cast(BinaryIO, stdout)) if stdout else None + + async def wait(self) -> int: + """Waits for exit by polling the Popen. + + A thread blocked in Popen.wait() cannot be cancelled by anyio, which + would defeat every timeout placed around this call. + """ + while (returncode := self.popen.poll()) is None: + await anyio.sleep(_EXIT_POLL_INTERVAL) + return returncode + + def terminate(self) -> None: + """Terminates the subprocess.""" + self.popen.terminate() + + def kill(self) -> None: + """Kills the subprocess (on Windows the same hard kill as terminate).""" + self.popen.kill() + + @property + def pid(self) -> int: + """Returns the process ID.""" + return self.popen.pid + + @property + def returncode(self) -> int | None: + """The exit code, or None while the process is still running. + + Polls the Popen so death is observable without anyone calling wait(). + """ + return self.popen.poll() + + +# The process handle stdio_client drives: anyio's Process, or the Popen-backed +# fallback used on Windows event loops without async subprocess support. +ServerProcess: TypeAlias = Process | FallbackProcess + + +async def create_windows_process( + command: str, + args: list[str], + env: dict[str, str] | None = None, + errlog: TextIO | None = sys.stderr, + cwd: Path | str | None = None, +) -> Process | FallbackProcess: + """Creates a subprocess with Job Object support for tree termination. + + Spawns via anyio's open_process; event loops without async subprocess + support (notably the SelectorEventLoop) raise NotImplementedError, in which + case the spawn falls back to a Popen-backed FallbackProcess. Either way the + process is then assigned to a Job Object so its children can be terminated + with it; children spawned before the assignment completes are not captured + (see the inline note below). + + Returns: + Process | FallbackProcess: The spawned process with async stdin/stdout streams. + """ + try: + process = await anyio.open_process( + [command, *args], + env=env, + # Ensure we don't create console windows for each process + creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0), + stderr=errlog, + cwd=cwd, + ) + except NotImplementedError: + # Windows event loops without async subprocess support (SelectorEventLoop) + process = await _create_windows_fallback_process(command, args, env, errlog, cwd) + + # Children spawned before the assignment completes land outside the job + # (membership is inherited at CreateProcess, never acquired retroactively); + # if that ever bites, the fix is a CREATE_SUSPENDED spawn -> assign -> resume. + job = _create_job_object() + _maybe_assign_process_to_job(process, job) + return process + + +async def _create_windows_fallback_process( + command: str, + args: list[str], + env: dict[str, str] | None = None, + errlog: TextIO | None = sys.stderr, + cwd: Path | str | None = None, +) -> FallbackProcess: + """Spawns via subprocess.Popen and wraps it in FallbackProcess.""" + popen_obj = subprocess.Popen( + [command, *args], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=errlog, + env=env, + cwd=cwd, + bufsize=0, # Unbuffered output + creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0), + ) + return FallbackProcess(popen_obj) + + +def _create_job_object() -> object | None: + """Creates a Windows Job Object configured to terminate all its processes when closed.""" + if sys.platform != "win32" or not win32api or not win32job: + return None + + job = None + try: + job = win32job.CreateJobObject(None, "") + extended_info = win32job.QueryInformationJobObject(job, win32job.JobObjectExtendedLimitInformation) + + extended_info["BasicLimitInformation"]["LimitFlags"] |= win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE + win32job.SetInformationJobObject(job, win32job.JobObjectExtendedLimitInformation, extended_info) + return job + except pywintypes.error: + logger.warning("Failed to create Job Object for process tree management", exc_info=True) + # If creation succeeded but configuration failed, close the handle now. + if job is not None: + _close_job_handle(job) + return None + + +def _maybe_assign_process_to_job(process: Process | FallbackProcess, job: object | None) -> None: + """Assigns the process to the job and records it for tree termination. + + On any failure the job handle is closed instead. + """ + if job is None: + return + + if sys.platform != "win32" or not win32api or not win32con or not win32job: + return + + try: + process_handle = win32api.OpenProcess( + win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, process.pid + ) + if not process_handle: + raise pywintypes.error(0, "OpenProcess", "Failed to open process handle") + + try: + win32job.AssignProcessToJobObject(job, process_handle) + finally: + win32api.CloseHandle(process_handle) + # Record only after the CloseHandle above succeeded: had it failed, the + # except below would close the job and KILL_ON_JOB_CLOSE takes the server. + _process_jobs[process] = job + except pywintypes.error: + logger.warning("Failed to assign process %d to Job Object", process.pid, exc_info=True) + _close_job_handle(job) + + +def close_process_job(process: Process | FallbackProcess) -> None: + """Closes the process's Job Object handle, if it still has one. + + KILL_ON_JOB_CLOSE makes the close also kill any members still alive, + deterministically rather than at GC time; a deliberate divergence from + POSIX, where a graceful server's children are left alive. + """ + if sys.platform != "win32": + return + + job = _process_jobs.pop(process, None) + if job is not None: + _close_job_handle(job) + + +async def terminate_windows_process_tree(process: Process | FallbackProcess) -> None: + """Terminates the process's job, or just the process if it has no job. + + Job termination is an immediate hard kill of every member. Windows has no + tree-wide SIGTERM; the stdin-close grace period is the server's chance to + exit cleanly. + """ + if sys.platform != "win32": + return + + job = _process_jobs.pop(process, None) + if job is not None and win32job: + try: + with suppress(pywintypes.error): # the job might already be terminated + win32job.TerminateJobObject(job, 1) + finally: + _close_job_handle(job) + + # The process may have no job (creation or assignment failed); kill it directly too. + try: + process.terminate() + except OSError: + pass + + +def _close_job_handle(job: object) -> None: + """Closes a Job Object handle, tolerating one that is already closed.""" + if win32api and pywintypes: + with suppress(pywintypes.error): + win32api.CloseHandle(job) diff --git a/src/mcp-client/mcp_client/py.typed b/src/mcp-client/mcp_client/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/mcp-client/mcp_client/shared/__init__.py b/src/mcp-client/mcp_client/shared/__init__.py new file mode 100644 index 0000000000..a9a2c5b3bb --- /dev/null +++ b/src/mcp-client/mcp_client/shared/__init__.py @@ -0,0 +1 @@ +__all__ = [] diff --git a/src/mcp-client/mcp_client/shared/_callable_inspection.py b/src/mcp-client/mcp_client/shared/_callable_inspection.py new file mode 100644 index 0000000000..0e89e446f8 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/_callable_inspection.py @@ -0,0 +1,33 @@ +"""Callable inspection utilities. + +Adapted from Starlette's `is_async_callable` implementation. +https://github.com/encode/starlette/blob/main/starlette/_utils.py +""" + +from __future__ import annotations + +import functools +import inspect +from collections.abc import Awaitable, Callable +from typing import Any, TypeGuard, TypeVar, overload + +T = TypeVar("T") + +AwaitableCallable = Callable[..., Awaitable[T]] + + +@overload +def is_async_callable(obj: AwaitableCallable[T]) -> TypeGuard[AwaitableCallable[T]]: ... + + +@overload +def is_async_callable(obj: Any) -> TypeGuard[AwaitableCallable[Any]]: ... + + +def is_async_callable(obj: Any) -> Any: + while isinstance(obj, functools.partial): # pragma: lax no cover + obj = obj.func + + return inspect.iscoroutinefunction(obj) or ( + callable(obj) and inspect.iscoroutinefunction(getattr(obj, "__call__", None)) + ) diff --git a/src/mcp-client/mcp_client/shared/_compat.py b/src/mcp-client/mcp_client/shared/_compat.py new file mode 100644 index 0000000000..88d50ba20a --- /dev/null +++ b/src/mcp-client/mcp_client/shared/_compat.py @@ -0,0 +1,19 @@ +"""Workarounds for CPython interpreter bugs the SDK papers over.""" + +import anyio.lowlevel + +__all__ = ["resync_tracer"] + + +async def resync_tracer() -> None: + """Resync coverage tracing after a cancelled task-group join. + + A cancel delivered at a join resumes the awaiting coroutine chain via + `coro.throw()`; on CPython 3.11 (python/cpython#106749) that drops the + `'call'` trace events for the outer frames and desyncs coverage's CTracer + until the chain next suspends and resumes normally. Yielding once here + resumes via `.send()`, which re-stamps the missing events. Shielded so a + pending outer cancel is not re-delivered at this point; behaviorally a + no-op. Delete this module when Python 3.11 support ends (EOL 2027-10). + """ + await anyio.lowlevel.cancel_shielded_checkpoint() diff --git a/src/mcp-client/mcp_client/shared/_context_streams.py b/src/mcp-client/mcp_client/shared/_context_streams.py new file mode 100644 index 0000000000..04c33306d9 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/_context_streams.py @@ -0,0 +1,119 @@ +"""Context-aware memory stream wrappers. + +anyio memory streams do not propagate ``contextvars.Context`` across task +boundaries. These thin wrappers capture the sender's context at ``send()`` +time and expose it on the receive side via ``last_context``, so consumers +can restore it with ``ctx.run(handler, item)``. + +The iteration interface is unchanged (yields ``T``, not tuples), keeping +these wrappers duck-type compatible with plain ``MemoryObjectSendStream`` +and ``MemoryObjectReceiveStream``. +""" + +from __future__ import annotations + +import contextvars +from types import TracebackType +from typing import Any, Generic, TypeVar + +import anyio +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream + +T = TypeVar("T") + +# Internal payload carried through the underlying raw stream. +_Envelope = tuple[contextvars.Context, T] + + +class ContextSendStream(Generic[T]): + """Send-side wrapper that snapshots ``contextvars.copy_context()`` on every ``send()``.""" + + __slots__ = ("_inner",) + + def __init__(self, inner: MemoryObjectSendStream[_Envelope[T]]) -> None: + self._inner = inner + + async def send(self, item: T) -> None: + await self._inner.send((contextvars.copy_context(), item)) + + def close(self) -> None: + self._inner.close() + + async def aclose(self) -> None: + await self._inner.aclose() + + def clone(self) -> ContextSendStream[T]: # pragma: no cover + return ContextSendStream(self._inner.clone()) + + async def __aenter__(self) -> ContextSendStream[T]: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> bool | None: + await self.aclose() + return None + + +class ContextReceiveStream(Generic[T]): + """Receive-side wrapper that yields ``T`` and stores the sender's context in ``last_context``.""" + + __slots__ = ("_inner", "last_context") + + def __init__(self, inner: MemoryObjectReceiveStream[_Envelope[T]]) -> None: + self._inner = inner + self.last_context: contextvars.Context | None = None + + async def receive(self) -> T: + ctx, item = await self._inner.receive() + self.last_context = ctx + return item + + def close(self) -> None: + self._inner.close() + + async def aclose(self) -> None: + await self._inner.aclose() + + def clone(self) -> ContextReceiveStream[T]: # pragma: no cover + return ContextReceiveStream(self._inner.clone()) + + def __aiter__(self) -> ContextReceiveStream[T]: + return self + + async def __anext__(self) -> T: + try: + return await self.receive() + except anyio.EndOfStream: + raise StopAsyncIteration + + async def __aenter__(self) -> ContextReceiveStream[T]: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> bool | None: + await self.aclose() + return None + + +class create_context_streams( + tuple[ContextSendStream[T], ContextReceiveStream[T]], +): + """Create context-aware memory object streams. + + Supports ``create_context_streams[T](n)`` bracket syntax, + matching anyio's ``create_memory_object_stream`` API style. + """ + + def __new__(cls, max_buffer_size: float = 0) -> tuple[ContextSendStream[T], ContextReceiveStream[T]]: # type: ignore[type-var] + raw_send: MemoryObjectSendStream[Any] + raw_receive: MemoryObjectReceiveStream[Any] + raw_send, raw_receive = anyio.create_memory_object_stream(max_buffer_size) + return (ContextSendStream(raw_send), ContextReceiveStream(raw_receive)) diff --git a/src/mcp-client/mcp_client/shared/_httpx_utils.py b/src/mcp-client/mcp_client/shared/_httpx_utils.py new file mode 100644 index 0000000000..940b9f08cc --- /dev/null +++ b/src/mcp-client/mcp_client/shared/_httpx_utils.py @@ -0,0 +1,218 @@ +"""Utilities for creating and using httpx2 AsyncClient instances in the MCP transports.""" + +from abc import ABC, abstractmethod +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Any, Protocol + +import httpx2 + +__all__ = ["create_mcp_http_client", "MCP_DEFAULT_TIMEOUT", "MCP_DEFAULT_SSE_READ_TIMEOUT"] + +# Default MCP timeout configuration +MCP_DEFAULT_TIMEOUT = 30.0 # General operations (seconds) +MCP_DEFAULT_SSE_READ_TIMEOUT = 300.0 # SSE streams - 5 minutes (seconds) + +# The headers httpx2.AsyncClient.sse() adds to an event-stream request. +_SSE_HEADERS = {"Accept": "text/event-stream", "Cache-Control": "no-store"} + +# How many redirects one auth-flow request may follow within its origin (see RedirectAwareAuth). +_AUTH_REDIRECT_LIMIT = 5 + + +class McpHttpClientFactory(Protocol): # pragma: no branch + def __call__( # pragma: no branch + self, + headers: dict[str, str] | None = None, + timeout: httpx2.Timeout | None = None, + auth: httpx2.Auth | None = None, + ) -> httpx2.AsyncClient: ... + + +def create_mcp_http_client( + headers: dict[str, str] | None = None, + timeout: httpx2.Timeout | None = None, + auth: httpx2.Auth | None = None, +) -> httpx2.AsyncClient: + """Create an httpx2 AsyncClient with the MCP transports' default timeouts. + + The client uses a 30-second timeout for connect/write/pool and a 300-second + read timeout, because a server may hold a response stream open. Redirect + following is left at the httpx2 default (off): the MCP transports follow + redirects within the endpoint's origin themselves, see `stream_within_origin`. + + Args: + headers: Optional headers to include with all requests. + timeout: Request timeout as httpx2.Timeout object. Defaults to 30s for + connect/write/pool and 300s for read (for long-lived SSE streams). + auth: Optional authentication handler. + + Returns: + Configured httpx2.AsyncClient instance. + + Note: + The returned AsyncClient must be used as a context manager to ensure + proper cleanup of connections. + """ + if timeout is None: + timeout = httpx2.Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT) + kwargs: dict[str, Any] = {"timeout": timeout} + if headers is not None: + kwargs["headers"] = headers + if auth is not None: # pragma: no cover + kwargs["auth"] = auth + return httpx2.AsyncClient(**kwargs) + + +def _within_origin(url: httpx2.URL, location: httpx2.URL) -> bool: + """Whether `location` is on `url`'s origin, or is its https upgrade on the default ports. + + httpx2 normalises a scheme's default port to None and lower-cases hosts, so + plain tuple comparison is exact. The upgrade rule is the one httpx2 itself + uses to decide a redirect has not left the origin (`_is_https_redirect`). + """ + if (url.scheme, url.host, url.port) == (location.scheme, location.host, location.port): + return True + return ( + url.host == location.host + and url.scheme == "http" + and url.port is None + and location.scheme == "https" + and location.port is None + ) + + +def next_request_within_origin(response: httpx2.Response) -> httpx2.Request | None: + """The request that follows `response`'s redirect, if it is one the MCP transports follow. + + That is when httpx2 built a next request for it (a redirect status with a + Location), the next request keeps the method (307/308, or any redirect of a + GET: httpx2 turns a POST into a body-less GET for 301/302/303, which would + drop the message), its URL stays within the origin of the request just sent + (same scheme, host and port, or http to https on the same host with default + ports), and the Location does not bring userinfo of its own (which httpx2 + would otherwise send as Basic auth; userinfo the configured URL already had + is kept by a relative Location and is fine). None for anything else, + including a non-redirect. + """ + next_request = response.next_request + if next_request is None: + return None + sent = response.request + if ( + next_request.method != sent.method + or (next_request.url.userinfo and next_request.url.userinfo != sent.url.userinfo) + or not _within_origin(sent.url, next_request.url) + ): + return None + return next_request + + +@asynccontextmanager +async def stream_within_origin( + client: httpx2.AsyncClient, method: str, url: httpx2.URL | str, **kwargs: Any +) -> AsyncGenerator[httpx2.Response]: + """`client.stream(...)`, following redirects only while they stay within the request's origin. + + An MCP transport talks to one configured endpoint, and everything on a request + (headers, auth, body) was configured for that endpoint. A redirect that + `next_request_within_origin` accepts, such as a 307/308 trailing-slash + normalisation, is followed, at most `client.max_redirects` times. Any other + redirect (or one past that budget) is not followed: the redirect response + itself is yielded, the way httpx2 hands one back when `follow_redirects` is + off, and the caller treats it as the non-success it is. The client's own + `follow_redirects` setting is not consulted. Requests an `httpx2.Auth` flow + makes during the call are sent without following either; the SDK's OAuth + providers apply the same rule to their own requests. + """ + request = client.build_request(method, url, **kwargs) + followed = 0 + while True: + response = await client.send(request, stream=True, follow_redirects=False) + next_request = next_request_within_origin(response) + if next_request is None or followed == client.max_redirects: + break + try: + # Drain the redirect body so the connection returns to the pool, as httpx2 does when it follows. + await response.aread() + finally: + await response.aclose() + request = next_request + followed += 1 + try: + yield response + finally: + await response.aclose() + + +async def request_within_origin( + client: httpx2.AsyncClient, method: str, url: httpx2.URL | str, **kwargs: Any +) -> httpx2.Response: + """`client.request(...)` with the redirect handling of `stream_within_origin`.""" + async with stream_within_origin(client, method, url, **kwargs) as response: + await response.aread() + return response + + +@asynccontextmanager +async def sse_within_origin( + client: httpx2.AsyncClient, url: httpx2.URL | str, *, headers: dict[str, str] | None = None +) -> AsyncGenerator[httpx2.EventSource]: + """`client.sse(url)` with the redirect handling of `stream_within_origin`.""" + merged = httpx2.Headers(_SSE_HEADERS) + merged.update(headers or {}) + async with stream_within_origin(client, "GET", url, headers=merged) as response: + yield httpx2.EventSource(response) + + +def redirect_location(response: httpx2.Response) -> httpx2.URL | None: + """Where `response` redirects to, for use in a message: without userinfo, query or fragment, + which can carry state that does not belong in an error or a log line. None if not a redirect.""" + if response.next_request is None: + return None + return response.next_request.url.copy_with(userinfo=b"", query=None, fragment=None) + + +def redirect_note(response: httpx2.Response) -> str: + """A suffix naming the location of a redirect response that was not followed, else empty.""" + location = redirect_location(response) + if location is None: + return "" + return f" (redirected to {location}; not followed)" + + +class RedirectAwareAuth(ABC, httpx2.Auth): + """An `httpx2.Auth` whose own requests follow redirects the way MCP transport requests do. + + The transports send every request with redirect following off and follow a + redirect themselves only within the endpoint's origin (`stream_within_origin`). + httpx2 applies that per-request setting to the requests an auth flow makes + too (metadata discovery, registration, token), so on their own those would + follow nothing. Subclasses write their flow as `_auth_flow`; this class + drives it and, for each request the flow makes other than the one being + authenticated, follows a redirect that `next_request_within_origin` accepts, + up to `_AUTH_REDIRECT_LIMIT` times. Any other redirect response is handed + to the flow as it is. + """ + + @abstractmethod + def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + """The subclass's flow, written as `httpx2.Auth.async_auth_flow` otherwise would be.""" + + async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + flow = self._auth_flow(request) + try: + outgoing = await flow.__anext__() + while True: + response = yield outgoing + if outgoing is not request: + for _ in range(_AUTH_REDIRECT_LIMIT): + follow = next_request_within_origin(response) + if follow is None: + break + response = yield follow + outgoing = await flow.asend(response) + except StopAsyncIteration: + return + finally: + await flow.aclose() diff --git a/src/mcp-client/mcp_client/shared/_otel.py b/src/mcp-client/mcp_client/shared/_otel.py new file mode 100644 index 0000000000..b7b05b11ab --- /dev/null +++ b/src/mcp-client/mcp_client/shared/_otel.py @@ -0,0 +1,60 @@ +"""OpenTelemetry helpers for MCP.""" + +from __future__ import annotations + +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from typing import Any + +from opentelemetry.context import Context +from opentelemetry.propagate import extract, inject +from opentelemetry.trace import SpanKind, get_current_span, get_tracer +from opentelemetry.trace.span import Span + +_tracer = get_tracer("mcp-python-sdk") + + +@contextmanager +def otel_span( + name: str, + *, + kind: SpanKind, + attributes: dict[str, Any] | None = None, + context: Context | None = None, + record_exception: bool = True, + set_status_on_exception: bool = True, +) -> Generator[Span]: + """Create an OTel span.""" + with _tracer.start_as_current_span( + name, + kind=kind, + attributes=attributes, + context=context, + record_exception=record_exception, + set_status_on_exception=set_status_on_exception, + ) as span: + yield span + + +def inject_trace_context(meta: dict[str, Any]) -> None: + """Inject W3C trace context (traceparent/tracestate) into a `_meta` dict.""" + inject(meta) + + +def extract_trace_context(meta: Mapping[str, Any] | None) -> Context | None: + """Extract W3C trace context from a `_meta` dict. + + Returns `None` when the carrier is absent, malformed, or carries no + valid `traceparent`, so callers fall through to ambient parenting; an + explicit empty `Context` would orphan the span instead of nesting under + the current one. + """ + if not meta: + return None + try: + ctx = extract(meta) + except (ValueError, TypeError): + return None + if not get_current_span(ctx).get_span_context().is_valid: + return None + return ctx diff --git a/src/mcp-client/mcp_client/shared/_stream_protocols.py b/src/mcp-client/mcp_client/shared/_stream_protocols.py new file mode 100644 index 0000000000..b799751329 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/_stream_protocols.py @@ -0,0 +1,49 @@ +"""Stream protocols for MCP transports. + +These are general-purpose protocols satisfied by both ``MemoryObjectSendStream``/ +``MemoryObjectReceiveStream`` and the context-aware wrappers in ``_context_streams``. +""" + +from __future__ import annotations + +from types import TracebackType +from typing import Protocol, TypeVar + +from typing_extensions import Self + +T_co = TypeVar("T_co", covariant=True) +T_contra = TypeVar("T_contra", contravariant=True) + + +class ReadStream(Protocol[T_co]): + """Protocol for reading items from a stream. + + Consumers that need the sender's context should use + ``getattr(stream, 'last_context', None)``. + """ + + async def receive(self) -> T_co: ... + async def aclose(self) -> None: ... + def __aiter__(self) -> ReadStream[T_co]: ... + async def __anext__(self) -> T_co: ... + async def __aenter__(self) -> Self: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> bool | None: ... + + +class WriteStream(Protocol[T_contra]): + """Protocol for writing items to a stream.""" + + async def send(self, item: T_contra, /) -> None: ... + async def aclose(self) -> None: ... + async def __aenter__(self) -> Self: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> bool | None: ... diff --git a/src/mcp-client/mcp_client/shared/auth.py b/src/mcp-client/mcp_client/shared/auth.py new file mode 100644 index 0000000000..881379d381 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/auth.py @@ -0,0 +1,258 @@ +from typing import Any, Literal, cast + +from pydantic import AnyHttpUrl, AnyUrl, BaseModel, ConfigDict, Field, field_validator, model_validator + +# RFC 7523 JWT bearer grant; SEP-990 leg 2 uses this to present the ID-JAG. +JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer" + +# Token-endpoint client authentication methods this SDK's clients request, and the set +# `OAuthContext.prepare_token_auth` recognizes on a registered client (`private_key_jwt` is +# applied by `PrivateKeyJWTOAuthProvider`; the rest send a client secret or nothing). +TokenEndpointAuthMethod = Literal["none", "client_secret_post", "client_secret_basic", "private_key_jwt"] + +# grant_types a client requests when it does not specify its own (RFC 7591 §2). +DEFAULT_GRANT_TYPES = ["authorization_code", "refresh_token"] + + +def _empty_str_to_none(v: object) -> object: + # RFC 7591 §2 marks these URL fields OPTIONAL; a "" placeholder means absent, so it + # must not fail AnyHttpUrl validation. (The registered-client record applies the same + # rule to every member; this coercion serves the request model.) + if v == "": + return None + return v + + +class OAuthToken(BaseModel): + """See https://datatracker.ietf.org/doc/html/rfc6749#section-5.1""" + + access_token: str + token_type: Literal["Bearer"] = "Bearer" + expires_in: int | None = None + scope: str | None = None + refresh_token: str | None = None + + @field_validator("token_type", mode="before") + @classmethod + def normalize_token_type(cls, v: str | None) -> str | None: + if isinstance(v, str): + # Bearer is title-cased in the spec, so we normalize it + # https://datatracker.ietf.org/doc/html/rfc6750#section-4 + return v.title() + return v # pragma: no cover + + +class AuthorizationCodeResult(BaseModel): + """Authorization-code-grant redirect parameters returned by a callback handler. + + `iss` carries the RFC 9207 authorization-response issuer when the authorization server + includes it in the redirect; the client validates it against the expected issuer. + """ + + code: str + state: str | None = None + iss: str | None = None + + +class InvalidScopeError(Exception): + def __init__(self, message: str): + self.message = message + + +class InvalidRedirectUriError(Exception): + def __init__(self, message: str): + self.message = message + + +class OAuthClientMetadataBase(BaseModel): + """RFC 7591 OAuth 2.0 Dynamic Client Registration metadata shared verbatim by the + registration request (`OAuthClientMetadata`) and the authorization server's record of a + registered client (`OAuthClientInformationFull`). Fields whose acceptable values differ + between the two - what this SDK sends versus what a third-party server may echo - are + declared on each model rather than here. + See https://datatracker.ietf.org/doc/html/rfc7591#section-2 + """ + + model_config = ConfigDict(url_preserve_empty_path=True) + + # The MCP spec requires the "code" response type, but OAuth + # servers may also return additional types they support + response_types: list[str] = ["code"] + scope: str | None = None + + # these fields are currently unused, but we support & store them for potential + # future use + client_name: str | None = None + client_uri: AnyHttpUrl | None = None + logo_uri: AnyHttpUrl | None = None + contacts: list[str] | None = None + tos_uri: AnyHttpUrl | None = None + policy_uri: AnyHttpUrl | None = None + jwks_uri: AnyHttpUrl | None = None + jwks: Any | None = None + software_id: str | None = None + software_version: str | None = None + + @field_validator( + "client_uri", + "logo_uri", + "tos_uri", + "policy_uri", + "jwks_uri", + mode="before", + ) + @classmethod + def _empty_string_optional_url_to_none(cls, v: object) -> object: + # These URL fields are OPTIONAL; an echoed "" would otherwise fail AnyHttpUrl + # and throw away an otherwise valid registration response. + return _empty_str_to_none(v) + + +class OAuthClientMetadata(OAuthClientMetadataBase): + """RFC 7591 OAuth 2.0 Dynamic Client Registration request metadata: what an MCP + client sends when it registers. Field values are narrowed to what this SDK will put + on the wire; parsing the authorization server's response is `OAuthClientInformationFull`'s + job. See https://datatracker.ietf.org/doc/html/rfc7591#section-2 + """ + + redirect_uris: list[AnyUrl] | None = Field(..., min_length=1) + # supported auth methods for the token endpoint + token_endpoint_auth_method: TokenEndpointAuthMethod | None = None + # supported grant_types of this implementation + grant_types: list[ + Literal["authorization_code", "refresh_token", "urn:ietf:params:oauth:grant-type:jwt-bearer"] | str + ] = list(DEFAULT_GRANT_TYPES) + # SEP-837: OIDC application_type. Defaults to "native" since MCP clients typically use + # loopback redirect URIs; set "web" for remote browser-based clients on a non-local host. + application_type: Literal["web", "native"] = "native" + + +class OAuthClientInformationFull(OAuthClientMetadataBase): + """RFC 7591 OAuth 2.0 Dynamic Client Registration client information response + (client information plus metadata) - the authorization server's record of a + registered client. See https://datatracker.ietf.org/doc/html/rfc7591#section-3.2.1 + + A third-party authorization server "MAY reject or replace any of the client's + requested metadata values submitted during the registration and substitute them with + suitable values", so `application_type`, `token_endpoint_auth_method`, and `grant_types` + are typed to accept any string the server echoes, and `redirect_uris` may be absent or + empty. A member the server serializes as a placeholder - an explicit `null`, or `""` - + is read as an omitted key, so the field's default applies rather than the parse failing. + Whether a substituted value is usable is decided where the value is used, not at parse. + `redirect_uris` elements are still parsed as URLs, as the authorization server compares + them against a client's requested `redirect_uri`. + """ + + redirect_uris: list[AnyUrl] | None = None + # RFC 7591 §3.2.1: the server may assign an auth method other than the one requested, + # including methods this SDK does not implement, or omit it. + token_endpoint_auth_method: str | None = None + grant_types: list[str] = list(DEFAULT_GRANT_TYPES) + # SEP-837: OIDC application_type. OIDC Registration §2 defines "web" and "native", but + # servers echo other strings or an explicit null; the value is informational here. + application_type: str | None = None + + # RFC 7591 §3.2.1: client_id is REQUIRED in a client information response - a body + # without one is not a registration, whatever else it echoes. + client_id: str + client_secret: str | None = None + client_id_issued_at: int | None = None + client_secret_expires_at: int | None = None + # SEP-2352: the issuer these credentials were registered with, recorded by the SDK (not an + # RFC 7591 field) to detect authorization-server migration and avoid cross-AS credential reuse. + issuer: str | None = None + + @model_validator(mode="before") + @classmethod + def _placeholder_members_read_as_omitted(cls, data: object) -> object: + # Servers dump unset members of their client record as null, or echo them as "", + # instead of omitting the keys. Either placeholder would otherwise fail the parse of a + # list field (or read "" as an unrecognized method) and discard an already-provisioned + # registration; a placeholder and an absent key mean the same thing. + if isinstance(data, dict): + members = cast(dict[str, Any], data) + return {key: value for key, value in members.items() if value is not None and value != ""} + return data + + def validate_scope(self, requested_scope: str | None) -> list[str] | None: + if requested_scope is None: + return None + requested_scopes = requested_scope.split(" ") + allowed_scopes = [] if self.scope is None else self.scope.split(" ") + for scope in requested_scopes: + if scope not in allowed_scopes: + raise InvalidScopeError(f"Client was not registered with scope {scope}") + return requested_scopes + + def validate_redirect_uri(self, redirect_uri: AnyUrl | None) -> AnyUrl: + if redirect_uri is not None: + # Validate redirect_uri against client's registered redirect URIs + if not self.redirect_uris or redirect_uri not in self.redirect_uris: + raise InvalidRedirectUriError(f"Redirect URI '{redirect_uri}' not registered for client") + return redirect_uri + elif self.redirect_uris and len(self.redirect_uris) == 1: + return self.redirect_uris[0] + else: + raise InvalidRedirectUriError( + "redirect_uri must be specified unless the client has exactly one registered URI" + ) + + +class OAuthMetadata(BaseModel): + """RFC 8414 OAuth 2.0 Authorization Server Metadata. + See https://datatracker.ietf.org/doc/html/rfc8414#section-2 + """ + + model_config = ConfigDict(url_preserve_empty_path=True) + + issuer: AnyHttpUrl + authorization_endpoint: AnyHttpUrl + token_endpoint: AnyHttpUrl + registration_endpoint: AnyHttpUrl | None = None + scopes_supported: list[str] | None = None + response_types_supported: list[str] = ["code"] + response_modes_supported: list[str] | None = None + grant_types_supported: list[str] | None = None + token_endpoint_auth_methods_supported: list[str] | None = None + token_endpoint_auth_signing_alg_values_supported: list[str] | None = None + service_documentation: AnyHttpUrl | None = None + ui_locales_supported: list[str] | None = None + op_policy_uri: AnyHttpUrl | None = None + op_tos_uri: AnyHttpUrl | None = None + revocation_endpoint: AnyHttpUrl | None = None + revocation_endpoint_auth_methods_supported: list[str] | None = None + revocation_endpoint_auth_signing_alg_values_supported: list[str] | None = None + introspection_endpoint: AnyHttpUrl | None = None + introspection_endpoint_auth_methods_supported: list[str] | None = None + introspection_endpoint_auth_signing_alg_values_supported: list[str] | None = None + code_challenge_methods_supported: list[str] | None = None + client_id_metadata_document_supported: bool | None = None + authorization_response_iss_parameter_supported: bool | None = None + # SEP-990 / draft-ietf-oauth-identity-assertion-authz-grant §7.2: profiles whose grants the + # authorization server supports, e.g. `urn:ietf:params:oauth:grant-profile:id-jag`. + authorization_grant_profiles_supported: list[str] | None = None + + +class ProtectedResourceMetadata(BaseModel): + """RFC 9728 OAuth 2.0 Protected Resource Metadata. + See https://datatracker.ietf.org/doc/html/rfc9728#section-2 + """ + + model_config = ConfigDict(url_preserve_empty_path=True) + + resource: AnyHttpUrl + authorization_servers: list[AnyHttpUrl] = Field(..., min_length=1) + jwks_uri: AnyHttpUrl | None = None + scopes_supported: list[str] | None = None + bearer_methods_supported: list[str] | None = Field(default=["header"]) # MCP only supports header method + resource_signing_alg_values_supported: list[str] | None = None + resource_name: str | None = None + resource_documentation: AnyHttpUrl | None = None + resource_policy_uri: AnyHttpUrl | None = None + resource_tos_uri: AnyHttpUrl | None = None + # tls_client_certificate_bound_access_tokens default is False, but omitted here for clarity + tls_client_certificate_bound_access_tokens: bool | None = None + authorization_details_types_supported: list[str] | None = None + dpop_signing_alg_values_supported: list[str] | None = None + # dpop_bound_access_tokens_required default is False, but omitted here for clarity + dpop_bound_access_tokens_required: bool | None = None diff --git a/src/mcp-client/mcp_client/shared/auth_utils.py b/src/mcp-client/mcp_client/shared/auth_utils.py new file mode 100644 index 0000000000..3ba880f40d --- /dev/null +++ b/src/mcp-client/mcp_client/shared/auth_utils.py @@ -0,0 +1,80 @@ +"""Utilities for OAuth 2.0 Resource Indicators (RFC 8707) and PKCE (RFC 7636).""" + +import time +from urllib.parse import urlparse, urlsplit, urlunsplit + +from pydantic import AnyUrl, HttpUrl + + +def resource_url_from_server_url(url: str | HttpUrl | AnyUrl) -> str: + """Convert server URL to canonical resource URL per RFC 8707. + + RFC 8707 section 2 states that resource URIs "MUST NOT include a fragment component". + Returns absolute URI with lowercase scheme/host for canonical form. + + Args: + url: Server URL to convert + + Returns: + Canonical resource URL string + """ + # Convert to string if needed + url_str = str(url) + + # Parse the URL and remove fragment, create canonical form + parsed = urlsplit(url_str) + canonical = urlunsplit(parsed._replace(scheme=parsed.scheme.lower(), netloc=parsed.netloc.lower(), fragment="")) + + return canonical + + +def check_resource_allowed(requested_resource: str, configured_resource: str) -> bool: + """Check if a requested resource URL matches a configured resource URL. + + A requested resource matches if it has the same scheme, domain, port, + and its path starts with the configured resource's path. This allows + hierarchical matching where a token for a parent resource can be used + for child resources. + + Args: + requested_resource: The resource URL being requested + configured_resource: The resource URL that has been configured + + Returns: + True if the requested resource matches the configured resource + """ + # Parse both URLs + requested = urlparse(requested_resource) + configured = urlparse(configured_resource) + + # Compare scheme, host, and port (origin) + if requested.scheme.lower() != configured.scheme.lower() or requested.netloc.lower() != configured.netloc.lower(): + return False + + # Normalize trailing slashes before comparison so that + # "/foo" and "/foo/" are treated as equivalent. + requested_path = requested.path + configured_path = configured.path + if not requested_path.endswith("/"): + requested_path += "/" + if not configured_path.endswith("/"): + configured_path += "/" + + # Check hierarchical match: requested must start with configured path. + # The trailing-slash normalization ensures "/api123/" won't match "/api/". + return requested_path.startswith(configured_path) + + +def calculate_token_expiry(expires_in: int | str | None) -> float | None: + """Calculate token expiry timestamp from expires_in seconds. + + Args: + expires_in: Seconds until token expiration (may be string from some servers) + + Returns: + Unix timestamp when token expires, or None if no expiry specified + """ + if expires_in is None: + return None # pragma: no cover + # Defensive: handle servers that return expires_in as string + return time.time() + int(expires_in) diff --git a/src/mcp-client/mcp_client/shared/context.py b/src/mcp-client/mcp_client/shared/context.py new file mode 100644 index 0000000000..a86ecbadee --- /dev/null +++ b/src/mcp-client/mcp_client/shared/context.py @@ -0,0 +1,85 @@ +"""`BaseContext` - the user-facing per-request context. + +Composition over a `DispatchContext`: forwards the transport metadata, the +back-channel (`send_raw_request`/`notify`), progress reporting, and the cancel +event. Adds `meta` (the inbound request's `_meta` field). + +Satisfies `Outbound`, so `ClientPeer` can wrap it. Shared between client and +server: the server's `Context` extends this with `lifespan`/`connection`; +`ClientContext` is just an alias. +""" + +from collections.abc import Mapping +from typing import Any, Generic + +import anyio +from mcp_types import RequestParamsMeta +from typing_extensions import TypeVar + +from mcp_client.shared.dispatcher import CallOptions, DispatchContext +from mcp_client.shared.transport_context import TransportContext + +__all__ = ["BaseContext"] + +TransportT = TypeVar("TransportT", bound=TransportContext, default=TransportContext, covariant=True) + + +class BaseContext(Generic[TransportT]): + """Per-request context wrapping a `DispatchContext`. + + `ServerRunner` constructs one per inbound request and passes it to the + user's handler. + """ + + def __init__(self, dctx: DispatchContext[TransportT], meta: RequestParamsMeta | None = None) -> None: + self._dctx = dctx + self._meta = meta + + @property + def transport(self) -> TransportT: + """Transport-specific metadata for this inbound request.""" + return self._dctx.transport + + @property + def cancel_requested(self) -> anyio.Event: + """Set when the peer sends `notifications/cancelled` for this request.""" + return self._dctx.cancel_requested + + @property + def can_send_request(self) -> bool: + """Whether the back-channel can currently deliver server-initiated requests. + + `False` when the transport has no back-channel, or when the underlying + dispatch context has been closed because the inbound request finished. + """ + return self._dctx.can_send_request + + @property + def meta(self) -> RequestParamsMeta | None: + """The inbound request's `_meta` field, if present.""" + return self._meta + + async def send_raw_request( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None = None, + ) -> dict[str, Any]: + """Send a request to the peer on the back-channel. + + Raises: + MCPError: The peer responded with an error. + NoBackChannelError: `can_send_request` is `False`. + """ + return await self._dctx.send_raw_request(method, params, opts) + + async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: + """Send a notification to the peer on the back-channel.""" + await self._dctx.notify(method, params, opts) + + async def report_progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: + """Report progress for this request, if the peer supplied a progress token. + + A no-op when no token was supplied. + """ + await self._dctx.progress(progress, total, message) diff --git a/src/mcp-client/mcp_client/shared/direct_dispatcher.py b/src/mcp-client/mcp_client/shared/direct_dispatcher.py new file mode 100644 index 0000000000..141cd0c0ae --- /dev/null +++ b/src/mcp-client/mcp_client/shared/direct_dispatcher.py @@ -0,0 +1,334 @@ +"""In-memory `Dispatcher` that wires two peers together with no transport. + +`DirectDispatcher` is the simplest possible `Dispatcher` implementation: a +request on one side directly invokes the other side's `on_request`. There is no +serialization, no JSON-RPC framing, and no streams. It exists to: + +* prove the `Dispatcher` Protocol is implementable without JSON-RPC +* provide a fast substrate for testing the layers above the dispatcher + (`ServerRunner`, `Context`, `Connection`) without wire-level moving parts +* embed a server in-process when the JSON-RPC overhead is unnecessary + +Like `JSONRPCDispatcher`, this is an exception-to-error boundary: a handler +exception surfaces to the caller as `MCPError`. The `raise_handler_exceptions` +knob controls whether unmapped exceptions are sanitized (matching the wire +path) or chained as ``__cause__`` for in-process debugging. +""" + +from __future__ import annotations + +import logging +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field +from typing import Any + +import anyio +import anyio.abc +from mcp_types import CONNECTION_CLOSED, INTERNAL_ERROR, INVALID_PARAMS, REQUEST_TIMEOUT, RequestId +from pydantic import ValidationError + +from mcp_client.shared._compat import resync_tracer +from mcp_client.shared.dispatcher import ( + CallOptions, + OnNotify, + OnNotifyIntercept, + OnRequest, + ProgressFnT, + coerce_request_id, + run_notify_intercept, +) +from mcp_client.shared.exceptions import MCPError, NoBackChannelError +from mcp_client.shared.message import MessageMetadata +from mcp_client.shared.transport_context import TransportContext + +logger = logging.getLogger("mcp.shared.direct_dispatcher") + +__all__ = ["DirectDispatcher", "create_direct_dispatcher_pair"] + +DIRECT_TRANSPORT_KIND = "direct" + + +_Request = Callable[[str, Mapping[str, Any] | None, CallOptions | None], Awaitable[dict[str, Any]]] +_Notify = Callable[[str, Mapping[str, Any] | None], Awaitable[None]] + + +@dataclass +class _DirectDispatchContext: + """`DispatchContext` for an inbound request on a `DirectDispatcher`. + + The back-channel callables target the *originating* side, so a handler's + `send_raw_request` reaches the peer that made the inbound request. + """ + + transport: TransportContext + _back_request: _Request + _back_notify: _Notify + request_id: RequestId | None = None + """The caller-supplied `CallOptions["request_id"]`, else a dispatcher-synthesized + id for requests; `None` for notifications.""" + message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework + """Always `None`: in-memory dispatch attaches no transport metadata.""" + _on_progress: ProgressFnT | None = None + cancel_requested: anyio.Event = field(default_factory=anyio.Event) + + @property + def can_send_request(self) -> bool: + return self.transport.can_send_request + + async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: + await self._back_notify(method, params) + + async def send_raw_request( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None = None, + ) -> dict[str, Any]: + if not self.can_send_request: + raise NoBackChannelError(method) + return await self._back_request(method, params, opts) + + async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: + if self._on_progress is not None: + await self._on_progress(progress, total, message) + + +class DirectDispatcher: + """A `Dispatcher` that calls a peer's handlers directly, in-process. + + Two instances are wired together with `create_direct_dispatcher_pair`; each + holds a reference to the other. `send_raw_request` on one awaits the peer's + `on_request`. `run` parks until `close` is called. + + Lifecycle mirrors `JSONRPCDispatcher`: `send_raw_request` requires `run()` + to have started, and once a side has closed - via `close()` or `run()` + ending - `send_raw_request` raises `MCPError` (`CONNECTION_CLOSED`) and + inbound requests fail the peer's call the same way instead of invoking the + handler. Notifications are fire-and-forget in both directions: after close + they are silently dropped. + """ + + def __init__(self, transport_ctx: TransportContext, *, raise_handler_exceptions: bool = True): + self._transport_ctx = transport_ctx + self._raise_handler_exceptions = raise_handler_exceptions + self._peer: DirectDispatcher | None = None + self._on_request: OnRequest | None = None + self._on_notify: OnNotify | None = None + self._on_notify_intercept: OnNotifyIntercept | None = None + self._next_id = 0 + self._in_flight_ids: set[RequestId] = set() + self._ready = anyio.Event() + self._close_event = anyio.Event() + self._running = False + self._closed = False + + def connect_to(self, peer: DirectDispatcher) -> None: + self._peer = peer + + async def send_raw_request( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None = None, + ) -> dict[str, Any]: + """Send a request by invoking the peer's `on_request` directly. + + Raises: + MCPError: The peer's handler raised; `REQUEST_TIMEOUT` if + `opts["timeout"]` elapsed; `CONNECTION_CLOSED` if either + side has closed. + RuntimeError: Called before `run()`. + """ + if self._peer is None: + raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()") + # Post-close sends get the same CONNECTION_CLOSED contract as JSONRPCDispatcher. + if self._closed: + raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") + if not self._running: + raise RuntimeError("DirectDispatcher.send_raw_request called before run()") + return await self._peer._dispatch_request(method, params, opts) + + async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: + """Send a notification by invoking the peer's `on_notify` directly. + + Fire-and-forget: usable before `run()` (delivery waits for the peer to + start), and after close it is silently dropped, matching + `JSONRPCDispatcher.notify`. `opts` is accepted for `Dispatcher` + conformance; there is no HTTP layer here so `headers` is ignored. + """ + if self._peer is None: + raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()") + if self._closed: + logger.debug("dropped notification %r on closed DirectDispatcher", method) + return + await self._peer._dispatch_notify(method, params) + + async def run( + self, + on_request: OnRequest, + on_notify: OnNotify, + on_notify_intercept: OnNotifyIntercept | None = None, + *, + task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, + ) -> None: + """Mark this side ready and park until `close()` is called. + + Single-shot, like `JSONRPCDispatcher.run`: once it returns the + dispatcher stays closed and cannot be restarted. + """ + try: + self._on_request = on_request + self._on_notify = on_notify + self._on_notify_intercept = on_notify_intercept + self._running = True + self._ready.set() + task_status.started() + await self._close_event.wait() + finally: + self._running = False + self._closed = True + # run() may end via cancellation without close() ever being + # called; setting the event wakes `_wait_ready` waiters so they + # observe the closed state instead of parking forever. + self._close_event.set() + + def close(self) -> None: + self._closed = True + self._close_event.set() + + def _make_context( + self, on_progress: ProgressFnT | None = None, request_id: RequestId | None = None + ) -> _DirectDispatchContext: + assert self._peer is not None + peer = self._peer + return _DirectDispatchContext( + transport=self._transport_ctx, + _back_request=lambda m, p, o: peer._dispatch_request(m, p, o), + _back_notify=lambda m, p: peer._dispatch_notify(m, p), + request_id=request_id, + _on_progress=on_progress, + ) + + async def _wait_ready(self) -> None: + """Park until `run()` has started, waking early if this side closes. + + Raises: + MCPError: `CONNECTION_CLOSED` if this side has closed. + """ + if not self._ready.is_set() and not self._close_event.is_set(): + async with anyio.create_task_group() as tg: + + async def wake_on(event: anyio.Event) -> None: + await event.wait() + tg.cancel_scope.cancel() + + tg.start_soon(wake_on, self._ready) + tg.start_soon(wake_on, self._close_event) + if self._closed: + raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") + + async def _dispatch_request( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None, + ) -> dict[str, Any]: + opts = opts or {} + try: + with anyio.fail_after(opts.get("timeout")): + # Inside the timeout scope, so a configured timeout also bounds + # waiting on a peer whose run() has not started yet. + await self._wait_ready() + assert self._on_request is not None + supplied_id = opts.get("request_id") + if supplied_id is not None: + request_id: RequestId = supplied_id + # Collisions use the same coerced domain as JSONRPCDispatcher's + # pending keys, so this in-memory stand-in raises for exactly + # the ids the wire dispatcher would; the context still sees + # the verbatim value. + in_flight_key = coerce_request_id(request_id) + if in_flight_key in self._in_flight_ids: + raise ValueError(f"request id {request_id!r} is already in flight") + else: + # Synthesize an id (the DispatchContext contract reserves None + # for notifications), minting past any key a supplied id + # occupies: the collision error is reserved for the caller + # who actually chose the id. + self._next_id += 1 + while self._next_id in self._in_flight_ids: + self._next_id += 1 + request_id = self._next_id + in_flight_key = request_id + self._in_flight_ids.add(in_flight_key) + dctx = self._make_context(on_progress=opts.get("on_progress"), request_id=request_id) + try: + return await self._on_request(dctx, method, params) + except MCPError: + raise + except ValidationError as e: + # Same shape JSONRPCDispatcher writes, so runner-over-direct + # tests see what runner-over-JSONRPC would. + raise MCPError(code=INVALID_PARAMS, message="Invalid request parameters", data="") from e + except Exception as e: + # Single owner of the in-proc exception-to-error policy (mirrors + # JSONRPCDispatcher / `_streamable_http_modern._to_jsonrpc_response` + # for the wire paths). True chains the original for in-process + # debugging; False sanitizes to match the wire path's leak guard. + if self._raise_handler_exceptions: + raise MCPError(code=INTERNAL_ERROR, message=str(e)) from e + logger.exception("request handler raised") + raise MCPError(code=INTERNAL_ERROR, message="Internal server error") from None + finally: + self._in_flight_ids.discard(in_flight_key) + except TimeoutError: + raise MCPError( + code=REQUEST_TIMEOUT, + message=f"Timed out after {opts.get('timeout')}s waiting for {method!r}", + ) from None + finally: + await resync_tracer() + + async def _dispatch_notify(self, method: str, params: Mapping[str, Any] | None) -> None: + try: + await self._wait_ready() + except MCPError: + # Notifications are fire-and-forget: a notify to a closed peer is + # dropped, not raised back into the sender's call. + logger.debug("dropped notification %r to closed DirectDispatcher", method) + return + if run_notify_intercept(self._on_notify_intercept, method, params): + return + assert self._on_notify is not None + dctx = self._make_context() + await self._on_notify(dctx, method, params) + + +def create_direct_dispatcher_pair( + *, + can_send_request: bool = True, + headers: Mapping[str, str] | None = None, + raise_handler_exceptions: bool = True, +) -> tuple[DirectDispatcher, DirectDispatcher]: + """Create two `DirectDispatcher` instances wired to each other. + + Args: + can_send_request: Sets `TransportContext.can_send_request` on both + sides. Pass `False` to simulate a transport with no back-channel. + headers: Sets `TransportContext.headers` on both sides. + raise_handler_exceptions: When `True` (the default - this is an + in-process debugging substrate), an unmapped handler exception + reaches the caller as `MCPError` with the original chained as + ``__cause__``. When `False` it is sanitized to an opaque + `INTERNAL_ERROR` so the in-process path matches the wire. + + Returns: + A `(client, server)` pair. The wiring is symmetric, so the roles + are conventional only. + """ + ctx = TransportContext(kind=DIRECT_TRANSPORT_KIND, can_send_request=can_send_request, headers=headers) + client = DirectDispatcher(ctx, raise_handler_exceptions=raise_handler_exceptions) + server = DirectDispatcher(ctx, raise_handler_exceptions=raise_handler_exceptions) + client.connect_to(server) + server.connect_to(client) + return client, server diff --git a/src/mcp-client/mcp_client/shared/dispatcher.py b/src/mcp-client/mcp_client/shared/dispatcher.py new file mode 100644 index 0000000000..113bbb1145 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/dispatcher.py @@ -0,0 +1,277 @@ +"""Dispatcher Protocol - the call/return boundary between transports and handlers. + +A Dispatcher turns a duplex message channel into two things: + +* an outbound API: `send_raw_request(method, params)` and `notify(method, params)` +* an inbound pump: `run(on_request, on_notify)` that drives the receive loop + and invokes the supplied handlers for each incoming request/notification + +It is deliberately *not* MCP-aware. Method names are strings, params and +results are `dict[str, Any]`. The MCP type layer (request/result models, +capability negotiation, `Context`) sits above this; the wire encoding +(JSON-RPC, gRPC, in-process direct calls) sits below it. + +See `JSONRPCDispatcher` for the production implementation and +`DirectDispatcher` for an in-memory implementation used in tests and for +embedding a server in-process. +""" + +import logging +from collections.abc import Awaitable, Callable, Mapping +from typing import Any, Protocol, TypedDict, TypeVar, runtime_checkable + +import anyio +import anyio.abc +from mcp_types import RequestId + +from mcp_client.shared.message import MessageMetadata +from mcp_client.shared.transport_context import TransportContext + +logger = logging.getLogger("mcp.shared.dispatcher") + +__all__ = [ + "CallOptions", + "DispatchContext", + "Dispatcher", + "OnNotify", + "OnNotifyIntercept", + "OnRequest", + "Outbound", + "ProgressFnT", + "as_request_id", + "coerce_request_id", + "run_notify_intercept", +] + +TransportT_co = TypeVar("TransportT_co", bound=TransportContext, covariant=True) + + +def as_request_id(value: object) -> RequestId | None: + """Narrow an untyped wire value to a `RequestId`, or None; rejects bool (True would alias request id 1).""" + if isinstance(value, str | int) and not isinstance(value, bool): + return value + return None + + +def coerce_request_id(request_id: RequestId) -> RequestId: + """Coerce a stringified int request id back to int so a peer-echoed id still correlates (matches the TS SDK). + + This is the collision/correlation domain dispatchers share: "7" and 7 are one + id for correlation purposes, even where the wire carries the verbatim value. + """ + if isinstance(request_id, str): + try: + return int(request_id) + except ValueError: + pass + return request_id + + +class ProgressFnT(Protocol): + """Callback invoked when a progress notification arrives for a pending request.""" + + async def __call__(self, progress: float, total: float | None, message: str | None) -> None: ... + + +class CallOptions(TypedDict, total=False): + """Per-call options for `Outbound.send_raw_request`. + + All keys are optional. Dispatchers ignore keys they do not understand. + """ + + request_id: RequestId + """Send the request under this caller-supplied id instead of a dispatcher-minted one. + + The peer sees the value verbatim ("7" stays a string). A value that collides + with one of the sender's own in-flight request ids raises `ValueError`. + Callers that need to know a request's id before its result arrives (a + `subscriptions/listen` stream is demultiplexed by it) mint their own ids + here; string ids that don't parse as integers can never collide with the + dispatcher's minted sequence. Per the class contract, dispatchers that + predate this key ignore it and mint as usual. + """ + + timeout: float + """Seconds to wait for a result before raising and sending `notifications/cancelled`.""" + + cancel_on_abandon: bool + """Whether abandoning this request (timeout or caller cancellation) sends `notifications/cancelled`. + + Defaults to `True`. Set `False` for requests the protocol forbids cancelling, such as `initialize`. + Also suppressed when resumption hints reach the transport, or when the request was never written. + """ + + on_progress: ProgressFnT + """Receive `notifications/progress` updates for this request.""" + + resumption_token: str + """Opaque token to resume a previously interrupted request. + + Client-side, streamable-HTTP only. Ignored by server dispatchers and other + transports, and also ignored (with a debug log) for requests sent from a + `DispatchContext`, where routing onto the inbound request's stream takes + precedence. Supports protocol version 2025-11-25 and earlier; SSE-stream + resumption is removed in the next protocol revision. + """ + + on_resumption_token: Callable[[str], Awaitable[None]] + """Receive a resumption token when the transport issues one for this request. + + Client-side, streamable-HTTP only. Ignored by server dispatchers and other + transports, and also ignored (with a debug log) for requests sent from a + `DispatchContext`, where routing onto the inbound request's stream takes + precedence. Supports protocol version 2025-11-25 and earlier; SSE-stream + resumption is removed in the next protocol revision. + """ + + headers: dict[str, str] + """Transport-layer hint: HTTP transports merge these onto the outgoing request; non-HTTP transports ignore.""" + + +@runtime_checkable +class Outbound(Protocol): + """Anything that can send requests and notifications to the peer. + + Both `Dispatcher` (top-level outbound) and `DispatchContext` (back-channel + during an inbound request) extend this. The MCP type layer (`ClientPeer`, + `Connection`) builds typed `send_request` / convenience methods on top of + this raw channel. + """ + + async def send_raw_request( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None = None, + ) -> dict[str, Any]: + """Send a request and await its raw result dict. + + Raises: + MCPError: If the peer responded with an error, or the handler + raised. Implementations normalize all handler exceptions to + `MCPError` so callers see a single exception type. + """ + ... + + async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: + """Send a fire-and-forget notification.""" + ... + + +class DispatchContext(Outbound, Protocol[TransportT_co]): + """Per-request context handed to `on_request` / `on_notify`. + + Carries the transport metadata for the inbound message and provides the + back-channel for sending requests/notifications to the peer while handling + it. `send_raw_request` raises `NoBackChannelError` if `can_send_request` + is `False`. + """ + + @property + def transport(self) -> TransportT_co: + """Transport-specific metadata for this inbound message.""" + ... + + @property + def can_send_request(self) -> bool: + """Whether the back-channel can currently deliver server-initiated requests. + + `False` when the transport has no back-channel, or when this context has + been closed (the inbound request finished). `send_raw_request` raises + `NoBackChannelError` exactly when this is `False`. + """ + ... + + @property + def request_id(self) -> RequestId | None: + """The id of the inbound request, or `None` for a notification. + + For JSON-RPC this is the wire `id` field. Handlers thread it through + as `related_request_id` on outbound notifications so HTTP transports + can route them onto the originating request's response stream. + """ + ... + + @property + def message_metadata(self) -> MessageMetadata: + """The metadata the transport attached to this inbound message, if any. + + This is `SessionMessage.metadata` passed through verbatim: HTTP + transports attach `ServerMessageMetadata` (the HTTP request, SSE + stream-close callbacks); stdio and in-memory dispatch attach nothing. + Tied to the `SessionMessage` wire format - goes away when transports + stop delivering messages that way. + """ + # TODO(maxisbey): remove for context rework + ... + + @property + def cancel_requested(self) -> anyio.Event: + """Set when the peer sends `notifications/cancelled` for this request.""" + ... + + async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: + """Report progress for the inbound request, if the peer supplied a progress token. + + A no-op when no token was supplied. + """ + ... + + +OnRequest = Callable[[DispatchContext[TransportContext], str, Mapping[str, Any] | None], Awaitable[dict[str, Any]]] +"""Handler for inbound requests: `(ctx, method, params) -> result`. Raise `MCPError` to send an error response.""" + +OnNotify = Callable[[DispatchContext[TransportContext], str, Mapping[str, Any] | None], Awaitable[None]] +"""Handler for inbound notifications: `(ctx, method, params)`.""" + +OnNotifyIntercept = Callable[[str, Mapping[str, Any] | None], bool] +"""Synchronous receive-order intercept for inbound notifications: `(method, params) -> consumed`. + +Runs before `on_notify` is scheduled so correlation state advances in wire order +relative to response resolution (the client's listen demux depends on this). +Returning True consumes the notification. Must not block the receive path. +""" + + +def run_notify_intercept(intercept: OnNotifyIntercept | None, method: str, params: Mapping[str, Any] | None) -> bool: + """Invoke `intercept`, containing a raise to that one notification (never the receive loop).""" + if intercept is None: + return False + try: + return intercept(method, params) + except Exception: + logger.exception("notification intercept raised; passing %r through", method) + return False + + +class Dispatcher(Outbound, Protocol[TransportT_co]): + """A duplex request/notification channel with call-return semantics. + + Implementations own correlation of outbound requests to inbound results, the + receive loop, per-request concurrency, and cancellation/progress wiring. + + The lifecycle surface is provisional; `run()` may change in a 2.x minor + release. + """ + + async def run( + self, + on_request: OnRequest, + on_notify: OnNotify, + on_notify_intercept: OnNotifyIntercept | None = None, + *, + task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, + ) -> None: + """Drive the receive loop until the underlying channel closes. + + Each inbound request is dispatched to `on_request` in its own task; + the returned dict (or raised `MCPError`) is sent back as the response. + Implementations MUST offer every inbound notification to + `on_notify_intercept` synchronously in receive order (via + `run_notify_intercept`), handing only unconsumed ones to `on_notify`. + + `task_status.started()` is called once the dispatcher is ready to + accept `send_request`/`notify` calls, so callers can use + `await tg.start(dispatcher.run, on_request, on_notify)`. + """ + ... diff --git a/src/mcp-client/mcp_client/shared/exceptions.py b/src/mcp-client/mcp_client/shared/exceptions.py new file mode 100644 index 0000000000..c2a7fd44e7 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/exceptions.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +from typing import Any, cast + +from mcp_types import INVALID_REQUEST, URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError + + +class MCPDeprecationWarning(UserWarning): + """A custom deprecation warning for the MCP SDK. + + Unlike the built-in `DeprecationWarning`, this inherits from `UserWarning` so + it is shown by default, helping users discover deprecated features without + enabling warnings explicitly. + + Reference: https://sethmlarson.dev/deprecations-via-warnings-dont-work-for-python-libraries + """ + + +class MCPError(Exception): + """Exception type raised when an error arrives over an MCP connection.""" + + error: ErrorData + + def __init__(self, code: int, message: str, data: Any = None): + super().__init__(code, message, data) + if data is not None: + self.error = ErrorData(code=code, message=message, data=data) + else: + self.error = ErrorData(code=code, message=message) + + @property + def code(self) -> int: + return self.error.code + + @property + def message(self) -> str: + return self.error.message + + @property + def data(self) -> Any: + return self.error.data + + @classmethod + def from_jsonrpc_error(cls, error: JSONRPCError) -> MCPError: + return cls.from_error_data(error.error) + + @classmethod + def from_error_data(cls, error: ErrorData) -> MCPError: + return cls(code=error.code, message=error.message, data=error.data) + + def __str__(self) -> str: + return self.message + + +class NoBackChannelError(MCPError): + """Raised when a server-initiated request has no channel that can deliver it. + + Raised by `DispatchContext.send_raw_request` when its request-scoped channel + reports `TransportContext.can_send_request` as `False` (the cases are + documented on that field), and by a connection's standalone channel when it + has none; serializes to an `INVALID_REQUEST` error response. + """ + + def __init__(self, method: str): + super().__init__( + code=INVALID_REQUEST, + message=( + f"Cannot send {method!r}: this transport context has no back-channel for server-initiated requests." + ), + ) + self.method = method + + +class UrlElicitationRequiredError(MCPError): + """Specialized error for when a tool requires URL mode elicitation(s) before proceeding. + + Servers can raise this error from tool handlers to indicate that the client + must complete one or more URL elicitations before the request can be processed. + + Example: + ```python + raise UrlElicitationRequiredError([ + ElicitRequestURLParams( + message="Authorization required for your files", + url="https://example.com/oauth/authorize", + elicitation_id="auth-001" + ) + ]) + ``` + """ + + def __init__(self, elicitations: list[ElicitRequestURLParams], message: str | None = None): + """Initialize UrlElicitationRequiredError.""" + if message is None: + message = f"URL elicitation{'s' if len(elicitations) > 1 else ''} required" + + self._elicitations = elicitations + + super().__init__( + code=URL_ELICITATION_REQUIRED, + message=message, + data={"elicitations": [e.model_dump(by_alias=True, exclude_none=True) for e in elicitations]}, + ) + + @property + def elicitations(self) -> list[ElicitRequestURLParams]: + """The list of URL elicitations required before the request can proceed.""" + return self._elicitations + + @classmethod + def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError: + """Reconstruct from an ErrorData received over the wire.""" + if error.code != URL_ELICITATION_REQUIRED: + raise ValueError(f"Expected error code {URL_ELICITATION_REQUIRED}, got {error.code}") + + data = cast(dict[str, Any], error.data or {}) + raw_elicitations = cast(list[dict[str, Any]], data.get("elicitations", [])) + elicitations = [ElicitRequestURLParams.model_validate(e) for e in raw_elicitations] + return cls(elicitations, error.message) diff --git a/src/mcp-client/mcp_client/shared/extension.py b/src/mcp-client/mcp_client/shared/extension.py new file mode 100644 index 0000000000..283e9ba89b --- /dev/null +++ b/src/mcp-client/mcp_client/shared/extension.py @@ -0,0 +1,28 @@ +"""Extension-identifier grammar shared by the server and client extension surfaces.""" + +from __future__ import annotations + +import re +from typing import Any + +__all__ = ["validate_extension_identifier"] + +# Extension identifiers follow the `_meta` key grammar with a mandatory prefix +# (SEP-2133 / basic/index.mdx): dot-separated labels, each starting with a +# letter and ending with a letter or digit (hyphens interior), then `/`, then a +# name that starts and ends alphanumeric (`.`/`_`/`-` interior). +_LABEL = r"[A-Za-z](?:[A-Za-z0-9-]*[A-Za-z0-9])?" +_NAME = r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?" +_IDENTIFIER_RE = re.compile(rf"{_LABEL}(?:\.{_LABEL})*/{_NAME}") + + +def validate_extension_identifier(identifier: Any, *, owner: str) -> None: + """Raise `TypeError` unless `identifier` is a `vendor-prefix/name` string. + + SEP-2133 requires extension identifiers to carry a reverse-DNS prefix. + """ + if not isinstance(identifier, str) or not _IDENTIFIER_RE.fullmatch(identifier): + raise TypeError( + f"{owner}.identifier must be a `vendor-prefix/name` string " + f"(reverse-DNS prefix required), got {identifier!r}" + ) diff --git a/src/mcp-client/mcp_client/shared/inbound.py b/src/mcp-client/mcp_client/shared/inbound.py new file mode 100644 index 0000000000..e33d6f502f --- /dev/null +++ b/src/mcp-client/mcp_client/shared/inbound.py @@ -0,0 +1,595 @@ +"""Inbound request classification for the modern per-request-envelope path. + +Pure module: no I/O, no transport, no `mcp.server` imports. Runs the +validation ladder against a decoded JSON-RPC body and returns either an +:class:`InboundModernRoute` (every rung passed) or an +:class:`InboundLadderRejection` (the first rung that failed). Callers map a +rejection's `code` through :data:`ERROR_CODE_HTTP_STATUS` to pick the HTTP +status. + +Also hosts the shared header-value codec and the `x-mcp-header` schema +validator so client emit and server validate read the same source of truth. +""" + +import base64 +import binascii +import re +from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Final, cast + +from mcp_types import ( + CLIENT_CAPABILITIES_META_KEY, + CLIENT_INFO_META_KEY, + PROTOCOL_VERSION_META_KEY, + UnsupportedProtocolVersionErrorData, +) +from mcp_types.jsonrpc import ( + HEADER_MISMATCH, + INVALID_PARAMS, + INVALID_REQUEST, + METHOD_NOT_FOUND, + MISSING_REQUIRED_CLIENT_CAPABILITY, + PARSE_ERROR, + UNSUPPORTED_PROTOCOL_VERSION, +) +from mcp_types.version import MODERN_PROTOCOL_VERSIONS + +__all__ = [ + "ERROR_CODE_HTTP_STATUS", + "InboundLadderRejection", + "InboundModernRoute", + "MCP_METHOD_HEADER", + "MCP_NAME_HEADER", + "MCP_PARAM_HEADER_PREFIX", + "MCP_PROTOCOL_VERSION_HEADER", + "NAME_BEARING_METHODS", + "X_MCP_HEADER_KEY", + "classify_inbound_request", + "decode_header_value", + "encode_header_value", + "find_duplicated_routing_header", + "find_invalid_x_mcp_header", + "mcp_param_headers", + "unsupported_protocol_version_rejection", + "validate_mcp_param_headers", + "x_mcp_header_map", +] + +MCP_PROTOCOL_VERSION_HEADER: Final = "mcp-protocol-version" +"""Canonical lowercase name of the HTTP header carrying the MCP protocol version.""" + +MCP_METHOD_HEADER: Final = "mcp-method" +"""Canonical lowercase name of the HTTP header carrying the JSON-RPC method.""" + +MCP_NAME_HEADER: Final = "mcp-name" +"""Canonical lowercase name of the HTTP header carrying the resource name (tool/prompt/resource URI).""" + +X_MCP_HEADER_KEY: Final = "x-mcp-header" +"""JSON-Schema property annotation that designates an `Mcp-Param-*` HTTP header.""" + +NAME_BEARING_METHODS: Final[Mapping[str, str]] = MappingProxyType( + { + "tools/call": "name", + "prompts/get": "name", + "resources/read": "uri", + } +) +"""Method → params key whose value is mirrored as the `Mcp-Name` HTTP header. + +Shared by client emit (which header to send) and server validate (which body +field to compare against), so both ends agree on the field by construction. +""" + +_B64_SENTINEL = re.compile(r"^=\?base64\?(?P.*)\?=$") +# RFC 7230 token chars minus DEL; visible ASCII 0x20-0x7E is the practical bound for a header value. +_HEADER_SAFE = re.compile(r"^[\x20-\x7E]*$") +# RFC 9110 §5.6.2 token: the only characters permitted in an HTTP field name. +_RFC9110_TOKEN = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") +# JSON-Schema types the spec permits to carry `x-mcp-header` (transports.mdx +# §Custom Headers). `number` is explicitly forbidden — float→str is not +# portable across implementations. +_X_MCP_HEADER_PRIMITIVE_TYPES: Final = frozenset({"string", "integer", "boolean"}) + +# JSON Schema 2020-12 applicator keywords whose values are themselves schema +# positions, grouped by value shape. `properties` is handled separately as the +# only keyword that preserves the statically-reachable chain; every keyword +# here drops the chain to None. Instance-data keywords (`default`, `examples`, +# `const`, `enum`) and `$ref`/`$dynamicRef` are deliberately absent so the +# walk never mistakes data for an annotation and never dereferences. +_SUBSCHEMA_SINGLE: Final = frozenset( + { + "items", + "contains", + "unevaluatedItems", + "additionalProperties", + "propertyNames", + "unevaluatedProperties", + "not", + "if", + "then", + "else", + "contentSchema", + } +) +_SUBSCHEMA_LIST: Final = frozenset({"allOf", "anyOf", "oneOf", "prefixItems"}) +_SUBSCHEMA_MAP: Final = frozenset({"patternProperties", "dependentSchemas", "$defs", "definitions"}) + + +def _walk_schema_positions(root: Any) -> Iterator[tuple[tuple[str, ...] | None, dict[str, Any]]]: + """Yield `(properties_path, schema)` for every schema position in `root`. + + `properties_path` is the chain of `properties` keys from the root to the + position, or `None` once any other applicator keyword has been crossed. + The root itself yields `()`. Only the JSON Schema 2020-12 applicators + listed above are entered; instance-data keywords are not, and `$ref` is + not dereferenced, so the walk terminates on any finite JSON value. An + explicit stack keeps the function total even on pathologically deep input. + """ + stack: list[tuple[tuple[str, ...] | None, Any]] = [((), root)] + while stack: + path, node = stack.pop() + if not isinstance(node, dict): + continue + schema = cast(dict[str, Any], node) + yield path, schema + for kw, val in schema.items(): + if kw == "properties" and isinstance(val, dict): + for name, sub in cast(dict[str, Any], val).items(): + stack.append(((*path, name) if path is not None else None, sub)) + elif kw in _SUBSCHEMA_SINGLE: + stack.append((None, val)) + elif kw in _SUBSCHEMA_LIST and isinstance(val, list): + stack.extend((None, sub) for sub in cast(list[Any], val)) + elif kw in _SUBSCHEMA_MAP and isinstance(val, dict): + stack.extend((None, sub) for sub in cast(dict[str, Any], val).values()) + + +def encode_header_value(value: str) -> str: + """Wrap `value` in the `=?base64?...?=` sentinel when it would not survive an HTTP field round-trip. + + Plain printable ASCII without leading/trailing whitespace passes verbatim; + anything else (control chars, non-ASCII, edge whitespace, or a value that + already looks like the sentinel) is base64-wrapped so the receiver can + recover the exact bytes. + """ + if _HEADER_SAFE.fullmatch(value) and value == value.strip() and not _B64_SENTINEL.fullmatch(value): + return value + return f"=?base64?{base64.b64encode(value.encode('utf-8')).decode('ascii')}?=" + + +def decode_header_value(value: str | None) -> str | None: + """Inverse of :func:`encode_header_value`. + + Returns the value verbatim unless it carries the `=?base64?...?=` sentinel, + in which case the payload is decoded as UTF-8. A malformed sentinel (bad + base64, non-canonical base64, or bad UTF-8) yields `None` so a corrupt + header never matches a body value by accident. `None` in → `None` out so + callers can pass `headers.get(...)` directly. + """ + if value is None: + return None + m = _B64_SENTINEL.fullmatch(value) + if m is None: + return value + payload = m.group("payload") + try: + decoded = base64.b64decode(payload, validate=True) + except binascii.Error: + return None + # Reject non-canonical base64 (e.g. non-zero trailing bits), which + # `validate=True` tolerates; the encoder only ever emits canonical form. + if base64.b64encode(decoded).decode("ascii") != payload: + return None + try: + return decoded.decode("utf-8") + except UnicodeDecodeError: + return None + + +def find_invalid_x_mcp_header(input_schema: Any) -> str | None: + """Return a reason string if any `x-mcp-header` annotation in `input_schema` is invalid; else `None`. + + Walks every JSON Schema 2020-12 schema position. An annotation is valid + only when it sits on a property statically reachable from the root via a + chain of pure `properties` keys, names a non-empty RFC 9110 token, is on + an integer/string/boolean property, and is case-insensitively unique + across the whole schema. A `None` / non-mapping schema has no schema + positions and returns `None`. + """ + seen: dict[str, str] = {} + for path, schema in _walk_schema_positions(input_schema): + if X_MCP_HEADER_KEY not in schema: + continue + if not path: # None (off the pure-properties chain) or () (the root itself) + return f"{X_MCP_HEADER_KEY} found at a schema position not reachable via a pure `properties` chain" + where = ".".join(path) + header = schema[X_MCP_HEADER_KEY] + # Wrong type and malformed value are distinct failures with distinct messages: the + # non-str arm returns before any interpolation, because `repr` of an arbitrary + # schema value is not total (a large `int` exceeds `sys.get_int_max_str_digits`). + if not isinstance(header, str): + return f"property {where!r}: {X_MCP_HEADER_KEY} must be a string, not {type(header).__name__}" + if not _RFC9110_TOKEN.fullmatch(header): + return f"property {where!r}: {X_MCP_HEADER_KEY} {header!r} is not an RFC 9110 token" + prop_type = schema.get("type") + if not isinstance(prop_type, str): + return ( + f"property {where!r}: {X_MCP_HEADER_KEY} is only permitted on " + f"integer/string/boolean properties (the type keyword is {type(prop_type).__name__}, not a string)" + ) + if prop_type not in _X_MCP_HEADER_PRIMITIVE_TYPES: + return ( + f"property {where!r}: {X_MCP_HEADER_KEY} is only permitted on " + f"integer/string/boolean properties (got {prop_type!r})" + ) + lower = header.lower() + if lower in seen: + return f"{X_MCP_HEADER_KEY} {header!r} on property {where!r} duplicates property {seen[lower]!r}" + seen[lower] = where + return None + + +MCP_PARAM_HEADER_PREFIX: Final = "Mcp-Param-" +"""Prefix the `x-mcp-header` token is joined to, forming the per-parameter HTTP header name.""" + + +def x_mcp_header_map(input_schema: Any) -> dict[tuple[str, ...], str]: + """Map each property carrying a valid `x-mcp-header` to its annotation token, keyed by property path. + + The key is the chain of `properties` keys from the schema root to the + annotated property; a top-level property has a one-element path, a nested + one a longer path. Call only on a schema that + :func:`find_invalid_x_mcp_header` accepts; an invalid schema yields an + undefined subset. + """ + return {path: token for path, token, _ in _annotated_positions(input_schema)} + + +def _annotated_positions(input_schema: Any) -> Iterator[tuple[tuple[str, ...], str, dict[str, Any]]]: + """Yield `(path, token, schema)` for every statically-reachable `x-mcp-header` annotation. + + Shared by client emit and server validate so both ends agree on what counts as a declared header. + """ + for path, schema in _walk_schema_positions(input_schema): + if path and isinstance(token := schema.get(X_MCP_HEADER_KEY), str): + yield path, token, schema + + +def _render_header_scalar(value: Any) -> str | None: + """Render `value` the way the client mirrors it into a header, or `None` when no rendering exists. + + Shared by emit and validate so both sides agree on what is mirrorable: + non-primitives and ints beyond CPython's int-to-str digit limit are not. + """ + if isinstance(value, bool): + return "true" if value else "false" + if not isinstance(value, str | int | float): + return None + try: + return str(value) + except ValueError: + return None + + +def mcp_param_headers(header_map: Mapping[tuple[str, ...], str], arguments: Mapping[str, Any]) -> dict[str, str]: + """Build the `Mcp-Param-*` headers a `tools/call` mirrors from its arguments. + + For each `(path, token)` in `header_map`, read the value at that property + path in `arguments` and, when it is present and not `None`, emit + `Mcp-Param-` carrying it: `bool` as `true`/`false`, other scalars via + `str`, each passed through :func:`encode_header_value` so a non-token value + is base64-wrapped. A path that hits a missing key or a non-mapping node is + skipped, matching the spec's "omit the header when no value is present", + as is a value with no header rendering. + """ + headers: dict[str, str] = {} + for path, token in header_map.items(): + value = _value_at_path(arguments, path) + if value is None or (rendered := _render_header_scalar(value)) is None: + continue + headers[f"{MCP_PARAM_HEADER_PREFIX}{token}"] = encode_header_value(rendered) + return headers + + +def _value_at_path(arguments: Mapping[str, Any], path: tuple[str, ...]) -> Any: + """Read the value at a `properties`-key path in `arguments`, or `None` if any step is missing or non-mapping.""" + node: Any = arguments + for key in path: + if not isinstance(node, Mapping): + return None + node = cast("Mapping[str, Any]", node).get(key) + return node + + +# INTERNAL_ERROR is deliberately unmapped (→ HTTP 200): the spec assigns no status to +# -32603, and whether handler-origin errors get 5xx is an open S4 question — see TODO(L66). +ERROR_CODE_HTTP_STATUS: Final[Mapping[int, int]] = MappingProxyType( + { + PARSE_ERROR: 400, + INVALID_REQUEST: 400, + INVALID_PARAMS: 400, + HEADER_MISMATCH: 400, + MISSING_REQUIRED_CLIENT_CAPABILITY: 400, + UNSUPPORTED_PROTOCOL_VERSION: 400, + METHOD_NOT_FOUND: 404, + } +) +"""HTTP status to send for a JSON-RPC `error.code`. + +Consulted for classifier-origin *and* handler-origin errors, so one table +decides the wire status regardless of where the error was produced. Unmapped +codes fall back to the caller's default (typically 200). +""" + + +@dataclass(frozen=True) +class InboundModernRoute: + """A modern-protocol request whose envelope passed every ladder rung. + + `client_info` and `client_capabilities` are the raw envelope values; the + classifier checks presence only, not shape, and `client_info` is `None` + when the (optional, SHOULD-include) key is absent. Method existence is not + a ladder rung — kernel dispatch is the single source of truth for that. + """ + + protocol_version: str + client_info: Any + client_capabilities: Any + + +@dataclass(frozen=True) +class InboundLadderRejection: + """The first ladder rung that failed, as JSON-RPC error fields.""" + + code: int + message: str + data: Any = None + + +_ROUTING_HEADER_NAMES: Final = frozenset({MCP_PROTOCOL_VERSION_HEADER, MCP_METHOD_HEADER, MCP_NAME_HEADER}) + + +def find_duplicated_routing_header(headers: Iterable[tuple[str, str]]) -> str | None: + """Name of a routing header supplied more than once in raw header lines, or `None`. + + Takes raw `(name, value)` pairs — a folded mapping hides duplicates. A + duplicate is rejected because first-copy and last-copy readers would + disagree. `Mcp-Param-*` duplicates are :func:`validate_mcp_param_headers`'s job. + """ + seen: set[str] = set() + for name, _ in headers: + key = name.lower() + if key in _ROUTING_HEADER_NAMES: + if key in seen: + return key + seen.add(key) + return None + + +def unsupported_protocol_version_rejection( + requested: str, supported_modern_versions: Sequence[str] = MODERN_PROTOCOL_VERSIONS +) -> InboundLadderRejection | None: + """The `UNSUPPORTED_PROTOCOL_VERSION` rejection for `requested`, or `None` if it is served. + + The request ladder's last rung, shared with the transport's notification arm + so both message kinds name the same `supported` list in the same words. + """ + if requested in supported_modern_versions: + return None + return InboundLadderRejection( + code=UNSUPPORTED_PROTOCOL_VERSION, + message="Unsupported protocol version", + data=UnsupportedProtocolVersionErrorData( + supported=list(supported_modern_versions), requested=requested + ).model_dump(mode="json"), + ) + + +def classify_inbound_request( + body: Mapping[str, Any], + *, + headers: Mapping[str, str] | None = None, + supported_modern_versions: Sequence[str] = MODERN_PROTOCOL_VERSIONS, +) -> InboundModernRoute | InboundLadderRejection: + """Run the modern-protocol validation ladder over a decoded JSON-RPC body. + + Rungs, in order — first failure wins: + + 1. `params._meta` is a mapping carrying the required envelope pair + (protocol version, client capabilities) → else + :data:`~mcp_types.jsonrpc.INVALID_PARAMS` naming the missing key(s) + (basic/index.mdx "Per-request protocol fields"). Client info is + optional (SHOULD-include, spec PR #3002); absent reads as `None`. + 2. When `headers` is given, `MCP-Protocol-Version` equals the envelope's + protocol version, `Mcp-Method` equals `body.method`, and — for the + methods in :data:`NAME_BEARING_METHODS` — `Mcp-Name` equals the named + body param → else :data:`~mcp_types.jsonrpc.HEADER_MISMATCH`. Runs + before the supported-version rung so a client that disagrees with itself + is told so, rather than told the body's version is unsupported. + 3. The envelope's protocol version is a string in + `supported_modern_versions` → non-string values are + :data:`~mcp_types.jsonrpc.INVALID_PARAMS` (a shape defect, not a + negotiation outcome), else + :data:`~mcp_types.jsonrpc.UNSUPPORTED_PROTOCOL_VERSION` with + `data = {"supported": [...], "requested": }`. + + Method existence is *not* a rung: kernel dispatch owns that decision so + custom-registered methods route and the answer lives in one place. + + Args: + body: The decoded JSON-RPC request mapping. Envelope shape + (`jsonrpc` / `id`) is not checked here. + headers: Transport headers keyed by lowercase name, or `None` to + skip the header rung (non-HTTP callers). + supported_modern_versions: Modern protocol revisions this server + accepts on the per-request-envelope path. + """ + try: + meta_value = body["params"]["_meta"] + except (KeyError, TypeError): + meta_value = None + if not isinstance(meta_value, Mapping): + return InboundLadderRejection( + code=INVALID_PARAMS, + message="params._meta must be an object carrying the required " + f"{PROTOCOL_VERSION_META_KEY!r} and {CLIENT_CAPABILITIES_META_KEY!r} envelope keys", + ) + meta = cast("Mapping[str, Any]", meta_value) + if missing := [key for key in (PROTOCOL_VERSION_META_KEY, CLIENT_CAPABILITIES_META_KEY) if key not in meta]: + return InboundLadderRejection( + code=INVALID_PARAMS, + message=f"params._meta is missing the required envelope key(s): {', '.join(missing)}", + ) + protocol_version: Any = meta[PROTOCOL_VERSION_META_KEY] + client_info: Any = meta.get(CLIENT_INFO_META_KEY) + client_capabilities: Any = meta[CLIENT_CAPABILITIES_META_KEY] + if headers is not None: + version_header = headers.get(MCP_PROTOCOL_VERSION_HEADER) + # Presence is checked explicitly: a null body version would otherwise + # slip the equality check (None == None) and mask the absent header. + if version_header is None or version_header != protocol_version: + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{MCP_PROTOCOL_VERSION_HEADER} header does not match the request envelope's protocol version", + ) + method: Any = body.get("method") + if headers.get(MCP_METHOD_HEADER) != method: + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{MCP_METHOD_HEADER} header does not match the request body's method", + ) + name_key = NAME_BEARING_METHODS.get(method) + if name_key is not None: + # Rung 1 already proved body["params"] is a mapping (its `_meta` is one). + body_value = cast("Mapping[str, Any]", body["params"]).get(name_key) + if body_value is not None and decode_header_value(headers.get(MCP_NAME_HEADER)) != body_value: + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{MCP_NAME_HEADER} header does not match the request body's {name_key!r} parameter", + ) + + if not isinstance(protocol_version, str): + # Rung 3's precondition: a shape defect, not a version-negotiation + # outcome - -32022 is the one code auto-negotiating clients do NOT + # fall back from, and the typed rung-3 payload itself requires a + # string `requested`. Sits after the header rung, which fires first + # for every header-bearing entry (an absent version header is a + # mismatch, and a present one is a string that can never equal a + # non-string body value) - so this rejection is reachable only on + # header-less transports. + return InboundLadderRejection( + code=INVALID_PARAMS, + message="the protocol-version envelope value must be a string", + ) + + if (unsupported := unsupported_protocol_version_rejection(protocol_version, supported_modern_versions)) is not None: + return unsupported + + return InboundModernRoute( + protocol_version=protocol_version, + client_info=client_info, + client_capabilities=client_capabilities, + ) + + +# Header values eligible for the spec's numeric-comparison SHOULD; scientific +# notation never compares numerically (matching the typescript-sdk's gate). +_CANONICAL_DECIMAL = re.compile(r"^-?[0-9]+(\.[0-9]+)?$") + + +def _mcp_param_value_matches(prop_type: Any, value: Any, rendered: str, decoded: str) -> bool: + """True when a decoded `Mcp-Param-*` header value agrees with the body argument. + + Integer-typed declarations with an integral body value compare numerically + (`42` matches `42.0`, the spec's SHOULD) for canonical-decimal headers — + exact, no float round-trip, so values beyond the IEEE754 safe range still + compare. Anything else compares against `rendered`, the emit-side rendering. + """ + if ( + prop_type == "integer" + and not isinstance(value, bool) + and (isinstance(value, int) or (isinstance(value, float) and value.is_integer())) + and _CANONICAL_DECIMAL.fullmatch(decoded) is not None + ): + whole, _, fraction = decoded.partition(".") + if fraction and set(fraction) != {"0"}: + return False + try: + return int(whole) == int(value) + except ValueError: + return False + return decoded == rendered + + +def validate_mcp_param_headers( + input_schema: Any, + arguments: Mapping[str, Any], + headers: Mapping[str, str], +) -> InboundLadderRejection | None: + """Compare a `tools/call` request's `Mcp-Param-*` headers against its body arguments. + + Each annotated property's header and argument must agree: present together + and equal after sentinel decoding, or absent together (`null` counts as + absent). Returns the first failure as a `HEADER_MISMATCH` rejection, else `None`. + + A header whose argument is absent or unrenderable is deliberately rejected: + the spec's purpose clause is exactly an intermediary routing on a value the + body never carried. A duplicated recognized header is rejected — first-copy + and last-copy readers would disagree. A schema :func:`find_invalid_x_mcp_header` + rejects validates nothing: conforming clients drop the tool and emit no headers. + """ + if find_invalid_x_mcp_header(input_schema) is not None: + return None + folded: dict[str, str] = {} + duplicated: set[str] = set() + for name, value in headers.items(): + key = name.lower() + if key in folded: + duplicated.add(key) + folded[key] = value + for path, token, schema in _annotated_positions(input_schema): + header_name = f"{MCP_PARAM_HEADER_PREFIX}{token}" + key = header_name.lower() + raw = folded.get(key) + value = _value_at_path(arguments, path) + argument = ".".join(path) + if raw is not None and key in duplicated: + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{header_name} header appears more than once", + ) + if value is None: + if raw is not None: + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{header_name} header is present but the request body's {argument!r} argument is absent", + ) + continue + rendered = _render_header_scalar(value) + if rendered is None: + # Unrenderable value: a conforming client omitted the header, so one claiming it can never match. + if raw is not None: + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{header_name} header does not match the request body's {argument!r} argument", + ) + continue + if raw is None: + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{header_name} header is missing but the request body's {argument!r} argument is present", + ) + decoded = decode_header_value(raw) + if decoded is None: + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{header_name} header carries a malformed base64 sentinel value", + ) + if not _mcp_param_value_matches(schema.get("type"), value, rendered, decoded): + return InboundLadderRejection( + code=HEADER_MISMATCH, + message=f"{header_name} header does not match the request body's {argument!r} argument", + ) + return None diff --git a/src/mcp-client/mcp_client/shared/jsonrpc_dispatcher.py b/src/mcp-client/mcp_client/shared/jsonrpc_dispatcher.py new file mode 100644 index 0000000000..be6b5a526f --- /dev/null +++ b/src/mcp-client/mcp_client/shared/jsonrpc_dispatcher.py @@ -0,0 +1,836 @@ +"""JSON-RPC `Dispatcher` over the `SessionMessage` stream contract all transports speak. + +Owns request-id correlation, the receive loop, per-request task isolation, +cancellation/progress wiring, and the single exception-to-wire boundary; +methods and params are otherwise opaque strings and dicts. +""" + +from __future__ import annotations + +import contextvars +import logging +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field +from functools import partial +from typing import Any, Generic, Literal, cast + +import anyio +import anyio.abc +import anyio.lowlevel +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream +from mcp_types import ( + CONNECTION_CLOSED, + INTERNAL_ERROR, + INVALID_PARAMS, + REQUEST_TIMEOUT, + ErrorData, + JSONRPCError, + JSONRPCMessage, + JSONRPCNotification, + JSONRPCRequest, + JSONRPCResponse, + ProgressToken, + RequestId, +) +from opentelemetry.trace import SpanKind +from pydantic import ValidationError +from typing_extensions import TypeVar + +from mcp_client.shared._compat import resync_tracer +from mcp_client.shared._otel import inject_trace_context, otel_span +from mcp_client.shared._stream_protocols import ReadStream, WriteStream +from mcp_client.shared.dispatcher import ( + CallOptions, + DispatchContext, + Dispatcher, + OnNotify, + OnNotifyIntercept, + OnRequest, + ProgressFnT, + as_request_id, + coerce_request_id, + run_notify_intercept, +) +from mcp_client.shared.exceptions import MCPError, NoBackChannelError +from mcp_client.shared.message import ( + ClientMessageMetadata, + MessageMetadata, + ServerMessageMetadata, + SessionMessage, +) +from mcp_client.shared.transport_context import TransportContext + +__all__ = [ + "JSONRPCDispatcher", + "cancelled_request_id_from_params", + "handler_exception_to_error_data", + "progress_token_from_params", +] + +logger = logging.getLogger("mcp.shared.jsonrpc_dispatcher") + +_ABANDON_WRITE_TIMEOUT: float = 5 +"""Bound for courtesy-cancel writes on the abandon paths; the caller-cancel +arm shields its write, so a wedged transport would otherwise hang it uncancellably.""" + +_SHUTDOWN_WRITE_TIMEOUT: float = 1 +"""Tighter bound for the shutdown-arm error write so a wedged transport can't hold session close.""" + +TransportT = TypeVar("TransportT", bound=TransportContext, default=TransportContext) + +PeerCancelMode = Literal["interrupt", "signal"] +"""How `notifications/cancelled` is applied: `"interrupt"` (default) cancels +the handler's scope; `"signal"` only sets `ctx.cancel_requested` and lets the +handler run to completion. Either way the cancelled request is never +answered - the handler's eventual result or error is dropped, not written.""" + + +def handler_exception_to_error_data(exc: BaseException) -> ErrorData | None: + """Map a handler-raised exception to its wire `ErrorData`. + + The two rungs every dispatcher shares: an `MCPError` carries its own + `ErrorData`; a pydantic `ValidationError` is the spec's INVALID_PARAMS + with empty ``data`` (no pydantic text on the wire). Returns ``None`` for + any other exception so each caller applies its own catch-all - + `JSONRPCDispatcher` currently pins ``code=0`` for v1 compat, + the modern HTTP entry uses `INTERNAL_ERROR`. + """ + if isinstance(exc, MCPError): + return exc.error + if isinstance(exc, ValidationError): + return ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data="") + return None + + +def progress_token_from_params(params: Mapping[str, Any] | None) -> ProgressToken | None: + """Read `params._meta.progressToken`; reject bool (bool subclasses int, so True would alias 1).""" + match params: + case {"_meta": {"progressToken": str() | int() as token}} if not isinstance(token, bool): + return token + case _: + return None + + +def cancelled_request_id_from_params(params: Mapping[str, Any] | None) -> RequestId | None: + """Read `params.requestId` from a `notifications/cancelled` (`as_request_id` shape rules).""" + return as_request_id((params or {}).get("requestId")) + + +@dataclass(slots=True) +class _Pending: + """An outbound request awaiting its response.""" + + send: MemoryObjectSendStream[dict[str, Any] | ErrorData] + receive: MemoryObjectReceiveStream[dict[str, Any] | ErrorData] + on_progress: ProgressFnT | None = None + + +@dataclass(slots=True) +class _InFlight(Generic[TransportT]): + """An inbound request currently being handled.""" + + scope: anyio.CancelScope + dctx: _JSONRPCDispatchContext[TransportT] + + +@dataclass +class _JSONRPCDispatchContext(Generic[TransportT]): + """Concrete `DispatchContext` produced for each inbound JSON-RPC message.""" + + transport: TransportT + _dispatcher: JSONRPCDispatcher[TransportT] + _request_id: RequestId | None + message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework + """Transport-attached `SessionMessage.metadata` that the server lifts onto its request context.""" + _progress_token: ProgressToken | None = None + _closed: bool = False + cancel_requested: anyio.Event = field(default_factory=anyio.Event) + + @property + def request_id(self) -> RequestId | None: + return self._request_id + + @property + def can_send_request(self) -> bool: + return self.transport.can_send_request and not self._closed + + async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: + if self._closed: + logger.debug("dropped %s: dispatch context closed", method) + return + await self._dispatcher.notify(method, params, opts, _related_request_id=self._request_id) + + async def send_raw_request( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None = None, + ) -> dict[str, Any]: + if not self.can_send_request: + raise NoBackChannelError(method) + return await self._dispatcher.send_raw_request(method, params, opts, _related_request_id=self._request_id) + + async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: + if self._progress_token is None: + return + params: dict[str, Any] = {"progressToken": self._progress_token, "progress": progress} + if total is not None: + params["total"] = total + if message is not None: + params["message"] = message + await self.notify("notifications/progress", params) + + def close(self) -> None: + self._closed = True + + +def _default_transport_builder(metadata: MessageMetadata) -> TransportContext: + """The `TransportContext` for a message, honoring the transport's own verdict when it stamps one. + + A message reads as riding a full duplex pipe (`can_send_request=True`) + unless the transport that framed it says otherwise on the metadata it + attached, so a transport whose response has no room for a server request + (streamable HTTP in JSON-response mode) needs no wiring from whoever drives + its streams. + """ + can_send_request = metadata.can_send_request if isinstance(metadata, ServerMessageMetadata) else True + return TransportContext(kind="jsonrpc", can_send_request=can_send_request) + + +def _shielded_progress(fn: ProgressFnT) -> ProgressFnT: + """Wrap a user progress callback so an exception can't cancel the dispatcher's task group.""" + + async def _wrapped(progress: float, total: float | None, message: str | None) -> None: + try: + await fn(progress, total, message) + except Exception: + logger.exception("progress callback raised") + + return _wrapped + + +def _contained_notify(fn: OnNotify) -> OnNotify: + """Wrap a notification handler so it can't crash the dispatcher (same boundary as `_shielded_progress`).""" + + async def _wrapped(dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None: + try: + await fn(dctx, method, params) + except Exception: + logger.exception("notification handler for %r raised", method) + + return _wrapped + + +@dataclass(slots=True, frozen=True) +class _OutboundPlan: + """Outbound metadata plus whether abandoning the request sends a courtesy `notifications/cancelled`.""" + + metadata: MessageMetadata + cancel_on_abandon: bool + + +def _plan_outbound(related_request_id: RequestId | None, opts: CallOptions | None) -> _OutboundPlan: + """Choose the outbound `SessionMessage.metadata` and the abandon-cancellation policy. + + `related_request_id` wins over resumption hints (they are dropped). Only + hints that actually reach the transport suppress the courtesy cancel - a + request that is neither resumable nor cancelled would leak the peer's work. + """ + opts = opts or {} + cancel_on_abandon = opts.get("cancel_on_abandon", True) + token = opts.get("resumption_token") + on_token = opts.get("on_resumption_token") + headers = opts.get("headers") + if related_request_id is not None: + if token is not None or on_token is not None: + logger.debug( + "dropping resumption hints: related_request_id %r takes precedence on metadata", related_request_id + ) + return _OutboundPlan(ServerMessageMetadata(related_request_id=related_request_id), cancel_on_abandon) + if token is not None or on_token is not None: + return _OutboundPlan( + ClientMessageMetadata(resumption_token=token, on_resumption_token_update=on_token, headers=headers), + cancel_on_abandon=False, + ) + if headers: + return _OutboundPlan(ClientMessageMetadata(headers=headers), cancel_on_abandon) + return _OutboundPlan(None, cancel_on_abandon) + + +class JSONRPCDispatcher(Dispatcher[TransportT]): + """`Dispatcher` over the `SessionMessage` stream contract. + + Explicit Protocol base so pyright checks conformance at the class definition. + """ + + def __init__( + self, + read_stream: ReadStream[SessionMessage | Exception], + write_stream: WriteStream[SessionMessage], + *, + transport_builder: Callable[[MessageMetadata], TransportT] | None = None, + peer_cancel_mode: PeerCancelMode = "interrupt", + raise_handler_exceptions: bool = False, + inline_methods: frozenset[str] = frozenset(), + on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None, + ) -> None: + """Wire a dispatcher over a transport's `SessionMessage` stream pair. + + Args: + transport_builder: Builds each message's `TransportContext` from + its `SessionMessage.metadata`. + raise_handler_exceptions: Re-raise handler exceptions out of + `run()` after the error response is written. + inline_methods: Methods awaited in the read loop before the next + message is dequeued (e.g. `initialize`); an inline handler + that awaits the peer deadlocks the parked loop. + on_stream_exception: Observer for `Exception` items on the read + stream; without it they are debug-logged and dropped. Awaited + inline in the read loop, so a slow observer stalls dispatch. + """ + self._read_stream = read_stream + self._write_stream = write_stream + # With transport_builder omitted, TransportT defaults to + # TransportContext; pyright can't connect the two, hence the cast. + self._transport_builder = cast( + "Callable[[MessageMetadata], TransportT]", + transport_builder or _default_transport_builder, + ) + self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode + self._raise_handler_exceptions = raise_handler_exceptions + self._inline_methods = inline_methods + self.on_stream_exception = on_stream_exception + """Observer for ``Exception`` items on the read stream. Mutable so a session can + bind it after the dispatcher is built (e.g. ``ClientSession`` routing into + ``message_handler``); only consulted inside ``run()`` so pre-enter assignment is safe.""" + + self._next_id = 0 + self._pending: dict[RequestId, _Pending] = {} + self._in_flight: dict[RequestId, _InFlight[TransportT]] = {} + self._on_notify_intercept: OnNotifyIntercept | None = None + self._tg: anyio.abc.TaskGroup | None = None + self._running = False + self._closed = False + + async def send_raw_request( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None = None, + *, + _related_request_id: RequestId | None = None, + ) -> dict[str, Any]: + """Send a JSON-RPC request and await its response. + + `_related_request_id` is set only by `_JSONRPCDispatchContext` so that + mid-handler requests route onto the inbound request's SSE stream. + + Raises: + MCPError: Peer error response; `REQUEST_TIMEOUT` if + `opts["timeout"]` elapsed; `CONNECTION_CLOSED` if the + transport closed or the dispatcher shut down. + RuntimeError: Called before `run()`. + """ + # Post-close sends get the same CONNECTION_CLOSED contract as in-flight waiters. + if self._closed: + raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") + if not self._running: + raise RuntimeError("JSONRPCDispatcher.send_raw_request called before run()") + opts = opts or {} + supplied_id = opts.get("request_id") + if supplied_id is not None: + request_id: RequestId = supplied_id + # The pending key gets the same coercion `_resolve_pending` applies + # to inbound response ids, so a supplied "7" still correlates + # whether the peer echoes "7" or 7. The wire id stays verbatim. + pending_key = coerce_request_id(request_id) + if pending_key in self._pending: + raise ValueError(f"request id {request_id!r} is already in flight") + else: + # Mint past any key a supplied id occupies: the collision error is + # reserved for the caller who actually chose the id. + request_id = self._allocate_id() + while request_id in self._pending: + request_id = self._allocate_id() + pending_key = request_id + out_params = dict(params) if params is not None else {} + out_meta = dict(out_params.get("_meta") or {}) + on_progress = opts.get("on_progress") + if on_progress is not None: + # The request id doubles as the progress token, so `_pending[token]` finds `on_progress` directly. + out_meta["progressToken"] = request_id + out_params["_meta"] = out_meta + + # buffer=1: a close signal can arrive before the waiter parks in receive(); + # a WouldBlock later just means the waiter already has its one outcome. + send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1) + pending = _Pending(send=send, receive=receive, on_progress=on_progress) + self._pending[pending_key] = pending + + plan = _plan_outbound(_related_request_id, opts) + # Spec MUST: only previously-issued requests may be cancelled. A write + # interrupted by cancellation may still have delivered (a memory-stream + # send can hand its item to the receiver and still raise), so a started + # write counts as issued: the peer ignores a cancel for an id it never + # saw, while skipping it would leak a delivered request's handler. + request_write_started = False + timeout_armed = False + + target = out_params.get("name") + span_name = f"MCP send {method}{f' {target}' if isinstance(target, str) else ''}" + # TODO(maxisbey): move the otel span + inject into an outbound + # middleware once that seam exists; the dispatcher should not own otel. + try: + with otel_span( + span_name, + kind=SpanKind.CLIENT, + attributes={"mcp.method.name": method, "jsonrpc.request.id": str(request_id)}, + ): + # SEP-414: inject W3C trace context; `_meta` stays on the wire even with a no-op tracer. + inject_trace_context(out_meta) + msg = JSONRPCRequest(jsonrpc="2.0", id=request_id, method=method, params=out_params) + # Surface a pre-existing cancellation while the request provably + # never started; past this point a cancelled write counts as issued. + await anyio.lowlevel.checkpoint_if_cancelled() + request_write_started = True + try: + await self._write(msg, plan.metadata) + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + # Transport tore down before run() noticed EOF; surface the documented contract. + raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") from None + with anyio.fail_after(opts.get("timeout")): + timeout_armed = True + outcome = await receive.receive() + except TimeoutError: + if not timeout_armed: + # `fail_after` arms only after the write, so this TimeoutError is the + # transport's own bounded send() failing - a transport error, not + # `opts["timeout"]` elapsing. Propagate it raw (v1 kept the write + # outside the timeout-catching try and did the same). + raise + # Courtesy cancel (spec-recommended, new vs v1) so the peer stops work; + # unshielded so an outer caller cancellation can still interrupt the write. + if plan.cancel_on_abandon: + await self._final_write( + partial( + self._cancel_outbound, + request_id, + f"timed out after {opts.get('timeout')}s", + _related_request_id, + ), + shield=False, + timeout=_ABANDON_WRITE_TIMEOUT, + describe=f"courtesy cancel for timed-out request {request_id!r}", + ) + raise MCPError(code=REQUEST_TIMEOUT, message=f"Request {method!r} timed out") from None + except anyio.get_cancelled_exc_class(): + # Caller cancelled: bare awaits re-raise here, so the shielded helper + # lets the courtesy cancel go out before we propagate. + if plan.cancel_on_abandon and request_write_started: + await self._final_write( + partial(self._cancel_outbound, request_id, "caller cancelled", _related_request_id), + shield=True, + timeout=_ABANDON_WRITE_TIMEOUT, + describe=f"courtesy cancel for caller-cancelled request {request_id!r}", + ) + raise + finally: + # Remove the waiter on every path so a late response is dropped, not leaked. + self._pending.pop(pending_key, None) + send.close() + receive.close() + + if isinstance(outcome, ErrorData): + raise MCPError(code=outcome.code, message=outcome.message, data=outcome.data) + return outcome + + async def notify( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None = None, + *, + _related_request_id: RequestId | None = None, + ) -> None: + """Send a fire-and-forget notification. + + Fire-and-forget all the way: a post-close send or a write onto a + torn-down transport drops the notification with a debug log instead + of raising (same policy as the response writes and `ctx.notify`). + """ + if self._closed: + logger.debug("dropped %s: dispatcher closed", method) + return + # Leave `params` unset when None: with `exclude_unset=True` an explicit + # None would serialize as `"params": null`, which JSON-RPC 2.0 forbids. + if params is not None: + msg = JSONRPCNotification(jsonrpc="2.0", method=method, params=dict(params)) + else: + msg = JSONRPCNotification(jsonrpc="2.0", method=method) + try: + await self._write(msg, _plan_outbound(_related_request_id, opts).metadata) + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + # Transport tore down before run() noticed EOF. + logger.debug("dropped %s: write stream closed", method) + + async def run( + self, + on_request: OnRequest, + on_notify: OnNotify, + on_notify_intercept: OnNotifyIntercept | None = None, + *, + task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, + ) -> None: + """Drive the receive loop until the read stream closes. + + `task_status.started()` fires once `send_raw_request` is usable. + Single-shot: once the loop ends the dispatcher stays closed and cannot be restarted. + """ + self._on_notify_intercept = on_notify_intercept + try: + # LIFO exits: the write stream closes only after the task-group join, so teardown writes still land. + async with self._write_stream: + async with anyio.create_task_group() as tg: + self._tg = tg + self._running = True + task_status.started() + try: + async with self._read_stream: + try: + async for item in self._read_stream: + # Duck-typed: only `ContextReceiveStream` carries the + # sender's per-message contextvars snapshot. + sender_ctx: contextvars.Context | None = getattr( + self._read_stream, "last_context", None + ) + await self._dispatch(item, on_request, on_notify, sender_ctx) + except anyio.ClosedResourceError: + # Receive end closed under us (stateless SHTTP teardown); same as EOF. + logger.debug("read stream closed by transport; treating as EOF") + # EOF: wake blocked `send_raw_request` waiters with CONNECTION_CLOSED. + self._running = False + self._closed = True + self._fan_out_closed() + finally: + # Cancel in-flight handlers; otherwise the task-group join + # waits on handlers whose callers are already gone. + tg.cancel_scope.cancel() + finally: + # Covers cancel/crash paths that skip the inline fan-out; idempotent. + self._running = False + self._closed = True + self._tg = None + self._fan_out_closed() + await resync_tracer() + + async def _dispatch( + self, + item: SessionMessage | Exception, + on_request: OnRequest, + on_notify: OnNotify, + sender_ctx: contextvars.Context | None, + ) -> None: + """Route one inbound item. + + Only `inline_methods` requests and the `on_stream_exception` observer + are awaited; any other `await` would head-of-line block the read loop. + """ + if isinstance(item, Exception): + if self.on_stream_exception is None: + logger.debug("transport yielded exception: %r", item) + return + try: + await self.on_stream_exception(item) + except Exception: + logger.exception("on_stream_exception observer raised") + return + metadata = item.metadata + msg = item.message + match msg: + case JSONRPCRequest(): + await self._dispatch_request(msg, metadata, on_request, sender_ctx) + case JSONRPCNotification(): + self._dispatch_notification(msg, metadata, on_notify, sender_ctx) + case JSONRPCResponse(): + self._resolve_pending(msg.id, msg.result) + case JSONRPCError(): # pragma: no branch + # Exhaustive over JSONRPCMessage, so the no-match arc is unreachable. + self._resolve_pending(msg.id, msg.error) + + async def _dispatch_request( + self, + req: JSONRPCRequest, + metadata: MessageMetadata, + on_request: OnRequest, + sender_ctx: contextvars.Context | None, + ) -> None: + progress_token = progress_token_from_params(req.params) + try: + transport_ctx = self._transport_builder(metadata) + except Exception: + # A raising builder must cost only this message, not the connection. + logger.exception("transport_builder raised; rejecting request %r", req.id) + self._spawn( + self._write_error, + req.id, + ErrorData(code=INTERNAL_ERROR, message="transport context unavailable"), + sender_ctx=sender_ctx, + ) + return + dctx = _JSONRPCDispatchContext( + transport=transport_ctx, + _dispatcher=self, + _request_id=req.id, + message_metadata=metadata, + _progress_token=progress_token, + ) + scope = anyio.CancelScope() + # TODO(maxisbey): duplicate ids blind-overwrite (v1/TS parity); revisit + # rejecting with INVALID_REQUEST. Key coerced so a stringified + # `notifications/cancelled` id still correlates. + self._in_flight[coerce_request_id(req.id)] = _InFlight(scope=scope, dctx=dctx) + if req.method in self._inline_methods: + # Spawn so `sender_ctx` applies, but park the read loop until the + # handler returns - that's the inline ordering guarantee. + done = anyio.Event() + + async def _run_inline() -> None: + try: + await self._handle_request(req, dctx, scope, on_request) + finally: + done.set() + + self._spawn(_run_inline, sender_ctx=sender_ctx) + await done.wait() + else: + self._spawn(self._handle_request, req, dctx, scope, on_request, sender_ctx=sender_ctx) + + def _dispatch_notification( + self, + msg: JSONRPCNotification, + metadata: MessageMetadata, + on_notify: OnNotify, + sender_ctx: contextvars.Context | None, + ) -> None: + """Route one inbound notification. + + `notifications/cancelled` and `notifications/progress` are intercepted + here (they correlate against the `_in_flight`/`_pending` tables this + layer owns) and still teed to `on_notify` afterwards. The caller's + `on_notify_intercept` then runs in receive order; only unconsumed + notifications reach the spawned `on_notify`. + """ + if msg.method == "notifications/cancelled": + rid = cancelled_request_id_from_params(msg.params) + if rid is not None and (in_flight := self._in_flight.get(coerce_request_id(rid))) is not None: + in_flight.dctx.cancel_requested.set() + if self._peer_cancel_mode == "interrupt": + in_flight.scope.cancel() + elif msg.method == "notifications/progress": + match msg.params: + case {"progressToken": str() | int() as token, "progress": int() | float() as progress} if ( + not isinstance(token, bool) + and not isinstance(progress, bool) + and (pending := self._pending.get(coerce_request_id(token))) is not None + and pending.on_progress is not None + ): + total = msg.params.get("total") + message = msg.params.get("message") + self._spawn( + _shielded_progress(pending.on_progress), + float(progress), + float(total) if isinstance(total, int | float) else None, + message if isinstance(message, str) else None, + sender_ctx=sender_ctx, + ) + case _: + pass + if run_notify_intercept(self._on_notify_intercept, msg.method, msg.params): + return + try: + transport_ctx = self._transport_builder(metadata) + except Exception: + # Same containment as `_dispatch_request`: drop the notification, keep the loop. + logger.exception("transport_builder raised; dropping notification %r", msg.method) + return + dctx = _JSONRPCDispatchContext( + transport=transport_ctx, _dispatcher=self, _request_id=None, message_metadata=metadata + ) + self._spawn(_contained_notify(on_notify), dctx, msg.method, msg.params, sender_ctx=sender_ctx) + + def _resolve_pending(self, request_id: RequestId | None, outcome: dict[str, Any] | ErrorData) -> None: + pending = self._pending.get(coerce_request_id(request_id)) if request_id is not None else None + if pending is None: + logger.debug("dropping response for unknown/late request id %r", request_id) + return + try: + pending.send.send_nowait(outcome) + except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError): + logger.debug("waiter for request id %r already gone", request_id) + + def _spawn( + self, + fn: Callable[..., Awaitable[Any]], + *args: object, + sender_ctx: contextvars.Context | None, + ) -> None: + """Schedule `fn(*args)` in the run() task group, propagating the sender's contextvars. + + ASGI middleware (auth, OTel) sets contextvars on the task that wrote the + message; `Context.run` makes the spawned handler inherit that context. + """ + assert self._tg is not None + if sender_ctx is not None: + sender_ctx.run(self._tg.start_soon, fn, *args) + else: + self._tg.start_soon(fn, *args) + + def _fan_out_closed(self) -> None: + """Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED`. + + Synchronous: callers may be inside a cancelled scope. Idempotent. + """ + closed = ErrorData(code=CONNECTION_CLOSED, message="Connection closed") + for pending in self._pending.values(): + try: + pending.send.send_nowait(closed) + except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError): + pass + self._pending.clear() + + async def _handle_request( + self, + req: JSONRPCRequest, + dctx: _JSONRPCDispatchContext[TransportT], + scope: anyio.CancelScope, + on_request: OnRequest, + ) -> None: + """Run `on_request` for one inbound request and write its response. + + The single exception-to-wire boundary: handler exceptions become + `JSONRPCError` here. A request the peer cancelled is never answered + (spec: MUST NOT send further messages for it) - it settles unanswered + instead, and `_settle_unanswered` tells the transport. + """ + answer_write_started = False + handler_failure: BaseException | None = None # re-raised once the request settles + try: + with scope: + try: + result = await on_request(dctx, req.method, req.params) + finally: + # Close the back-channel and drop from `_in_flight`; no checkpoint + # since handler return, so a peer cancel can't interleave. + # Identity guard: don't evict a duplicate id's newer entry. + dctx.close() + key = coerce_request_id(req.id) + if (entry := self._in_flight.get(key)) is not None and entry.dctx is dctx: + del self._in_flight[key] + if not dctx.cancel_requested.is_set(): + # A write interrupted by cancellation may still have delivered + # (a memory-stream send can hand its item to the receiver and + # still raise), so a started answer write counts as sent below: + # peers drop late responses, while a second answer for one id + # would break JSON-RPC. + answer_write_started = True + await self._write_result(req.id, result) + except anyio.get_cancelled_exc_class(): + # Shutdown: answer the request so the peer isn't left waiting - unless + # an answer write already started (it may have reached the transport; + # prefer possibly-zero answers over possibly-two), or the peer already + # cancelled it and stopped waiting. The shielded helper is needed + # because bare awaits re-raise here. + if not answer_write_started and not dctx.cancel_requested.is_set(): + await self._final_write( + partial(self._write_error, req.id, ErrorData(code=CONNECTION_CLOSED, message="Connection closed")), + shield=True, + timeout=_SHUTDOWN_WRITE_TIMEOUT, + describe=f"shutdown error response for request {req.id!r}", + ) + raise + except Exception as e: + error = handler_exception_to_error_data(e) + if error is None: + logger.exception("handler for %r raised", req.method) + # TODO(L58): code=0 pins existing-server compat; JSON-RPC says + # INTERNAL_ERROR. Revisit per the suite's divergence entry. + error = ErrorData(code=0, message=str(e)) + if self._raise_handler_exceptions: + handler_failure = e + # A cancel silences only the wire; the failure stays as visible as before. + if not dctx.cancel_requested.is_set(): + answer_write_started = True + await self._write_error(req.id, error) + # The one place a cancelled request settles: the handler is done (any + # mode) with nothing written. A peer-interrupt cancel is absorbed at + # scope __exit__ and lands here too. + if not answer_write_started: + await self._settle_unanswered(dctx) + if handler_failure is not None: + raise handler_failure + # No `_in_flight` pop here: the inner finally covers every path, and a late pop could evict a reused id. + + def _allocate_id(self) -> int: + self._next_id += 1 + return self._next_id + + async def _write(self, message: JSONRPCMessage, metadata: MessageMetadata = None) -> None: + await self._write_stream.send(SessionMessage(message=message, metadata=metadata)) + + async def _write_result(self, request_id: RequestId, result: dict[str, Any]) -> None: + try: + await self._write(JSONRPCResponse(jsonrpc="2.0", id=request_id, result=result)) + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + logger.debug("dropped result for %r: write stream closed", request_id) + + async def _write_error(self, request_id: RequestId, error: ErrorData) -> None: + try: + await self._write(JSONRPCError(jsonrpc="2.0", id=request_id, error=error)) + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + logger.debug("dropped error for %r: write stream closed", request_id) + + async def _settle_unanswered(self, dctx: _JSONRPCDispatchContext[TransportT]) -> None: + """Run the transport's `on_request_unanswered` hook: this request settled with no response. + + The dispatcher writes nothing for it; a transport whose wire must still + end the request (2025-era streamable HTTP) does so from this hook. A + raising hook is contained here, like the other callback boundaries. + """ + metadata = dctx.message_metadata + if not isinstance(metadata, ServerMessageMetadata) or metadata.on_request_unanswered is None: + return + try: + await metadata.on_request_unanswered() + except (anyio.BrokenResourceError, anyio.ClosedResourceError): + logger.debug("on_request_unanswered dropped: connection closing") + except Exception: + logger.exception("on_request_unanswered hook raised") + + async def _final_write( + self, + write: Callable[[], Awaitable[None]], + *, + shield: bool, + timeout: float, + describe: str, + ) -> None: + """Attempt one last write under the shared abandon/teardown policy. + + `shield=True` is for arms already inside a cancelled scope (a bare + `await` would re-raise); the bound keeps a wedged transport write + from becoming an uncancellable hang. + """ + with anyio.move_on_after(timeout, shield=shield) as scope: + await write() + if scope.cancelled_caught: + logger.warning("%s gave up: transport write blocked", describe) + + async def _cancel_outbound(self, request_id: RequestId, reason: str, related_request_id: RequestId | None) -> None: + # Thread `related_request_id` so streamable HTTP routes the cancel onto + # the request's own SSE stream instead of a possibly-absent GET stream. + # `notify` swallows connection-state errors itself, so no guard here. + await self.notify( + "notifications/cancelled", + {"requestId": request_id, "reason": reason}, + _related_request_id=related_request_id, + ) diff --git a/src/mcp-client/mcp_client/shared/memory.py b/src/mcp-client/mcp_client/shared/memory.py new file mode 100644 index 0000000000..4fc2dd0309 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/memory.py @@ -0,0 +1,33 @@ +"""In-memory transports""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager + +from mcp_client.shared._compat import resync_tracer +from mcp_client.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams +from mcp_client.shared.message import SessionMessage + +MessageStream = tuple[ContextReceiveStream[SessionMessage | Exception], ContextSendStream[SessionMessage | Exception]] + + +@asynccontextmanager +async def create_client_server_memory_streams() -> AsyncGenerator[tuple[MessageStream, MessageStream], None]: + """Creates a pair of bidirectional memory streams for client-server communication. + + Yields: + A tuple of (client_streams, server_streams) where each is a tuple of + (read_stream, write_stream) + """ + # Create streams for both directions + server_to_client_send, server_to_client_receive = create_context_streams[SessionMessage | Exception](1) + client_to_server_send, client_to_server_receive = create_context_streams[SessionMessage | Exception](1) + + client_streams = (server_to_client_receive, client_to_server_send) + server_streams = (client_to_server_receive, server_to_client_send) + + async with server_to_client_receive, client_to_server_send, client_to_server_receive, server_to_client_send: + yield client_streams, server_streams + # Heals caller-driven cancels; closing memory streams never suspends. + await resync_tracer() diff --git a/src/mcp-client/mcp_client/shared/message.py b/src/mcp-client/mcp_client/shared/message.py new file mode 100644 index 0000000000..31e51e7128 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/message.py @@ -0,0 +1,63 @@ +"""Message wrapper with metadata support. + +This module defines a wrapper type that combines JSONRPCMessage with metadata +to support transport-specific features like resumability. +""" + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any + +from mcp_types import JSONRPCMessage, RequestId + +ResumptionToken = str + +ResumptionTokenUpdateCallback = Callable[[ResumptionToken], Awaitable[None]] + +# Callback type for closing SSE streams without terminating +CloseSSEStreamCallback = Callable[[], Awaitable[None]] + + +@dataclass +class ClientMessageMetadata: + """Metadata specific to client messages.""" + + resumption_token: ResumptionToken | None = None + on_resumption_token_update: Callable[[ResumptionToken], Awaitable[None]] | None = None + # Per-message HTTP headers (e.g. MCP-Protocol-Version, Mcp-Method) the transport should set. + headers: dict[str, str] | None = None + + +@dataclass +class ServerMessageMetadata: + """Metadata specific to server messages.""" + + related_request_id: RequestId | None = None + # Transport-specific request context (e.g. starlette Request for HTTP + # transports, None for stdio). Typed as Any because the server layer is + # transport-agnostic. + request_context: Any = None + # Callback to close SSE stream for the current request without terminating + close_sse_stream: CloseSSEStreamCallback | None = None + # Callback to close the standalone GET SSE stream (for unsolicited notifications) + close_standalone_sse_stream: CloseSSEStreamCallback | None = None + # Callback the dispatcher runs when this request settles without a response + # (e.g. it was cancelled), for a transport whose wire must still end the + # request even though no response is written. + on_request_unanswered: Callable[[], Awaitable[None]] | None = None + # The transport's verdict on whether this message's request-scoped channel + # can deliver a server-initiated request (see + # `TransportContext.can_send_request`); a transport that says nothing leaves + # it True. + can_send_request: bool = True + + +MessageMetadata = ClientMessageMetadata | ServerMessageMetadata | None + + +@dataclass +class SessionMessage: + """A message with specific metadata for transport-specific features.""" + + message: JSONRPCMessage + metadata: MessageMetadata = None diff --git a/src/mcp-client/mcp_client/shared/metadata_utils.py b/src/mcp-client/mcp_client/shared/metadata_utils.py new file mode 100644 index 0000000000..b646133477 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/metadata_utils.py @@ -0,0 +1,46 @@ +"""Utility functions for working with metadata in MCP types. + +These utilities are primarily intended for client-side usage to properly display +human-readable names in user interfaces in a spec-compliant way. +""" + +from mcp_types import Implementation, Prompt, Resource, ResourceTemplate, Tool + + +def get_display_name(obj: Tool | Resource | Prompt | ResourceTemplate | Implementation) -> str: + """Get the display name for an MCP object with proper precedence. + + This is a client-side utility function designed to help MCP clients display + human-readable names in their user interfaces. When servers provide a 'title' + field, it should be preferred over the programmatic 'name' field for display. + + For tools: title > annotations.title > name + For other objects: title > name + + Example: + ```python + # In a client displaying available tools + tools = await session.list_tools() + for tool in tools.tools: + display_name = get_display_name(tool) + print(f"Available tool: {display_name}") + ``` + + Args: + obj: An MCP object with name and optional title fields + + Returns: + The display name to use for UI presentation + """ + if isinstance(obj, Tool): + # Tools have special precedence: title > annotations.title > name + if hasattr(obj, "title") and obj.title is not None: + return obj.title + if obj.annotations and hasattr(obj.annotations, "title") and obj.annotations.title is not None: + return obj.annotations.title + return obj.name + else: + # All other objects: title > name + if hasattr(obj, "title") and obj.title is not None: + return obj.title + return obj.name diff --git a/src/mcp-client/mcp_client/shared/path_security.py b/src/mcp-client/mcp_client/shared/path_security.py new file mode 100644 index 0000000000..f663ee5bd2 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/path_security.py @@ -0,0 +1,176 @@ +"""Filesystem path safety primitives for resource handlers. + +These functions help MCP servers reject paths that would resolve +outside the served root when extracted URI template parameters are +used in filesystem operations. They are standalone utilities usable from both the +high-level :class:`~mcp.server.mcpserver.MCPServer` and lowlevel server +implementations. + +The canonical safe pattern:: + + from mcp_client.shared.path_security import safe_join + + @mcp.resource("file://docs/{+path}") + def read_doc(path: str) -> str: + return safe_join("/data/docs", path).read_text(encoding="utf-8") +""" + +import string +from pathlib import Path + +__all__ = ["PathEscapeError", "contains_path_traversal", "is_absolute_path", "safe_join"] + + +class PathEscapeError(ValueError): + """Raised by :func:`safe_join` when the resolved path escapes the base.""" + + +def contains_path_traversal(value: str) -> bool: + r"""Check whether a value, treated as a relative path, escapes its origin. + + This is a **base-free** check: it does not know the sandbox root, so + it detects only whether ``..`` components would move above the + starting point. Use :func:`safe_join` when you know the root — it + additionally catches symlink escapes and absolute-path injection. + + Note: + This is a string-level check on the value as supplied. It does + not model platform-specific filesystem normalisation (e.g. Win32 + stripping of trailing dots and spaces from the final path + component). For filesystem access, use :func:`safe_join`, which + resolves through the OS and verifies containment. + + The check is component-based: ``..`` is dangerous only as a + standalone path segment, not as a substring. Both ``/`` and ``\`` + are treated as separators. + + Example:: + + >>> contains_path_traversal("a/b/c") + False + >>> contains_path_traversal("../etc") + True + >>> contains_path_traversal("a/../../b") + True + >>> contains_path_traversal("a/../b") + False + >>> contains_path_traversal("1.0..2.0") + False + >>> contains_path_traversal("..") + True + + Args: + value: A string that may be used as a filesystem path. + + Returns: + ``True`` if the path would escape its starting directory. + """ + depth = 0 + for part in value.replace("\\", "/").split("/"): + if part == "..": + depth -= 1 + if depth < 0: + return True + elif part and part != ".": + depth += 1 + return False + + +def is_absolute_path(value: str) -> bool: + r"""Check whether a value is an absolute filesystem path. + + Absolute paths are dangerous when joined onto a base: in Python, + ``Path("/data") / "/etc/passwd"`` yields ``/etc/passwd`` — the + absolute right-hand side silently discards the base. + + Detects POSIX absolute (``/foo``), Windows drive-absolute + (``C:\foo``) and drive-relative (``C:foo``), and Windows + UNC/root-relative (``\\server\share``, ``\foo``). + + Example:: + + >>> is_absolute_path("relative/path") + False + >>> is_absolute_path("/etc/passwd") + True + >>> is_absolute_path("C:\\Windows") + True + >>> is_absolute_path("") + False + + Args: + value: A string that may be used as a filesystem path. + + Returns: + ``True`` if the path is absolute on any common platform. + """ + if not value: + return False + if value[0] in ("/", "\\"): + return True + # Windows drive form: C:, C:\, C:foo (drive-relative). A drive- + # relative right-hand side discards the join base when drives + # differ, so flag it even though PureWindowsPath.is_absolute() + # is False. This means single-letter-prefixed identifiers like + # "x:y" also match — opt out via ResourceSecurity(exempt_params=). + if len(value) >= 2 and value[1] == ":" and value[0] in string.ascii_letters: + return True + return False + + +def safe_join(base: str | Path, *parts: str) -> Path: + """Join path components onto a base, rejecting escapes. + + Resolves the joined path and verifies it remains within ``base``. + This is the **gold-standard** check: it catches ``..`` traversal, + absolute-path injection, and symlink escapes that the base-free + checks cannot. + + The symlink check is point-in-time: a directory swapped for a + symlink between this call and the caller's subsequent open would not + be re-checked. Handlers serving a tree that may be modified + concurrently should additionally open with ``O_NOFOLLOW`` or use + platform path-confinement primitives. + + Example:: + + >>> safe_join("/data/docs", "readme.txt") + PosixPath('/data/docs/readme.txt') + >>> safe_join("/data/docs", "../../../etc/passwd") + Traceback (most recent call last): + ... + PathEscapeError: ... + + Args: + base: The sandbox root. May be relative; it will be resolved. + parts: Path components to join. Each is checked for null bytes + and absolute form before joining. + + Returns: + The resolved path, verified to be within ``base`` at resolution + time. + + Raises: + PathEscapeError: If any part contains a null byte, any part is + absolute, or the resolved path is not contained within the + resolved base. + """ + base_resolved = Path(base).resolve() + + for part in parts: + # Null bytes pass through Path construction but fail at the + # syscall boundary with a cryptic error. Reject here so callers + # get a clear PathEscapeError instead. + if "\0" in part: + raise PathEscapeError(f"Path component contains a null byte; refusing to join onto {base_resolved}") + # Absolute parts would silently discard everything to the left + # in Path's / operator. + if is_absolute_path(part): + raise PathEscapeError(f"Path component {part!r} is absolute; refusing to join onto {base_resolved}") + + target = base_resolved.joinpath(*parts).resolve() + + if not target.is_relative_to(base_resolved): + raise PathEscapeError(f"Path {target} escapes base {base_resolved}") + + return target diff --git a/src/mcp-client/mcp_client/shared/peer.py b/src/mcp-client/mcp_client/shared/peer.py new file mode 100644 index 0000000000..3896096511 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/peer.py @@ -0,0 +1,240 @@ +"""Typed MCP request sugar over an `Outbound`. + +`ClientPeer` wraps any `Outbound` (anything with `send_raw_request` and +`notify`) and exposes the server-to-client request methods (sampling, +elicitation, roots, ping) as typed methods. + +`ClientPeer` does no capability gating: it builds the params, calls +`send_raw_request(method, params)`, and parses the result into the typed +model. Gating (and `NoBackChannelError`) is the wrapped `Outbound`'s job. +""" + +from collections.abc import Mapping +from typing import Any, cast, overload + +from mcp_types import ( + CreateMessageRequestParams, + CreateMessageResult, + CreateMessageResultWithTools, + ElicitRequestedSchema, + ElicitRequestFormParams, + ElicitRequestURLParams, + ElicitResult, + IncludeContext, + ListRootsResult, + ModelPreferences, + RequestParams, + RequestParamsMeta, + SamplingMessage, + Tool, + ToolChoice, +) +from pydantic import BaseModel +from typing_extensions import deprecated + +from mcp_client.shared.dispatcher import CallOptions, Outbound +from mcp_client.shared.exceptions import MCPDeprecationWarning + +__all__ = ["ClientPeer", "Meta"] + +Meta = dict[str, Any] +"""Type alias for the `_meta` field carried on request/notification params.""" + + +def dump_params(model: BaseModel | None, meta: Meta | None = None) -> dict[str, Any] | None: + """Serialize a params model to a wire dict, merging `meta` into `_meta`. + + Shared by `ClientPeer` and `Connection` so every typed convenience method + gets the same `_meta` handling. `meta` keys take precedence over any + `_meta` already present on the model. + + `meta` is serialized through `RequestParams` so Python field names emit + their wire aliases: an inbound `ctx.meta` carries `progress_token` (the + key `_extract_meta` validation produces), and forwarding it outbound via + `meta=ctx.meta` must put `progressToken` back on the wire. Keys not + declared on `RequestParamsMeta` pass through unchanged. + """ + out = model.model_dump(by_alias=True, mode="json", exclude_none=True) if model is not None else None + if meta: + wire_meta = RequestParams(_meta=cast(RequestParamsMeta, meta)).model_dump(by_alias=True, mode="json")["_meta"] + out = dict(out or {}) + out["_meta"] = {**out.get("_meta", {}), **wire_meta} + return out + + +class ClientPeer: + """Typed server-to-client request methods over a wrapped `Outbound`. + + Use this when you have a bare dispatcher (or any `Outbound`) and want the + typed methods (`sample`, `elicit_form`, `elicit_url`, `list_roots`, + `ping`) without writing your own host class. + """ + + def __init__(self, outbound: Outbound) -> None: + self._outbound = outbound + + async def send_raw_request( + self, + method: str, + params: Mapping[str, Any] | None, + opts: CallOptions | None = None, + ) -> dict[str, Any]: + return await self._outbound.send_raw_request(method, params, opts) + + async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: + await self._outbound.notify(method, params, opts) + + @overload + @deprecated("The sampling capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def sample( + self, + messages: list[SamplingMessage], + *, + max_tokens: int, + system_prompt: str | None = None, + include_context: IncludeContext | None = None, + temperature: float | None = None, + stop_sequences: list[str] | None = None, + metadata: dict[str, Any] | None = None, + model_preferences: ModelPreferences | None = None, + tools: None = None, + tool_choice: None = None, + meta: Meta | None = None, + opts: CallOptions | None = None, + ) -> CreateMessageResult: ... + @overload + @deprecated("The sampling capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def sample( + self, + messages: list[SamplingMessage], + *, + max_tokens: int, + system_prompt: str | None = None, + include_context: IncludeContext | None = None, + temperature: float | None = None, + stop_sequences: list[str] | None = None, + metadata: dict[str, Any] | None = None, + model_preferences: ModelPreferences | None = None, + tools: list[Tool], + tool_choice: ToolChoice | None = None, + meta: Meta | None = None, + opts: CallOptions | None = None, + ) -> CreateMessageResultWithTools: ... + @overload + @deprecated("The sampling capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def sample( + self, + messages: list[SamplingMessage], + *, + max_tokens: int, + system_prompt: str | None = None, + include_context: IncludeContext | None = None, + temperature: float | None = None, + stop_sequences: list[str] | None = None, + metadata: dict[str, Any] | None = None, + model_preferences: ModelPreferences | None = None, + tools: list[Tool] | None = None, + tool_choice: ToolChoice, + meta: Meta | None = None, + opts: CallOptions | None = None, + ) -> CreateMessageResultWithTools: ... + @deprecated("The sampling capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def sample( + self, + messages: list[SamplingMessage], + *, + max_tokens: int, + system_prompt: str | None = None, + include_context: IncludeContext | None = None, + temperature: float | None = None, + stop_sequences: list[str] | None = None, + metadata: dict[str, Any] | None = None, + model_preferences: ModelPreferences | None = None, + tools: list[Tool] | None = None, + tool_choice: ToolChoice | None = None, + meta: Meta | None = None, + opts: CallOptions | None = None, + ) -> CreateMessageResult | CreateMessageResultWithTools: + """Send a `sampling/createMessage` request to the peer. + + Raises: + MCPError: The peer responded with an error. + NoBackChannelError: No back-channel for server-initiated requests. + pydantic.ValidationError: The peer's result does not match the expected result type. + """ + params = CreateMessageRequestParams( + messages=messages, + system_prompt=system_prompt, + include_context=include_context, + temperature=temperature, + max_tokens=max_tokens, + stop_sequences=stop_sequences, + metadata=metadata, + model_preferences=model_preferences, + tools=tools, + tool_choice=tool_choice, + ) + result = await self.send_raw_request("sampling/createMessage", dump_params(params, meta), opts) + if tools is not None or tool_choice is not None: + return CreateMessageResultWithTools.model_validate(result, by_name=False) + return CreateMessageResult.model_validate(result, by_name=False) + + async def elicit_form( + self, + message: str, + requested_schema: ElicitRequestedSchema, + *, + meta: Meta | None = None, + opts: CallOptions | None = None, + ) -> ElicitResult: + """Send a form-mode `elicitation/create` request. + + Raises: + MCPError: The peer responded with an error. + NoBackChannelError: No back-channel for server-initiated requests. + pydantic.ValidationError: The peer's result does not match the expected result type. + """ + params = ElicitRequestFormParams(message=message, requested_schema=requested_schema) + result = await self.send_raw_request("elicitation/create", dump_params(params, meta), opts) + return ElicitResult.model_validate(result, by_name=False) + + async def elicit_url( + self, + message: str, + url: str, + elicitation_id: str, + *, + meta: Meta | None = None, + opts: CallOptions | None = None, + ) -> ElicitResult: + """Send a URL-mode `elicitation/create` request. + + Raises: + MCPError: The peer responded with an error. + NoBackChannelError: No back-channel for server-initiated requests. + pydantic.ValidationError: The peer's result does not match the expected result type. + """ + params = ElicitRequestURLParams(message=message, url=url, elicitation_id=elicitation_id) + result = await self.send_raw_request("elicitation/create", dump_params(params, meta), opts) + return ElicitResult.model_validate(result, by_name=False) + + @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) + async def list_roots(self, *, meta: Meta | None = None, opts: CallOptions | None = None) -> ListRootsResult: + """Send a `roots/list` request. + + Raises: + MCPError: The peer responded with an error. + NoBackChannelError: No back-channel for server-initiated requests. + pydantic.ValidationError: The peer's result does not match the expected result type. + """ + result = await self.send_raw_request("roots/list", dump_params(None, meta), opts) + return ListRootsResult.model_validate(result, by_name=False) + + async def ping(self, *, meta: Meta | None = None, opts: CallOptions | None = None) -> None: + """Send a `ping` request and ignore the result. + + Raises: + MCPError: The peer responded with an error. + NoBackChannelError: No back-channel for server-initiated requests. + """ + await self.send_raw_request("ping", dump_params(None, meta), opts) diff --git a/src/mcp-client/mcp_client/shared/subscriptions.py b/src/mcp-client/mcp_client/shared/subscriptions.py new file mode 100644 index 0000000000..30449a82ff --- /dev/null +++ b/src/mcp-client/mcp_client/shared/subscriptions.py @@ -0,0 +1,111 @@ +"""Typed event vocabulary for `subscriptions/listen` (2026-07-28, SEP-2575), shared by server and client. + +Every event is a level trigger ("this changed, refetch if you care"), so both sides bound buffers by dedupe. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from mcp_types import ( + NotificationParams, + PromptListChangedNotification, + ResourceListChangedNotification, + ResourceUpdatedNotification, + ResourceUpdatedNotificationParams, + ServerNotification, + SubscriptionFilter, + ToolListChangedNotification, +) + +__all__ = [ + "LISTEN_STREAM_METHODS", + "SUBSCRIPTION_ID_META_KEY", + "PromptsListChanged", + "ResourceUpdated", + "ResourcesListChanged", + "ServerEvent", + "ToolsListChanged", + "event_from_wire", + "event_matches", + "event_to_notification", +] + +SUBSCRIPTION_ID_META_KEY = "io.modelcontextprotocol/subscriptionId" +"""The `_meta` key on every listen-stream frame; the value is the `subscriptions/listen` request's JSON-RPC id.""" + + +@dataclass(frozen=True) +class ToolsListChanged: + """The server's tool list changed.""" + + +@dataclass(frozen=True) +class PromptsListChanged: + """The server's prompt list changed.""" + + +@dataclass(frozen=True) +class ResourcesListChanged: + """The server's resource list changed.""" + + +@dataclass(frozen=True) +class ResourceUpdated: + """The resource at `uri` changed and may need to be read again.""" + + uri: str + + +ServerEvent = ToolsListChanged | PromptsListChanged | ResourcesListChanged | ResourceUpdated +"""An event a server publishes for delivery to listen subscribers.""" + + +def event_to_notification(event: ServerEvent, meta: dict[str, Any]) -> ServerNotification: + """Build the stamped wire notification for `event` (the server's direction).""" + if isinstance(event, ToolsListChanged): + return ToolListChangedNotification(params=NotificationParams(_meta=meta)) + if isinstance(event, PromptsListChanged): + return PromptListChangedNotification(params=NotificationParams(_meta=meta)) + if isinstance(event, ResourcesListChanged): + return ResourceListChangedNotification(params=NotificationParams(_meta=meta)) + return ResourceUpdatedNotification(params=ResourceUpdatedNotificationParams(uri=event.uri, _meta=meta)) + + +_LIST_CHANGED_EVENTS: dict[str, ServerEvent] = { + "notifications/tools/list_changed": ToolsListChanged(), + "notifications/prompts/list_changed": PromptsListChanged(), + "notifications/resources/list_changed": ResourcesListChanged(), +} + +LISTEN_STREAM_METHODS: frozenset[str] = frozenset({*_LIST_CHANGED_EVENTS, "notifications/resources/updated"}) +"""The notification methods that ride `subscriptions/listen` streams at 2026-07-28 +(and, at that era, nowhere else): the change-notification vocabulary.""" + + +def event_from_wire(method: str, params: Mapping[str, Any] | None) -> ServerEvent | None: + """The event a raw listen-stream frame announces, or None if it carries none. + + Takes the raw wire dict: the client demultiplexes before the typed notification parse.""" + if (event := _LIST_CHANGED_EVENTS.get(method)) is not None: + return event + if method == "notifications/resources/updated": + uri = (params or {}).get("uri") + if isinstance(uri, str): + return ResourceUpdated(uri=uri) + return None + + +def event_matches(honored: SubscriptionFilter, uris: frozenset[str], event: ServerEvent) -> bool: + """Whether `event` is within the stream's honored filter (`uris`: the honored resource subscriptions as a set). + + The admission predicate both sides share: server delivery and client intake honor only what was acknowledged.""" + if isinstance(event, ToolsListChanged): + return honored.tools_list_changed is True + if isinstance(event, PromptsListChanged): + return honored.prompts_list_changed is True + if isinstance(event, ResourcesListChanged): + return honored.resources_list_changed is True + return event.uri in uris diff --git a/src/mcp-client/mcp_client/shared/tool_name_validation.py b/src/mcp-client/mcp_client/shared/tool_name_validation.py new file mode 100644 index 0000000000..bd2b7a0294 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/tool_name_validation.py @@ -0,0 +1,129 @@ +"""Tool name validation utilities according to SEP-986. + +Tool names SHOULD be between 1 and 128 characters in length (inclusive). +Tool names are case-sensitive. +Allowed characters: uppercase and lowercase ASCII letters (A-Z, a-z), +digits (0-9), underscore (_), dash (-), and dot (.). +Tool names SHOULD NOT contain spaces, commas, or other special characters. + +See: https://modelcontextprotocol.io/specification/2025-11-25/server/tools#tool-names +""" + +from __future__ import annotations + +import logging +import re +from dataclasses import dataclass, field + +logger = logging.getLogger("mcp.shared.tool_name_validation") + +# Regular expression for valid tool names according to SEP-986 specification +TOOL_NAME_REGEX = re.compile(r"^[A-Za-z0-9._-]{1,128}$") + +# SEP reference URL for warning messages +SEP_986_URL = "https://modelcontextprotocol.io/specification/2025-11-25/server/tools#tool-names" + + +@dataclass +class ToolNameValidationResult: + """Result of tool name validation. + + Attributes: + is_valid: Whether the tool name conforms to SEP-986 requirements. + warnings: List of warning messages for non-conforming aspects. + """ + + is_valid: bool + warnings: list[str] = field(default_factory=lambda: []) + + +def validate_tool_name(name: str) -> ToolNameValidationResult: + """Validate a tool name according to the SEP-986 specification. + + Args: + name: The tool name to validate. + + Returns: + ToolNameValidationResult containing validation status and any warnings. + """ + warnings: list[str] = [] + + # Check for empty name + if not name: + return ToolNameValidationResult( + is_valid=False, + warnings=["Tool name cannot be empty"], + ) + + # Check length + if len(name) > 128: + return ToolNameValidationResult( + is_valid=False, + warnings=[f"Tool name exceeds maximum length of 128 characters (current: {len(name)})"], + ) + + # Check for problematic patterns (warnings, not validation failures) + if " " in name: + warnings.append("Tool name contains spaces, which may cause parsing issues") + + if "," in name: + warnings.append("Tool name contains commas, which may cause parsing issues") + + # Check for potentially confusing leading/trailing characters + if name.startswith("-") or name.endswith("-"): + warnings.append("Tool name starts or ends with a dash, which may cause parsing issues in some contexts") + + if name.startswith(".") or name.endswith("."): + warnings.append("Tool name starts or ends with a dot, which may cause parsing issues in some contexts") + + # Check for invalid characters + if not TOOL_NAME_REGEX.fullmatch(name): + # Find all invalid characters (unique, preserving order) + invalid_chars: list[str] = [] + seen: set[str] = set() + for char in name: + if not re.match(r"[A-Za-z0-9._-]", char) and char not in seen: + invalid_chars.append(char) + seen.add(char) + + warnings.append(f"Tool name contains invalid characters: {', '.join(repr(c) for c in invalid_chars)}") + warnings.append("Allowed characters are: A-Z, a-z, 0-9, underscore (_), dash (-), and dot (.)") + + return ToolNameValidationResult(is_valid=False, warnings=warnings) + + return ToolNameValidationResult(is_valid=True, warnings=warnings) + + +def issue_tool_name_warning(name: str, warnings: list[str]) -> None: + """Log warnings for non-conforming tool names. + + Args: + name: The tool name that triggered the warnings. + warnings: List of warning messages to log. + """ + if not warnings: + return + + logger.warning(f'Tool name validation warning for "{name}":') + for warning in warnings: + logger.warning(f" - {warning}") + logger.warning("Tool registration will proceed, but this may cause compatibility issues.") + logger.warning("Consider updating the tool name to conform to the MCP tool naming standard.") + logger.warning(f"See SEP-986 ({SEP_986_URL}) for more details.") + + +def validate_and_warn_tool_name(name: str) -> bool: + """Validate a tool name and issue warnings for non-conforming names. + + This is the primary entry point for tool name validation. It validates + the name and logs any warnings via the logging module. + + Args: + name: The tool name to validate. + + Returns: + True if the name is valid, False otherwise. + """ + result = validate_tool_name(name) + issue_tool_name_warning(name, result.warnings) + return result.is_valid diff --git a/src/mcp-client/mcp_client/shared/transport_context.py b/src/mcp-client/mcp_client/shared/transport_context.py new file mode 100644 index 0000000000..8d15a2eaa2 --- /dev/null +++ b/src/mcp-client/mcp_client/shared/transport_context.py @@ -0,0 +1,45 @@ +"""Transport-specific metadata attached to each inbound message. + +`TransportContext` is the base; each transport defines its own subclass with +whatever fields make sense (HTTP request id, ASGI scope, stdio process handle, +etc.). The dispatcher passes it through opaquely; only the layers above the +dispatcher (`ServerRunner`, `Context`, user handlers) read its concrete fields. +""" + +from collections.abc import Mapping +from dataclasses import dataclass + +__all__ = ["TransportContext"] + + +@dataclass(kw_only=True, frozen=True) +class TransportContext: + """Base transport metadata for an inbound message. + + Subclass per transport and add fields as needed. Instances are immutable. + """ + + kind: str + """Short identifier for the transport (e.g. `"stdio"`, `"streamable-http"`).""" + + can_send_request: bool + """Whether this message's request-scoped channel can deliver a server-initiated request. + + `False` for any of three reasons: the response has no room (streamable + HTTP in JSON-response mode and the 2026-07-28 single-exchange entry answer + with one JSON-RPC reply), the client's reply has nowhere to land (stateless + HTTP, no session), or the protocol forbids server-initiated requests (any + 2026-07-28 connection, whose dispatch masks the flag off). `True` for a + plain duplex pipe (stdio, SSE) and stateful streamable HTTP with SSE + responses, all pre-2026-07-28. When `False`, + `DispatchContext.send_raw_request` raises `NoBackChannelError` instead of + parking a waiter no reply can reach. Says nothing about the connection's + standalone channel, which refuses separately. + """ + + headers: Mapping[str, str] | None = None + """Request headers carried by this message, when the transport has them. + + Populated by HTTP-based transports; `None` on stdio. Handlers should + None-check before use. + """ diff --git a/src/mcp-client/mcp_client/shared/uri_template.py b/src/mcp-client/mcp_client/shared/uri_template.py new file mode 100644 index 0000000000..399a23e84d --- /dev/null +++ b/src/mcp-client/mcp_client/shared/uri_template.py @@ -0,0 +1,1116 @@ +"""RFC 6570 URI Templates with bidirectional support. + +Provides both expansion (template + variables → URI) and matching +(URI → variables). RFC 6570 only specifies expansion; matching is the +inverse operation needed by MCP servers to route ``resources/read`` +requests to handlers. + +Supports Levels 1-3 fully, plus Level 4 explode modifier for path-like +operators (``{/var*}``, ``{.var*}``, ``{;var*}``). The Level 4 prefix +modifier (``{var:N}``) and query-explode (``{?var*}``) are not supported. + +Matching semantics +------------------ + +Matching is not specified by RFC 6570 (§1.4 explicitly defers to regex +languages). This implementation uses a two-ended scan that never +backtracks: match time is O(n·v) where n is URI length and v is the +number of template variables. Realistic templates have v < 10, making +this effectively linear; there is no input that produces +superpolynomial time. + +A template may contain **at most one multi-segment variable** — +``{+var}``, ``{#var}``, or an explode-modified variable (``{/var*}``, +``{.var*}``, ``{;var*}``). This variable greedily consumes whatever the +surrounding bounded variables and literals do not. Two such variables +in one template are inherently ambiguous (which one gets the extra +segment?) and are rejected at parse time. So are any two variables +adjacent with no literal between them — including a variable adjacent +to the multi-segment variable: the scan has nothing to anchor the +boundary on. Operators that emit their own lead character supply that +literal themselves, so ``{+path}{.ext}`` and ``{a}{.b}`` are fine +while ``{+path}{ext}`` and ``{a}{b}`` are not. + +Bounded variables before the multi-segment variable match **lazily** +(first occurrence of the following literal); those after match +**greedily** (last occurrence of the preceding literal). Templates +without a multi-segment variable match greedily throughout, identical +to regex semantics. + +Reserved expansion ``{+var}`` leaves ``?`` and ``#`` unencoded, but +the scan stops at those characters so ``{+path}{?q}`` can separate path +from query. A value containing a literal ``?`` or ``#`` expands fine +but will not round-trip through ``match()``. +""" + +from __future__ import annotations + +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Literal, TypeAlias, cast +from urllib.parse import quote, unquote + +__all__ = [ + "DEFAULT_MAX_TEMPLATE_LENGTH", + "DEFAULT_MAX_VARIABLES", + "DEFAULT_MAX_URI_LENGTH", + "InvalidUriTemplate", + "Operator", + "UriTemplate", + "Variable", +] + +Operator = Literal["", "+", "#", ".", "/", ";", "?", "&"] + +_OPERATORS: frozenset[str] = frozenset({"+", "#", ".", "/", ";", "?", "&"}) + +# RFC 6570 §2.3: varname = varchar *(["."] varchar), varchar = ALPHA / DIGIT / "_" +# Dots appear only between varchar groups — not consecutive, not trailing. +# (Percent-encoded varchars are technically allowed but unseen in practice.) +_VARNAME_RE = re.compile(r"^[A-Za-z0-9_]+(?:\.[A-Za-z0-9_]+)*$") + +DEFAULT_MAX_TEMPLATE_LENGTH = 8_192 +DEFAULT_MAX_VARIABLES = 256 +DEFAULT_MAX_URI_LENGTH = 65_536 + +# RFC 3986 reserved characters, kept unencoded by {+var} and {#var}. +_RESERVED = ":/?#[]@!$&'()*+,;=" + + +@dataclass(frozen=True) +class _OperatorSpec: + """Expansion behavior for a single operator (RFC 6570 §3.2, Table in §A).""" + + prefix: str + """Leading character emitted before the first variable.""" + separator: str + """Character between variables (and between exploded list items).""" + named: bool + """Emit ``name=value`` pairs (query/path-param style) rather than bare values.""" + allow_reserved: bool + """Keep reserved characters unencoded ({+var}, {#var}).""" + ifemp: str + """Suffix after a named variable whose expanded value is empty (RFC §A): '' for ;, '=' for ?/&.""" + + +_OPERATOR_SPECS: dict[Operator, _OperatorSpec] = { + "": _OperatorSpec(prefix="", separator=",", named=False, allow_reserved=False, ifemp=""), + "+": _OperatorSpec(prefix="", separator=",", named=False, allow_reserved=True, ifemp=""), + "#": _OperatorSpec(prefix="#", separator=",", named=False, allow_reserved=True, ifemp=""), + ".": _OperatorSpec(prefix=".", separator=".", named=False, allow_reserved=False, ifemp=""), + "/": _OperatorSpec(prefix="/", separator="/", named=False, allow_reserved=False, ifemp=""), + ";": _OperatorSpec(prefix=";", separator=";", named=True, allow_reserved=False, ifemp=""), + "?": _OperatorSpec(prefix="?", separator="&", named=True, allow_reserved=False, ifemp="="), + "&": _OperatorSpec(prefix="&", separator="&", named=True, allow_reserved=False, ifemp="="), +} + +# Per-operator stop characters for the linear scan. A bounded variable's +# value ends at the first occurrence of any character in its stop set, +# mirroring the character-class boundaries a regex would use but without +# the backtracking. +_STOP_CHARS: dict[Operator, str] = { + "": "/?#&,", # simple: everything structural is pct-encoded + "+": "?#", # reserved: / allowed, stop at query/fragment + "#": "", # fragment: tail of URI, nothing stops it + ".": "./?#", # label: stop at next . + "/": "/?#", # path segment: stop at next / + ";": ";/?#", # path-param value (may be empty: ;name) + "?": "&#", # query value (may be empty: ?name=) + "&": "&#", # query-cont value +} + + +class InvalidUriTemplate(ValueError): + """Raised when a URI template string is malformed or unsupported. + + Attributes: + template: The template string that failed to parse. + position: Character offset where the error was detected, or None + if the error is not tied to a specific position. + """ + + def __init__(self, message: str, *, template: str, position: int | None = None) -> None: + super().__init__(message) + self.template = template + self.position = position + + +@dataclass(frozen=True) +class Variable: + """A single variable within a URI template expression.""" + + name: str + operator: Operator + explode: bool = False + + +@dataclass +class _Expression: + """A parsed ``{...}`` expression: one operator, one or more variables.""" + + operator: Operator + variables: list[Variable] + + +_Part = str | _Expression + + +@dataclass(frozen=True) +class _Lit: + """A literal run in the flattened match-atom sequence.""" + + text: str + + +@dataclass(frozen=True) +class _Cap: + """A single-variable capture in the flattened match-atom sequence. + + ``ifemp`` marks the ``;`` operator's optional-equals quirk: ``{;id}`` + expands to ``;id=value`` or bare ``;id`` when the value is empty, so + the scan must accept both forms. + """ + + var: Variable + ifemp: bool = False + + +_Atom: TypeAlias = _Lit | _Cap + + +def _is_greedy(var: Variable) -> bool: + """Return True if this variable can span multiple path segments. + + Reserved/fragment expansion and explode variables are the only + constructs whose match range is not bounded by a single structural + delimiter. A template may contain at most one such variable. + """ + return var.explode or var.operator in ("+", "#") + + +def _is_str_sequence(value: object) -> bool: + """Check if value is a non-string sequence whose items are all strings.""" + if isinstance(value, str) or not isinstance(value, Sequence): + return False + seq = cast(Sequence[object], value) + return all(isinstance(item, str) for item in seq) + + +_PCT_TRIPLET_RE = re.compile(r"%[0-9A-Fa-f]{2}") + + +def _encode(value: str, *, allow_reserved: bool) -> str: + """Percent-encode a value per RFC 6570 §3.2.1. + + Simple expansion encodes everything except unreserved characters. + Reserved expansion (``{+var}``, ``{#var}``) additionally keeps + RFC 3986 reserved characters intact and passes through existing + ``%XX`` pct-triplets unchanged (RFC 6570 §3.2.3). A bare ``%`` not + followed by two hex digits is still encoded to ``%25``. + """ + if not allow_reserved: + return quote(value, safe="") + + # Reserved expansion: walk the string, pass through triplets as-is, + # quote the gaps between them. A bare % with no triplet lands in a + # gap and gets encoded normally. + out: list[str] = [] + last = 0 + for m in _PCT_TRIPLET_RE.finditer(value): + out.append(quote(value[last : m.start()], safe=_RESERVED)) + out.append(m.group()) + last = m.end() + out.append(quote(value[last:], safe=_RESERVED)) + return "".join(out) + + +def _expand_expression(expr: _Expression, variables: Mapping[str, str | Sequence[str]]) -> str: + """Expand a single ``{...}`` expression into its URI fragment. + + Walks the expression's variables, encoding and joining defined ones + according to the operator's spec. Undefined variables are skipped + (RFC 6570 §2.3); if all are undefined, the expression contributes + nothing (no prefix is emitted). + """ + spec = _OPERATOR_SPECS[expr.operator] + rendered: list[str] = [] + + for var in expr.variables: + if var.name not in variables: + # Undefined: skip entirely, no placeholder. + continue + + value = variables[var.name] + + # Explicit type guard: reject non-str scalars with a clear message + # rather than a confusing "not iterable" from the sequence branch. + if not isinstance(value, str) and not _is_str_sequence(value): + raise TypeError(f"Variable {var.name!r} must be str or a sequence of str, got {type(value).__name__}") + + if isinstance(value, str): + encoded = _encode(value, allow_reserved=spec.allow_reserved) + if spec.named: + rendered.append(f"{var.name}{spec.ifemp}" if value == "" else f"{var.name}={encoded}") + else: + rendered.append(encoded) + else: + # Sequence value. + items = [_encode(v, allow_reserved=spec.allow_reserved) for v in value] + if not items: + continue + if var.explode: + # Each item gets the operator's separator; named ops repeat the key. + if spec.named: + rendered.append( + spec.separator.join(f"{var.name}{spec.ifemp}" if v == "" else f"{var.name}={v}" for v in items) + ) + else: + rendered.append(spec.separator.join(items)) + else: + # Non-explode: comma-join into a single value, then apply + # ifemp to the joined result (RFC §3.2.1: behaves as if the + # value were the joined string). + joined = ",".join(items) + if spec.named: + rendered.append(f"{var.name}{spec.ifemp}" if joined == "" else f"{var.name}={joined}") + else: + rendered.append(joined) + + if not rendered: + return "" + return spec.prefix + spec.separator.join(rendered) + + +@dataclass(frozen=True) +class UriTemplate: + """A parsed RFC 6570 URI template. + + Construct via :meth:`parse`. Instances are immutable and hashable; + equality is based on the template string alone. + """ + + template: str + _parts: list[_Part] = field(repr=False, compare=False) + _variables: list[Variable] = field(repr=False, compare=False) + _prefix: list[_Atom] = field(repr=False, compare=False) + _greedy: Variable | None = field(repr=False, compare=False) + _suffix: list[_Atom] = field(repr=False, compare=False) + _query_variables: list[Variable] = field(repr=False, compare=False) + + @staticmethod + def is_template(value: str) -> bool: + """Check whether a string contains URI template expressions. + + A cheap heuristic for distinguishing concrete URIs from templates + without the cost of full parsing. Returns ``True`` if the string + contains at least one ``{...}`` pair. + + Example:: + + >>> UriTemplate.is_template("file://docs/{name}") + True + >>> UriTemplate.is_template("file://docs/readme.txt") + False + + Note: + This does not validate the template. A ``True`` result does + not guarantee :meth:`parse` will succeed. + """ + open_i = value.find("{") + return open_i != -1 and value.find("}", open_i) != -1 + + @classmethod + def parse( + cls, + template: str, + *, + max_length: int = DEFAULT_MAX_TEMPLATE_LENGTH, + max_variables: int = DEFAULT_MAX_VARIABLES, + ) -> UriTemplate: + """Parse a URI template string. + + Args: + template: An RFC 6570 URI template. + max_length: Maximum permitted length of the template string. + Guards against resource exhaustion. + max_variables: Maximum number of variables permitted across + all expressions. Counting variables rather than + ``{...}`` expressions closes the gap where a single + ``{v0,v1,...,vN}`` expression packs arbitrarily many + variables under one expression count. + + Raises: + InvalidUriTemplate: If the template is malformed, exceeds the + size limits, or uses unsupported RFC 6570 features. + """ + if len(template) > max_length: + raise InvalidUriTemplate( + f"Template exceeds maximum length of {max_length}", + template=template, + ) + + parts, variables = _parse(template, max_variables=max_variables) + + # Trailing {?...}/{&...} expressions are split off and matched as + # a query string (order-agnostic, partial, extras ignored) rather + # than via the linear scan. + path_parts, query_vars = _split_query_tail(parts) + atoms = _flatten(path_parts) + prefix, greedy, suffix = _partition_greedy(atoms, template) + + return cls( + template=template, + _parts=parts, + _variables=variables, + _prefix=prefix, + _greedy=greedy, + _suffix=suffix, + _query_variables=query_vars, + ) + + @property + def variables(self) -> list[Variable]: + """All variables in the template, in order of appearance.""" + return list(self._variables) + + @property + def variable_names(self) -> list[str]: + """All variable names in the template, in order of appearance.""" + return [v.name for v in self._variables] + + @property + def query_variable_names(self) -> frozenset[str]: + """Names of variables that :meth:`match` treats as optional query parameters. + + These are the variables in a trailing run of ``{?...}``/``{&...}`` + expressions, which are matched leniently: a URI that omits some + (or all) of them still matches, and the omitted names are simply + absent from the result. Any value bound to such a name therefore + needs a fallback for the omitted case. + + Every other variable is bound on every successful :meth:`match` + (possibly to an empty string) and is *not* in this set. That + includes a ``{&...}`` expression with no preceding ``{?...}``: it + never emits the ``?`` the lenient query split keys on, so it is + matched strictly. + """ + return frozenset(v.name for v in self._query_variables) + + def expand(self, variables: Mapping[str, str | Sequence[str]]) -> str: + """Expand the template by substituting variable values. + + String values are percent-encoded according to their operator: + simple ``{var}`` encodes reserved characters; ``{+var}`` and + ``{#var}`` leave them intact. Sequence values are joined with + commas for non-explode variables, or with the operator's + separator for explode variables. + + Example:: + + >>> t = UriTemplate.parse("file://docs/{name}") + >>> t.expand({"name": "hello world.txt"}) + 'file://docs/hello%20world.txt' + + >>> t = UriTemplate.parse("file://docs/{+path}") + >>> t.expand({"path": "src/main.py"}) + 'file://docs/src/main.py' + + >>> t = UriTemplate.parse("/search{?q,lang}") + >>> t.expand({"q": "mcp", "lang": "en"}) + '/search?q=mcp&lang=en' + + >>> t = UriTemplate.parse("/files{/path*}") + >>> t.expand({"path": ["a", "b", "c"]}) + '/files/a/b/c' + + Args: + variables: Values for each template variable. Keys must be + strings; values must be ``str`` or a sequence of ``str``. + + Returns: + The expanded URI string. + + Note: + Per RFC 6570, variables absent from the mapping are + **silently omitted**. This is the correct behavior for + optional query parameters (``{?page}`` with no page yields + no ``?page=``), but for required path segments it produces + a structurally incomplete URI. If you need all variables + present, validate before calling:: + + missing = set(t.variable_names) - variables.keys() + if missing: + raise ValueError(f"Missing: {missing}") + + Raises: + TypeError: If a value is neither ``str`` nor an iterable of + ``str``. Non-string scalars (``int``, ``None``) are not + coerced. + """ + out: list[str] = [] + for part in self._parts: + if isinstance(part, str): + out.append(part) + else: + out.append(_expand_expression(part, variables)) + return "".join(out) + + def match(self, uri: str, *, max_uri_length: int = DEFAULT_MAX_URI_LENGTH) -> dict[str, str | list[str]] | None: + """Match a concrete URI against this template and extract variables. + + This is the inverse of :meth:`expand`. The URI is matched via a + linear scan of the template and captured values are + percent-decoded. The round-trip ``match(expand({k: v})) == {k: v}`` + holds when ``v`` does not contain its operator's separator + unencoded: ``{.ext}`` with ``ext="tar.gz"`` expands to + ``.tar.gz`` but does not match — the scan stops ``ext`` at the + first ``.`` and the trailing ``.gz`` has nothing to consume it. + RFC 6570 §1.4 notes this is an inherent reversal limitation. + + Matching is structural at the URI level only: a simple ``{name}`` + will not match across a literal ``/`` in the URI (the scan stops + there), but a percent-encoded ``%2F`` that decodes to ``/`` is + accepted as part of the value. Path-safety validation belongs at + a higher layer; see :mod:`mcp_client.shared.path_security`. + + Example:: + + >>> t = UriTemplate.parse("file://docs/{name}") + >>> t.match("file://docs/readme.txt") + {'name': 'readme.txt'} + >>> t.match("file://docs/hello%20world.txt") + {'name': 'hello world.txt'} + + >>> t = UriTemplate.parse("file://docs/{+path}") + >>> t.match("file://docs/src/main.py") + {'path': 'src/main.py'} + + >>> t = UriTemplate.parse("/files{/path*}") + >>> t.match("/files/a/b/c") + {'path': ['a', 'b', 'c']} + + **Query parameters** (``{?q,lang}`` at the end of a template) + are matched leniently: order-agnostic, partial, and unrecognized + params are ignored. Absent params are omitted from the result so + downstream function defaults can apply:: + + >>> t = UriTemplate.parse("logs://{service}{?since,level}") + >>> t.match("logs://api") + {'service': 'api'} + >>> t.match("logs://api?level=error") + {'service': 'api', 'level': 'error'} + >>> t.match("logs://api?level=error&since=5m&utm=x") + {'service': 'api', 'since': '5m', 'level': 'error'} + + Args: + uri: A concrete URI string. + max_uri_length: Maximum permitted length of the input URI. + Oversized inputs return ``None`` without scanning, + guarding against resource exhaustion. + + Returns: + A mapping from variable names to decoded values (``str`` for + scalar variables, ``list[str]`` for explode variables), or + ``None`` if the URI does not match the template or exceeds + ``max_uri_length``. + """ + if len(uri) > max_uri_length: + return None + + if self._query_variables: + # Two-phase: scan matches the path, the query is split and + # decoded manually. Query params may be partial, reordered, + # or include extras; absent params stay absent so downstream + # defaults can apply. Fragment is stripped first since the + # template's {?...} tail never describes a fragment. + before_fragment, _, _ = uri.partition("#") + path, _, query = before_fragment.partition("?") + result = self._scan(path) + if result is None: + return None + if query: + parsed = _parse_query(query) + for var in self._query_variables: + if var.name in parsed: + result[var.name] = parsed[var.name] + return result + + return self._scan(uri) + + def _scan(self, uri: str) -> dict[str, str | list[str]] | None: + """Run the two-ended linear scan against the path portion of a URI.""" + n = len(uri) + + if self._greedy is None: + # No greedy var: the suffix IS the whole template, scanned + # right-to-left and anchored so atoms[0] matches at position 0. + suffix = _scan_suffix(self._suffix, uri, n, anchored=True) + if suffix is None: + return None + suffix_result, suffix_start = suffix + return suffix_result if suffix_start == 0 else None + + # Greedy var present. The parser rejects a capture adjacent to + # the greedy slot, so a non-empty suffix begins with a _Lit whose + # rfind-derived anchor does not depend on how far the prefix + # scans. Scan the suffix first, then give the prefix that exact + # position as its ceiling so it cannot consume past the anchor. + suffix = _scan_suffix(self._suffix, uri, n, anchored=False) + if suffix is None: + return None + suffix_result, suffix_start = suffix + prefix = _scan_prefix(self._prefix, uri, 0, suffix_start) + if prefix is None: + return None + prefix_result, prefix_end = prefix + + # Prefix consumed [0, prefix_end); suffix consumed [suffix_start, n); + # the greedy var takes the gap. The prefix scan is bounded by + # suffix_start, so this holds by construction; guard explicitly + # rather than asserting so a future regression surfaces as a + # non-match, not an exception. + if suffix_start < prefix_end: + return None # pragma: no cover - unreachable while bounds hold + middle = uri[prefix_end:suffix_start] + greedy_value = _extract_greedy(self._greedy, middle) + if greedy_value is None: + return None + + return {**prefix_result, self._greedy.name: greedy_value, **suffix_result} + + def __str__(self) -> str: + return self.template + + +def _parse_query(query: str) -> dict[str, str]: + """Parse a query string into a name→value mapping. + + Unlike ``urllib.parse.parse_qs``, this follows RFC 3986 semantics: + ``+`` is a literal sub-delim, not a space. Form-urlencoding treats + ``+`` as space for HTML form submissions, but RFC 6570 and MCP + resource URIs follow RFC 3986 where only ``%20`` encodes a space. + + Parameter names are **not** percent-decoded. RFC 6570 expansion + never encodes variable names, so a legitimate match will always + have the name in literal form. Decoding names would let + ``%74oken=evil&token=real`` shadow the real ``token`` parameter + via first-wins. + + Duplicate keys keep the first value. Pairs without ``=`` are + treated as empty-valued. + """ + result: dict[str, str] = {} + for pair in query.split("&"): + name, _, value = pair.partition("=") + if name and name not in result: + result[name] = unquote(value) + return result + + +def _extract_greedy(var: Variable, raw: str) -> str | list[str] | None: + """Decode the greedy variable's isolated middle span. + + For scalar greedy (``{+var}``, ``{#var}``) this is a stop-char + validation and a single ``unquote``. For explode variables the span + is a run of separator-delimited segments (``/a/b/c`` or + ``;keys=a;keys=b``) that is split, validated, and decoded per item. + """ + spec = _OPERATOR_SPECS[var.operator] + stops = _STOP_CHARS[var.operator] + + if not var.explode: + if any(c in stops for c in raw): + return None + return unquote(raw) + + sep = spec.separator + if not raw: + return [] + # A non-empty explode span must begin with the separator: {/a*} + # expands to "/x/y", never "x/y". The scan does not consume the + # separator itself, so it must be the first character here. + if raw[0] != sep: + return None + # Segments must not contain the operator's non-separator stop + # characters (e.g. {/path*} segments may contain neither ? nor #). + body_stops = set(stops) - {sep} + if any(c in body_stops for c in raw): + return None + + segments: list[str] = [] + prefix = f"{var.name}=" + # split()[0] is always "" because raw starts with the separator; + # subsequent empties are legitimate values ({/path*} with + # ["a","","c"] expands to /a//c). + for seg in raw.split(sep)[1:]: + if spec.named: + # Named explode emits name=value per item (or bare name + # under ; with empty value). Validate the name and strip + # the prefix before decoding. + if seg.startswith(prefix): + seg = seg[len(prefix) :] + elif seg == var.name: + seg = "" + else: + return None + segments.append(unquote(seg)) + return segments + + +def _split_query_tail(parts: list[_Part]) -> tuple[list[_Part], list[Variable]]: + """Separate trailing ``?``/``&`` expressions from the path portion. + + Lenient query matching (order-agnostic, partial, ignores extras) + applies when a template ends with one or more consecutive ``?``/``&`` + expressions and the preceding path portion contains no literal + ``?``. If the path has a literal ``?`` (e.g., ``?fixed=1{&page}``), + the URI's ``?`` split won't align with the template's expression + boundary, so the strict scan is used instead. + + Returns: + A pair ``(path_parts, query_vars)``. If lenient matching does + not apply, ``query_vars`` is empty and ``path_parts`` is the + full input. + """ + split = len(parts) + for i in range(len(parts) - 1, -1, -1): + part = parts[i] + if isinstance(part, _Expression) and part.operator in ("?", "&"): + split = i + else: + break + + if split == len(parts): + return parts, [] + + # The tail must start with a {?...} expression so that expand() + # emits a ? the URI can split on. A standalone {&page} expands + # with an & prefix, which partition("?") won't find. + first = parts[split] + assert isinstance(first, _Expression) + if first.operator != "?": + return parts, [] + + # If the path portion contains a literal ?/# or a {?...}/{#...} + # expression, lenient matching's partition("#") then partition("?") + # would strip content the path scan expects to see. Fall back to + # the strict scan. + for part in parts[:split]: + if isinstance(part, str): + if "?" in part or "#" in part: + return parts, [] + elif part.operator in ("?", "#"): + return parts, [] + + query_vars: list[Variable] = [] + for part in parts[split:]: + assert isinstance(part, _Expression) + query_vars.extend(part.variables) + + return parts[:split], query_vars + + +def _parse(template: str, *, max_variables: int) -> tuple[list[_Part], list[Variable]]: + """Split a template into an ordered sequence of literals and expressions. + + Walks the string, alternating between collecting literal runs and + parsing ``{...}`` expressions. The resulting ``parts`` sequence + preserves positional interleaving so ``match()`` and ``expand()`` can + walk it in order. + + Raises: + InvalidUriTemplate: On unclosed braces, too many expressions, or + any error surfaced by :func:`_parse_expression`. + """ + parts: list[_Part] = [] + variables: list[Variable] = [] + i = 0 + n = len(template) + + while i < n: + # Find the next expression opener from the current cursor. + brace = template.find("{", i) + + if brace == -1: + # No more expressions; everything left is a trailing literal. + parts.append(template[i:]) + break + + if brace > i: + # Literal text between cursor and the brace. + parts.append(template[i:brace]) + + end = template.find("}", brace) + if end == -1: + raise InvalidUriTemplate( + f"Unclosed expression at position {brace}", + template=template, + position=brace, + ) + + # Delegate body (between braces, exclusive) to the expression parser. + expr = _parse_expression(template, template[brace + 1 : end], brace) + parts.append(expr) + variables.extend(expr.variables) + + if len(variables) > max_variables: + raise InvalidUriTemplate( + f"Template exceeds maximum of {max_variables} variables", + template=template, + ) + + # Advance past the closing brace. + i = end + 1 + + _check_duplicate_variables(template, variables) + _check_single_query_expression(template, parts) + return parts, variables + + +def _parse_expression(template: str, body: str, pos: int) -> _Expression: + """Parse the body of a single ``{...}`` expression. + + The body is everything between the braces. It consists of an optional + leading operator character followed by one or more comma-separated + variable specifiers. Each specifier is a name with an optional + trailing ``*`` (explode modifier). + + Args: + template: The full template string, for error reporting. + body: The expression body, braces excluded. + pos: Character offset of the opening brace, for error reporting. + + Raises: + InvalidUriTemplate: On empty body, invalid variable names, or + unsupported modifiers. + """ + if not body: + raise InvalidUriTemplate(f"Empty expression at position {pos}", template=template, position=pos) + + # Peel off the operator, if any. Membership check justifies the cast. + operator: Operator = "" + if body[0] in _OPERATORS: + operator = cast(Operator, body[0]) + body = body[1:] + if not body: + raise InvalidUriTemplate( + f"Expression has operator but no variables at position {pos}", + template=template, + position=pos, + ) + + # Remaining body is comma-separated variable specs: name[*] + variables: list[Variable] = [] + for spec in body.split(","): + if ":" in spec: + raise InvalidUriTemplate( + f"Prefix modifier {{var:N}} is not supported (in {spec!r} at position {pos})", + template=template, + position=pos, + ) + + explode = spec.endswith("*") + name = spec[:-1] if explode else spec + + if not _VARNAME_RE.fullmatch(name): + raise InvalidUriTemplate( + f"Invalid variable name {name!r} at position {pos}", + template=template, + position=pos, + ) + + # Explode only makes sense for operators that repeat a separator. + # Simple/reserved/fragment have no per-item separator; query-explode + # needs order-agnostic dict matching which we don't support yet. + if explode and operator in ("", "+", "#", "?", "&"): + raise InvalidUriTemplate( + f"Explode modifier on {{{operator}{name}*}} is not supported for matching", + template=template, + position=pos, + ) + + variables.append(Variable(name=name, operator=operator, explode=explode)) + + return _Expression(operator=operator, variables=variables) + + +def _check_duplicate_variables(template: str, variables: list[Variable]) -> None: + """Reject templates that use the same variable name more than once. + + RFC 6570 requires repeated variables to expand to the same value, + which would require backreference matching with potentially + exponential cost. Rather than silently returning only the last + captured value, we reject at parse time. + + Raises: + InvalidUriTemplate: If any variable name appears more than once. + """ + seen: set[str] = set() + for var in variables: + if var.name in seen: + raise InvalidUriTemplate( + f"Variable {var.name!r} appears more than once; repeated variables are not supported", + template=template, + ) + seen.add(var.name) + + +def _check_single_query_expression(template: str, parts: list[_Part]) -> None: + """Reject templates with more than one ``{?...}`` expression. + + The ``?`` operator emits a leading ``?``, so two such expressions + expand to a URI with two ``?`` characters — malformed per RFC 3986 + §3.4. Use ``{?a,b}`` or ``{?a}{&b}`` for multiple query parameters. + """ + seen = False + for part in parts: + if isinstance(part, _Expression) and part.operator == "?": + if seen: + raise InvalidUriTemplate( + "Template contains more than one {?...} expression; " + "use {?a,b} or {?a}{&b} for multiple query parameters", + template=template, + ) + seen = True + + +def _flatten(parts: list[_Part]) -> list[_Atom]: + """Lower expressions into a flat sequence of literals and single-variable captures. + + Operator prefixes and separators become explicit ``_Lit`` atoms so + the scan only ever sees two atom kinds. Adjacent literals are + coalesced so that anchor-finding (``find``/``rfind``) operates on + the longest possible literal, reducing false matches. + + Explode variables emit no lead literal: the explode capture + includes its own separator-prefixed repetitions (``{/a*}`` → + ``/x/y/z``, not ``/`` then ``x/y/z``). + """ + atoms: list[_Atom] = [] + + def push_lit(text: str) -> None: + if not text: + return + if atoms and isinstance(atoms[-1], _Lit): + atoms[-1] = _Lit(atoms[-1].text + text) + else: + atoms.append(_Lit(text)) + + for part in parts: + if isinstance(part, str): + push_lit(part) + continue + spec = _OPERATOR_SPECS[part.operator] + for i, var in enumerate(part.variables): + lead = spec.prefix if i == 0 else spec.separator + if var.explode: + atoms.append(_Cap(var)) + elif spec.named: + # ; uses ifemp (bare name when empty); ? and & always + # emit name= so the equals is part of the literal. + if part.operator == ";": + push_lit(f"{lead}{var.name}") + atoms.append(_Cap(var, ifemp=True)) + else: + push_lit(f"{lead}{var.name}=") + atoms.append(_Cap(var)) + else: + push_lit(lead) + atoms.append(_Cap(var)) + return atoms + + +def _partition_greedy(atoms: list[_Atom], template: str) -> tuple[list[_Atom], Variable | None, list[_Atom]]: + """Split atoms at the single greedy variable, if any. + + Returns ``(prefix, greedy_var, suffix)``. If there is no greedy + variable the entire atom list is returned as the suffix so that + the right-to-left scan (which matches regex-greedy semantics) + handles it. + + Raises: + InvalidUriTemplate: If two variables are adjacent with no + literal between them — whether or not one is the + multi-segment variable, the scan has nothing to anchor the + boundary on — or if more than one multi-segment variable + is present (two are inherently ambiguous: there is no + principled way to decide which one absorbs an extra + segment). + """ + greedy_idx: int | None = None + prev: _Atom | None = None + for i, atom in enumerate(atoms): + if isinstance(atom, _Cap): + if isinstance(prev, _Cap): + raise InvalidUriTemplate( + f"Variables {prev.var.name!r} and {atom.var.name!r} are adjacent " + "with no literal separator; matching cannot determine where one " + "ends and the other begins. Add a literal between them or use a " + "single variable.", + template=template, + ) + if _is_greedy(atom.var): + if greedy_idx is not None: + raise InvalidUriTemplate( + "Template contains more than one multi-segment variable " + "({+var}, {#var}, or explode modifier); matching would be ambiguous", + template=template, + ) + greedy_idx = i + prev = atom + if greedy_idx is None: + return [], None, atoms + greedy = atoms[greedy_idx] + assert isinstance(greedy, _Cap) + return atoms[:greedy_idx], greedy.var, atoms[greedy_idx + 1 :] + + +def _scan_suffix( + atoms: Sequence[_Atom], uri: str, end: int, *, anchored: bool +) -> tuple[dict[str, str | list[str]], int] | None: + """Scan atoms right-to-left from ``end``, returning captures and start position. + + Each bounded variable takes the minimum span that lets its + preceding literal match (found via ``rfind``), which makes the + *first* variable in template order greedy — identical to Python + regex semantics for a sequence of greedy groups. + + When ``anchored`` is true the atom sequence is the entire template + (no greedy variable), so ``atoms[0]`` must match at URI position 0 + rather than at its rightmost occurrence. + """ + result: dict[str, str | list[str]] = {} + pos = end + i = len(atoms) - 1 + while i >= 0: + atom = atoms[i] + if isinstance(atom, _Lit): + n = len(atom.text) + if pos < n or uri[pos - n : pos] != atom.text: + return None + pos -= n + i -= 1 + continue + + var = atom.var + stops = _STOP_CHARS[var.operator] + prev = atoms[i - 1] if i > 0 else None + + if atom.ifemp: + # ;name or ;name=value. The preceding _Lit is ";name". + # Try empty first: if the lit ends at pos the value is + # absent (RFC ifemp). Otherwise require =value. + assert isinstance(prev, _Lit) + if uri.endswith(prev.text, 0, pos): + result[var.name] = "" + i -= 1 + continue + earliest = pos + while earliest > 0 and uri[earliest - 1] not in stops: + earliest -= 1 + eq = uri.find("=", earliest, pos) + if eq == -1: + return None + result[var.name] = unquote(uri[eq + 1 : pos]) + pos = eq + i -= 1 + continue + + # Earliest valid start: the var cannot extend left past any + # stop-char, so scan backward to find that boundary. + earliest = pos + while earliest > 0 and uri[earliest - 1] not in stops: + earliest -= 1 + + if prev is None: + start = earliest + else: + # prev is a _Lit: the parser rejects two adjacent captures, + # so the only possible neighbour kind is a literal. + assert isinstance(prev, _Lit) + if anchored and i - 1 == 0: + # First atom of the whole template: positionally fixed at + # 0, not rightmost occurrence. rfind would land inside the + # value when the literal repeats there (e.g. "prefix-{id}" + # against "prefix-prefix-123"). + start = len(prev.text) + if start < earliest or start > pos: + return None + else: + # Rightmost occurrence of the preceding literal whose end + # falls within the var's valid range. + idx = uri.rfind(prev.text, 0, pos) + if idx == -1 or idx + len(prev.text) < earliest: + return None + start = idx + len(prev.text) + + result[var.name] = unquote(uri[start:pos]) + pos = start + i -= 1 + return result, pos + + +def _scan_prefix( + atoms: Sequence[_Atom], uri: str, start: int, limit: int +) -> tuple[dict[str, str | list[str]], int] | None: + """Scan atoms left-to-right from ``start``, not exceeding ``limit``. + + Each bounded variable takes the minimum span that lets its + following literal match (found via ``find``), leaving the + greedy variable as much of the URI as possible. + """ + result: dict[str, str | list[str]] = {} + pos = start + for i, atom in enumerate(atoms): + if isinstance(atom, _Lit): + end = pos + len(atom.text) + if end > limit or uri[pos:end] != atom.text: + return None + pos = end + continue + + var = atom.var + stops = _STOP_CHARS[var.operator] + # Every capture here is followed by a literal: the parser rejects + # two adjacent captures, and a capture at the END of the prefix + # would be adjacent to the greedy variable. + nxt = atoms[i + 1] + assert isinstance(nxt, _Lit) + + if atom.ifemp: + # RFC §3.2.7 ifemp: ;name=val for non-empty, bare ;name for + # empty. Decide which form is present without falling through + # to the stop-char scan when the value is empty. + if uri.startswith(nxt.text, pos): + # Following literal begins immediately: value is empty. + # Checked before '=' so a literal that itself starts + # with '=' is not mistaken for the ifemp separator. + result[var.name] = "" + continue + if pos < limit and uri[pos] == "=": + pos += 1 # value follows; fall through to the scan + else: + # The following literal does not start here and there is + # no '=': the URI's name continued past the template's + # (e.g. ;keys vs ;key) — no parse. + return None + + # Latest valid end: the var stops at the first stop-char or + # the scan limit, whichever comes first. + latest = pos + while latest < limit and uri[latest] not in stops: + latest += 1 + + # First occurrence of the following literal: the capture takes + # the minimum span, leaving the greedy variable as much of the + # URI as possible. The search window's upper bound already + # forces any hit to start at or before ``latest``, so the var + # never extends past a stop-char. + end = uri.find(nxt.text, pos, latest + len(nxt.text)) + if end == -1: + return None + + result[var.name] = unquote(uri[pos:end]) + pos = end + return result, pos diff --git a/src/mcp-client/pyproject.toml b/src/mcp-client/pyproject.toml new file mode 100644 index 0000000000..be46159e63 --- /dev/null +++ b/src/mcp-client/pyproject.toml @@ -0,0 +1,64 @@ +[project] +name = "mcp-client" +dynamic = ["version", "dependencies"] +description = "Model Context Protocol client SDK" +readme = "README.md" +requires-python = ">=3.10" +authors = [{ name = "Model Context Protocol a Series of LF Projects, LLC." }] +maintainers = [ + { name = "David Soria Parra", email = "davidsp@anthropic.com" }, + { name = "Marcelo Trylesinski", email = "marcelotryle@gmail.com" }, + { name = "Max Isbey", email = "maxisbey@anthropic.com" }, + { name = "Felix Weinberger", email = "fweinberger@anthropic.com" }, +] +keywords = ["mcp", "llm", "automation"] +license = { text = "MIT" } +classifiers = [ + "Development Status :: 5 - Production/Stable", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", +] + +[project.urls] +Homepage = "https://modelcontextprotocol.io" +Documentation = "https://py.sdk.modelcontextprotocol.io/" +Repository = "https://github.com/modelcontextprotocol/python-sdk" +Issues = "https://github.com/modelcontextprotocol/python-sdk/issues" + +[build-system] +requires = ["hatchling", "uv-dynamic-versioning"] +build-backend = "hatchling.build" + +[tool.hatch.version] +source = "uv-dynamic-versioning" + +[tool.uv-dynamic-versioning] +vcs = "git" +style = "pep440" +bump = true + +[tool.hatch.metadata.hooks.uv-dynamic-versioning] +dependencies = [ + "anyio>=4.10; python_version >= '3.14'", + "anyio>=4.9; python_version < '3.14'", + "httpx2>=2.5.0", + "mcp-types=={{ version }}", + "pydantic>=2.12.0", + "jsonschema>=4.20.0", + "pywin32>=311; sys_platform == 'win32'", + "pyjwt[crypto]>=2.10.1", + "typing-extensions>=4.13.0", + "opentelemetry-api>=1.28.0", +] + +[tool.hatch.build.targets.sdist.force-include] +"../../LICENSE" = "LICENSE" + +[tool.hatch.build.targets.wheel] +packages = ["mcp_client"] diff --git a/src/mcp/client/__init__.py b/src/mcp/client/__init__.py index d6b07045ce..fc7c410de2 100644 --- a/src/mcp/client/__init__.py +++ b/src/mcp/client/__init__.py @@ -1,26 +1,100 @@ -"""MCP Client module.""" +from mcp_client.client import ( + CacheConfig as CacheConfig, +) +from mcp_client.client import ( + CacheEntry as CacheEntry, +) +from mcp_client.client import ( + CacheKey as CacheKey, +) +from mcp_client.client import ( + CacheMode as CacheMode, +) +from mcp_client.client import ( + ClaimContext as ClaimContext, +) +from mcp_client.client import ( + Client as Client, +) +from mcp_client.client import ( + ClientExtension as ClientExtension, +) +from mcp_client.client import ( + ClientRequestContext as ClientRequestContext, +) +from mcp_client.client import ( + ClientSession as ClientSession, +) +from mcp_client.client import ( + IncomingMessage as IncomingMessage, +) +from mcp_client.client import ( + InMemoryResponseCacheStore as InMemoryResponseCacheStore, +) +from mcp_client.client import ( + InputRequiredRoundsExceededError as InputRequiredRoundsExceededError, +) +from mcp_client.client import ( + NotificationBinding as NotificationBinding, +) +from mcp_client.client import ( + ResponseCacheStore as ResponseCacheStore, +) +from mcp_client.client import ( + ResultClaim as ResultClaim, +) +from mcp_client.client import ( + Transport as Transport, +) +from mcp_client.client import ( + UnexpectedClaimedResult as UnexpectedClaimedResult, +) +from mcp_client.client import ( + advertise as advertise, +) -from mcp.client._input_required import InputRequiredRoundsExceededError -from mcp.client._transport import Transport -from mcp.client.caching import ( - CacheConfig, - CacheEntry, - CacheKey, - CacheMode, - InMemoryResponseCacheStore, - ResponseCacheStore, -) -from mcp.client.client import Client -from mcp.client.context import ClientRequestContext -from mcp.client.extension import ( - ClaimContext, - ClientExtension, - NotificationBinding, - ResultClaim, - UnexpectedClaimedResult, - advertise, -) -from mcp.client.session import ClientSession, IncomingMessage +from . import ( + _input_required as _input_required, +) +from . import ( + _memory as _memory, +) +from . import ( + _probe as _probe, +) +from . import ( + _transport as _transport, +) +from . import ( + caching as caching, +) +from . import ( + client as client, +) +from . import ( + context as context, +) +from . import ( + extension as extension, +) +from . import ( + session as session, +) +from . import ( + session_group as session_group, +) +from . import ( + sse as sse, +) +from . import ( + stdio as stdio, +) +from . import ( + streamable_http as streamable_http, +) +from . import ( + subscriptions as subscriptions, +) __all__ = [ "CacheConfig", diff --git a/src/mcp/client/__main__.py b/src/mcp/client/__main__.py index 60e3b02390..a8e8005074 100644 --- a/src/mcp/client/__main__.py +++ b/src/mcp/client/__main__.py @@ -1,81 +1,4 @@ -import argparse -import logging -import sys -import warnings -from functools import partial -from urllib.parse import urlparse - -import anyio -import mcp_types as types - -from mcp.client._transport import ReadStream, WriteStream -from mcp.client.session import ClientSession, IncomingMessage -from mcp.client.sse import sse_client -from mcp.client.stdio import StdioServerParameters, stdio_client -from mcp.shared.message import SessionMessage - -if not sys.warnoptions: - warnings.simplefilter("ignore") - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger("client") - - -async def message_handler(message: IncomingMessage) -> None: - if isinstance(message, Exception): - logger.error("Error: %s", message) - return - - logger.info("Received message from server: %s", message) - - -async def run_session( - read_stream: ReadStream[SessionMessage | Exception], - write_stream: WriteStream[SessionMessage], - client_info: types.Implementation | None = None, -): - async with ClientSession( - read_stream, - write_stream, - message_handler=message_handler, - client_info=client_info, - ) as session: - logger.info("Initializing session") - await session.initialize() - logger.info("Initialized") - - -async def main(command_or_url: str, args: list[str], env: list[tuple[str, str]]): - env_dict = dict(env) - - if urlparse(command_or_url).scheme in ("http", "https"): - # Use SSE client for HTTP(S) URLs - async with sse_client(command_or_url) as streams: - await run_session(*streams) - else: - # Use stdio client for commands - server_parameters = StdioServerParameters(command=command_or_url, args=args, env=env_dict) - async with stdio_client(server_parameters) as streams: - await run_session(*streams) - - -def cli(): - parser = argparse.ArgumentParser() - parser.add_argument("command_or_url", help="Command or URL to connect to") - parser.add_argument("args", nargs="*", help="Additional arguments") - parser.add_argument( - "-e", - "--env", - nargs=2, - action="append", - metavar=("KEY", "VALUE"), - help="Environment variables to set. Can be used multiple times.", - default=[], - ) - - args = parser.parse_args() - anyio.run(partial(main, args.command_or_url, args.args, args.env), backend="trio") - +from mcp_client.client.__main__ import cli if __name__ == "__main__": cli() diff --git a/src/mcp/client/_input_required.py b/src/mcp/client/_input_required.py index fe3f59e175..d74becc731 100644 --- a/src/mcp/client/_input_required.py +++ b/src/mcp/client/_input_required.py @@ -1,127 +1,26 @@ -"""SEP-2322 client-side multi-round-trip driver. - -When a server returns `InputRequiredResult` instead of the normal result of a -`tools/call` / `prompts/get` / `resources/read`, the client fulfils the -embedded `input_requests` (sampling, elicitation, roots) and retries the -original request carrying the responses and the echoed opaque `request_state`. -This module implements that retry loop as a pure function so it can drive any -of the three methods identically; `Client` builds the `dispatch` and `retry` -closures, `ClientSession` stays mechanics-only. -""" - -from __future__ import annotations - -from collections.abc import Awaitable, Callable -from typing import TypeVar - -import anyio -import anyio.abc -from mcp_types import ErrorData, InputRequest, InputRequiredResult, InputResponse, InputResponses - -from mcp.shared.exceptions import MCPError - -DEFAULT_INPUT_REQUIRED_MAX_ROUNDS = 10 -"""Default cap on `InputRequiredResult` retry rounds before the driver gives up. - -Matches the typescript-sdk default; csharp-sdk and go-sdk use the same value -as a hard constant. -""" - -_STATE_ONLY_BACKOFF_INITIAL_SECONDS = 0.05 -"""First sleep when an `InputRequiredResult` carries only `request_state` (no input requests).""" - -_STATE_ONLY_BACKOFF_CAP_SECONDS = 0.25 -"""Upper bound on the state-only backoff sleep; reached after three consecutive state-only legs.""" - - -ResultT = TypeVar("ResultT") - - -class InputRequiredRoundsExceededError(RuntimeError): - """The server kept returning `InputRequiredResult` past the configured `max_rounds`.""" - - def __init__(self, max_rounds: int) -> None: - super().__init__( - f"Server returned InputRequiredResult for more than {max_rounds} rounds; " - "raise input_required_max_rounds on the Client, or use " - "client.session.(..., allow_input_required=True) to drive the loop manually." - ) - self.max_rounds = max_rounds - - -async def run_input_required_driver( - first: InputRequiredResult, - *, - dispatch: Callable[[str, InputRequest], Awaitable[InputResponse | ErrorData]], - retry: Callable[[InputResponses | None, str | None], Awaitable[ResultT | InputRequiredResult]], - max_rounds: int = DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, -) -> ResultT: - """Resolve an `InputRequiredResult` to its terminal result. - - Loops until `retry` returns a non-`InputRequiredResult`, or `max_rounds` is - exhausted. Each round either dispatches all `input_requests` concurrently - and retries with the collected responses, or — when the server sent only - `request_state` — sleeps with exponential backoff (50ms doubling to a 250ms - cap, reset by any leg that carries input requests) and retries empty. - `request_state` is passed through byte-exact and never inspected. - - Args: - first: The `InputRequiredResult` the original call returned. - dispatch: Runs one embedded `InputRequest` through the client's - sampling / elicitation / roots callbacks. Called concurrently per - request key. An `ErrorData` return aborts the loop as an `MCPError`. - retry: Re-issues the original request with the collected responses and - the latest `request_state`. Each call mints a fresh JSON-RPC id. - max_rounds: Cap on retry rounds. - - Raises: - InputRequiredRoundsExceededError: `max_rounds` exhausted. - MCPError: A `dispatch` call returned `ErrorData`. - """ - rounds = 0 - state_only_delay = _STATE_ONLY_BACKOFF_INITIAL_SECONDS - current: ResultT | InputRequiredResult = first - while isinstance(current, InputRequiredResult): - rounds += 1 - if rounds > max_rounds: - raise InputRequiredRoundsExceededError(max_rounds) - if current.input_requests: - state_only_delay = _STATE_ONLY_BACKOFF_INITIAL_SECONDS - responses: InputResponses | None = await _dispatch_all(current.input_requests, dispatch) - else: - await anyio.sleep(state_only_delay) - state_only_delay = min(state_only_delay * 2, _STATE_ONLY_BACKOFF_CAP_SECONDS) - responses = None - current = await retry(responses, current.request_state) - return current - - -async def _dispatch_all( - requests: dict[str, InputRequest], - dispatch: Callable[[str, InputRequest], Awaitable[InputResponse | ErrorData]], -) -> InputResponses: - """Run `dispatch` concurrently for every key, raising `MCPError` on the first `ErrorData`. - - The first task to return `ErrorData` cancels its siblings via the task - group's cancel scope, so a refused input does not wait on a slow peer. - A callback that *raises* propagates as an `ExceptionGroup` like any other - task-group failure. - """ - responses: InputResponses = {} - refused: ErrorData | None = None - - async def run_one(tg: anyio.abc.TaskGroup, key: str, req: InputRequest) -> None: - nonlocal refused - result = await dispatch(key, req) - if isinstance(result, ErrorData): - refused = result - tg.cancel_scope.cancel() - else: - responses[key] = result - - async with anyio.create_task_group() as tg: - for key, req in requests.items(): - tg.start_soon(run_one, tg, key, req) - if refused is not None: - raise MCPError.from_error_data(refused) - return responses +import sys + +import mcp_client.client._input_required as _implementation +from mcp_client.client._input_required import ( + _STATE_ONLY_BACKOFF_CAP_SECONDS as _STATE_ONLY_BACKOFF_CAP_SECONDS, +) +from mcp_client.client._input_required import ( + _STATE_ONLY_BACKOFF_INITIAL_SECONDS as _STATE_ONLY_BACKOFF_INITIAL_SECONDS, +) +from mcp_client.client._input_required import ( + DEFAULT_INPUT_REQUIRED_MAX_ROUNDS as DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, +) +from mcp_client.client._input_required import ( + InputRequiredRoundsExceededError as InputRequiredRoundsExceededError, +) +from mcp_client.client._input_required import ( + ResultT as ResultT, +) +from mcp_client.client._input_required import ( + _dispatch_all as _dispatch_all, +) +from mcp_client.client._input_required import ( + run_input_required_driver as run_input_required_driver, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/_memory.py b/src/mcp/client/_memory.py index 187131e380..c887f935c2 100644 --- a/src/mcp/client/_memory.py +++ b/src/mcp/client/_memory.py @@ -8,11 +8,11 @@ from typing import Any import anyio +from mcp_client.client._transport import TransportStreams +from mcp_client.shared.memory import create_client_server_memory_streams -from mcp.client._transport import TransportStreams from mcp.server import Server from mcp.server.mcpserver import MCPServer -from mcp.shared.memory import create_client_server_memory_streams SERVER_SHUTDOWN_GRACE = 2.0 """Seconds to wait for the in-process server to exit on EOF before cancelling.""" diff --git a/src/mcp/client/_probe.py b/src/mcp/client/_probe.py index 0e46ae57d9..5f52dca29a 100644 --- a/src/mcp/client/_probe.py +++ b/src/mcp/client/_probe.py @@ -1,114 +1,11 @@ -"""Connect-time era negotiation for ``mode='auto'``. +import sys -The ``server/discover`` probe is sent at the newest modern version. Anything -that is not positive evidence the peer is a modern MCP server falls back to -the legacy ``initialize`` handshake — a *denylist* (only the disjoint-modern -case raises) rather than an allowlist of fallback codes. - -Every ``MCPError`` falls back except ``-32022`` with a disjoint modern-only -``supported`` list. The streamable-HTTP transport already maps HTTP-layer -4xx rejections (no JSON-RPC body) into ``MCPError`` codes, so those reach -the same path. Any non-``MCPError`` exception (network/connection errors, -anyio cancellation) propagates to the caller; an outage or in-process bug -is never an era verdict. - -A successful ``DiscoverResult`` whose ``supportedVersions`` shares no modern -version with this client is treated the same way: the server speaks discover -but advertises only handshake-era versions, which is a legacy advertisement, -not an incompatibility. - -The fallback handshake itself can be answered with ``-32022`` — e.g. a probe -that timed out client-side but succeeded on a slow-starting server locked the -connection modern before the pipelined ``initialize`` arrived. That code is -itself positive modern evidence (it names the server's versions), so it -triggers one re-probe at a mutual version instead of failing the connect. -""" - -from __future__ import annotations - -from typing import Any - -import mcp_types as types -from mcp_types import UNSUPPORTED_PROTOCOL_VERSION -from mcp_types.version import ( - HANDSHAKE_PROTOCOL_VERSIONS, - LATEST_MODERN_VERSION, - MODERN_PROTOCOL_VERSIONS, +import mcp_client.client._probe as _implementation +from mcp_client.client._probe import ( + _parse_supported as _parse_supported, +) +from mcp_client.client._probe import ( + negotiate_auto as negotiate_auto, ) -from pydantic import ValidationError - -from mcp.client.session import ClientSession -from mcp.shared.exceptions import MCPError - - -def _parse_supported(data: Any) -> list[str] | None: - """Pull ``data.supported`` off a -32022 error, or ``None`` if not actionable.""" - try: - return types.UnsupportedProtocolVersionErrorData.model_validate(data).supported - except ValidationError: - return None - - -async def negotiate_auto(session: ClientSession) -> None: - """Drive the ``mode='auto'`` connect-time policy on ``session``. - - Probes ``server/discover`` once (twice if the server names a mutual - modern version via -32022), then either ``adopt()``s the result or falls - back to ``initialize()``. Idempotent only in the sense that one of - ``session.discover_result`` / ``session.initialize_result`` is set on - return. - Raises: - MCPError: The server is modern-only and shares no version with this - client (-32022 with a disjoint ``supported`` list), or the - fallback handshake failed and one corrective re-probe did too. - Exception: Any transport/network error from the probe propagates as-is. - """ - version = LATEST_MODERN_VERSION - for attempt in range(2): - try: - raw = await session.send_discover(version) - except MCPError as e: - if e.code == UNSUPPORTED_PROTOCOL_VERSION: - supported = _parse_supported(e.error.data) - mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in (supported or ())] - if mutual and attempt == 0: - version = mutual[-1] - continue - if supported is not None and not any(v in HANDSHAKE_PROTOCOL_VERSIONS for v in supported): - raise # server is modern-only and disjoint — real incompatibility - try: - await session.initialize() # every other rpc-error → legacy (the denylist) - except MCPError as handshake_exc: - if handshake_exc.code != UNSUPPORTED_PROTOCOL_VERSION or attempt != 0: - raise - # -32022 from the handshake is itself modern evidence: a probe - # that timed out client-side but succeeded on the server locked - # the connection modern before this initialize arrived. Re-probe - # once at a version the server names; the era is already - # settled, so the second probe answers without the slow start. - supported = _parse_supported(handshake_exc.error.data) - mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in (supported or ())] - if not mutual: - raise - version = mutual[-1] - continue - return - # any other exception (httpx2.TransportError, ConnectionError, - # anyio errors) → propagate - try: - result = types.DiscoverResult.model_validate(raw) - except ValidationError: - await session.initialize() # unparseable result → not modern evidence - return - if not any(v in result.supported_versions for v in MODERN_PROTOCOL_VERSIONS): - # A discover-answering server that advertises no modern version - # (go-sdk's stateful streamable default does this) is an explicit - # legacy advertisement: fall back like the -32022 branch above - # instead of letting `adopt()` raise. The ts and go clients fall - # back here too. - await session.initialize() - return - session.adopt(result) - return - raise AssertionError("unreachable") # pragma: no cover — loop body always returns or raises +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/_transport.py b/src/mcp/client/_transport.py index 0163fef950..d927b59ac1 100644 --- a/src/mcp/client/_transport.py +++ b/src/mcp/client/_transport.py @@ -1,21 +1,17 @@ -"""Transport protocol for MCP clients.""" - -from __future__ import annotations - -from contextlib import AbstractAsyncContextManager -from typing import Protocol - -from mcp.shared._stream_protocols import ReadStream, WriteStream -from mcp.shared.message import SessionMessage - -__all__ = ["ReadStream", "WriteStream", "Transport", "TransportStreams"] - -TransportStreams = tuple[ReadStream[SessionMessage | Exception], WriteStream[SessionMessage]] - - -class Transport(AbstractAsyncContextManager[TransportStreams], Protocol): - """Protocol for MCP transports. - - A transport is an async context manager that yields read and write streams - for bidirectional communication with an MCP server. - """ +import sys + +import mcp_client.client._transport as _implementation +from mcp_client.client._transport import ( + ReadStream as ReadStream, +) +from mcp_client.client._transport import ( + Transport as Transport, +) +from mcp_client.client._transport import ( + TransportStreams as TransportStreams, +) +from mcp_client.client._transport import ( + WriteStream as WriteStream, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/auth/__init__.py b/src/mcp/client/auth/__init__.py index a6d67093ae..f1d6b8e1fc 100644 --- a/src/mcp/client/auth/__init__.py +++ b/src/mcp/client/auth/__init__.py @@ -1,15 +1,28 @@ -"""OAuth2 Authentication implementation for httpx2. - -Implements authorization code flow with PKCE and automatic token refresh. -""" - -from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError -from mcp.client.auth.oauth2 import ( - OAuthClientProvider, - PKCEParameters, - TokenStorage, +from mcp_client.client.auth import ( + AuthorizationCodeResult as AuthorizationCodeResult, +) +from mcp_client.client.auth import ( + OAuthClientProvider as OAuthClientProvider, +) +from mcp_client.client.auth import ( + OAuthFlowError as OAuthFlowError, +) +from mcp_client.client.auth import ( + OAuthRegistrationError as OAuthRegistrationError, ) -from mcp.shared.auth import AuthorizationCodeResult +from mcp_client.client.auth import ( + OAuthTokenError as OAuthTokenError, +) +from mcp_client.client.auth import ( + PKCEParameters as PKCEParameters, +) +from mcp_client.client.auth import ( + TokenStorage as TokenStorage, +) + +from . import exceptions as exceptions +from . import oauth2 as oauth2 +from . import utils as utils __all__ = [ "AuthorizationCodeResult", diff --git a/src/mcp/client/auth/exceptions.py b/src/mcp/client/auth/exceptions.py index 5ce8777b86..a304799d69 100644 --- a/src/mcp/client/auth/exceptions.py +++ b/src/mcp/client/auth/exceptions.py @@ -1,10 +1,14 @@ -class OAuthFlowError(Exception): - """Base exception for OAuth flow errors.""" - - -class OAuthTokenError(OAuthFlowError): - """Raised when token operations fail.""" - - -class OAuthRegistrationError(OAuthFlowError): - """Raised when client registration fails.""" +import sys + +import mcp_client.client.auth.exceptions as _implementation +from mcp_client.client.auth.exceptions import ( + OAuthFlowError as OAuthFlowError, +) +from mcp_client.client.auth.exceptions import ( + OAuthRegistrationError as OAuthRegistrationError, +) +from mcp_client.client.auth.exceptions import ( + OAuthTokenError as OAuthTokenError, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/auth/extensions/__init__.py b/src/mcp/client/auth/extensions/__init__.py index e69de29bb2..a9a2c5b3bb 100644 --- a/src/mcp/client/auth/extensions/__init__.py +++ b/src/mcp/client/auth/extensions/__init__.py @@ -0,0 +1 @@ +__all__ = [] diff --git a/src/mcp/client/auth/extensions/client_credentials.py b/src/mcp/client/auth/extensions/client_credentials.py index 091f9e39ee..c123908dca 100644 --- a/src/mcp/client/auth/extensions/client_credentials.py +++ b/src/mcp/client/auth/extensions/client_credentials.py @@ -1,416 +1,26 @@ -"""OAuth client credential extensions for MCP. - -Provides OAuth providers for machine-to-machine authentication flows: -- ClientCredentialsOAuthProvider: For client_credentials with client_id + client_secret -- PrivateKeyJWTOAuthProvider: For client_credentials with private_key_jwt authentication - (typically using a pre-built JWT from workload identity federation) -""" - -import time -import warnings -from collections.abc import Awaitable, Callable -from typing import Any, Literal -from urllib.parse import urlparse -from uuid import uuid4 - -import httpx2 -import jwt -from pydantic import BaseModel, Field - -from mcp.client.auth import OAuthClientProvider, OAuthFlowError, TokenStorage -from mcp.client.auth.oauth2 import OAuthContext -from mcp.client.auth.utils import issuers_match -from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata -from mcp.shared.exceptions import MCPDeprecationWarning - - -def _checked_issuer(issuer: str | None) -> str | None: - if issuer is None: - warnings.warn( - "Omitting `issuer` is deprecated and it will be required in 3.0. Without it, the MCP server " - "decides which authorization server receives this client's credentials; pass " - "issuer= so they are only ever sent there.", - MCPDeprecationWarning, - stacklevel=3, - ) - return None - if urlparse(issuer).scheme not in ("http", "https"): - raise ValueError(f"issuer must be the authorization server's http(s) issuer URL, got {issuer!r}") - return issuer - - -def _preferred_authorization_server(advertised: list[str], issuer: str | None) -> str: - """The advertised server matching the configured issuer if there is one, else the first.""" - return next( - (server for server in advertised if issuer is not None and issuers_match(server, issuer)), advertised[0] - ) - - -def _require_metadata_for_configured_issuer(context: OAuthContext, issuer: str | None) -> None: - """With an issuer configured, a token request is only built from metadata discovered for that issuer. - - Anything else held is dropped along with the tokens, so the next request starts discovery afresh - rather than refreshing against it. - """ - if issuer is None: - return - metadata = context.oauth_metadata - if metadata is not None and issuers_match(str(metadata.issuer), issuer): - return - context.oauth_metadata = None - context.clear_tokens() - if metadata is None: - raise OAuthFlowError(f"No authorization server metadata discovered for configured issuer {issuer}") - raise OAuthFlowError(f"Authorization server metadata issuer mismatch: {metadata.issuer} != {issuer}") - - -class ClientCredentialsOAuthProvider(OAuthClientProvider): - """OAuth provider for client_credentials grant with client_id + client_secret. - - This provider sets client_info directly, bypassing dynamic client registration. - Use this when you already have client credentials (client_id and client_secret). - Pass `issuer` to name the authorization server those credentials belong to: token - requests are then only built from authorization server metadata for that issuer, and - the flow stops if the MCP server leads anywhere else. - - Example: - ```python - provider = ClientCredentialsOAuthProvider( - server_url="https://api.example.com", - storage=my_token_storage, - client_id="my-client-id", - client_secret="my-client-secret", - issuer="https://auth.example.com", - ) - ``` - """ - - def __init__( - self, - server_url: str, - storage: TokenStorage, - client_id: str, - client_secret: str, - token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_basic", - scope: str | None = None, - issuer: str | None = None, - ) -> None: - """Initialize client_credentials OAuth provider. - - Args: - server_url: The MCP server URL. - storage: Token storage implementation. - client_id: The OAuth client ID. - client_secret: The OAuth client secret. - token_endpoint_auth_method: Authentication method for token endpoint. - Either "client_secret_basic" (default) or "client_secret_post". - scope: Optional space-separated list of scopes to request. - issuer: The issuer identifier of the authorization server that issued - `client_id` and `client_secret`. When set, token requests are only built from - discovered authorization server metadata whose `issuer` is exactly this string; - otherwise the flow stops with `OAuthFlowError`. Omitting it is deprecated - (`MCPDeprecationWarning`) and it will be required in 3.0; until then, whichever - authorization server discovery yields is used. - """ - # Build minimal client_metadata for the base class - client_metadata = OAuthClientMetadata( - redirect_uris=None, - grant_types=["client_credentials"], - token_endpoint_auth_method=token_endpoint_auth_method, - scope=scope, - ) - super().__init__(server_url, client_metadata, storage, None, None) - self._issuer = _checked_issuer(issuer) - # Store client_info to be set during _initialize - no dynamic registration needed - self._fixed_client_info = OAuthClientInformationFull( - redirect_uris=None, - client_id=client_id, - client_secret=client_secret, - grant_types=["client_credentials"], - token_endpoint_auth_method=token_endpoint_auth_method, - scope=scope, - ) - - async def _initialize(self) -> None: - """Load stored tokens and set pre-configured client_info.""" - self.context.current_tokens = await self.context.storage.get_tokens() - self.context.client_info = self._fixed_client_info - self._initialized = True - - def _select_authorization_server(self, advertised: list[str]) -> str: - return _preferred_authorization_server(advertised, self._issuer) - - async def _perform_authorization(self) -> httpx2.Request: - """Perform client_credentials authorization.""" - return await self._exchange_token_client_credentials() - - async def _exchange_token_client_credentials(self) -> httpx2.Request: - """Build token exchange request for client_credentials grant.""" - _require_metadata_for_configured_issuer(self.context, self._issuer) - - token_data: dict[str, Any] = { - "grant_type": "client_credentials", - } - - headers: dict[str, str] = {"Content-Type": "application/x-www-form-urlencoded"} - - # Use standard auth methods (client_secret_basic, client_secret_post, none) - token_data, headers = self.context.prepare_token_auth(token_data, headers) - - if self.context.should_include_resource_param(self.context.protocol_version): - token_data["resource"] = self.context.get_resource_url() - - if self.context.client_metadata.scope: - token_data["scope"] = self.context.client_metadata.scope - - token_url = self._get_token_endpoint() - return httpx2.Request("POST", token_url, data=token_data, headers=headers) - - -def static_assertion_provider(token: str) -> Callable[[str], Awaitable[str]]: - """Create an assertion provider that returns a static JWT token. - - Use this when you have a pre-built JWT (e.g., from workload identity federation) - that doesn't need the audience parameter. - - Example: - ```python - provider = PrivateKeyJWTOAuthProvider( - server_url="https://api.example.com", - storage=my_token_storage, - client_id="my-client-id", - assertion_provider=static_assertion_provider(my_prebuilt_jwt), - issuer="https://auth.example.com", - ) - ``` - - Args: - token: The pre-built JWT assertion string. - - Returns: - An async callback suitable for use as an assertion_provider. - """ - - async def provider(audience: str) -> str: - return token - - return provider - - -class SignedJWTParameters(BaseModel): - """Parameters for creating SDK-signed JWT assertions. - - Use `create_assertion_provider()` to create an assertion provider callback - for use with `PrivateKeyJWTOAuthProvider`. - - Example: - ```python - jwt_params = SignedJWTParameters( - issuer="my-client-id", - subject="my-client-id", - signing_key=private_key_pem, - ) - provider = PrivateKeyJWTOAuthProvider( - server_url="https://api.example.com", - storage=my_token_storage, - client_id="my-client-id", - assertion_provider=jwt_params.create_assertion_provider(), - issuer="https://auth.example.com", - ) - ``` - """ - - issuer: str = Field(description="Issuer for JWT assertions (typically client_id).") - subject: str = Field(description="Subject identifier for JWT assertions (typically client_id).") - signing_key: str = Field(description="Private key for JWT signing (PEM format).") - signing_algorithm: str = Field(default="RS256", description="Algorithm for signing JWT assertions.") - lifetime_seconds: int = Field(default=300, description="Lifetime of generated JWT in seconds.") - additional_claims: dict[str, Any] | None = Field(default=None, description="Additional claims.") - - def create_assertion_provider(self) -> Callable[[str], Awaitable[str]]: - """Create an assertion provider callback for use with PrivateKeyJWTOAuthProvider. - - Returns: - An async callback that takes the audience (authorization server issuer URL) - and returns a signed JWT assertion. - """ - - async def provider(audience: str) -> str: - now = int(time.time()) - claims: dict[str, Any] = { - "iss": self.issuer, - "sub": self.subject, - "aud": audience, - "exp": now + self.lifetime_seconds, - "iat": now, - "jti": str(uuid4()), - } - if self.additional_claims: - claims.update(self.additional_claims) - - return jwt.encode(claims, self.signing_key, algorithm=self.signing_algorithm) - - return provider - - -class PrivateKeyJWTOAuthProvider(OAuthClientProvider): - """OAuth provider for client_credentials grant with private_key_jwt authentication. - - Uses RFC 7523 Section 2.2 for client authentication via JWT assertion. - - The JWT assertion's audience MUST be the authorization server's issuer identifier - (per RFC 7523bis security updates). The `assertion_provider` callback receives - this audience value and must return a JWT with that audience. Pass `issuer` to name - the authorization server this client is registered with: an assertion is then only - minted once metadata for that issuer has been discovered, and token requests are only - built from that metadata. - - **Option 1: Pre-built JWT via Workload Identity Federation** - - In production scenarios, the JWT assertion is typically obtained from a workload - identity provider (e.g., GCP, AWS IAM, Azure AD): - - ```python - async def get_workload_identity_token(audience: str) -> str: - # Fetch JWT from your identity provider - # The JWT's audience must match the provided audience parameter - return await fetch_token_from_identity_provider(audience=audience) - - provider = PrivateKeyJWTOAuthProvider( - server_url="https://api.example.com", - storage=my_token_storage, - client_id="my-client-id", - assertion_provider=get_workload_identity_token, - issuer="https://auth.example.com", - ) - ``` - - **Option 2: Static pre-built JWT** - - If you have a static JWT that doesn't need the audience parameter: - - ```python - provider = PrivateKeyJWTOAuthProvider( - server_url="https://api.example.com", - storage=my_token_storage, - client_id="my-client-id", - assertion_provider=static_assertion_provider(my_prebuilt_jwt), - issuer="https://auth.example.com", - ) - ``` - - **Option 3: SDK-signed JWT (for testing/simple setups)** - - For testing or simple deployments, use `SignedJWTParameters.create_assertion_provider()`: - - ```python - jwt_params = SignedJWTParameters( - issuer="my-client-id", - subject="my-client-id", - signing_key=private_key_pem, - ) - provider = PrivateKeyJWTOAuthProvider( - server_url="https://api.example.com", - storage=my_token_storage, - client_id="my-client-id", - assertion_provider=jwt_params.create_assertion_provider(), - issuer="https://auth.example.com", - ) - ``` - """ - - def __init__( - self, - server_url: str, - storage: TokenStorage, - client_id: str, - assertion_provider: Callable[[str], Awaitable[str]], - scope: str | None = None, - issuer: str | None = None, - ) -> None: - """Initialize private_key_jwt OAuth provider. - - Args: - server_url: The MCP server URL. - storage: Token storage implementation. - client_id: The OAuth client ID. - assertion_provider: Async callback that takes the audience (authorization - server's issuer identifier) and returns a JWT assertion. Use - `SignedJWTParameters.create_assertion_provider()` for SDK-signed JWTs, - `static_assertion_provider()` for pre-built JWTs, or provide your own - callback for workload identity federation. - scope: Optional space-separated list of scopes to request. - issuer: The issuer identifier of the authorization server `client_id` is - registered with. When set, an assertion is only minted, and token requests - are only built, once authorization server metadata whose `issuer` is exactly this - string has been discovered; otherwise the flow stops with `OAuthFlowError`. - Omitting it is deprecated (`MCPDeprecationWarning`) and it will be required in - 3.0; until then, whichever authorization server discovery yields is used. - """ - # Build minimal client_metadata for the base class - client_metadata = OAuthClientMetadata( - redirect_uris=None, - grant_types=["client_credentials"], - token_endpoint_auth_method="private_key_jwt", - scope=scope, - ) - super().__init__(server_url, client_metadata, storage, None, None) - self._assertion_provider = assertion_provider - self._issuer = _checked_issuer(issuer) - # Store client_info to be set during _initialize - no dynamic registration needed - self._fixed_client_info = OAuthClientInformationFull( - redirect_uris=None, - client_id=client_id, - grant_types=["client_credentials"], - token_endpoint_auth_method="private_key_jwt", - scope=scope, - ) - - async def _initialize(self) -> None: - """Load stored tokens and set pre-configured client_info.""" - self.context.current_tokens = await self.context.storage.get_tokens() - self.context.client_info = self._fixed_client_info - self._initialized = True - - def _select_authorization_server(self, advertised: list[str]) -> str: - return _preferred_authorization_server(advertised, self._issuer) - - async def _perform_authorization(self) -> httpx2.Request: - """Perform client_credentials authorization with private_key_jwt.""" - return await self._exchange_token_client_credentials() - - async def _add_client_authentication_jwt(self, *, token_data: dict[str, Any]) -> None: - """Add JWT assertion for client authentication to token endpoint parameters.""" - if not self.context.oauth_metadata: - raise OAuthFlowError("Missing OAuth metadata for private_key_jwt flow") # pragma: no cover - - # Audience MUST be the issuer identifier of the authorization server - # https://datatracker.ietf.org/doc/html/draft-ietf-oauth-rfc7523bis-01 - audience = str(self.context.oauth_metadata.issuer) - assertion = await self._assertion_provider(audience) - - # RFC 7523 Section 2.2: client authentication via JWT - token_data["client_assertion"] = assertion - token_data["client_assertion_type"] = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer" - - async def _exchange_token_client_credentials(self) -> httpx2.Request: - """Build token exchange request for client_credentials grant with private_key_jwt.""" - _require_metadata_for_configured_issuer(self.context, self._issuer) - - token_data: dict[str, Any] = { - "grant_type": "client_credentials", - } - - headers: dict[str, str] = {"Content-Type": "application/x-www-form-urlencoded"} - - # Add JWT client authentication (RFC 7523 Section 2.2) - await self._add_client_authentication_jwt(token_data=token_data) - - if self.context.should_include_resource_param(self.context.protocol_version): - token_data["resource"] = self.context.get_resource_url() - - if self.context.client_metadata.scope: - token_data["scope"] = self.context.client_metadata.scope - - token_url = self._get_token_endpoint() - return httpx2.Request("POST", token_url, data=token_data, headers=headers) +import sys + +import mcp_client.client.auth.extensions.client_credentials as _implementation +from mcp_client.client.auth.extensions.client_credentials import ( + ClientCredentialsOAuthProvider as ClientCredentialsOAuthProvider, +) +from mcp_client.client.auth.extensions.client_credentials import ( + PrivateKeyJWTOAuthProvider as PrivateKeyJWTOAuthProvider, +) +from mcp_client.client.auth.extensions.client_credentials import ( + SignedJWTParameters as SignedJWTParameters, +) +from mcp_client.client.auth.extensions.client_credentials import ( + _checked_issuer as _checked_issuer, +) +from mcp_client.client.auth.extensions.client_credentials import ( + _preferred_authorization_server as _preferred_authorization_server, +) +from mcp_client.client.auth.extensions.client_credentials import ( + _require_metadata_for_configured_issuer as _require_metadata_for_configured_issuer, +) +from mcp_client.client.auth.extensions.client_credentials import ( + static_assertion_provider as static_assertion_provider, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/auth/extensions/identity_assertion.py b/src/mcp/client/auth/extensions/identity_assertion.py index 35e48c03fa..606ba19e1b 100644 --- a/src/mcp/client/auth/extensions/identity_assertion.py +++ b/src/mcp/client/auth/extensions/identity_assertion.py @@ -1,216 +1,14 @@ -"""SEP-990 Identity Assertion Authorization Grant (RFC 7523 jwt-bearer) client provider. +import sys -`IdentityAssertionOAuthProvider` is the client side of SEP-990 leg 2: it presents an Identity -Assertion Authorization Grant (ID-JAG) - a signed JWT issued by the enterprise identity provider - -to the MCP authorization server's token endpoint using the RFC 7523 jwt-bearer grant -(`grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer`, ID-JAG as `assertion`), and receives an -MCP access token. - -The authorization server is configuration, not discovery. SEP-990's trust model is the inverse of -the default OAuth client's: the AS issuer is supplied at construction, authorization-server metadata -is fetched from that issuer's own RFC 8414 well-known, and the resource server is never asked which -AS to use - so it cannot redirect the ID-JAG or client secret elsewhere. There is no protected -resource metadata fetch, no dynamic client registration, and no server-driven scope selection. - -Obtaining the ID-JAG (logging into the IdP and the leg-1 token exchange against it) is -deployment-specific and out of scope for the SDK. The caller supplies it through the -`assertion_provider` callback, which receives the configured issuer (the `aud` the ID-JAG must -carry) and the MCP server's resource identifier (the `resource` claim it must carry, per ext-auth -section 4.3), and returns the ID-JAG. -""" - -import base64 -import time -from collections.abc import AsyncGenerator, Awaitable, Callable -from typing import Literal -from urllib.parse import quote, urlsplit - -import anyio -import httpx2 - -from mcp.client.auth import OAuthFlowError, OAuthTokenError, TokenStorage -from mcp.client.auth.utils import ( - build_oauth_authorization_server_metadata_discovery_urls, - create_oauth_metadata_request, - extract_field_from_www_auth, - extract_scope_from_www_auth, - handle_auth_metadata_response, - handle_token_response_scopes, - union_scopes, - validate_metadata_issuer, +import mcp_client.client.auth.extensions.identity_assertion as _implementation +from mcp_client.client.auth.extensions.identity_assertion import ( + _DEFAULT_PORTS as _DEFAULT_PORTS, +) +from mcp_client.client.auth.extensions.identity_assertion import ( + IdentityAssertionOAuthProvider as IdentityAssertionOAuthProvider, +) +from mcp_client.client.auth.extensions.identity_assertion import ( + _origin as _origin, ) -from mcp.shared._httpx_utils import RedirectAwareAuth, redirect_note -from mcp.shared.auth import JWT_BEARER_GRANT_TYPE, OAuthClientInformationFull, OAuthToken -from mcp.shared.auth_utils import calculate_token_expiry, resource_url_from_server_url - -_DEFAULT_PORTS = {"https": 443, "http": 80} - - -def _origin(url: str) -> tuple[str, str, int | None]: - """Return the (scheme, host, port) origin of a URL for same-origin comparison. - - The port is normalized to the scheme's default so an explicit `:443`/`:80` compares equal to the - same origin written without a port. - """ - parsed = urlsplit(url) - port = parsed.port if parsed.port is not None else _DEFAULT_PORTS.get(parsed.scheme) - return (parsed.scheme, parsed.hostname or "", port) - - -class IdentityAssertionOAuthProvider(RedirectAwareAuth): - """`httpx2.Auth` for the SEP-990 ID-JAG flow (RFC 7523 jwt-bearer grant) against a configured AS. - - The authorization server `issuer` is fixed at construction; metadata is fetched from its - RFC 8414 well-known and the ID-JAG and client secret are sent only to that issuer's token - endpoint. The resource server is never consulted for AS selection. The ID-JAG is fetched lazily - from `assertion_provider` so a fresh assertion is used on each exchange. - - Example: - ```python - async def fetch_id_jag(audience: str, resource: str) -> str: - # `audience` is the configured issuer (the ID-JAG `aud`); `resource` is the MCP - # server's identifier (the ID-JAG `resource` claim). Obtaining the ID-JAG from the - # enterprise IdP is deployment-specific and not handled by the SDK. - return await my_idp.issue_id_jag(audience=audience, resource=resource) - - - provider = IdentityAssertionOAuthProvider( - server_url="https://mcp.example.com/mcp", - storage=my_token_storage, - client_id="my-client-id", - client_secret="my-client-secret", - issuer="https://auth.example.com", - assertion_provider=fetch_id_jag, - ) - ``` - """ - - requires_response_body = True - - def __init__( - self, - server_url: str, - storage: TokenStorage, - client_id: str, - client_secret: str, - issuer: str, - assertion_provider: Callable[[str, str], Awaitable[str]], - scope: str | None = None, - token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_post", - ) -> None: - """Initialize the identity-assertion OAuth provider. - - Args: - server_url: The MCP server URL. - storage: Token storage implementation. - client_id: The OAuth client ID registered with the MCP authorization server. - client_secret: The client secret. SEP-990 section 5.1 requires a confidential client. - issuer: The issuer identifier of the MCP authorization server this client is provisioned - for. Authorization-server metadata is fetched from this issuer's well-known and the - ID-JAG and secret are sent only to its token endpoint. - assertion_provider: Async callback taking `(audience, resource)` - the configured issuer - and the MCP server's resource identifier - and returning the ID-JAG. - scope: Optional space-separated list of scopes to request. - token_endpoint_auth_method: Confidential-client auth method, either `client_secret_post` - (default) or `client_secret_basic`. - """ - if not client_secret: - raise ValueError("client_secret is required: SEP-990 mandates a confidential client") - if not issuer: - raise ValueError("issuer is required: the authorization server is configuration, not discovery") - self._resource = resource_url_from_server_url(server_url) - self._storage = storage - self._issuer = issuer - self._assertion_provider = assertion_provider - self._scope = scope - self._client = OAuthClientInformationFull( - client_id=client_id, - client_secret=client_secret, - redirect_uris=None, - grant_types=[JWT_BEARER_GRANT_TYPE], - token_endpoint_auth_method=token_endpoint_auth_method, - issuer=issuer, - ) - self._token_endpoint: str | None = None - self._tokens: OAuthToken | None = None - self._expiry: float | None = None - self._lock = anyio.Lock() - self._initialized = False - - def _build_token_request(self, scope: str | None, assertion: str) -> httpx2.Request: - """Build the RFC 7523 jwt-bearer token request, applying confidential-client auth.""" - assert self._token_endpoint is not None - assert self._client.client_id is not None and self._client.client_secret is not None - data: dict[str, str] = { - "grant_type": JWT_BEARER_GRANT_TYPE, - "assertion": assertion, - "client_id": self._client.client_id, - "resource": self._resource, - } - if scope: - data["scope"] = scope - headers = {"Content-Type": "application/x-www-form-urlencoded"} - if self._client.token_endpoint_auth_method == "client_secret_basic": - # RFC 6749 section 2.3.1: URL-encode each part, then base64 the colon-joined pair. - encoded_id = quote(self._client.client_id, safe="") - encoded_secret = quote(self._client.client_secret, safe="") - credentials = base64.b64encode(f"{encoded_id}:{encoded_secret}".encode()).decode() - headers["Authorization"] = f"Basic {credentials}" - else: - data["client_secret"] = self._client.client_secret - return httpx2.Request("POST", self._token_endpoint, data=data, headers=headers) - - async def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: - async with self._lock: - if not self._initialized: - self._tokens = await self._storage.get_tokens() - self._expiry = calculate_token_expiry(self._tokens.expires_in) if self._tokens else None - self._initialized = True - - if self._tokens and (self._expiry is None or time.time() <= self._expiry): - request.headers["Authorization"] = f"Bearer {self._tokens.access_token}" - response = yield request - - if response.status_code == 401: - scope_to_request = self._scope - elif response.status_code == 403 and extract_field_from_www_auth(response, "error") == "insufficient_scope": - scope_to_request = union_scopes(self._scope, extract_scope_from_www_auth(response)) - else: - return - - # Discover ASM from the configured issuer's well-known. The RS is not consulted: both - # arguments are the issuer, so even the helper's legacy fallback resolves there. - if self._token_endpoint is None: - for url in build_oauth_authorization_server_metadata_discovery_urls(self._issuer, self._issuer): - asm_response = yield create_oauth_metadata_request(url) - ok, asm = await handle_auth_metadata_response(asm_response) - if not ok: - break - if asm is not None: - validate_metadata_issuer(asm, self._issuer) - token_endpoint = str(asm.token_endpoint) - if _origin(token_endpoint) != _origin(self._issuer): - raise OAuthFlowError( - f"Token endpoint {token_endpoint} is not on the configured issuer origin {self._issuer}" - ) - self._token_endpoint = token_endpoint - break - if self._token_endpoint is None: - raise OAuthFlowError(f"No authorization server metadata at configured issuer {self._issuer}") - - assertion = await self._assertion_provider(self._issuer, self._resource) - token_response = yield self._build_token_request(scope_to_request, assertion) - if token_response.status_code != 200: - body = (await token_response.aread()).decode(errors="replace") - raise OAuthTokenError( - f"Token exchange failed ({token_response.status_code}){redirect_note(token_response)}: {body}" - ) - tokens = await handle_token_response_scopes(token_response) - if tokens.scope is None: - tokens.scope = scope_to_request - self._tokens = tokens - self._expiry = calculate_token_expiry(tokens.expires_in) - await self._storage.set_tokens(tokens) - request.headers["Authorization"] = f"Bearer {tokens.access_token}" - yield request +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index 8588208924..2d01c0d6ef 100644 --- a/src/mcp/client/auth/oauth2.py +++ b/src/mcp/client/auth/oauth2.py @@ -1,793 +1,38 @@ -"""OAuth2 Authentication implementation for httpx2. +import sys -Implements authorization code flow with PKCE and automatic token refresh. -""" - -import base64 -import hashlib -import logging -import secrets -import string -import time -from collections.abc import AsyncGenerator, Awaitable, Callable -from dataclasses import dataclass, field -from typing import Any, Protocol, get_args -from urllib.parse import quote, urlencode, urljoin, urlparse - -import anyio -import httpx2 -from mcp_types.version import is_version_at_least -from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, TypeAdapter, ValidationError - -from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError -from mcp.client.auth.utils import ( - build_oauth_authorization_server_metadata_discovery_urls, - build_protected_resource_metadata_discovery_urls, - create_client_info_from_metadata_url, - create_client_registration_request, - create_oauth_metadata_request, - credentials_match_issuer, - extract_field_from_www_auth, - extract_resource_metadata_from_www_auth, - extract_scope_from_www_auth, - get_client_metadata_scopes, - handle_auth_metadata_response, - handle_protected_resource_response, - handle_registration_response, - handle_token_response_scopes, - is_valid_client_metadata_url, - issuers_match, - should_use_client_metadata_url, - union_scopes, - validate_authorization_response_iss, - validate_metadata_issuer, +import mcp_client.client.auth.oauth2 as _implementation +from mcp_client.client.auth.oauth2 import ( + _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS as _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS, ) -from mcp.shared._httpx_utils import RedirectAwareAuth, redirect_note -from mcp.shared.auth import ( - AuthorizationCodeResult, - OAuthClientInformationFull, - OAuthClientMetadata, - OAuthMetadata, - OAuthToken, - ProtectedResourceMetadata, - TokenEndpointAuthMethod, +from mcp_client.client.auth.oauth2 import ( + _ORIGIN_URL as _ORIGIN_URL, ) -from mcp.shared.auth_utils import ( - calculate_token_expiry, - check_resource_allowed, - resource_url_from_server_url, +from mcp_client.client.auth.oauth2 import ( + _REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS as _REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS, ) -from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER - -logger = logging.getLogger(__name__) - -# Methods a registered client's record may carry without a token request being an error, -# derived from the set the SDK is willing to request so the two cannot drift. `None`/"none" -# send no client secret. `private_key_jwt` sends none from here either: only -# `PrivateKeyJWTOAuthProvider` signs the assertion, and only in its client-credentials -# exchange, so its inherited refresh path must pass through here without raising - a refresh -# the server then rejects falls back to a fresh client-credentials exchange, which signs. -# Anything else is a method no client here can apply. -_KNOWN_TOKEN_ENDPOINT_AUTH_METHODS: tuple[str | None, ...] = (None, *get_args(TokenEndpointAuthMethod)) - -# Methods that authenticate the token request with the minted `client_secret`; a -# registration assigning one is only usable if the server issued that secret. -_SECRET_TOKEN_ENDPOINT_AUTH_METHODS = ("client_secret_post", "client_secret_basic") - -# Methods a registration completed by the authorization-code flow can act on. That flow -# authenticates the token request with the minted client secret (or nothing); it holds no key -# to sign a `private_key_jwt` assertion, so a server assigning that method has registered a -# client this flow cannot use. `PrivateKeyJWTOAuthProvider` never registers dynamically. -_REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS: tuple[str | None, ...] = tuple( - method for method in _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS if method != "private_key_jwt" +from mcp_client.client.auth.oauth2 import ( + _SECRET_TOKEN_ENDPOINT_AUTH_METHODS as _SECRET_TOKEN_ENDPOINT_AUTH_METHODS, +) +from mcp_client.client.auth.oauth2 import ( + OAuthClientProvider as OAuthClientProvider, +) +from mcp_client.client.auth.oauth2 import ( + OAuthContext as OAuthContext, +) +from mcp_client.client.auth.oauth2 import ( + PKCEParameters as PKCEParameters, +) +from mcp_client.client.auth.oauth2 import ( + TokenStorage as TokenStorage, +) +from mcp_client.client.auth.oauth2 import ( + _origin_issuer as _origin_issuer, +) +from mcp_client.client.auth.oauth2 import ( + check_registration_usable as check_registration_usable, +) +from mcp_client.client.auth.oauth2 import ( + logger as logger, ) - -def check_registration_usable(client_info: OAuthClientInformationFull) -> None: - """Confirm a registration this flow completed is one it can act on. - - RFC 7591 §3.2.1 lets the authorization server replace requested metadata and leaves it to - the client to "check the values in the response to determine if the registration is - sufficient for use". Two substitutions make the minted credentials unusable, and both are - judged here - before the record is persisted or any interactive authorization begins - - rather than surfacing later as an opaque failure at the token endpoint: a token-endpoint - auth method the authorization-code flow cannot apply (one it does not implement, or - `private_key_jwt`, whose assertion this flow has no key to sign), and a secret-based - method the flow could apply but for which the server issued no `client_secret`. - - Raises: - OAuthRegistrationError: The server registered the client with a - `token_endpoint_auth_method` this flow cannot apply, or with a secret-based - method but no `client_secret`. - """ - method = client_info.token_endpoint_auth_method - if method not in _REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS: - raise OAuthRegistrationError( - f"Authorization server registered the client with unsupported token_endpoint_auth_method {method!r}" - ) - if method in _SECRET_TOKEN_ENDPOINT_AUTH_METHODS and client_info.client_secret is None: - raise OAuthRegistrationError( - f"Authorization server registered the client for {method!r} but issued no client_secret" - ) - - -class PKCEParameters(BaseModel): - """PKCE (Proof Key for Code Exchange) parameters.""" - - code_verifier: str = Field(..., min_length=43, max_length=128) - code_challenge: str = Field(..., min_length=43, max_length=128) - - @classmethod - def generate(cls) -> "PKCEParameters": - """Generate new PKCE parameters.""" - code_verifier = "".join(secrets.choice(string.ascii_letters + string.digits + "-._~") for _ in range(128)) - digest = hashlib.sha256(code_verifier.encode()).digest() - code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=") - return cls(code_verifier=code_verifier, code_challenge=code_challenge) - - -class TokenStorage(Protocol): - """Protocol for token storage implementations.""" - - async def get_tokens(self) -> OAuthToken | None: - """Get stored tokens.""" - ... - - async def set_tokens(self, tokens: OAuthToken) -> None: - """Store tokens.""" - ... - - async def get_client_info(self) -> OAuthClientInformationFull | None: - """Get stored client information.""" - ... - - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: - """Store client information.""" - ... - - -@dataclass -class OAuthContext: - """OAuth flow context.""" - - server_url: str - client_metadata: OAuthClientMetadata - storage: TokenStorage - redirect_handler: Callable[[str], Awaitable[None]] | None - callback_handler: Callable[[], Awaitable[AuthorizationCodeResult]] | None - client_metadata_url: str | None = None - - # Discovered metadata - protected_resource_metadata: ProtectedResourceMetadata | None = None - oauth_metadata: OAuthMetadata | None = None - auth_server_url: str | None = None - protocol_version: str | None = None - - # Client registration - client_info: OAuthClientInformationFull | None = None - - # Token management - current_tokens: OAuthToken | None = None - token_expiry_time: float | None = None - - # State - lock: anyio.Lock = field(default_factory=anyio.Lock) - - def get_authorization_base_url(self, server_url: str) -> str: - """Extract base URL by removing path component.""" - parsed = urlparse(server_url) - return f"{parsed.scheme}://{parsed.netloc}" - - def update_token_expiry(self, token: OAuthToken) -> None: - """Update token expiry time using shared util function.""" - self.token_expiry_time = calculate_token_expiry(token.expires_in) - - def is_token_valid(self) -> bool: - """Check if current token is valid.""" - return bool( - self.current_tokens - and self.current_tokens.access_token - and (not self.token_expiry_time or time.time() <= self.token_expiry_time) - ) - - def can_refresh_token(self) -> bool: - """Check if token can be refreshed.""" - return bool(self.current_tokens and self.current_tokens.refresh_token and self.client_info) - - def clear_tokens(self) -> None: - """Clear current tokens.""" - self.current_tokens = None - self.token_expiry_time = None - - def get_resource_url(self) -> str: - """Get resource URL for RFC 8707. - - Uses PRM resource if it's a valid parent, otherwise uses canonical server URL. - """ - resource = resource_url_from_server_url(self.server_url) - - # If PRM provides a resource that's a valid parent, use it - if self.protected_resource_metadata and self.protected_resource_metadata.resource: - prm_resource = str(self.protected_resource_metadata.resource) - if check_resource_allowed(requested_resource=resource, configured_resource=prm_resource): - resource = prm_resource - - return resource - - def should_include_resource_param(self, protocol_version: str | None = None) -> bool: - """Determine if the resource parameter should be included in OAuth requests. - - Returns True if: - - Protected resource metadata is available, OR - - MCP-Protocol-Version header is 2025-06-18 or later - """ - # If we have protected resource metadata, include the resource param - if self.protected_resource_metadata is not None: - return True - - # If no protocol version provided, don't include resource param - if not protocol_version: - return False - - return is_version_at_least(protocol_version, "2025-06-18") - - def prepare_token_auth( - self, data: dict[str, str], headers: dict[str, str] | None = None - ) -> tuple[dict[str, str], dict[str, str]]: - """Prepare authentication for token requests. - - Args: - data: The form data to send - headers: Optional headers dict to update - - Returns: - Tuple of (updated_data, updated_headers) - - Raises: - OAuthTokenError: The client record carries a `token_endpoint_auth_method` this - client does not know. A dynamic registration assigning an unusable method is - rejected earlier, by `check_registration_usable`; this fires for a stored or - pre-registered record that reaches a token request with such a method. - """ - if headers is None: - headers = {} # pragma: no cover - - if not self.client_info: - return data, headers - - auth_method = self.client_info.token_endpoint_auth_method - - if auth_method == "client_secret_basic" and self.client_info.client_secret: - # URL-encode client ID and secret per RFC 6749 Section 2.3.1 - encoded_id = quote(self.client_info.client_id, safe="") - encoded_secret = quote(self.client_info.client_secret, safe="") - credentials = f"{encoded_id}:{encoded_secret}" - encoded_credentials = base64.b64encode(credentials.encode()).decode() - headers["Authorization"] = f"Basic {encoded_credentials}" - # Don't include client_secret in body for basic auth - data = {k: v for k, v in data.items() if k != "client_secret"} - elif auth_method == "client_secret_post" and self.client_info.client_secret: - # Include client_id and client_secret in request body (RFC 6749 §2.3.1) - data["client_id"] = self.client_info.client_id - data["client_secret"] = self.client_info.client_secret - elif auth_method not in _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS: - raise OAuthTokenError(f"Registered client uses unsupported token_endpoint_auth_method {auth_method!r}") - # For "none" (or absent), don't add any client_secret; "private_key_jwt" adds its - # assertion in the provider that implements it, not here. - - return data, headers - - -_ORIGIN_URL = TypeAdapter(AnyHttpUrl, config=ConfigDict(url_preserve_empty_path=True)) - - -def _origin_issuer(server_url: str) -> str: - """The resource server's origin as an issuer identifier: `scheme://authority`, rendered the way - `OAuthMetadata.issuer` renders URLs (host case, default ports) so the two compare as strings.""" - parsed = urlparse(server_url) - return str(_ORIGIN_URL.validate_python(f"{parsed.scheme}://{parsed.netloc}")) - - -class OAuthClientProvider(RedirectAwareAuth): - """OAuth2 authentication for httpx2. - - Handles OAuth flow with automatic client registration and token storage. - """ - - requires_response_body = True - - def __init__( - self, - server_url: str, - client_metadata: OAuthClientMetadata, - storage: TokenStorage, - redirect_handler: Callable[[str], Awaitable[None]] | None = None, - callback_handler: Callable[[], Awaitable[AuthorizationCodeResult]] | None = None, - client_metadata_url: str | None = None, - validate_resource_url: Callable[[str, str | None], Awaitable[None]] | None = None, - ): - """Initialize OAuth2 authentication. - - Args: - server_url: The MCP server URL. - client_metadata: OAuth client metadata for registration. - storage: Token storage implementation. - redirect_handler: Handler for authorization redirects. - callback_handler: Handler for authorization callbacks. - client_metadata_url: URL-based client ID. When provided and the server - advertises client_id_metadata_document_supported=True, this URL will be - used as the client_id instead of performing dynamic client registration. - Must be a valid HTTPS URL with a non-root pathname. - validate_resource_url: Optional callback to override resource URL validation. - Called with (server_url, prm_resource) where prm_resource is the resource - from Protected Resource Metadata (or None if not present). If not provided, - default validation rejects mismatched resources per RFC 8707. - - Raises: - ValueError: If client_metadata_url is provided but not a valid HTTPS URL - with a non-root pathname. - """ - # Validate client_metadata_url if provided - if client_metadata_url is not None and not is_valid_client_metadata_url(client_metadata_url): - raise ValueError( - f"client_metadata_url must be a valid HTTPS URL with a non-root pathname, got: {client_metadata_url}" - ) - - self.context = OAuthContext( - server_url=server_url, - client_metadata=client_metadata, - storage=storage, - redirect_handler=redirect_handler, - callback_handler=callback_handler, - client_metadata_url=client_metadata_url, - ) - self._validate_resource_url_callback = validate_resource_url - self._initialized = False - - async def _handle_protected_resource_response(self, response: httpx2.Response) -> bool: - """Handle protected resource metadata discovery response. - - Per SEP-985, supports fallback when discovery fails at one URL. - - Returns: - True if metadata was successfully discovered, False if we should try next URL - """ - if response.status_code == 200: - try: - content = await response.aread() - metadata = ProtectedResourceMetadata.model_validate_json(content) - self.context.protected_resource_metadata = metadata - if metadata.authorization_servers: # pragma: no branch - self.context.auth_server_url = str(metadata.authorization_servers[0]) - return True - - except ValidationError: # pragma: no cover - # Invalid metadata - try next URL - logger.warning(f"Invalid protected resource metadata at {response.request.url}") - return False - elif response.status_code == 404: # pragma: no cover - # Not found - try next URL in fallback chain - logger.debug(f"Protected resource metadata not found at {response.request.url}, trying next URL") - return False - else: - # Other error - fail immediately - raise OAuthFlowError( - f"Protected Resource Metadata request failed: {response.status_code}" - ) # pragma: no cover - - async def _perform_authorization(self) -> httpx2.Request: - """Perform the authorization flow.""" - auth_code, code_verifier = await self._perform_authorization_code_grant() - token_request = await self._exchange_token_authorization_code(auth_code, code_verifier) - return token_request - - async def _perform_authorization_code_grant(self) -> tuple[str, str]: - """Perform the authorization redirect and get auth code.""" - if self.context.client_metadata.redirect_uris is None: - raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover - if not self.context.redirect_handler: - raise OAuthFlowError("No redirect handler provided for authorization code grant") # pragma: no cover - if not self.context.callback_handler: - raise OAuthFlowError("No callback handler provided for authorization code grant") # pragma: no cover - - if self.context.oauth_metadata and self.context.oauth_metadata.authorization_endpoint: - auth_endpoint = str(self.context.oauth_metadata.authorization_endpoint) - else: - auth_base_url = self.context.get_authorization_base_url(self.context.server_url) - auth_endpoint = urljoin(auth_base_url, "/authorize") - - if not self.context.client_info: - raise OAuthFlowError("No client info available for authorization") # pragma: no cover - - # Generate PKCE parameters - pkce_params = PKCEParameters.generate() - state = secrets.token_urlsafe(32) - - auth_params = { - "response_type": "code", - "client_id": self.context.client_info.client_id, - "redirect_uri": str(self.context.client_metadata.redirect_uris[0]), - "state": state, - "code_challenge": pkce_params.code_challenge, - "code_challenge_method": "S256", - } - - # Only include resource param if conditions are met - if self.context.should_include_resource_param(self.context.protocol_version): - auth_params["resource"] = self.context.get_resource_url() # RFC 8707 - - if self.context.client_metadata.scope: # pragma: no branch - auth_params["scope"] = self.context.client_metadata.scope - - # OIDC requires prompt=consent when offline_access is requested - # https://openid.net/specs/openid-connect-core-1_0.html#OfflineAccess - if "offline_access" in self.context.client_metadata.scope.split(): - auth_params["prompt"] = "consent" - - authorization_url = f"{auth_endpoint}?{urlencode(auth_params)}" - await self.context.redirect_handler(authorization_url) - - # Wait for callback - result = await self.context.callback_handler() - - if result.state is None or not secrets.compare_digest(result.state, state): - raise OAuthFlowError(f"State parameter mismatch: {result.state} != {state}") - - # RFC 9207: validate the authorization-response issuer - validate_authorization_response_iss(result.iss, self.context.oauth_metadata) - - if not result.code: - raise OAuthFlowError("No authorization code received") - - # Return auth code and code verifier for token exchange - return result.code, pkce_params.code_verifier - - def _get_token_endpoint(self) -> str: - if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint: - token_url = str(self.context.oauth_metadata.token_endpoint) - else: - auth_base_url = self.context.get_authorization_base_url(self.context.server_url) - token_url = urljoin(auth_base_url, "/token") - return token_url - - async def _exchange_token_authorization_code(self, auth_code: str, code_verifier: str) -> httpx2.Request: - """Build token exchange request for authorization_code flow.""" - if self.context.client_metadata.redirect_uris is None: - raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover - if not self.context.client_info: - raise OAuthFlowError("Missing client info") # pragma: no cover - - token_url = self._get_token_endpoint() - token_data: dict[str, Any] = { - "grant_type": "authorization_code", - "code": auth_code, - "redirect_uri": str(self.context.client_metadata.redirect_uris[0]), - "client_id": self.context.client_info.client_id, - "code_verifier": code_verifier, - } - - # Only include resource param if conditions are met - if self.context.should_include_resource_param(self.context.protocol_version): - token_data["resource"] = self.context.get_resource_url() # RFC 8707 - - # Prepare authentication based on preferred method - headers = {"Content-Type": "application/x-www-form-urlencoded"} - token_data, headers = self.context.prepare_token_auth(token_data, headers) - - return httpx2.Request("POST", token_url, data=token_data, headers=headers) - - async def _handle_token_response(self, response: httpx2.Response) -> None: - """Handle token exchange response.""" - if response.status_code not in {200, 201}: - body = await response.aread() - body_text = body.decode("utf-8") - raise OAuthTokenError( - f"Token exchange failed ({response.status_code}){redirect_note(response)}: {body_text}" - ) - - # Parse and validate response with scope validation - token_response = await handle_token_response_scopes(response) - - # RFC 6749 §5.1: an omitted scope means the granted scope equals the requested - # scope. Record it explicitly so the persisted token is self-describing — the - # SEP-2350 step-up union reads it after a restart, when client_metadata.scope - # has reverted to its constructor value. - if token_response.scope is None: - token_response.scope = self.context.client_metadata.scope - - # Store tokens in context - self.context.current_tokens = token_response - self.context.update_token_expiry(token_response) - await self.context.storage.set_tokens(token_response) - - async def _refresh_token(self) -> httpx2.Request: - """Build token refresh request.""" - if not self.context.current_tokens or not self.context.current_tokens.refresh_token: - raise OAuthTokenError("No refresh token available") # pragma: no cover - - if not self.context.client_info or not self.context.client_info.client_id: - raise OAuthTokenError("No client info available") # pragma: no cover - - if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint: - token_url = str(self.context.oauth_metadata.token_endpoint) - else: - auth_base_url = self.context.get_authorization_base_url(self.context.server_url) - token_url = urljoin(auth_base_url, "/token") - - refresh_data: dict[str, str] = { - "grant_type": "refresh_token", - "refresh_token": self.context.current_tokens.refresh_token, - "client_id": self.context.client_info.client_id, - } - - # Only include resource param if conditions are met - if self.context.should_include_resource_param(self.context.protocol_version): - refresh_data["resource"] = self.context.get_resource_url() # RFC 8707 - - # Prepare authentication based on preferred method - headers = {"Content-Type": "application/x-www-form-urlencoded"} - refresh_data, headers = self.context.prepare_token_auth(refresh_data, headers) - - return httpx2.Request("POST", token_url, data=refresh_data, headers=headers) - - async def _handle_refresh_response(self, response: httpx2.Response) -> bool: - """Handle token refresh response. Returns True if successful.""" - if response.status_code != 200: - logger.warning(f"Token refresh failed: {response.status_code}{redirect_note(response)}") - self.context.clear_tokens() - return False - - try: - content = await response.aread() - token_response = OAuthToken.model_validate_json(content) - - # RFC 6749 §6: a refresh response may omit scope (unchanged) and refresh_token - # (the AS does not rotate). Carry both forward so the persisted token stays - # self-describing for the SEP-2350 step-up union and the next expiry can - # still refresh instead of forcing a full re-authorization. - prior = self.context.current_tokens - if token_response.scope is None and prior is not None: - token_response.scope = prior.scope - if token_response.refresh_token is None and prior is not None: - token_response.refresh_token = prior.refresh_token - - self.context.current_tokens = token_response - self.context.update_token_expiry(token_response) - await self.context.storage.set_tokens(token_response) - - return True - except ValidationError: # pragma: no cover - logger.exception("Invalid refresh response") - self.context.clear_tokens() - return False - - async def _initialize(self) -> None: - """Load stored tokens and client info.""" - self.context.current_tokens = await self.context.storage.get_tokens() - self.context.client_info = await self.context.storage.get_client_info() - self._initialized = True - - def _add_auth_header(self, request: httpx2.Request) -> None: - """Add authorization header to request if we have valid tokens.""" - if self.context.current_tokens and self.context.current_tokens.access_token: # pragma: no branch - request.headers["Authorization"] = f"Bearer {self.context.current_tokens.access_token}" - - async def _handle_oauth_metadata_response(self, response: httpx2.Response) -> None: - content = await response.aread() - metadata = OAuthMetadata.model_validate_json(content) - self.context.oauth_metadata = metadata - - async def _validate_resource_match(self, prm: ProtectedResourceMetadata) -> None: - """Validate that PRM resource matches the server URL per RFC 8707.""" - prm_resource = str(prm.resource) if prm.resource else None - - if self._validate_resource_url_callback is not None: - await self._validate_resource_url_callback(self.context.server_url, prm_resource) - return - - if not prm_resource: - return # pragma: no cover - default_resource = resource_url_from_server_url(self.context.server_url) - if not check_resource_allowed(requested_resource=default_resource, configured_resource=prm_resource): - raise OAuthFlowError(f"Protected resource {prm_resource} does not match expected {default_resource}") - - def _select_authorization_server(self, advertised: list[str]) -> str: - """Which of the servers listed in protected resource metadata to use: the first (the list is never empty).""" - return advertised[0] - - def _expected_issuer(self) -> str: - """The issuer that authorization server metadata and client credentials must belong to: the - PRM-advertised server, or on the legacy no-PRM path the resource server's origin, which is what - the 2025-03-26 well-known URL is built from (RFC 8414 §3.3).""" - return self.context.auth_server_url or _origin_issuer(self.context.server_url) - - async def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: - """The OAuth flow proper; `async_auth_flow` drives it (see `RedirectAwareAuth`).""" - async with self.context.lock: - if not self._initialized: - await self._initialize() - - # Capture protocol version from request headers - self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER) - - if not self.context.is_token_valid() and self.context.can_refresh_token(): - # Try to refresh token - refresh_request = await self._refresh_token() - refresh_response = yield refresh_request - - if not await self._handle_refresh_response(refresh_response): - # Refresh failed, need full re-authentication - self._initialized = False - - if self.context.is_token_valid(): - self._add_auth_header(request) - - response = yield request - - step_up = ( - response.status_code == 403 and extract_field_from_www_auth(response, "error") == "insufficient_scope" - ) - - if response.status_code == 401 or step_up: - # Perform full OAuth flow - try: - # Read before discovery, which may clear the tokens: on a restart the stored - # token's scope is the only record of what was granted (see Step 3). - granted_scope = self.context.current_tokens.scope if self.context.current_tokens else None - - # OAuth flow must be inline due to generator constraints. - # Steps 1-2 run on every 401. A scope step-up reuses the metadata discovered earlier - # in this process, and discovers it first when none is held yet (for example when - # tokens were loaded from storage), so re-authorization targets the right server. - if response.status_code == 401 or self.context.oauth_metadata is None: - www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response) - - # Step 1: Discover protected resource metadata (SEP-985 with fallback support) - prm_discovery_urls = build_protected_resource_metadata_discovery_urls( - www_auth_resource_metadata_url, self.context.server_url - ) - - prm_request_failed: int | None = None - for url in prm_discovery_urls: - discovery_request = create_oauth_metadata_request(url) - - discovery_response = yield discovery_request # sending request - - if discovery_response.status_code >= 500 or discovery_response.status_code == 429: - prm_request_failed = discovery_response.status_code - prm = await handle_protected_resource_response(discovery_response) - if prm: - # Validate PRM resource matches server URL (RFC 8707) - await self._validate_resource_match(prm) - self.context.protected_resource_metadata = prm - - self.context.auth_server_url = self._select_authorization_server( - [str(url) for url in prm.authorization_servers] - ) - break - else: - logger.debug(f"Protected resource metadata discovery failed: {url}") - else: - if prm_request_failed is not None: - # A server error says nothing about whether the resource publishes - # metadata, so it must not send the flow down the legacy path. - raise OAuthFlowError( - f"Protected resource metadata request failed: HTTP {prm_request_failed}" - ) - - expected_issuer = self._expected_issuer() - - # SEP-2352: stored credentials are bound to the issuer that registered them. - # Decided before any metadata is fetched: if the expected issuer is a different - # server, drop them (and the old tokens) so the flow re-registers instead of - # presenting another server's credentials. - if self.context.client_info is not None and not credentials_match_issuer( - self.context.client_info, expected_issuer, self.context.client_metadata_url - ): - logger.debug( - "Authorization server changed; discarding bound credentials and re-registering" - ) - self.context.client_info = None - self.context.clear_tokens() - # Any cached AS metadata is for the old server; drop it so a failed - # rediscovery cannot leak the old registration/token endpoints into Step 4. - self.context.oauth_metadata = None - - asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls( - self.context.auth_server_url, self.context.server_url - ) - - # Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers) - for url in asm_discovery_urls: # pragma: no branch - oauth_metadata_request = create_oauth_metadata_request(url) - oauth_metadata_response = yield oauth_metadata_request - - ok, asm = await handle_auth_metadata_response(oauth_metadata_response) - if not ok: - break - if ok and asm: - # SEP-2468 / RFC 8414 §3.3: the metadata must name the expected issuer. - # On the legacy path a root issuer rendered with its trailing slash - # names the same origin. - if self.context.auth_server_url is None and issuers_match( - str(asm.issuer), expected_issuer - ): - expected_issuer = str(asm.issuer) - validate_metadata_issuer(asm, expected_issuer) - self.context.oauth_metadata = asm - break - else: - logger.debug(f"OAuth metadata discovery failed: {url}") - - # Step 3: Apply scope selection strategy - challenged_scope = get_client_metadata_scopes( - extract_scope_from_www_auth(response), - self.context.protected_resource_metadata, - self.context.oauth_metadata, - self.context.client_metadata.grant_types, - ) - if step_up: - # SEP-2350: union previously requested scopes with the newly challenged ones so - # escalating one operation keeps the others' grants, folding in the granted - # scope read above since client_metadata.scope is not reloaded on a restart. - prior_scope = union_scopes(self.context.client_metadata.scope, granted_scope) - self.context.client_metadata.scope = union_scopes(prior_scope, challenged_scope) - else: - self.context.client_metadata.scope = challenged_scope - - # Step 4: Register client or use URL-based client ID (CIMD) - if not self.context.client_info: - # SEP-2352: the issuer to bind these credentials to, once metadata for it - # was actually found. - discovered_issuer = self._expected_issuer() if self.context.oauth_metadata is not None else None - - if should_use_client_metadata_url( - self.context.oauth_metadata, self.context.client_metadata_url - ): - # Use URL-based client ID (CIMD). CIMD records are portable across - # authorization servers, so the issuer stamp is informational. - logger.debug(f"Using URL-based client ID (CIMD): {self.context.client_metadata_url}") - client_information = create_client_info_from_metadata_url( - self.context.client_metadata_url, # type: ignore[arg-type] - redirect_uris=self.context.client_metadata.redirect_uris, - ) - client_information.issuer = discovered_issuer - self.context.client_info = client_information - await self.context.storage.set_client_info(client_information) - else: - # Fallback to Dynamic Client Registration - fallback_base = self.context.get_authorization_base_url(self.context.server_url) - registration_request = create_client_registration_request( - self.context.oauth_metadata, self.context.client_metadata, fallback_base - ) - registration_response = yield registration_request - client_information = await handle_registration_response(registration_response) - check_registration_usable(client_information) - # Only record the issuer when the registration above actually targeted - # the discovered AS — either via its published registration_endpoint, - # or because the resource-origin /register fallback is on the issuer's - # own host (legacy same-origin embedded AS). Otherwise the fallback hit - # a different server and recording a binding to the PRM-advertised AS - # would persist a binding that was never established. - if ( - self.context.oauth_metadata is not None - and discovered_issuer is not None - and ( - self.context.oauth_metadata.registration_endpoint is not None - or self.context.get_authorization_base_url(discovered_issuer) == fallback_base - ) - ): - client_information.issuer = discovered_issuer - self.context.client_info = client_information - await self.context.storage.set_client_info(client_information) - - # Step 5: Perform authorization and complete token exchange - token_response = yield await self._perform_authorization() - await self._handle_token_response(token_response) - except Exception: - logger.exception("OAuth flow error") - raise - - # Retry with new tokens - self._add_auth_header(request) - yield request +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/auth/utils.py b/src/mcp/client/auth/utils.py index 0a2ba80aec..74528bc1f9 100644 --- a/src/mcp/client/auth/utils.py +++ b/src/mcp/client/auth/utils.py @@ -1,442 +1,50 @@ -import re -from typing import Any, cast -from urllib.parse import urljoin, urlparse - -from httpx2 import Request, Response -from mcp_types import LATEST_PROTOCOL_VERSION -from pydantic import AnyUrl, ValidationError -from pydantic_core import from_json - -from mcp.client.auth import OAuthFlowError, OAuthRegistrationError, OAuthTokenError -from mcp.shared._httpx_utils import redirect_note -from mcp.shared.auth import ( - OAuthClientInformationFull, - OAuthClientMetadata, - OAuthMetadata, - OAuthToken, - ProtectedResourceMetadata, +import sys + +import mcp_client.client.auth.utils as _implementation +from mcp_client.client.auth.utils import ( + build_oauth_authorization_server_metadata_discovery_urls, + build_protected_resource_metadata_discovery_urls, + create_client_info_from_metadata_url, + create_client_registration_request, + create_oauth_metadata_request, + credentials_match_issuer, + extract_field_from_www_auth, + extract_resource_metadata_from_www_auth, + extract_scope_from_www_auth, + get_client_metadata_scopes, + handle_auth_metadata_response, + handle_protected_resource_response, + handle_registration_response, + handle_token_response_scopes, + is_valid_client_metadata_url, + issuers_match, + should_use_client_metadata_url, + union_scopes, + validate_authorization_response_iss, + validate_metadata_issuer, ) -from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER - - -def extract_field_from_www_auth(response: Response, field_name: str) -> str | None: - """Extract field from WWW-Authenticate header. - - Returns: - Field value if found in WWW-Authenticate header, None otherwise - """ - www_auth_header = response.headers.get("WWW-Authenticate") - if not www_auth_header: - return None - - # Pattern matches: field_name="value" or field_name=value (unquoted) - pattern = rf'{field_name}=(?:"([^"]+)"|([^\s,]+))' - match = re.search(pattern, www_auth_header) - - if match: - # Return quoted value if present, otherwise unquoted value - return match.group(1) or match.group(2) - - return None - - -def extract_scope_from_www_auth(response: Response) -> str | None: - """Extract scope parameter from WWW-Authenticate header as per RFC 6750. - - Returns: - Scope string if found in WWW-Authenticate header, None otherwise - """ - return extract_field_from_www_auth(response, "scope") - - -def extract_resource_metadata_from_www_auth(response: Response) -> str | None: - """Extract protected resource metadata URL from WWW-Authenticate header as per RFC 9728. - - Returns: - Resource metadata URL if found in WWW-Authenticate header, None otherwise - """ - if not response or response.status_code not in (401, 403): - return None # pragma: no cover - - return extract_field_from_www_auth(response, "resource_metadata") - - -def build_protected_resource_metadata_discovery_urls(www_auth_url: str | None, server_url: str) -> list[str]: - """Build ordered list of URLs to try for protected resource metadata discovery. - - Per SEP-985, the client MUST: - 1. Try resource_metadata from WWW-Authenticate header (if present) - 2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path} - 3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource - - Args: - www_auth_url: Optional resource_metadata URL extracted from the WWW-Authenticate header - server_url: Server URL - - Returns: - Ordered list of URLs to try for discovery - """ - urls: list[str] = [] - - # Priority 1: WWW-Authenticate header with resource_metadata parameter - if www_auth_url: - urls.append(www_auth_url) - - # Priority 2-3: Well-known URIs (RFC 9728) - parsed = urlparse(server_url) - base_url = f"{parsed.scheme}://{parsed.netloc}" - - # Priority 2: Path-based well-known URI (if server has a path component) - if parsed.path and parsed.path != "/": - path_based_url = urljoin(base_url, f"/.well-known/oauth-protected-resource{parsed.path}") - urls.append(path_based_url) - - # Priority 3: Root-based well-known URI - root_based_url = urljoin(base_url, "/.well-known/oauth-protected-resource") - urls.append(root_based_url) - - return urls - - -def get_client_metadata_scopes( - www_authenticate_scope: str | None, - protected_resource_metadata: ProtectedResourceMetadata | None, - authorization_server_metadata: OAuthMetadata | None = None, - client_grant_types: list[str] | None = None, -) -> str | None: - """Select effective scopes and augment for refresh token support.""" - selected_scope: str | None = None - - # MCP spec scope selection priority: - # 1. WWW-Authenticate header scope - # 2. PRM scopes_supported - # 3. AS scopes_supported (SDK fallback) - # 4. Omit scope parameter - if www_authenticate_scope is not None: - selected_scope = www_authenticate_scope - elif protected_resource_metadata is not None and protected_resource_metadata.scopes_supported is not None: - selected_scope = " ".join(protected_resource_metadata.scopes_supported) - elif authorization_server_metadata is not None and authorization_server_metadata.scopes_supported is not None: - selected_scope = " ".join(authorization_server_metadata.scopes_supported) - - # SEP-2207: append offline_access when the AS supports it and the client can use refresh tokens - if ( - selected_scope is not None - and authorization_server_metadata is not None - and authorization_server_metadata.scopes_supported is not None - and "offline_access" in authorization_server_metadata.scopes_supported - and client_grant_types is not None - and "refresh_token" in client_grant_types - and "offline_access" not in selected_scope.split() - ): - selected_scope = f"{selected_scope} offline_access" - - return selected_scope - - -def union_scopes(previous_scope: str | None, new_scope: str | None) -> str | None: - """Merge two space-delimited scope strings, preserving order and dropping duplicates. - - SEP-2350: on step-up re-authorization the client requests the union of previously requested - scopes and the newly challenged scopes, so escalating one operation does not drop the - permissions granted for another. Previously requested scopes come first; new scopes are - appended in order. - """ - if not previous_scope: - return new_scope - if not new_scope: - return previous_scope - - merged = previous_scope.split() - seen = set(merged) - for scope in new_scope.split(): - if scope not in seen: - merged.append(scope) - seen.add(scope) - return " ".join(merged) - - -def build_oauth_authorization_server_metadata_discovery_urls(auth_server_url: str | None, server_url: str) -> list[str]: - """Generate an ordered list of URLs for authorization server metadata discovery. - - Args: - auth_server_url: OAuth Authorization Server Metadata URL if found, otherwise None - server_url: URL for the MCP server, used as a fallback if auth_server_url is None - """ - - if not auth_server_url: - # Legacy path using the 2025-03-26 spec: - # link: https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization - parsed = urlparse(server_url) - return [f"{parsed.scheme}://{parsed.netloc}/.well-known/oauth-authorization-server"] - - urls: list[str] = [] - parsed = urlparse(auth_server_url) - base_url = f"{parsed.scheme}://{parsed.netloc}" - - # RFC 8414: Path-aware OAuth discovery - if parsed.path and parsed.path != "/": - oauth_path = f"/.well-known/oauth-authorization-server{parsed.path.rstrip('/')}" - urls.append(urljoin(base_url, oauth_path)) - - # RFC 8414 section 5: Path-aware OIDC discovery - # See https://www.rfc-editor.org/rfc/rfc8414.html#section-5 - oidc_path = f"/.well-known/openid-configuration{parsed.path.rstrip('/')}" - urls.append(urljoin(base_url, oidc_path)) - - # https://openid.net/specs/openid-connect-discovery-1_0.html - oidc_path = f"{parsed.path.rstrip('/')}/.well-known/openid-configuration" - urls.append(urljoin(base_url, oidc_path)) - return urls - - # OAuth root - urls.append(urljoin(base_url, "/.well-known/oauth-authorization-server")) - - # OIDC 1.0 fallback (appends to full URL per OIDC spec) - # https://openid.net/specs/openid-connect-discovery-1_0.html - urls.append(urljoin(base_url, "/.well-known/openid-configuration")) - - return urls - - -async def handle_protected_resource_response( - response: Response, -) -> ProtectedResourceMetadata | None: - """Handle protected resource metadata discovery response. - - Per SEP-985, supports fallback when discovery fails at one URL. - - Returns: - ProtectedResourceMetadata if successfully discovered, None if we should try next URL - """ - if response.status_code == 200: - try: - content = await response.aread() - metadata = ProtectedResourceMetadata.model_validate_json(content) - return metadata - - except ValidationError: # pragma: no cover - # Invalid metadata - try next URL - return None - else: - # Not found - try next URL in fallback chain - return None - - -async def handle_auth_metadata_response(response: Response) -> tuple[bool, OAuthMetadata | None]: - if response.status_code == 200: - try: - content = await response.aread() - asm = OAuthMetadata.model_validate_json(content) - return True, asm - except ValidationError: # pragma: no cover - return True, None - elif 300 <= response.status_code < 500: - return True, None # Not served at this URL (redirects are not followed) - try the next candidate - return False, None # Server error or unexpected status, stop trying - - -def validate_authorization_response_iss(iss: str | None, oauth_metadata: OAuthMetadata | None) -> None: - """Validate the RFC 9207 `iss` authorization-response parameter. - - Per RFC 9207 section 2.4, the client compares `iss` against the issuer of the - authorization server the request was sent to, using simple string comparison - (RFC 3986 section 6.2.1, i.e. without URL normalization), and rejects on mismatch. - A response that omits `iss` is rejected only when the server advertised support via - `authorization_response_iss_parameter_supported`. - - Raises: - OAuthFlowError: If `iss` is present and does not match, or is absent when the - authorization server advertised support. - """ - expected = str(oauth_metadata.issuer) if oauth_metadata else None - - if iss is not None: - if iss != expected: - raise OAuthFlowError(f"Authorization response iss mismatch: {iss} != {expected}") - return - - if oauth_metadata is not None and oauth_metadata.authorization_response_iss_parameter_supported: - raise OAuthFlowError("Authorization response missing iss parameter advertised by the authorization server") - - -def validate_metadata_issuer(oauth_metadata: OAuthMetadata, expected_issuer: str) -> None: - """Validate that authorization server metadata `issuer` matches the discovery issuer. - - Per RFC 8414 section 3.3 / SEP-2468, the `issuer` in the metadata must match the issuer - used to construct the well-known URL, compared as a simple string (RFC 3986 section 6.2.1). - - Raises: - OAuthFlowError: If the metadata issuer does not match `expected_issuer`. - """ - if str(oauth_metadata.issuer) != expected_issuer: - raise OAuthFlowError( - f"Authorization server metadata issuer mismatch: {oauth_metadata.issuer} != {expected_issuer}" - ) - - -def create_oauth_metadata_request(url: str) -> Request: - return Request("GET", url, headers={MCP_PROTOCOL_VERSION_HEADER: LATEST_PROTOCOL_VERSION}) - - -def create_client_registration_request( - auth_server_metadata: OAuthMetadata | None, client_metadata: OAuthClientMetadata, auth_base_url: str -) -> Request: - """Build a client registration request.""" - - if auth_server_metadata and auth_server_metadata.registration_endpoint: - registration_url = str(auth_server_metadata.registration_endpoint) - else: - registration_url = urljoin(auth_base_url, "/register") - - registration_data = client_metadata.model_dump(by_alias=True, mode="json", exclude_none=True) - - return Request("POST", registration_url, json=registration_data, headers={"Content-Type": "application/json"}) - - -async def handle_registration_response(response: Response) -> OAuthClientInformationFull: - """Handle registration response.""" - if response.status_code not in (200, 201): - await response.aread() - raise OAuthRegistrationError( - f"Registration failed: {response.status_code}{redirect_note(response)} {response.text}" - ) - - try: - content = await response.aread() - body = from_json(content) - # `issuer` is the SDK's own binding of these credentials to the server they were - # registered with (SEP-2352), stamped by the auth flow - never sourced from the - # wire, so it is dropped before the body is parsed rather than trusted or cleared. - if isinstance(body, dict): - cast(dict[str, Any], body).pop("issuer", None) - return OAuthClientInformationFull.model_validate(body) - except ValueError as e: - # `from_json` reports malformed bytes/JSON as ValueError, and pydantic's - # ValidationError is itself a ValueError, so both parse layers surface here. - raise OAuthRegistrationError(f"Invalid registration response: {e}") from e - - -def is_valid_client_metadata_url(url: str | None) -> bool: - """Validate that a URL is suitable for use as a client_id (CIMD). - - The URL must be HTTPS with a non-root pathname. - - Args: - url: The URL to validate - - Returns: - True if the URL is a valid HTTPS URL with a non-root pathname - """ - if not url: - return False - try: - parsed = urlparse(url) - return parsed.scheme == "https" and parsed.path not in ("", "/") - except Exception: - return False - - -def credentials_match_issuer( - client_info: OAuthClientInformationFull, issuer: str, client_metadata_url: str | None -) -> bool: - """Whether stored client credentials may be reused against `issuer` (SEP-2352). - - A URL-based client ID (CIMD) is portable across authorization servers — the same self-hosted - document is resolved by whichever server is in use — so it always matches; CIMD is identified - by the client ID being the configured `client_metadata_url`, not by URL shape (a registration - server may also issue URL-shaped IDs that are bound to it). Credentials with a recorded issuer - match only when it equals `issuer` (simple string comparison; a root issuer with and without - its trailing slash count as equal). Credentials with no recorded - issuer (pre-registered, or stored before issuer binding existed) carry no binding to enforce - and are left as-is. - """ - if client_metadata_url is not None and client_info.client_id == client_metadata_url: - return True - if client_info.issuer is None: - return True - return issuers_match(client_info.issuer, issuer) - - -def issuers_match(a: str, b: str) -> bool: - """Simple string comparison of two issuer identifiers (RFC 8414 section 3.3), except that a root - issuer with and without its trailing slash (`scheme://authority` and `scheme://authority/`) name - the same server.""" - if a == b: - return True - shorter, longer = sorted((a, b), key=len) - parsed = urlparse(shorter) - return longer == f"{shorter}/" and shorter == f"{parsed.scheme}://{parsed.netloc}" - - -def should_use_client_metadata_url( - oauth_metadata: OAuthMetadata | None, - client_metadata_url: str | None, -) -> bool: - """Determine if URL-based client ID (CIMD) should be used instead of DCR. - - URL-based client IDs should be used when: - 1. The server advertises client_id_metadata_document_supported=True - 2. The client has a valid client_metadata_url configured - - Args: - oauth_metadata: OAuth authorization server metadata - client_metadata_url: URL-based client ID (already validated) - - Returns: - True if CIMD should be used, False if DCR should be used - """ - if not client_metadata_url: - return False - - if not oauth_metadata: - return False - - return oauth_metadata.client_id_metadata_document_supported is True - - -def create_client_info_from_metadata_url( - client_metadata_url: str, redirect_uris: list[AnyUrl] | None = None -) -> OAuthClientInformationFull: - """Create client information using a URL-based client ID (CIMD). - - When using URL-based client IDs, the URL itself becomes the client_id - and no client_secret is used (token_endpoint_auth_method="none"). - - Args: - client_metadata_url: The URL to use as the client_id - redirect_uris: The redirect URIs from the client metadata, recorded on the client - information alongside the client_id - - Returns: - OAuthClientInformationFull with the URL as client_id - """ - return OAuthClientInformationFull( - client_id=client_metadata_url, - token_endpoint_auth_method="none", - redirect_uris=redirect_uris, - ) - - -async def handle_token_response_scopes( - response: Response, -) -> OAuthToken: - """Parse and validate a token response. - - Parses token response JSON. Callers should check response.status_code before calling. - - Args: - response: HTTP response from token endpoint (status already checked by caller) - - Returns: - Validated OAuthToken model - Raises: - OAuthTokenError: If response JSON is invalid - """ - try: - content = await response.aread() - token_response = OAuthToken.model_validate_json(content) - return token_response - except ValidationError as e: # pragma: no cover - raise OAuthTokenError(f"Invalid token response: {e}") +__all__ = [ + "build_oauth_authorization_server_metadata_discovery_urls", + "build_protected_resource_metadata_discovery_urls", + "create_client_info_from_metadata_url", + "create_client_registration_request", + "create_oauth_metadata_request", + "credentials_match_issuer", + "extract_field_from_www_auth", + "extract_resource_metadata_from_www_auth", + "extract_scope_from_www_auth", + "get_client_metadata_scopes", + "handle_auth_metadata_response", + "handle_protected_resource_response", + "handle_registration_response", + "handle_token_response_scopes", + "is_valid_client_metadata_url", + "issuers_match", + "should_use_client_metadata_url", + "union_scopes", + "validate_authorization_response_iss", + "validate_metadata_issuer", +] + +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/caching.py b/src/mcp/client/caching.py index a464accd16..a4e55e93bf 100644 --- a/src/mcp/client/caching.py +++ b/src/mcp/client/caching.py @@ -1,387 +1,38 @@ -"""Client-side response caching primitives (SEP-2549, protocol revision 2026-07-28).""" +import sys -from __future__ import annotations - -import json -import logging -import time -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any, Final, Literal, Protocol - -import anyio -import anyio.lowlevel -from mcp_types import ( - CacheableResult, - PromptListChangedNotification, - ResourceListChangedNotification, - ResourceUpdatedNotification, - ServerNotification, - ToolListChangedNotification, +import mcp_client.client.caching as _implementation +from mcp_client.client.caching import ( + _GENERATION_MAP_CAP as _GENERATION_MAP_CAP, +) +from mcp_client.client.caching import ( + _STORE_CLEANUP_TIMEOUT as _STORE_CLEANUP_TIMEOUT, +) +from mcp_client.client.caching import ( + MAX_TTL_MS as MAX_TTL_MS, +) +from mcp_client.client.caching import ( + CacheConfig as CacheConfig, +) +from mcp_client.client.caching import ( + CacheEntry as CacheEntry, +) +from mcp_client.client.caching import ( + CacheKey as CacheKey, +) +from mcp_client.client.caching import ( + CacheMode as CacheMode, +) +from mcp_client.client.caching import ( + ClientResponseCache as ClientResponseCache, +) +from mcp_client.client.caching import ( + InMemoryResponseCacheStore as InMemoryResponseCacheStore, +) +from mcp_client.client.caching import ( + ResponseCacheStore as ResponseCacheStore, +) +from mcp_client.client.caching import ( + logger as logger, ) -from mcp_types.version import MODERN_PROTOCOL_VERSIONS - -__all__ = [ - "MAX_TTL_MS", - "CacheConfig", - "CacheEntry", - "CacheKey", - "CacheMode", - "InMemoryResponseCacheStore", - "ResponseCacheStore", -] - -logger = logging.getLogger(__name__) - -CacheMode = Literal["use", "refresh", "bypass"] -"""Per-call cache behavior: `"use"` serves and stores, `"refresh"` stores -without serving, `"bypass"` skips the cache entirely.""" - -MAX_TTL_MS: Final[int] = 24 * 60 * 60 * 1000 -"""Cap on any entry's time-to-live (24 hours, in milliseconds); larger `ttlMs` values are clamped down.""" - - -@dataclass(frozen=True, slots=True) -class CacheKey: - """Identity of one cached response; compare as the field tuple, never a flattened string (collision hazard).""" - - method: str - - params_key: str = "" - """Result-affecting params discriminator: the uri for `resources/read`, `""` for the list methods.""" - - partition: str = "" - """Coordinator-computed arm identifier; opaque to stores.""" - - -@dataclass(frozen=True, slots=True) -class CacheEntry: - """One cached response with its freshness and sharing metadata.""" - - value: Any - """The cached result; the SDK deep-copies on write and on serve, so a store may hold it as-is.""" - - scope: Literal["public", "private"] - """Server-asserted `cacheScope`: only `"public"` entries may be shared across authorization contexts.""" - - expires_at: float | None - """Epoch seconds after which the entry is stale; `None` is never fresh.""" - - -class ResponseCacheStore(Protocol): - """Storage contract for the client response cache. - - Each `Client` calls its store from a single event loop; per-operation - atomicity is the implementation's responsibility. Operations may raise - - the SDK degrades to a miss rather than failing the call. A serializing - store must round-trip `value` back to the result model object (a - wrong-shape entry is a miss, never an error). A lookup may issue two - sequential `get` calls (private arm, then public). - """ - - async def get(self, key: CacheKey) -> CacheEntry | None: ... - - async def set(self, key: CacheKey, entry: CacheEntry) -> None: ... - - async def delete(self, key: CacheKey) -> None: ... - - async def clear(self) -> None: ... - - -@dataclass(frozen=True, slots=True) -class CacheConfig: - """Configuration for a `Client`'s response cache. - - Raises: - ValueError: On a custom `store` without `partition`, an empty `target_id`, or a negative `default_ttl_ms`. - """ - - store: ResponseCacheStore | None = None - """Backing store; `None` means a per-client `InMemoryResponseCacheStore`. - A custom store requires an explicit `partition`.""" - - partition: str = "" - """Authorization-context identifier isolating `"private"`-scoped entries - within a shared store. Derive it from a verified credential - never from - request-supplied data or the server URL. Fixed for the `Client`'s - lifetime: construct a new `Client` when the principal changes.""" - - target_id: str | None = None - """Server-identity override for custom transports and proxies where the - SDK cannot derive one from a URL; must be non-empty when provided.""" - - default_ttl_ms: int = 0 - """TTL in milliseconds for results carrying no `ttlMs` hint; the default `0` leaves them uncached.""" - - clock: Callable[[], float] = time.time - """Wall-clock source returning epoch seconds; injectable for expiry tests.""" - - share_public: bool = False - """Serve server-marked `"public"` entries across every partition in the store. - - WARNING: this trusts the server's `"public"` classification for every - principal sharing the store - a mislabeled response leaks across tenants. - Constructor-level only: the per-call `cache_mode` can never widen sharing.""" - - def __post_init__(self) -> None: - if self.store is not None and not self.partition: - raise ValueError("a custom store requires an explicit partition") - if self.target_id == "": - raise ValueError("target_id must be a non-empty string or omitted") - if self.default_ttl_ms < 0: - raise ValueError(f"default_ttl_ms must be >= 0, got {self.default_ttl_ms}") - - -class InMemoryResponseCacheStore: - """Default in-process `ResponseCacheStore`. - - Method bodies are synchronous, so concurrent tasks never observe a torn - write. `max_entries` caps the whole store, evicting least-recently-used - at the cap (`0` disables it); `get` and `set` both refresh recency, so a - hot entry survives churn from other keys. - - Raises: - ValueError: If `max_entries` is negative. - """ - - def __init__(self, *, max_entries: int = 1024) -> None: - if max_entries < 0: - raise ValueError(f"max_entries must be >= 0, got {max_entries}") - self._max_entries = max_entries - self._entries: dict[CacheKey, CacheEntry] = {} - - async def get(self, key: CacheKey) -> CacheEntry | None: - entry = self._entries.get(key) - if entry is not None: - # Pop-and-reinsert moves the key to the back: the dict's insertion order is the LRU ledger. - self._entries[key] = self._entries.pop(key) - return entry - - async def set(self, key: CacheKey, entry: CacheEntry) -> None: - self._entries.pop(key, None) - self._entries[key] = entry - if self._max_entries and len(self._entries) > self._max_entries: - del self._entries[next(iter(self._entries))] - - async def delete(self, key: CacheKey) -> None: - self._entries.pop(key, None) - - async def clear(self) -> None: - self._entries.clear() - - -_GENERATION_MAP_CAP: Final[int] = 4096 -"""Cap on the generation map; at the cap the oldest key's eviction-race guard is dropped (FIFO).""" - -_STORE_CLEANUP_TIMEOUT: Final[float] = 5 -"""Bound for must-complete store cleanup deletes (mirrors the dispatcher's final-write bound); -a wedged store delete must not hold client teardown uncancellably.""" - - -class ClientResponseCache: - """Coordinates the `Client` caching verbs with a `ResponseCacheStore`: keys, era gate, TTL/scope, eviction.""" - - def __init__( - self, - *, - store: ResponseCacheStore, - partition: str, - arm_id: str, - default_ttl_ms: int, - clock: Callable[[], float], - share_public: bool, - negotiated_version: Callable[[], str | None], - generation_map_cap: int = _GENERATION_MAP_CAP, - store_cleanup_timeout: float = _STORE_CLEANUP_TIMEOUT, - ) -> None: - self._store = store - self._partition = partition - self._arm_id = arm_id - self._share_public = share_public - self._default_ttl_ms = default_ttl_ms - self._clock = clock - self._negotiated_version = negotiated_version - # A key is eviction-race-guarded iff registered here. - self._generations: dict[tuple[str, str], int] = {} - self._generation_map_cap = generation_map_cap - self._store_cleanup_timeout = store_cleanup_timeout - self._warned_store_ops: set[str] = set() - - def _arm(self, scope: Literal["public", "private"]) -> str: - # JSON arrays so crafted arm_id/partition values cannot collide across field boundaries. - # The negotiated version era-scopes every arm: a session never serves an entry written - # under a different protocol era (its content differs - sieve-stripped fields, header - # filtering). Every caller runs post-connect; were that ever untrue, the supplier's - # None still partitions harmlessly. - fields: list[str | None] = [scope, self._negotiated_version(), self._arm_id] - if scope == "private" or not self._share_public: - fields.append(self._partition) - return json.dumps(fields) - - async def read(self, method: str, params_key: str) -> CacheableResult | None: - """Serve a fresh entry for the key, or `None`; the served result is a deep copy.""" - # A hit completes without any other yielding await, so checkpoint here: a poll - # loop over a fresh entry must not starve spawned tasks (eviction dispatch). - await anyio.lowlevel.checkpoint() - # A wrong-shape entry raises as late as the copy, so the boundary wraps the whole read path. - try: - entry = await self._get_fresh(CacheKey(method, params_key, self._arm("private"))) - if entry is None: - # After a scope flip, a stale private entry must not shadow a fresh public one. - entry = await self._get_fresh(CacheKey(method, params_key, self._arm("public"))) - if entry is not None and entry.scope != "public": - # Never serve an entry the server scoped "private" out of the shared arm. - entry = None - copied: CacheableResult | None = None if entry is None else entry.value.model_copy(deep=True) - except Exception: # boundary around user store code: any read-path failure is a miss, never a failed call - self._warn_store_failure("get") - return None - self._warned_store_ops.discard("get") - return copied - - async def _get_fresh(self, key: CacheKey) -> CacheEntry | None: - entry = await self._store.get(key) - if entry is None or entry.expires_at is None or entry.expires_at <= self._clock(): - return None - return entry - - def capture(self, method: str, params_key: str) -> int: - """Register the key for eviction-race detection before the fetch; `write` takes the returned generation.""" - gen_key = (method, params_key) - if gen_key not in self._generations: - if len(self._generations) >= self._generation_map_cap: - # FIFO overflow: the dropped key's race guard degrades to the accepted co-tenant class. - del self._generations[next(iter(self._generations))] - self._generations[gen_key] = 0 - return self._generations[gen_key] - - async def write( - self, - method: str, - params_key: str, - result: CacheableResult, - gen_at_capture: int, - mode: Literal["use", "refresh"], - ) -> None: - """Store a fetched result under the arm its resolved scope selects.""" - gen_key = (method, params_key) - if self._generation_moved(gen_key, gen_at_capture): - return # the key was evicted while the fetch was in flight - ttl_ms, scope = self._resolve(result) - private_key = CacheKey(method, params_key, self._arm("private")) - public_key = CacheKey(method, params_key, self._arm("public")) - if ttl_ms <= 0: - if mode == "refresh": - # The refetch superseded the warm entry, which a cancellation must not leave serving. - await self._cleanup_delete(private_key, public_key) - return - own, opposite = (public_key, private_key) if scope == "public" else (private_key, public_key) - # Opposite arm first: a failed delete aborts before the set - never two arms answering for one key. - if not await self._delete(opposite): - # The own arm's entry is superseded too: best-effort delete, degrading to a full miss. - await self._cleanup_delete(own) - return - entry = CacheEntry(value=result.model_copy(deep=True), scope=scope, expires_at=self._clock() + ttl_ms / 1000) - try: - if not await self._set(own, entry): - # The fetch superseded any pre-existing own-arm entry, and the failed set - # left it in place: purge it (mirrors the opposite-arm-failure path). - await self._cleanup_delete(own) - finally: - # An eviction can land while the set commits - even when the await - # is cancelled - so re-check on every exit; the delete must complete - # so the pending cancellation cannot resurrect the evicted entry. - if self._generation_moved(gen_key, gen_at_capture): - await self._cleanup_delete(own) - - async def evict_method(self, method: str) -> None: - """Evict the method's cursor-less entry.""" - await self.evict_key(method, "") - - async def evict_key(self, method: str, params_key: str) -> None: - """Evict one key from both arms. - - Only the current era's arms are touched; other-era entries in a persistent store age out by TTL. - """ - gen_key = (method, params_key) - # Bump first so an in-flight fetch cannot write the evicted entry back. - # Unregistered keys skip the bump (uris must not grow the map) but not - # the deletes - a persistent store may hold uncaptured entries. - if gen_key in self._generations: - self._generations[gen_key] += 1 - # Must complete: a cancellation between the deletes would leave one arm serving the evicted entry. - await self._cleanup_delete( - CacheKey(method, params_key, self._arm("private")), - CacheKey(method, params_key, self._arm("public")), - ) - - async def evict_for_notification(self, notification: ServerNotification) -> None: - """Map a server notification to the entries it makes stale. - - Eviction is eventual (spawned-task dispatch): the generation bump closes - the write-back race; a racing read may briefly serve the old entry. - """ - match notification: - case ToolListChangedNotification(): - await self.evict_method("tools/list") - case PromptListChangedNotification(): - await self.evict_method("prompts/list") - case ResourceListChangedNotification(): - # Templates enumerate the same changed resource space. - await self.evict_method("resources/list") - await self.evict_method("resources/templates/list") - case ResourceUpdatedNotification(): - await self.evict_key("resources/read", notification.params.uri) - case _: - pass - - def _resolve(self, result: CacheableResult) -> tuple[int, Literal["public", "private"]]: - # A legacy peer can also put `ttlMs`/`cacheScope` keys on the wire, so - # wire presence is not a peer-era signal - hints count only when modern. - modern = self._negotiated_version() in MODERN_PROTOCOL_VERSIONS - if modern and "ttl_ms" in result.model_fields_set: - # An explicit `ttlMs: 0` stays 0, and negatives are unconstructible - # upstream (model ge=0, parse-seam floor) - only the cap applies. - ttl_ms = result.ttl_ms - else: - ttl_ms = self._default_ttl_ms - scope: Literal["public", "private"] = "public" if modern and result.cache_scope == "public" else "private" - return min(ttl_ms, MAX_TTL_MS), scope - - def _generation_moved(self, gen_key: tuple[str, str], gen_at_capture: int) -> bool: - # A FIFO-dropped key fails open (the accepted co-tenant race) rather than discarding the fetch. - return self._generations.get(gen_key, gen_at_capture) != gen_at_capture - - async def _set(self, key: CacheKey, entry: CacheEntry) -> bool: - try: - await self._store.set(key, entry) - except Exception: # boundary around user store code: nothing cached, the fetch already succeeded - self._warn_store_failure("set") - return False - self._warned_store_ops.discard("set") - return True - - async def _cleanup_delete(self, *keys: CacheKey) -> None: - # Must-complete cleanup: shielded so a pending cancellation cannot skip the deletes, - # bounded so a wedged store delete cannot hold client teardown uncancellably. - with anyio.move_on_after(self._store_cleanup_timeout, shield=True) as scope: - for key in keys: - await self._delete(key) - if scope.cancelled_caught: - logger.warning("Response cache store delete timed out; the entry will age out by TTL") - - async def _delete(self, key: CacheKey) -> bool: - try: - await self._store.delete(key) - except Exception: # boundary around user store code: callers decide whether a failed delete aborts - self._warn_store_failure("delete") - return False - self._warned_store_ops.discard("delete") - return True - def _warn_store_failure(self, kind: Literal["get", "set", "delete"]) -> None: - # One warning per failure burst, per op kind; re-armed only when that - # same kind succeeds, so a healthy delete cannot re-arm a broken set. - if kind not in self._warned_store_ops: - self._warned_store_ops.add(kind) - logger.warning("Response cache store operation failed; continuing without the cache", exc_info=True) +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/client.py b/src/mcp/client/client.py index f921c7e30b..62f317ed20 100644 --- a/src/mcp/client/client.py +++ b/src/mcp/client/client.py @@ -1,949 +1,53 @@ -"""Unified MCP Client that wraps ClientSession with transport management.""" +import sys +from typing import Any -from __future__ import annotations - -import hashlib -import logging -import uuid -from collections.abc import Awaitable, Callable, Mapping, Sequence -from contextlib import AbstractAsyncContextManager, AsyncExitStack -from dataclasses import KW_ONLY, dataclass, field -from typing import Any, Literal, TypeVar, cast - -import anyio -import anyio.lowlevel -import mcp_types as types -from mcp_types import ( - INVALID_PARAMS, - CacheableResult, - CallToolResult, - CompleteResult, - EmptyResult, - ErrorData, - GetPromptResult, - Implementation, - InputRequest, - InputRequiredResult, - InputResponse, - InputResponses, - ListPromptsResult, - ListResourcesResult, - ListResourceTemplatesResult, - ListToolsResult, - LoggingLevel, - PaginatedRequestParams, - PromptReference, - ReadResourceResult, - RequestParamsMeta, - ResourceTemplateReference, - Result, - ServerCapabilities, +import mcp_client.client.client as _implementation +from mcp_client.client.client import ( + _T as _T, ) -from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS -from typing_extensions import deprecated - -from mcp.client._input_required import DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, run_input_required_driver -from mcp.client._memory import InMemoryTransport -from mcp.client._probe import negotiate_auto -from mcp.client._transport import Transport -from mcp.client.caching import CacheConfig, CacheMode, ClientResponseCache, InMemoryResponseCacheStore -from mcp.client.extension import ClaimContext, ClientExtension, NotificationBinding, ResultClaim -from mcp.client.session import ( - ClientRequestContext, - ClientSession, - ElicitationFnT, - IncomingMessage, - ListRootsFnT, - LoggingFnT, - MessageHandlerFnT, - SamplingFnT, +from mcp_client.client.client import ( + Client as Client, +) +from mcp_client.client.client import ( + ConnectMode as ConnectMode, +) +from mcp_client.client.client import ( + _CacheableT as _CacheableT, +) +from mcp_client.client.client import ( + _connect_transport as _connect_transport, +) +from mcp_client.client.client import ( + _connected as _connected, +) +from mcp_client.client.client import ( + _Connector as _Connector, +) +from mcp_client.client.client import ( + _evicting_message_handler as _evicting_message_handler, ) -from mcp.client.stdio import StdioServerParameters, stdio_client -from mcp.client.streamable_http import streamable_http_client -from mcp.client.subscriptions import ServerEvent, Subscription -from mcp.client.subscriptions import listen as _listen +from mcp_client.client.client import ( + _fold_extensions as _fold_extensions, +) +from mcp_client.client.client import ( + _FoldedExtensions as _FoldedExtensions, +) +from mcp_client.client.client import ( + _ResultT as _ResultT, +) +from mcp_client.client.client import ( + _strip_userinfo as _strip_userinfo, +) +from mcp_client.client.client import ( + _synthesize_discover as _synthesize_discover, +) +from mcp_client.client.client import ( + logger as logger, +) + from mcp.server import Server from mcp.server.mcpserver import MCPServer -from mcp.server.runner import modern_on_request -from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair -from mcp.shared.dispatcher import Dispatcher, ProgressFnT -from mcp.shared.exceptions import MCPDeprecationWarning, MCPError -from mcp.shared.extension import validate_extension_identifier -from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher -from mcp.shared.subscriptions import event_to_notification - -logger = logging.getLogger(__name__) - -ConnectMode = Literal["legacy", "auto"] | str -"""``mode=`` value: ``"legacy"`` (initialize handshake), ``"auto"`` (discover, fall back to -initialize), or a modern protocol-version string (adopt directly). The ``str`` arm is for -forward-compat; ``Client.__post_init__`` rejects anything outside that set at construction.""" - -_T = TypeVar("_T") -_ResultT = TypeVar("_ResultT") -_CacheableT = TypeVar("_CacheableT", bound=CacheableResult) - -_Connector = Callable[[AsyncExitStack, ConnectMode, bool], Awaitable["Dispatcher[Any]"]] -"""Resolved at ``__post_init__`` from the shape of ``server`` alone: enter whatever resources -are needed onto the exit stack and hand back the ``Dispatcher`` ``ClientSession`` will drive. -``mode`` and ``raise_exceptions`` are passed at call time so they're read at the same moment -``__aenter__`` reads them for the handshake step.""" - - -def _connect_transport(transport: Transport) -> _Connector: - """Connector for the stream-backed paths (URL, user-supplied ``Transport``).""" - - async def connect(exit_stack: AsyncExitStack, _mode: ConnectMode, _raise_exceptions: bool) -> Dispatcher[Any]: - read_stream, write_stream = await exit_stack.enter_async_context(transport) - return JSONRPCDispatcher(read_stream, write_stream) - - return connect - - -def _connect_inproc(server: Server[Any]) -> _Connector: - """Connector for an in-process ``Server``: legacy mode drives the stream loop via - ``InMemoryTransport``; any other mode drives the modern per-request path through a - ``DirectDispatcher`` peer pair (no streams, no JSON-RPC framing, no initialize handshake).""" - - async def connect(exit_stack: AsyncExitStack, mode: ConnectMode, raise_exceptions: bool) -> Dispatcher[Any]: - if mode == "legacy": - transport = InMemoryTransport(server, raise_exceptions=raise_exceptions) - read_stream, write_stream = await exit_stack.enter_async_context(transport) - return JSONRPCDispatcher(read_stream, write_stream) - lifespan_state = await exit_stack.enter_async_context(server.lifespan(server)) - client_disp, server_disp = create_direct_dispatcher_pair(raise_handler_exceptions=raise_exceptions) - tg = await exit_stack.enter_async_context(anyio.create_task_group()) - exit_stack.callback(server_disp.close) - on_request = modern_on_request(server, lifespan_state) - await tg.start(server_disp.run, on_request, _no_inbound_client_notifications) - return client_disp - - return connect - - -def _connected(value: _T | None) -> _T: - """Narrow a post-handshake session attribute from ``T | None`` to ``T``. - - ``Client.__aenter__`` only assigns ``_session`` after the handshake succeeds, so inside - ``async with Client(...)`` these attributes are always populated; the ``.session`` gate - raises before this is reached otherwise. The guard exists for pyright, not runtime. - """ - if value is None: # pragma: no cover - raise RuntimeError("Client must be used within an async context manager") - return value - - -def _strip_userinfo(url: str) -> str: - """Drop any userinfo from the URL's authority component; byte-exact otherwise. - - Credentials must not enter cache-key material; any further normalization could merge distinct servers. - """ - # Pure text, no urlsplit: it strips embedded tab/CR/LF before parsing, which would misalign slices. - sep = url.find("//") - if sep == -1: - return url - start = sep + 2 - end = len(url) - for delimiter in "/?#": - if (found := url.find(delimiter, start)) != -1: - end = min(end, found) - authority = url[start:end] - if "@" not in authority: - return url - return url[:start] + authority.rpartition("@")[2] + url[end:] - - -def _evicting_message_handler(cache: ClientResponseCache, user_handler: MessageHandlerFnT | None) -> MessageHandlerFnT: - """Wrap the session message handler with cache eviction on server notifications.""" - - async def handler(message: IncomingMessage) -> None: - if isinstance(message, types.ServerNotification): - try: - await cache.evict_for_notification(message) - except Exception: # boundary: eviction reaches user store code; a cache fault must not block delivery - logger.exception("Response cache eviction failed; the notification is still delivered") - if user_handler is not None: - await user_handler(message) - else: - # Mirrors ClientSession's default handler (session._default_message_handler). - await anyio.lowlevel.checkpoint() - - return handler - - -def _synthesize_discover(protocol_version: str) -> types.DiscoverResult: - return types.DiscoverResult( - supported_versions=[protocol_version], - capabilities=types.ServerCapabilities(), - result_type="complete", - ttl_ms=0, - cache_scope="public", - ) - - -async def _no_inbound_client_notifications(_dctx: Any, _method: str, _params: Mapping[str, Any] | None) -> None: - """Server-side inbound ``OnNotify`` for the modern in-process path — receives nothing. - - At 2026-07-28 the spec defines no client→server notifications: ``initialized`` and - ``roots/list_changed`` are removed, and cancellation is structural (anyio scope cancel - through the direct await, not a notify). Server→client notifications (progress, log - messages) flow the other way via the per-request ``DispatchContext`` into the client's - callbacks, and are not seen here. - """ - - -@dataclass(frozen=True) -class _FoldedExtensions: - """`Client.extensions` instances folded into the shapes `ClientSession` consumes.""" - - ad: dict[str, dict[str, Any]] | None - claims: dict[str, tuple[ResultClaim[Any], ...]] | None - bindings: tuple[NotificationBinding[Any], ...] | None - by_model: Mapping[type[Result], ResultClaim[Any]] - - -def _fold_extensions(extensions: Sequence[ClientExtension] | None) -> _FoldedExtensions: - """Fold extension contributions at construction, naming both owners on duplicate tags or methods.""" - if isinstance(extensions, Mapping): - raise TypeError( - "extensions= takes a sequence of ClientExtension instances. The mapping form was " - "replaced: use advertise(identifier, settings) for advertise-only entries" - ) - if not extensions: - return _FoldedExtensions(ad=None, claims=None, bindings=None, by_model={}) - ad: dict[str, dict[str, Any]] = {} - claims: dict[str, tuple[ResultClaim[Any], ...]] = {} - bindings: list[NotificationBinding[Any]] = [] - by_model: dict[type[Result], ResultClaim[Any]] = {} - claim_owners: dict[str, str] = {} - binding_owners: dict[str, str] = {} - for extension in extensions: - identifier = getattr(extension, "identifier", None) - if identifier is None: - raise ValueError( - f"{type(extension).__name__} has no `identifier`; a ClientExtension must set the " - "`identifier` class attribute (or assign one in `__init__`) before it can be used" - ) - validate_extension_identifier(identifier, owner=type(extension).__name__) - if identifier in ad: - raise ValueError(f"extension identifier {identifier!r} is passed more than once") - ad[identifier] = extension.settings() - extension_claims = tuple(extension.claims()) - for claim in extension_claims: - tag = claim.result_type - if tag in claim_owners: - owner = claim_owners[tag] - both = ( - f"extension {identifier!r} claims" - if owner == identifier - else (f"extensions {owner!r} and {identifier!r} both claim") - ) - raise ValueError(f"{both} resultType {tag!r}; a wire tag can have only one resolver") - claim_owners[tag] = identifier - # Each model pins its result_type Literal to one tag, so this index cannot collide. - by_model[claim.model] = claim - if extension_claims: - claims[identifier] = extension_claims - for binding in extension.notifications(): - if binding.method in binding_owners: - owner = binding_owners[binding.method] - both = ( - f"extension {identifier!r} binds" - if owner == identifier - else (f"extensions {owner!r} and {identifier!r} both bind") - ) - raise ValueError(f"{both} notification method {binding.method!r}; a method can have only one observer") - binding_owners[binding.method] = identifier - bindings.append(binding) - return _FoldedExtensions(ad=ad, claims=claims or None, bindings=tuple(bindings) or None, by_model=by_model) - - -@dataclass -class Client: - """A high-level MCP client for connecting to MCP servers. - - Pass a URL string (Streamable HTTP), a `StdioServerParameters` (launch the command as a - subprocess and talk over its stdin/stdout), any `Transport`, or - in tests - a `Server` or - `MCPServer` instance to connect to it in-process. - - Example: - ```python - import asyncio - - from mcp import Client - - async def main(): - async with Client("http://localhost:8000/mcp") as client: - result = await client.call_tool("add", {"a": 1, "b": 2}) - - asyncio.run(main()) - ``` - """ - - server: Server[Any] | MCPServer | Transport | StdioServerParameters | str - """The MCP server to connect to. - - If the server is a URL string, it will be used as the URL for a `streamable_http_client` transport. - If the server is a `StdioServerParameters`, the command is launched with `stdio_client`. - If the server is a `Transport` instance, it will be used directly. - If the server is a `Server` or `MCPServer` instance, it will be connected in-process. - """ - - _: KW_ONLY - - # TODO(Marcelo): When do `raise_exceptions=True` actually raises? - raise_exceptions: bool = False - """Whether to raise exceptions from the server.""" - - read_timeout_seconds: float | None = None - """Timeout for read operations.""" - - sampling_callback: SamplingFnT | None = None - """Callback for handling sampling requests.""" - - sampling_capabilities: types.SamplingCapability | None = None - """Sampling sub-capabilities (e.g. tools) declared alongside `sampling_callback`; no effect without it.""" - - list_roots_callback: ListRootsFnT | None = None - """Callback for handling list roots requests.""" - - logging_callback: LoggingFnT | None = None - """Callback for handling logging notifications.""" - - log_level: LoggingLevel | None = None - """The log level to opt in to on 2026-07-28+ connections (deprecated logging feature, SEP-2577). - - Modern (2026-07-28+) servers send `notifications/message` only for requests that opt in by - carrying `io.modelcontextprotocol/logLevel` in `_meta`, and only at or above that level. Setting - this stamps that opt-in on every request; `None` (the default) means no opt-in, so no log - messages arrive - a `logging_callback` alone is not an opt-in. No effect on handshake-era - connections, where the deprecated `logging/setLevel` request governs delivery instead. A - per-request `_meta` entry with the same key overrides this default.""" - - # TODO(Marcelo): Why do we have both "callback" and "handler"? - message_handler: MessageHandlerFnT | None = None - """Callback for handling raw messages.""" - - client_info: Implementation | None = None - """Client implementation info to send to server.""" - - mode: ConnectMode = "auto" - """How to negotiate the protocol version. - - 'auto' (the default) probes `server/discover` and falls back to the initialize handshake on legacy servers; - for an in-process `Server`/`MCPServer` it dispatches directly without JSON-RPC framing. 'legacy' forces the - initialize handshake (byte-identical pre-2026 behavior). A modern protocol-version string (e.g. '2026-07-28') - adopts that version directly without a probe — supply `prior_discover` to reuse a known DiscoverResult, or - omit it to synthesize a minimal one.""" - - prior_discover: types.DiscoverResult | None = None - """A previously-obtained DiscoverResult to install via .adopt() when mode is a version pin. - Ignored when mode='legacy'.""" - - elicitation_callback: ElicitationFnT | None = None - """Callback for handling elicitation requests.""" - - input_required_max_rounds: int = DEFAULT_INPUT_REQUIRED_MAX_ROUNDS - """Cap on `InputRequiredResult` retry rounds before `call_tool` / `get_prompt` / - `read_resource` give up. Use `client.session.(..., allow_input_required=True)` - to drive the loop manually instead.""" - - extensions: Sequence[ClientExtension] | None = None - """Opt-in client extensions (SEP-2133). - - Each instance contributes its capability ad, its result claims (resolved - transparently by `call_tool`), and its notification bindings. For an - ad-only entry use `mcp.client.advertise(identifier, settings)`.""" - - cache: CacheConfig | None = field(default_factory=CacheConfig) - """Client-side response caching for the SEP-2549 cacheable methods (2026-07-28). - - The default `CacheConfig()` honors server `ttlMs`/`cacheScope` hints with a - per-client in-memory store; pass a customized `CacheConfig`, or `None` to - disable. The cacheable verbs take a per-call `cache_mode` (see `CacheMode`); - calls carrying `meta` always reach the server. A `CacheConfig` with a custom - `store` requires `target_id` when the server is not a URL (no identity can be - derived).""" - - _entered: bool = field(init=False, default=False) - _session: ClientSession | None = field(init=False, default=None) - _exit_stack: AsyncExitStack | None = field(init=False, default=None) - _connect: _Connector = field(init=False, repr=False, compare=False) - _response_cache: ClientResponseCache | None = field(init=False, default=None, repr=False, compare=False) - _folded_extensions: _FoldedExtensions = field(init=False, repr=False, compare=False) - - def __post_init__(self) -> None: - if self.mode not in ("legacy", "auto") and self.mode not in MODERN_PROTOCOL_VERSIONS: - hint = ( - f" ({self.mode!r} is a handshake-era version; use mode='legacy')" - if self.mode in HANDSHAKE_PROTOCOL_VERSIONS - else "" - ) - raise ValueError( - f"mode must be 'legacy', 'auto', or one of {list(MODERN_PROTOCOL_VERSIONS)}; got {self.mode!r}{hint}" - ) - - self._folded_extensions = _fold_extensions(self.extensions) - - srv = self.server - if isinstance(srv, MCPServer): - srv = srv._lowlevel_server # pyright: ignore[reportPrivateUsage] - if isinstance(srv, Server): - self._connect = _connect_inproc(srv) - elif isinstance(srv, str): - self._connect = _connect_transport(streamable_http_client(srv)) - elif isinstance(srv, StdioServerParameters): - self._connect = _connect_transport(stdio_client(srv)) - else: - self._connect = _connect_transport(srv) - - if self.cache is not None: - config = self.cache - # Only the hash below leaves this scope - the raw identity may carry credentials; never log or store it. - target_id = config.target_id - if target_id is None and isinstance(self.server, str): - target_id = _strip_userinfo(self.server) - if target_id is None: - if config.store is not None: - raise ValueError( - "a custom cache store requires CacheConfig.target_id when the server is not a URL: " - "in-process servers and Transport instances get a random per-client identity, so " - "their entries in a shared store could never be served to another client" - ) - target_id = uuid.uuid4().hex - self._response_cache = ClientResponseCache( - store=config.store if config.store is not None else InMemoryResponseCacheStore(), - partition=config.partition, - arm_id=hashlib.sha256(target_id.encode()).hexdigest(), - default_ttl_ms=config.default_ttl_ms, - clock=config.clock, - share_public=config.share_public, - # Lazy: the negotiated version is unknown until __aenter__'s handshake. - negotiated_version=lambda: self._session.protocol_version if self._session is not None else None, - ) - - async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession: - """Enter the resolved connector and return an un-entered ClientSession.""" - dispatcher = await self._connect(exit_stack, self.mode, self.raise_exceptions) - message_handler = self.message_handler - if self._response_cache is not None: - message_handler = _evicting_message_handler(self._response_cache, self.message_handler) - return ClientSession( - dispatcher=dispatcher, - read_timeout_seconds=self.read_timeout_seconds, - sampling_callback=self.sampling_callback, - sampling_capabilities=self.sampling_capabilities, - list_roots_callback=self.list_roots_callback, - logging_callback=self.logging_callback, - log_level=self.log_level, - message_handler=message_handler, - client_info=self.client_info, - elicitation_callback=self.elicitation_callback, - extensions=self._folded_extensions.ad, - result_claims=self._folded_extensions.claims, - notification_bindings=self._folded_extensions.bindings, - ) - - async def __aenter__(self) -> Client: - """Enter the async context manager.""" - if self._entered: - raise RuntimeError("Client is already entered; cannot reenter") - self._entered = True - - async with AsyncExitStack() as exit_stack: - session = await self._build_session(exit_stack) - session = await exit_stack.enter_async_context(session) - - if self.mode == "legacy": - await session.initialize() - elif self.mode == "auto": - await negotiate_auto(session) - else: - session.adopt(self.prior_discover or _synthesize_discover(self.mode)) - - # Only publish the session after the handshake succeeds, so `_session is not None` - # implies the protocol_version/server_capabilities are populated (server_info - # stays optional: 2026-era servers may not identify themselves). If the - # handshake raised above, the local exit_stack unwinds the transport for us. - self._session = session - self._exit_stack = exit_stack.pop_all() - return self - - async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: - """Exit the async context manager.""" - if self._exit_stack: # pragma: no branch - await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb) - self._session = None - - @property - def session(self) -> ClientSession: - """Get the underlying ClientSession. - - This provides access to the full ClientSession API for advanced use cases. - - Raises: - RuntimeError: If accessed before entering the context manager. - """ - if self._session is None: - raise RuntimeError("Client must be used within an async context manager") - return self._session - - # TODO(maxisbey): the by-construction shape is for __aenter__ to return a connected-view - # type whose protocol_version/server_capabilities are non-Optional fields, - # eliminating these guards (and the one in .session). Same family as resolving the - # transport/connector at __post_init__ so the Optional internal fields disappear. - # (server_info stays Optional even connected: the 2026-era stamp is optional.) - @property - def protocol_version(self) -> str: - """Negotiated protocol version (set by initialize/discover/adopt during ``__aenter__``).""" - return _connected(self.session.protocol_version) - - @property - def server_info(self) -> Implementation | None: - """Server name/version, or `None` when the server did not identify itself. - - Legacy connections always carry it (`InitializeResult.serverInfo` is - required); on 2026-era connections the `_meta` `serverInfo` stamp is - optional, so an anonymous server reads as `None`. - """ - return self.session.server_info - - @property - def server_capabilities(self) -> ServerCapabilities: - """Server capabilities (set by initialize/discover/adopt during ``__aenter__``).""" - return _connected(self.session.server_capabilities) - - @property - def instructions(self) -> str | None: - """Server-provided instructions text, if any.""" - return self.session.instructions - - @deprecated( - "ping is removed as of 2026-07-28; the method only works under mode='legacy'.", - category=MCPDeprecationWarning, - ) - async def send_ping(self, *, meta: RequestParamsMeta | None = None) -> EmptyResult: - """Send a ping request to the server.""" - return await self.session.send_ping(meta=meta) - - @deprecated( - "Client-to-server progress is deprecated as of 2026-07-28; progress is server-to-client only.", - category=MCPDeprecationWarning, - ) - async def send_progress_notification( - self, - progress_token: str | int, - progress: float, - total: float | None = None, - message: str | None = None, - ) -> None: - """Send a progress notification to the server.""" - await self.session.send_progress_notification( # pyright: ignore[reportDeprecated] - progress_token=progress_token, - progress=progress, - total=total, - message=message, - ) - - @deprecated("The logging capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def set_logging_level(self, level: LoggingLevel, *, meta: RequestParamsMeta | None = None) -> EmptyResult: - """Set the logging level on the server.""" - return await self.session.set_logging_level(level=level, meta=meta) # pyright: ignore[reportDeprecated] - - async def _cached_fetch( - self, - method: str, - *, - cursor: str | None, - meta: RequestParamsMeta | None, - cache_mode: CacheMode, - send: Callable[[], Awaitable[_CacheableT]], - absorb: Callable[[_CacheableT], _CacheableT] | None = None, - ) -> _CacheableT: - """Serve one of the four list verbs through the response cache. - - `absorb` (tools/list only) re-applies session-side derived state to a served cache hit. - """ - cache = self._response_cache - if cache is None or cache_mode == "bypass": - return await send() - # A closed (or never-entered) client must raise, never serve cached entries. - _ = self.session - if meta is not None and cache_mode == "use": - # meta (a progress token, tracing fields) expects a wire request; fetch and replace the entry. - cache_mode = "refresh" - if cursor is not None: - # Continuation pages skip the cache, but an expired cursor means the listing changed (spec SHOULD evict). - try: - return await send() - except MCPError as e: - if e.code == INVALID_PARAMS: - await cache.evict_method(method) - raise - if cache_mode == "use" and (hit := await cache.read(method, "")) is not None: - # The hit is a private deep copy, so absorption may mutate it freely. - served = cast(_CacheableT, hit) - return served if absorb is None else absorb(served) - gen = cache.capture(method, "") - result = await send() - await cache.write(method, "", result, gen, cache_mode) - return result - - async def list_resources( - self, - *, - cursor: str | None = None, - meta: RequestParamsMeta | None = None, - cache_mode: CacheMode = "use", - ) -> ListResourcesResult: - """List available resources from the server.""" - return await self._cached_fetch( - "resources/list", - cursor=cursor, - meta=meta, - cache_mode=cache_mode, - send=lambda: self.session.list_resources(params=PaginatedRequestParams(cursor=cursor, _meta=meta)), - ) - - async def list_resource_templates( - self, - *, - cursor: str | None = None, - meta: RequestParamsMeta | None = None, - cache_mode: CacheMode = "use", - ) -> ListResourceTemplatesResult: - """List available resource templates from the server.""" - return await self._cached_fetch( - "resources/templates/list", - cursor=cursor, - meta=meta, - cache_mode=cache_mode, - send=lambda: self.session.list_resource_templates(params=PaginatedRequestParams(cursor=cursor, _meta=meta)), - ) - - async def read_resource( - self, - uri: str, - *, - input_responses: InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - cache_mode: CacheMode = "use", - ) -> ReadResourceResult: - """Read a resource from the server. - - If the server returns an `InputRequiredResult`, the embedded input - requests are dispatched to this client's sampling / elicitation / roots - callbacks and the read is retried automatically (up to - `input_required_max_rounds`). - - Args: - uri: The URI of the resource to read. - input_responses: Responses to seed the first call with (e.g. when - resuming from a persisted `InputRequiredResult`). - request_state: Opaque state to seed the first call with. - meta: Additional metadata for the request. - cache_mode: Cache behavior for this call (see `CacheMode`); seeded - calls (`input_responses` or `request_state` set) ignore it. - - Returns: - The resource content. - - Raises: - InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted. - MCPError: A callback returned `ErrorData` for an embedded input request. - pydantic.ValidationError: The server returned a result that does not - conform to the negotiated protocol version. - """ - - async def retry(r: InputResponses | None, s: str | None) -> ReadResourceResult | InputRequiredResult: - return await self.session.read_resource( - uri, input_responses=r, request_state=s, meta=meta, allow_input_required=True - ) - - # Seeded calls resume a specific exchange and must never be cached (spec MUST). - seeded = input_responses is not None or request_state is not None - cache = None if seeded else self._response_cache - if cache is None or cache_mode == "bypass": - return await self._drive_input_required(await retry(input_responses, request_state), retry) - # A closed (or never-entered) client must raise, never serve cached entries. - _ = self.session - if meta is not None and cache_mode == "use": - # Calls carrying meta always reach the server (mirrors `_cached_fetch`). - cache_mode = "refresh" - if cache_mode == "use" and (hit := await cache.read("resources/read", uri)) is not None: - # Only terminal first-round results are stored, so a hit legitimately skips the driver. - return cast(ReadResourceResult, hit) - gen = cache.capture("resources/read", uri) - first = await retry(None, None) - if not isinstance(first, InputRequiredResult): - await cache.write("resources/read", uri, first, gen, cache_mode) - elif cache_mode == "refresh": - # The refresh superseded whatever was cached, but an input_required resolution - # cannot be stored: purge the warm entry so it cannot be served again. - await cache.evict_key("resources/read", uri) - # Driver rounds carry inputResponses, so a terminal result reached through them is never cached (spec MUST). - return await self._drive_input_required(first, retry) - - def listen( - self, - *, - tools_list_changed: bool = False, - prompts_list_changed: bool = False, - resources_list_changed: bool = False, - resource_subscriptions: Sequence[str] = (), - ) -> AbstractAsyncContextManager[Subscription]: - """Open a `subscriptions/listen` stream of typed change events (2026-07-28 only). - - Keyword args mirror the wire `SubscriptionFilter`; entering waits for the ack (honored subset: `sub.honored`): - - async with client.listen(tools_list_changed=True) as sub: - async for event in sub: - tools = await client.list_tools() # refetch on change - - A graceful close ends the loop; an abrupt drop raises `SubscriptionLost`. No replay: re-listen and refetch. - - Raises: - ListenNotSupportedError: The negotiated protocol version predates 2026-07-28. - MCPError: The server rejected the request or the connection failed first. - SubscriptionLost: The stream ended before it was acknowledged. - TimeoutError: The read timeout elapsed before the acknowledgment. - """ - return _listen( - self.session, - tools_list_changed=tools_list_changed, - prompts_list_changed=prompts_list_changed, - resources_list_changed=resources_list_changed, - resource_subscriptions=resource_subscriptions, - on_event=self._evict_for_listen_event if self._response_cache is not None else None, - ) - - async def _evict_for_listen_event(self, event: ServerEvent) -> None: - """Finish response-cache eviction before a listen consumer can refetch. - - Without it the iterator wakes first and refetches a still-warm entry, with no - corrective wake (events are deduplicated level triggers). The tee path repeats - the eviction; deliberate: idempotent, and it covers non-iterating consumers. - """ - cache = self._response_cache - assert cache is not None # installed as the event barrier only when a cache exists - try: - await cache.evict_for_notification(event_to_notification(event, {})) - except Exception: # boundary: eviction reaches user store code; a cache fault must not block delivery - logger.exception("Response cache eviction failed; the event is still delivered") - - @deprecated( - "resources/subscribe is removed as of 2026-07-28; use Client.listen() instead.", - category=MCPDeprecationWarning, - ) - async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> EmptyResult: - """Subscribe to resource updates (2025-era servers only).""" - return await self.session.subscribe_resource(uri, meta=meta) # pyright: ignore[reportDeprecated] - - @deprecated( - "resources/unsubscribe is removed as of 2026-07-28; use Client.listen() instead.", - category=MCPDeprecationWarning, - ) - async def unsubscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> EmptyResult: - """Unsubscribe from resource updates (2025-era servers only).""" - return await self.session.unsubscribe_resource(uri, meta=meta) # pyright: ignore[reportDeprecated] - - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - ) -> CallToolResult: - """Call a tool on the server. - - If the server returns an `InputRequiredResult`, the embedded input - requests are dispatched to this client's sampling / elicitation / roots - callbacks and the call is retried automatically (up to - `input_required_max_rounds`). To drive the loop yourself — e.g. to - persist `request_state` across process restarts — use - `client.session.call_tool(..., allow_input_required=True)`. Persisted - state is still subject to the server's TTL, request binding, and key - lifetime; a server on the default process-local key rejects it after a restart. - - Result shapes claimed by this client's `extensions` are finished by the - owning claim's resolver, whose `CallToolResult` is returned; resolver - exceptions propagate as-is. To receive the claimed shape yourself, use - `client.session.call_tool(..., allow_claimed=True)`. - - Args: - name: The name of the tool to call. - arguments: Arguments to pass to the tool. - read_timeout_seconds: Timeout for each underlying `tools/call` round. - progress_callback: Callback for progress updates. - input_responses: Responses to seed the first call with (e.g. when - resuming from a persisted `InputRequiredResult`). - request_state: Opaque state to seed the first call with. - meta: Additional metadata for the request. - - Returns: - The tool result. - - Raises: - InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted. - MCPError: A callback returned `ErrorData` for an embedded input request. - pydantic.ValidationError: The server returned a result that does not - conform to the negotiated protocol version. - """ - - async def retry(r: InputResponses | None, s: str | None) -> CallToolResult | InputRequiredResult | Result: - return await self.session.call_tool( - name, - arguments, - read_timeout_seconds=read_timeout_seconds, - progress_callback=progress_callback, - input_responses=r, - request_state=s, - meta=meta, - allow_input_required=True, - # Input rounds resolve before a claimed result, so a claim may end any round. - allow_claimed=True, - ) - - result = await self._drive_input_required(await retry(input_responses, request_state), retry) - if isinstance(result, CallToolResult): - return result - # Only claimed shapes reach this point, so the lookup is total. - claim = self._folded_extensions.by_model[type(result)] - final = await claim.resolve( - result, - ClaimContext(session=self.session, tool_name=name, read_timeout_seconds=read_timeout_seconds), - ) - if not final.is_error: - # Match the direct path: revalidate the output schema, but never for isError results. - await self.session.validate_tool_result(name, final) - return final - - async def list_prompts( - self, - *, - cursor: str | None = None, - meta: RequestParamsMeta | None = None, - cache_mode: CacheMode = "use", - ) -> ListPromptsResult: - """List available prompts from the server.""" - return await self._cached_fetch( - "prompts/list", - cursor=cursor, - meta=meta, - cache_mode=cache_mode, - send=lambda: self.session.list_prompts(params=PaginatedRequestParams(cursor=cursor, _meta=meta)), - ) - - async def get_prompt( - self, - name: str, - arguments: dict[str, str] | None = None, - *, - input_responses: InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - ) -> GetPromptResult: - """Get a prompt from the server. - - If the server returns an `InputRequiredResult`, the embedded input - requests are dispatched to this client's sampling / elicitation / roots - callbacks and the get is retried automatically (up to - `input_required_max_rounds`). - - Args: - name: The name of the prompt. - arguments: Arguments to pass to the prompt. - input_responses: Responses to seed the first call with (e.g. when - resuming from a persisted `InputRequiredResult`). - request_state: Opaque state to seed the first call with. - meta: Additional metadata for the request. - - Returns: - The prompt content. - - Raises: - InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted. - MCPError: A callback returned `ErrorData` for an embedded input request. - pydantic.ValidationError: The server returned a result that does not - conform to the negotiated protocol version. - """ - - async def retry(r: InputResponses | None, s: str | None) -> GetPromptResult | InputRequiredResult: - return await self.session.get_prompt( - name, arguments, input_responses=r, request_state=s, meta=meta, allow_input_required=True - ) - - return await self._drive_input_required(await retry(input_responses, request_state), retry) - - async def _drive_input_required( - self, - first: _ResultT | InputRequiredResult, - retry: Callable[[InputResponses | None, str | None], Awaitable[_ResultT | InputRequiredResult]], - ) -> _ResultT: - """Hand an `InputRequiredResult` to the SEP-2322 driver, or pass a terminal result through. - - `dispatch` routes each embedded request through the same callback table - that serves legacy server→client RPCs, so the two paths stay - behaviourally identical by construction. - """ - if not isinstance(first, InputRequiredResult): - return first - session = self.session - - async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData: - ctx = ClientRequestContext(session=session, request_id=key, meta=req.params.meta if req.params else None) - return await session.dispatch_input_request(ctx, req) - - return await run_input_required_driver( - first, dispatch=dispatch, retry=retry, max_rounds=self.input_required_max_rounds - ) - - async def complete( - self, - ref: ResourceTemplateReference | PromptReference, - argument: dict[str, str], - context_arguments: dict[str, str] | None = None, - ) -> CompleteResult: - """Get completions for a prompt or resource template argument. - - Args: - ref: Reference to the prompt or resource template - argument: The argument to complete - context_arguments: Additional context arguments - - Returns: - Completion suggestions. - """ - return await self.session.complete(ref=ref, argument=argument, context_arguments=context_arguments) - - async def list_tools( - self, - *, - cursor: str | None = None, - meta: RequestParamsMeta | None = None, - cache_mode: CacheMode = "use", - ) -> ListToolsResult: - """List available tools from the server.""" - return await self._cached_fetch( - "tools/list", - cursor=cursor, - meta=meta, - cache_mode=cache_mode, - send=lambda: self.session.list_tools(params=PaginatedRequestParams(cursor=cursor, _meta=meta)), - # A cache hit skips session.list_tools, so the session re-absorbs the served - # listing to rebuild its derived per-tool state. Hits are cursorless, but a - # cached page 1 can carry next_cursor - never prune on a partial listing. - absorb=lambda hit: self.session._absorb_tool_listing( # pyright: ignore[reportPrivateUsage] - hit, complete=hit.next_cursor is None - ), - ) - @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def send_roots_list_changed(self) -> None: - """Send a notification that the roots list has changed.""" - await self.session.send_roots_list_changed() # pyright: ignore[reportDeprecated] +# Keep runtime annotations resolvable through the full SDK's existing import path. +_implementation._InProcessServer = Server[Any] | MCPServer +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/context.py b/src/mcp/client/context.py index aecd29527f..413e0c64b4 100644 --- a/src/mcp/client/context.py +++ b/src/mcp/client/context.py @@ -1,5 +1,8 @@ -"""Request context for MCP client handlers.""" +import sys -from mcp.client.session import ClientRequestContext +import mcp_client.client.context as _implementation +from mcp_client.client.context import ( + ClientRequestContext as ClientRequestContext, +) -__all__ = ["ClientRequestContext"] +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/extension.py b/src/mcp/client/extension.py index a813475e5e..e71aa4e825 100644 --- a/src/mcp/client/extension.py +++ b/src/mcp/client/extension.py @@ -1,196 +1,41 @@ -"""Opt-in extension interface for MCP clients. - -Subclass `ClientExtension`, set `identifier`, override the hooks you need, and -pass instances to `Client(extensions=[...])`. For an identifier-only -capability ad, use `advertise()`. -""" - -from __future__ import annotations - -from collections.abc import Awaitable, Callable, Sequence -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeVar, get_args - -from mcp_types import CORE_RESULT_TYPES, CallToolResult, InputRequiredResult, Result -from mcp_types.version import MODERN_PROTOCOL_VERSIONS -from pydantic import AliasChoices, AliasPath, BaseModel -from pydantic.fields import FieldInfo - -from mcp.shared.extension import validate_extension_identifier - -if TYPE_CHECKING: - from mcp.client.session import ClientSession - -__all__ = [ - "ClaimContext", - "ClientExtension", - "NotificationBinding", - "ResultClaim", - "UnexpectedClaimedResult", - "advertise", -] - -_CLAIM_METHODS: Final[frozenset[str]] = frozenset({"tools/call"}) -"""The closed set of verbs a claim may attach to; widen together with the `method` Literal.""" - -_RESERVED_WIRE_ALIASES: Final[frozenset[str]] = frozenset({"requestState", "inputRequests"}) -"""Typed optional fields of the core result surface that pre-validates every inbound result.""" - - -def _wire_keys(name: str, field: FieldInfo) -> frozenset[str]: - """Every top-level wire key this field can read from or write to.""" - keys = {field.alias or name} - if field.serialization_alias: - keys.add(field.serialization_alias) - validation_alias = field.validation_alias - choices = validation_alias.choices if isinstance(validation_alias, AliasChoices) else [validation_alias] - for choice in choices: - if isinstance(choice, AliasPath): - choice = choice.path[0] - if isinstance(choice, str): - keys.add(choice) - return frozenset(keys) - - -ClaimedT = TypeVar("ClaimedT", bound=Result) -NotifyParamsT = TypeVar("NotifyParamsT", bound=BaseModel) - - -@dataclass(frozen=True, kw_only=True) -class ClaimContext: - """Host-injected context for one `ResultClaim.resolve` call.""" - - session: ClientSession - tool_name: str - read_timeout_seconds: float | None - - -@dataclass(frozen=True, kw_only=True) -class ResultClaim(Generic[ClaimedT]): - """One extra result shape on one spec verb, keyed by the wire `resultType`. - - Active only while the declaring extension is constructed into the client and - the negotiated protocol version admits it. `resolve` finishes a claimed - result, may send follow-ups through `ctx.session`, and must return the - verb's ordinary result. All field constraints are enforced at construction. - """ - - result_type: str - model: type[ClaimedT] - resolve: Callable[[ClaimedT, ClaimContext], Awaitable[CallToolResult]] - method: Literal["tools/call"] = "tools/call" - protocol_versions: frozenset[str] | None = None - - def __post_init__(self) -> None: - if self.method not in _CLAIM_METHODS: - raise ValueError(f"claims attach to {sorted(_CLAIM_METHODS)} only; got method {self.method!r}") - if self.result_type in CORE_RESULT_TYPES: - raise ValueError(f"resultType {self.result_type!r} is core protocol vocabulary") - if Result not in self.model.__mro__: # runtime guard; the ClaimedT bound only constrains checked callers - raise ValueError(f"{self.model.__name__} must subclass mcp_types.Result") - if issubclass(self.model, CallToolResult | InputRequiredResult): - raise ValueError("claim models must not subclass core result types") - for name, model_field in self.model.model_fields.items(): - for clash in sorted(_wire_keys(name, model_field) & _RESERVED_WIRE_ALIASES): - raise ValueError( - f"{self.model.__name__}.{name} aliases {clash!r}, a typed field of the core " - "result surface; a colliding value would fail core validation before the " - "claim adapter runs" - ) - field = self.model.model_fields.get("result_type") - if field is None or get_args(field.annotation) != (self.result_type,): - raise ValueError(f"{self.model.__name__}.result_type must be Literal[{self.result_type!r}]") - if self.protocol_versions is not None and not self.protocol_versions: - raise ValueError("empty protocol_versions could never activate; use None for all") - if self.protocol_versions is not None and not self.protocol_versions.issubset(MODERN_PROTOCOL_VERSIONS): - unrecognized = sorted(self.protocol_versions.difference(MODERN_PROTOCOL_VERSIONS)) - raise ValueError( - f"protocol_versions {unrecognized} are not modern protocol revisions; claimed shapes " - "cannot be delivered on a legacy wire (None means every modern version)" - ) - - -class UnexpectedClaimedResult(RuntimeError): - """A claimed (extension) result arrived on a `call_tool` that did not opt in. - - The parsed value is carried as `result`; the server may already hold state it - references. Opt in via `Client(extensions=[...])` or `allow_claimed=True`. - """ - - def __init__(self, result: Result) -> None: - super().__init__( - f"Server returned a claimed result ({type(result).__name__}); pass the owning extension to " - "Client(extensions=[...]) for transparent resolution, or call with allow_claimed=True " - "and handle the shape. The carried result may reference server-side state needing cleanup." - ) - self.result = result - - -@dataclass(frozen=True, kw_only=True) -class NotificationBinding(Generic[NotifyParamsT]): - """Deliver server notifications for `method` (the bare wire name) to `handler`. - - Observation-only: validated params arrive one at a time per binding, in - dispatch order, through a bounded queue that drops the oldest with a warning - on overflow. Stream transports dispatch each notification independently, so - near-simultaneous notifications may be dispatched out of wire order. Methods - the negotiated version's core tables handle are never delivered to bindings. - """ - - method: str - params_type: type[NotifyParamsT] - handler: Callable[[NotifyParamsT], Awaitable[None]] - - -class ClientExtension: - """Base class for an opt-in client extension; override only what you need. - - The surface is declarative, fixed at construction, and never receives the client. - """ - - #: Reverse-DNS extension identifier, advertised under `ClientCapabilities.extensions`. - identifier: str - - def __init_subclass__(cls, **kwargs: Any) -> None: - super().__init_subclass__(**kwargs) - # Per-instance identifiers (assigned in __init__) are validated at consumption instead. - if (identifier := cls.__dict__.get("identifier")) is not None: - validate_extension_identifier(identifier, owner=cls.__name__) - - def settings(self) -> dict[str, Any]: - """Per-extension settings advertised at `ClientCapabilities.extensions[identifier]`. - - Read once at `Client` construction. A claim-bearing extension is - advertised only at protocol versions where at least one of its claims - is active. - """ - return {} - - def claims(self) -> Sequence[ResultClaim[Any]]: - """Extra result shapes this extension claims, with their resolvers.""" - return () - - def notifications(self) -> Sequence[NotificationBinding[Any]]: - """Server notifications this extension observes.""" - return () - - -class _AdvertiseOnly(ClientExtension): - """Ad-only extension returned by `advertise()`.""" - - def __init__(self, identifier: str, settings: dict[str, Any]) -> None: - self.identifier = identifier - self._settings = settings - - def settings(self) -> dict[str, Any]: - return self._settings - - -def advertise(identifier: str, settings: dict[str, Any] | None = None) -> ClientExtension: - """Advertise an extension identifier (with optional settings) and nothing else. - - Advertising an extension you do not implement asserts wire support you do - not have; for behavioral extensions construct the real extension instead. - """ - validate_extension_identifier(identifier, owner="advertise") - return _AdvertiseOnly(identifier, {} if settings is None else settings) +import sys + +import mcp_client.client.extension as _implementation +from mcp_client.client.extension import ( + _CLAIM_METHODS as _CLAIM_METHODS, +) +from mcp_client.client.extension import ( + _RESERVED_WIRE_ALIASES as _RESERVED_WIRE_ALIASES, +) +from mcp_client.client.extension import ( + ClaimContext as ClaimContext, +) +from mcp_client.client.extension import ( + ClaimedT as ClaimedT, +) +from mcp_client.client.extension import ( + ClientExtension as ClientExtension, +) +from mcp_client.client.extension import ( + NotificationBinding as NotificationBinding, +) +from mcp_client.client.extension import ( + NotifyParamsT as NotifyParamsT, +) +from mcp_client.client.extension import ( + ResultClaim as ResultClaim, +) +from mcp_client.client.extension import ( + UnexpectedClaimedResult as UnexpectedClaimedResult, +) +from mcp_client.client.extension import ( + _AdvertiseOnly as _AdvertiseOnly, +) +from mcp_client.client.extension import ( + _wire_keys as _wire_keys, +) +from mcp_client.client.extension import ( + advertise as advertise, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index a618112153..0ba4f1b464 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -1,1506 +1,113 @@ -from __future__ import annotations +import sys -import json -import logging -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass -from functools import cache, reduce -from operator import or_ -from types import TracebackType, UnionType -from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, cast, get_args, overload - -import anyio -import anyio.abc -import anyio.lowlevel -import mcp_types as types -from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream -from mcp_types import ( - CLIENT_CAPABILITIES_META_KEY, - CLIENT_INFO_META_KEY, - CONNECTION_CLOSED, - INTERNAL_ERROR, - LOG_LEVEL_META_KEY, - METHOD_NOT_FOUND, - PROTOCOL_VERSION_META_KEY, - SERVER_INFO_META_KEY, - UNSUPPORTED_PROTOCOL_VERSION, - RequestId, - RequestParamsMeta, +import mcp_client.client.session as _implementation +from mcp_client.client.session import ( + _NOTIFICATION_QUEUE_SIZE as _NOTIFICATION_QUEUE_SIZE, ) -from mcp_types import methods as _methods -from mcp_types.version import ( - HANDSHAKE_PROTOCOL_VERSIONS, - KNOWN_PROTOCOL_VERSIONS, - LATEST_HANDSHAKE_VERSION, - LATEST_MODERN_VERSION, - MODERN_PROTOCOL_VERSIONS, +from mcp_client.client.session import ( + DEFAULT_CLIENT_INFO as DEFAULT_CLIENT_INFO, ) -from pydantic import BaseModel, Discriminator, Tag, TypeAdapter, ValidationError -from typing_extensions import Self, TypeVar, deprecated - -from mcp.client._transport import ReadStream, WriteStream -from mcp.client.extension import NotificationBinding, ResultClaim, UnexpectedClaimedResult -from mcp.client.subscriptions import ListenRoute -from mcp.shared._compat import resync_tracer -from mcp.shared.dispatcher import CallOptions, DispatchContext, Dispatcher, ProgressFnT, as_request_id -from mcp.shared.exceptions import MCPDeprecationWarning, MCPError -from mcp.shared.inbound import ( - MCP_METHOD_HEADER, - MCP_NAME_HEADER, - MCP_PROTOCOL_VERSION_HEADER, - NAME_BEARING_METHODS, - encode_header_value, - find_invalid_x_mcp_header, - mcp_param_headers, - x_mcp_header_map, +from mcp_client.client.session import ( + DISCOVER_TIMEOUT_SECONDS as DISCOVER_TIMEOUT_SECONDS, ) -from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher, cancelled_request_id_from_params -from mcp.shared.message import ClientMessageMetadata, SessionMessage -from mcp.shared.subscriptions import SUBSCRIPTION_ID_META_KEY, event_from_wire -from mcp.shared.transport_context import TransportContext - -if TYPE_CHECKING: - # `jsonschema` is imported lazily inside `validate_tool_result`: pulling it (and its - # `attrs`/`referencing` tree) in at module scope costs every client that never validates. - from jsonschema.protocols import Validator - -DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0") -DISCOVER_TIMEOUT_SECONDS = 10.0 -_NOTIFICATION_QUEUE_SIZE: Final = 256 - -logger = logging.getLogger("client") - - -def _clamp_inbound_ttl(raw: dict[str, Any]) -> None: - """Floor a negative inbound `ttlMs` to 0 before `ge=0` validation fails the call (2026-07-28 caching SHOULD).""" - ttl = raw.get("ttlMs") - if isinstance(ttl, int | float) and not isinstance(ttl, bool) and ttl < 0: - raw["ttlMs"] = 0 - - -@cache -def _wire_fields(target: type[BaseModel] | UnionType) -> frozenset[str]: - """Top-level wire keys `target` declares (its members', for a union).""" - members: tuple[Any, ...] = get_args(target) if isinstance(target, UnionType) else (target,) - models = [m for m in members if isinstance(m, type) and issubclass(m, BaseModel)] - fields: set[str] = set() - for model in models: - fields.update(field.alias or name for name, field in model.model_fields.items()) - return frozenset(fields) - - -@cache -def _later_revision_fields(method: str, version: str) -> frozenset[str]: - """Result keys a revision newer than `version` declares for `method` but `version` doesn't. - - The version-free result types carry every revision's fields, so such a key - (e.g. 2026-07-28 `ttlMs`/`cacheScope` on a pre-2026 session) is outside the - negotiated contract yet would still parse into the model and trip that later - revision's constraints. Empty at the newest known revision. - """ - current = _methods.SERVER_RESULTS.get((method, version)) - if current is None or version not in KNOWN_PROTOCOL_VERSIONS: - return frozenset() - newer = KNOWN_PROTOCOL_VERSIONS[KNOWN_PROTOCOL_VERSIONS.index(version) + 1 :] - later: set[str] = set() - for revision in newer: - row = _methods.SERVER_RESULTS.get((method, revision)) - if row is not None: - later |= _wire_fields(row) - return frozenset(later) - _wire_fields(current) - - -def _same_schema(a: dict[str, Any] | None, b: dict[str, Any] | None) -> bool: - """JSON equality for two output schemas. - - Python `==` is not JSON equality: it conflates `True`/`1` and `False`/`0`, which JSON - Schema keeps distinct (`const: true` vs `const: 1`). Canonical serialization compares as - JSON does; where it is stricter (`1` vs `1.0`), erring toward "changed" only costs a - recompile, never a stale validator. - """ - return json.dumps(a, sort_keys=True) == json.dumps(b, sort_keys=True) - - -def _preconnect_stamp(data: dict[str, Any], opts: CallOptions) -> None: - # initialize/discover forbid cancellation; other pre-handshake requests (lowlevel - # ClientSession callers may skip the handshake entirely) keep the courtesy cancel. - if data["method"] in ("initialize", "server/discover"): - opts["cancel_on_abandon"] = False - - -def _parse_server_info_stamp(result: types.DiscoverResult) -> types.Implementation | None: - """The typed identity from a discover result's `_meta` serverInfo stamp. - - The stamp is display-only per the spec, so absent and malformed both read - as `None` rather than failing the connection. - """ - raw = (result.meta or {}).get(SERVER_INFO_META_KEY) - if raw is None: - return None - try: - return types.Implementation.model_validate(raw) - except ValidationError: - return None - - -def _make_handshake_stamp(protocol_version: str) -> Callable[[dict[str, Any], CallOptions], None]: - def stamp(data: dict[str, Any], opts: CallOptions) -> None: - opts.setdefault("headers", {})[MCP_PROTOCOL_VERSION_HEADER] = protocol_version - - return stamp - - -def _make_modern_stamp( - protocol_version: str, - client_info: dict[str, Any], - capabilities: dict[str, Any], - resolve_param_headers: Callable[[str, Mapping[str, Any]], dict[str, str]], - *, - log_level: types.LoggingLevel | None = None, -) -> Callable[[dict[str, Any], CallOptions], None]: - def stamp(data: dict[str, Any], opts: CallOptions) -> None: - params = data.setdefault("params", {}) - meta = params.setdefault("_meta", {}) - meta[PROTOCOL_VERSION_META_KEY] = protocol_version - meta[CLIENT_INFO_META_KEY] = client_info - meta[CLIENT_CAPABILITIES_META_KEY] = capabilities - # The per-request log-delivery opt-in (2026 logging is opt-in per - # request). A default the caller can override on any single call by - # supplying the key in that request's `_meta`, hence setdefault. - if log_level is not None: - meta.setdefault(LOG_LEVEL_META_KEY, log_level) - # `cancel_on_abandon` stays at the dispatcher default (True): the - # courtesy `notifications/cancelled` is the abandon signal. On the - # stream transports it is the 2026 wire's cancellation spelling; the - # streamable-HTTP transport translates it into aborting the request's - # own POST instead of writing it (the 2026 HTTP wire has no - # client-to-server notifications - closing the stream is the signal). - # The negotiation methods still opt out, mirroring `_preconnect_stamp`: - # the spec forbids cancelling them. - if data["method"] in ("initialize", "server/discover"): - opts["cancel_on_abandon"] = False - headers = opts.setdefault("headers", {}) - headers[MCP_PROTOCOL_VERSION_HEADER] = protocol_version - headers[MCP_METHOD_HEADER] = data["method"] - name_key = NAME_BEARING_METHODS.get(data["method"]) - if name_key is not None and isinstance(name := params.get(name_key), str): - headers[MCP_NAME_HEADER] = encode_header_value(name) - if data["method"] == "tools/call" and isinstance(name := params.get("name"), str): - headers.update(resolve_param_headers(name, params.get("arguments") or {})) - - return stamp - - -ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel) - - -@dataclass(kw_only=True) -class ClientRequestContext: - """Context for a server-initiated request, passed to the sampling/elicitation/list-roots callbacks.""" - - session: ClientSession - request_id: RequestId - meta: RequestParamsMeta | None = None - - -class SamplingFnT(Protocol): - async def __call__( - self, - context: ClientRequestContext, - params: types.CreateMessageRequestParams, - ) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: ... # pragma: no branch - - -class ElicitationFnT(Protocol): - async def __call__( - self, - context: ClientRequestContext, - params: types.ElicitRequestParams, - ) -> types.ElicitResult | types.ErrorData: ... # pragma: no branch - - -class ListRootsFnT(Protocol): - async def __call__( - self, context: ClientRequestContext - ) -> types.ListRootsResult | types.ErrorData: ... # pragma: no branch - - -class LoggingFnT(Protocol): - async def __call__(self, params: types.LoggingMessageNotificationParams) -> None: ... # pragma: no branch - - -IncomingMessage: TypeAlias = types.ServerNotification | Exception -"""What `message_handler` receives: the server notifications the session surfaces, plus transport-level exceptions. - -`notifications/cancelled` is applied by the dispatcher and never surfaced, and a -`notifications/subscriptions/acknowledged` for a live `listen()` stream is consumed by that -stream, so neither reaches the handler. -""" - - -class MessageHandlerFnT(Protocol): - async def __call__(self, message: IncomingMessage) -> None: ... # pragma: no branch - - -async def _default_message_handler(message: IncomingMessage) -> None: - await anyio.lowlevel.checkpoint() - - -async def _default_sampling_callback( - context: ClientRequestContext, - params: types.CreateMessageRequestParams, -) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: - return types.ErrorData( - code=types.INVALID_REQUEST, - message="Sampling not supported", - ) - - -async def _default_elicitation_callback( - context: ClientRequestContext, - params: types.ElicitRequestParams, -) -> types.ElicitResult | types.ErrorData: - return types.ErrorData( - code=types.INVALID_REQUEST, - message="Elicitation not supported", - ) - - -async def _default_list_roots_callback( - context: ClientRequestContext, -) -> types.ListRootsResult | types.ErrorData: - return types.ErrorData( - code=types.INVALID_REQUEST, - message="List roots not supported", - ) - - -async def _default_logging_callback( - params: types.LoggingMessageNotificationParams, -) -> None: - pass - - -ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(types.ClientResult | types.ErrorData) - -# Typed against the wide parse union so adopt-built claim adapters share this attribute type. -_CallToolResultAdapter: TypeAdapter[types.CallToolResult | types.InputRequiredResult | types.Result] = TypeAdapter( - types.CallToolResult | types.InputRequiredResult +from mcp_client.client.session import ( + ClientRequestContext as ClientRequestContext, ) -_GetPromptResultAdapter: TypeAdapter[types.GetPromptResult | types.InputRequiredResult] = TypeAdapter( - types.GetPromptResult | types.InputRequiredResult +from mcp_client.client.session import ( + ClientResponse as ClientResponse, ) -_ReadResourceResultAdapter: TypeAdapter[types.ReadResourceResult | types.InputRequiredResult] = TypeAdapter( - types.ReadResourceResult | types.InputRequiredResult +from mcp_client.client.session import ( + ClientSession as ClientSession, +) +from mcp_client.client.session import ( + ElicitationFnT as ElicitationFnT, +) +from mcp_client.client.session import ( + IncomingMessage as IncomingMessage, +) +from mcp_client.client.session import ( + ListRootsFnT as ListRootsFnT, +) +from mcp_client.client.session import ( + LoggingFnT as LoggingFnT, +) +from mcp_client.client.session import ( + MessageHandlerFnT as MessageHandlerFnT, +) +from mcp_client.client.session import ( + ReceiveResultT as ReceiveResultT, +) +from mcp_client.client.session import ( + SamplingFnT as SamplingFnT, +) +from mcp_client.client.session import ( + _active_claims_at as _active_claims_at, +) +from mcp_client.client.session import ( + _build_call_tool_adapter as _build_call_tool_adapter, +) +from mcp_client.client.session import ( + _CallToolResultAdapter as _CallToolResultAdapter, +) +from mcp_client.client.session import ( + _claim_active as _claim_active, +) +from mcp_client.client.session import ( + _clamp_inbound_ttl as _clamp_inbound_ttl, +) +from mcp_client.client.session import ( + _default_elicitation_callback as _default_elicitation_callback, +) +from mcp_client.client.session import ( + _default_list_roots_callback as _default_list_roots_callback, +) +from mcp_client.client.session import ( + _default_logging_callback as _default_logging_callback, +) +from mcp_client.client.session import ( + _default_message_handler as _default_message_handler, +) +from mcp_client.client.session import ( + _default_sampling_callback as _default_sampling_callback, +) +from mcp_client.client.session import ( + _GetPromptResultAdapter as _GetPromptResultAdapter, +) +from mcp_client.client.session import ( + _index_bindings as _index_bindings, +) +from mcp_client.client.session import ( + _index_claims as _index_claims, +) +from mcp_client.client.session import ( + _input_required_unexpected as _input_required_unexpected, +) +from mcp_client.client.session import ( + _later_revision_fields as _later_revision_fields, +) +from mcp_client.client.session import ( + _make_handshake_stamp as _make_handshake_stamp, +) +from mcp_client.client.session import ( + _make_modern_stamp as _make_modern_stamp, +) +from mcp_client.client.session import ( + _parse_server_info_stamp as _parse_server_info_stamp, +) +from mcp_client.client.session import ( + _preconnect_stamp as _preconnect_stamp, +) +from mcp_client.client.session import ( + _ReadResourceResultAdapter as _ReadResourceResultAdapter, +) +from mcp_client.client.session import ( + _same_schema as _same_schema, +) +from mcp_client.client.session import ( + _wire_fields as _wire_fields, +) +from mcp_client.client.session import ( + logger as logger, ) - -def _claim_active(claim: ResultClaim[Any], version: str) -> bool: - """A claim is active at modern versions only, narrowed by its optional version subset.""" - return version in MODERN_PROTOCOL_VERSIONS and ( - claim.protocol_versions is None or version in claim.protocol_versions - ) - - -def _active_claims_at( - claims_by_extension: Mapping[str, tuple[ResultClaim[Any], ...]], version: str -) -> dict[str, ResultClaim[Any]]: - """Claims active at `version`, keyed by wire tag; empty at any legacy version.""" - return { - claim.result_type: claim - for claims in claims_by_extension.values() - for claim in claims - if _claim_active(claim, version) - } - - -def _build_call_tool_adapter( - active: Mapping[str, ResultClaim[Any]], -) -> TypeAdapter[types.CallToolResult | types.InputRequiredResult | types.Result]: - """Build a discriminated tools/call adapter: a core arm plus one arm per active claim.""" - if not active: - return _CallToolResultAdapter - tags = frozenset(active) - core_arm = "core" - while core_arm in tags: # the routing sentinel must never collide with a claimed tag - core_arm += "-" - - def _route(value: Any) -> str: - # pydantic hands the discriminator either the raw dict or an already-built model. - # Unknown or non-string tags route to the core arm and fail core validation there. - if isinstance(value, dict): - tag = cast("dict[str, Any]", value).get("resultType") - else: - tag = getattr(value, "result_type", None) - return tag if isinstance(tag, str) and tag in tags else core_arm - - arms: list[Any] = [Annotated[types.CallToolResult | types.InputRequiredResult, Tag(core_arm)]] - arms += [Annotated[claim.model, Tag(tag)] for tag, claim in active.items()] - # reduce(or_) rather than Union star-unpack, which needs py3.11+. - return TypeAdapter(Annotated[reduce(or_, arms), Discriminator(_route)]) - - -def _index_claims( - result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None, - extensions: dict[str, dict[str, Any]] | None, -) -> dict[str, tuple[ResultClaim[Any], ...]]: - """Validate and copy the claims-by-extension mapping.""" - indexed: dict[str, tuple[ResultClaim[Any], ...]] = {} - seen: set[str] = set() - for identifier, claims in (result_claims or {}).items(): - if extensions is None or identifier not in extensions: - raise ValueError( - f"result_claims key {identifier!r} has no extensions entry; a claim is only " - "advertised through its extension's capability ad" - ) - if not claims: - raise ValueError( - f"result_claims[{identifier!r}] is empty and would drop the extension from " - "the capability ad at every version. Omit the key instead" - ) - for claim in claims: - if claim.result_type in seen: - raise ValueError(f"duplicate result claim for resultType {claim.result_type!r}") - seen.add(claim.result_type) - indexed[identifier] = tuple(claims) - return indexed - - -def _index_bindings( - notification_bindings: Sequence[NotificationBinding[Any]] | None, -) -> dict[str, NotificationBinding[Any]]: - """Index bindings by wire method, rejecting duplicates.""" - indexed: dict[str, NotificationBinding[Any]] = {} - for binding in notification_bindings or (): - if binding.method in indexed: - raise ValueError(f"duplicate notification binding for method {binding.method!r}") - indexed[binding.method] = binding - return indexed - - -def _input_required_unexpected(method: str) -> RuntimeError: - return RuntimeError( - "Server returned InputRequiredResult; pass allow_input_required=True to receive it " - f"and retry {method}(..., input_responses=..., request_state=result.request_state)." - ) - - -class ClientSession: - """Client half of an MCP connection, running on a `Dispatcher`. - - Construct it over a transport's stream pair (or pass a pre-built - `dispatcher=`), enter as an async context manager, then call - `initialize()`. The dispatcher owns the receive loop and request - correlation; this class owns the typed MCP layer and the constructor - callbacks. Transport `Exception` items reach `message_handler` on any - stream-backed dispatcher (`JSONRPCDispatcher`), whether built here from a - stream pair or supplied without a stream-exception hook of its own; an - in-process `DirectDispatcher` carries none. - - Extension `result_claims` fold into tools/call parsing at `adopt()`; - `notification_bindings` observe vendor notifications via bounded FIFOs. - """ - - def __init__( - self, - read_stream: ReadStream[SessionMessage | Exception] | None = None, - write_stream: WriteStream[SessionMessage] | None = None, - read_timeout_seconds: float | None = None, - sampling_callback: SamplingFnT | None = None, - elicitation_callback: ElicitationFnT | None = None, - list_roots_callback: ListRootsFnT | None = None, - logging_callback: LoggingFnT | None = None, - message_handler: MessageHandlerFnT | None = None, - client_info: types.Implementation | None = None, - *, - log_level: types.LoggingLevel | None = None, - sampling_capabilities: types.SamplingCapability | None = None, - extensions: dict[str, dict[str, Any]] | None = None, - result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None = None, - notification_bindings: Sequence[NotificationBinding[Any]] | None = None, - dispatcher: Dispatcher[Any] | None = None, - ) -> None: - self._session_read_timeout_seconds = read_timeout_seconds - self._client_info = client_info or DEFAULT_CLIENT_INFO - self._sampling_callback = sampling_callback or _default_sampling_callback - self._sampling_capabilities = sampling_capabilities - self._extensions = dict(extensions) if extensions is not None else None - self._result_claims = _index_claims(result_claims, extensions) - self._notification_bindings = _index_bindings(notification_bindings) - self._active_claims: dict[str, ResultClaim[Any]] = {} - self._call_tool_adapter = _CallToolResultAdapter - self._binding_queues: dict[ - str, tuple[MemoryObjectSendStream[BaseModel], MemoryObjectReceiveStream[BaseModel]] - ] = {} - self._elicitation_callback = elicitation_callback or _default_elicitation_callback - self._list_roots_callback = list_roots_callback or _default_list_roots_callback - self._logging_callback = logging_callback or _default_logging_callback - self._log_level: types.LoggingLevel | None = log_level - self._message_handler = message_handler or _default_message_handler - self._tool_output_schemas: dict[str, dict[str, Any] | None] = {} - # Compiled output-schema validators, derived from `_tool_output_schemas` and owned by - # `_absorb_tool_listing`, which evicts a tool's entry whenever its schema changes. - self._tool_output_validators: dict[str, Validator] = {} - self._x_mcp_header_maps: dict[str, dict[tuple[str, ...], str]] = {} - self._initialize_result: types.InitializeResult | None = None - self._discover_result: types.DiscoverResult | None = None - self._discover_server_info: types.Implementation | None = None - self._negotiated_version: str | None = None - self._stamp: Callable[[dict[str, Any], CallOptions], None] = _preconnect_stamp - self._task_group: anyio.abc.TaskGroup | None = None - # subscriptions/listen demux routes; membership decides ack consumption (raw listens are never registered) - self._listen_routes: dict[RequestId, ListenRoute] = {} - if dispatcher is not None: - if read_stream is not None or write_stream is not None: - raise ValueError("pass read_stream/write_stream or dispatcher, not both") - self._dispatcher: Dispatcher[Any] = dispatcher - if isinstance(dispatcher, JSONRPCDispatcher) and dispatcher.on_stream_exception is None: - # Route transport-level Exception items into message_handler — only - # stream-backed dispatchers carry these; DirectDispatcher has none. - # Don't clobber a caller-supplied hook. - # TODO(L78): this leaves a bound-method ref on the dispatcher after the - # session exits (memory pin) and a second wrap of the same dispatcher would - # skip install. The Transport-as-Dispatcher rework (L77) removes this seam. - dispatcher.on_stream_exception = self._on_stream_exception - else: - if read_stream is None or write_stream is None: - raise ValueError("read_stream and write_stream are required when no dispatcher is given") - # Built eagerly so notifications can be sent before entering the context manager. - self._dispatcher = JSONRPCDispatcher( - read_stream, write_stream, on_stream_exception=self._on_stream_exception - ) - - async def __aenter__(self) -> Self: - self._task_group = anyio.create_task_group() - await self._task_group.__aenter__() - try: - # Queues must exist before the dispatcher starts: _on_notify enqueues into this dict. - for binding in self._notification_bindings.values(): - send, receive = anyio.create_memory_object_stream[BaseModel](_NOTIFICATION_QUEUE_SIZE) - self._binding_queues[binding.method] = (send, receive) - await self._task_group.start( - self._dispatcher.run, self._on_request, self._on_notify, self._intercept_notification - ) - for binding in self._notification_bindings.values(): - _, receive = self._binding_queues[binding.method] - self._task_group.start_soon(self._deliver_bound_notifications, binding, receive) - except BaseException: - # Unwind the entered task group before propagating: a cancellation - # landing here (e.g. `move_on_after` around connect) would abandon - # it and anyio would later raise "exited non-innermost cancel scope". - task_group = self._task_group - self._task_group = None - task_group.cancel_scope.cancel() - # Shield the group's own scope (a new one would break LIFO exit) - # so a pending outer cancellation cannot re-fire inside __aexit__. - task_group.cancel_scope.shield = True - try: - await task_group.__aexit__(None, None, None) - finally: - self._close_binding_queues() - raise - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: TracebackType | None, - ) -> bool | None: - # Exit must not block: cancel the dispatcher, binding consumers, and in-flight callbacks. - assert self._task_group is not None - self._task_group.cancel_scope.cancel() - try: - result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb) - finally: - self._close_binding_queues() - self._settle_listen_routes_closed() - await resync_tracer() - return result - - def _close_binding_queues(self) -> None: - # Unclosed memory object streams warn at garbage collection; close is idempotent. - for send, receive in self._binding_queues.values(): - send.close() - receive.close() - self._binding_queues.clear() - - async def _deliver_bound_notifications( - self, binding: NotificationBinding[Any], receive: MemoryObjectReceiveStream[BaseModel] - ) -> None: - """Consume one binding's FIFO, decoupled from the dispatcher so handlers can do session I/O.""" - while True: - params = await receive.receive() - try: - await binding.handler(params) - except Exception: - # A raising handler costs only that delivery, as in _on_notify. - logger.exception("notification binding handler for %r raised", binding.method) - - async def send_request( - self, - request: types.ClientRequest | types.Request[Any, Any], - result_type: type[ReceiveResultT] | TypeAdapter[ReceiveResultT], - request_read_timeout_seconds: float | None = None, - metadata: ClientMessageMetadata | None = None, - progress_callback: ProgressFnT | None = None, - ) -> ReceiveResultT: - """Send a request and wait for its typed result. - - Args: - metadata: Streamable HTTP resumption hints. - - Raises: - MCPError: Error response, read timeout, or connection closed. - RuntimeError: Called before entering the context manager. - ValueError: The request declares `name_param` but its params carry no string name. - pydantic.ValidationError: The server returned a result that does not - conform to the negotiated protocol version. - """ - data = request.model_dump(by_alias=True, mode="json", exclude_none=True) - method: str = data["method"] - opts: CallOptions = {} - self._stamp(data, opts) - # The stamp runs first, so its NAME_BEARING_METHODS rows win; a missing name fails loud. - headers = opts.setdefault("headers", {}) - if (key := type(request).name_param) is not None and MCP_NAME_HEADER not in headers: - params_data: dict[str, Any] = data.get("params") or {} - name = params_data.get(key) - if not isinstance(name, str): - raise ValueError(f"{method} requires params[{key!r}] for Mcp-Name") - headers[MCP_NAME_HEADER] = encode_header_value(name) - timeout = ( - request_read_timeout_seconds - if request_read_timeout_seconds is not None - else self._session_read_timeout_seconds - ) - if timeout is not None: - opts["timeout"] = timeout - if progress_callback is not None: - opts["on_progress"] = progress_callback - if metadata is not None: - if metadata.resumption_token is not None: - opts["resumption_token"] = metadata.resumption_token - if metadata.on_resumption_token_update is not None: - opts["on_resumption_token"] = metadata.on_resumption_token_update - raw = await self._dispatcher.send_raw_request(method, data.get("params"), opts) - _clamp_inbound_ttl(raw) - # Literal fallback covers pre-handshake and stateless; matches runner.py. - version = self._negotiated_version or "2025-11-25" - try: - _methods.validate_server_result(method, version, raw) - except KeyError: - pass - # Drop a later revision's fields (e.g. 2026-07-28 cache hints on a pre-2026 - # session): they are outside the negotiated contract, and the version-free - # result type would otherwise apply that revision's constraints to them. - if not (foreign := _later_revision_fields(method, version)).isdisjoint(raw): - raw = {key: value for key, value in raw.items() if key not in foreign} - if isinstance(result_type, TypeAdapter): - return result_type.validate_python(raw, by_name=False) - return result_type.model_validate(raw, by_name=False) - - async def send_notification(self, notification: types.ClientNotification) -> None: - """Send a one-way notification. Usable before entering the context manager. - - Fire-and-forget: after the connection has closed, the notification is - dropped with a debug log instead of raising. - """ - data = notification.model_dump(by_alias=True, mode="json", exclude_none=True) - opts: CallOptions = {} - self._stamp(data, opts) - await self._dispatcher.notify(data["method"], data.get("params"), opts) - - def _build_capabilities(self, version: str) -> types.ClientCapabilities: - """Build the capability ad for a wire speaking `version`. - - Claim-bearing identifiers whose claims are all inactive at `version` drop, so - the client never advertises result shapes it would reject; claim-less - identifiers always advertise. - """ - extensions = self._extensions - if extensions is not None and self._result_claims: - extensions = { - identifier: settings - for identifier, settings in extensions.items() - if identifier not in self._result_claims - or any(_claim_active(claim, version) for claim in self._result_claims[identifier]) - } or None - sampling = ( - (self._sampling_capabilities or types.SamplingCapability()) - if self._sampling_callback is not _default_sampling_callback - else None - ) - elicitation = ( - types.ElicitationCapability(form=types.FormElicitationCapability(), url=types.UrlElicitationCapability()) - if self._elicitation_callback is not _default_elicitation_callback - else None - ) - roots = ( - # TODO: Should this be based on whether we - # _will_ send notifications, or only whether - # they're supported? - types.RootsCapability(list_changed=True) - if self._list_roots_callback is not _default_list_roots_callback - else None - ) - return types.ClientCapabilities( - sampling=sampling, elicitation=elicitation, experimental=None, extensions=extensions, roots=roots - ) - - async def initialize(self) -> types.InitializeResult: - if self._initialize_result is not None: - return self._initialize_result - result = await self.send_request( - types.InitializeRequest( - params=types.InitializeRequestParams( - protocol_version=LATEST_HANDSHAKE_VERSION, - # The handshake negotiates only legacy versions, where no claim is active. - capabilities=self._build_capabilities(LATEST_HANDSHAKE_VERSION), - client_info=self._client_info, - ), - ), - types.InitializeResult, - ) - - if result.protocol_version not in HANDSHAKE_PROTOCOL_VERSIONS: - raise RuntimeError(f"Unsupported protocol version from the server: {result.protocol_version}") - - self.adopt(result) - - await self.send_notification(types.InitializedNotification()) - - return result - - def adopt(self, result: types.InitializeResult | types.DiscoverResult) -> None: - """Install negotiated state from a result the caller already holds (no wire traffic). - - Clears the opposite slot, so at most one of `initialize_result` / - `discover_result` is ever non-None. - - Raises: - RuntimeError: `result` is a `DiscoverResult` whose `supported_versions` - shares nothing with this client's `MODERN_PROTOCOL_VERSIONS`. - """ - if isinstance(result, types.DiscoverResult): - # ordered oldest→newest via MODERN_PROTOCOL_VERSIONS - mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in result.supported_versions] - if not mutual: - raise RuntimeError( - f"No mutually supported modern protocol version " - f"(server: {result.supported_versions}, client: {list(MODERN_PROTOCOL_VERSIONS)})" - ) - version = mutual[-1] - client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True) - capabilities = self._build_capabilities(version).model_dump(by_alias=True, mode="json", exclude_none=True) - self._stamp = _make_modern_stamp( - version, client_info, capabilities, self._resolve_param_headers, log_level=self._log_level - ) - self._discover_result = result - self._discover_server_info = _parse_server_info_stamp(result) - self._initialize_result = None - else: - version = result.protocol_version - self._stamp = _make_handshake_stamp(version) - self._initialize_result = result - self._discover_result = None - self._discover_server_info = None - self._negotiated_version = version - # Both arms reach here, so re-adoption resets cleanly; legacy versions activate no claims. - # Core-vocabulary tags are unconstructible (ResultClaim.__post_init__), so no exclusion needed. - self._active_claims = _active_claims_at(self._result_claims, version) - self._call_tool_adapter = _build_call_tool_adapter(self._active_claims) - for method in self._notification_bindings: - # Bindings are consulted only for methods core does not know, so this one can never fire. - if (method, version) in _methods.SERVER_NOTIFICATIONS: - logger.warning( - "notification binding for %r will never fire at %s: the core protocol defines this method", - method, - version, - ) - - async def send_discover(self, version: str) -> dict[str, Any]: - """Send a single ``server/discover`` at ``version`` and return the raw result dict. - - No retry, no ``adopt()``. The ``_meta`` envelope and the - ``Mcp-Protocol-Version`` header are stamped at ``version`` so the - server-side era router sees a coherent probe. Used by ``discover()`` and - the connect-time auto-negotiation policy. - - Raises: - MCPError: The server returned a JSON-RPC error, or the transport - bounced the request at its own layer (a bare HTTP 4xx is - synthesized into a JSON-RPC error by the transport). - """ - client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True) - capabilities = self._build_capabilities(version).model_dump(by_alias=True, mode="json", exclude_none=True) - request = types.DiscoverRequest( - params=types.RequestParams( - _meta={ - PROTOCOL_VERSION_META_KEY: version, - CLIENT_INFO_META_KEY: client_info, - CLIENT_CAPABILITIES_META_KEY: capabilities, - } - ) - ) - data = request.model_dump(by_alias=True, mode="json", exclude_none=True) - opts: CallOptions = { - "timeout": DISCOVER_TIMEOUT_SECONDS, - "cancel_on_abandon": False, - "headers": {MCP_PROTOCOL_VERSION_HEADER: version, MCP_METHOD_HEADER: data["method"]}, - } - raw = await self._dispatcher.send_raw_request(data["method"], data.get("params"), opts) - # Un-floored, a negative ttl fails the mode='auto' probe's validation and silently downgrades the handshake. - _clamp_inbound_ttl(raw) - return raw - - async def discover(self) -> types.DiscoverResult: - """Probe `server/discover` and adopt the result. - - Sends a single `server/discover` proposing the newest modern protocol - version. On `UNSUPPORTED_PROTOCOL_VERSION` (-32022) the server's - `supported` list is intersected with `MODERN_PROTOCOL_VERSIONS` and the - probe is retried once at the highest mutual version. Any other error — - including `METHOD_NOT_FOUND` (-32601) and `REQUEST_TIMEOUT` (-32001) — - propagates; the legacy `initialize()` fallback is the caller's policy. - - Raises: - MCPError: The server rejected `server/discover`, the probe timed - out, or the -32022 retry found no mutual version / failed again. - RuntimeError: `adopt()` found no mutual version in the returned - `supported_versions`. - """ - if self._discover_result is not None: - return self._discover_result - - try: - raw = await self.send_discover(LATEST_MODERN_VERSION) - except MCPError as e: - if e.code != UNSUPPORTED_PROTOCOL_VERSION: - raise - try: - data = types.UnsupportedProtocolVersionErrorData.model_validate(e.error.data) - except ValidationError: - raise e from None - # ordered oldest→newest via MODERN_PROTOCOL_VERSIONS - mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in data.supported] - if not mutual: - raise - raw = await self.send_discover(mutual[-1]) - - result = types.DiscoverResult.model_validate(raw) - self.adopt(result) - return result - - @property - def initialize_result(self) -> types.InitializeResult | None: - """The server's InitializeResult. None unless `initialize()` ran (or was adopted).""" - return self._initialize_result - - @property - def discover_result(self) -> types.DiscoverResult | None: - """The server's DiscoverResult. None unless `discover()` ran (or was adopted). - - Retained intact (supported_versions, ttl_ms, cache_scope) so callers - can round-trip it as ``prior_discover=``. - """ - return self._discover_result - - @property - def protocol_version(self) -> str | None: - """Negotiated protocol version. None until `initialize()`, `discover()`, or `adopt()`.""" - return self._negotiated_version - - @property - def server_info(self) -> types.Implementation | None: - """Server name/version. None until `initialize()`, `discover()`, or `adopt()`. - - On 2026-era connections this is the discover result's optional `_meta` - `serverInfo` stamp, parsed once at adopt time; `None` when the server - did not identify itself. The stamp is display-only per the spec, so a - malformed value reads as absent rather than failing the connection. - """ - if self._discover_result is not None: - return self._discover_server_info - if self._initialize_result is not None: - return self._initialize_result.server_info - return None - - @property - def server_capabilities(self) -> types.ServerCapabilities | None: - """Server capabilities. None until `initialize()`, `discover()`, or `adopt()`.""" - if self._discover_result is not None: - return self._discover_result.capabilities - if self._initialize_result is not None: - return self._initialize_result.capabilities - return None - - @property - def instructions(self) -> str | None: - """Server-provided instructions text, if any.""" - if self._discover_result is not None: - return self._discover_result.instructions - if self._initialize_result is not None: - return self._initialize_result.instructions - return None - - async def send_ping(self, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: - """Send a ping request.""" - return await self.send_request(types.PingRequest(params=types.RequestParams(_meta=meta)), types.EmptyResult) - - @deprecated( - "Client-to-server progress is deprecated as of 2026-07-28; progress is server-to-client only.", - category=MCPDeprecationWarning, - ) - async def send_progress_notification( - self, - progress_token: str | int, - progress: float, - total: float | None = None, - message: str | None = None, - *, - meta: RequestParamsMeta | None = None, - ) -> None: - """Send a progress notification.""" - await self.send_notification( - types.ProgressNotification( - params=types.ProgressNotificationParams( - progress_token=progress_token, - progress=progress, - total=total, - message=message, - _meta=meta, - ), - ) - ) - - @deprecated("The logging capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def set_logging_level( - self, - level: types.LoggingLevel, - *, - meta: RequestParamsMeta | None = None, - ) -> types.EmptyResult: - """Send a logging/setLevel request.""" - return await self.send_request( - types.SetLevelRequest(params=types.SetLevelRequestParams(level=level, _meta=meta)), - types.EmptyResult, - ) - - async def list_resources(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListResourcesResult: - """Send a resources/list request. - - Args: - params: Full pagination parameters including cursor and any future fields - """ - return await self.send_request(types.ListResourcesRequest(params=params), types.ListResourcesResult) - - async def list_resource_templates( - self, *, params: types.PaginatedRequestParams | None = None - ) -> types.ListResourceTemplatesResult: - """Send a resources/templates/list request. - - Args: - params: Full pagination parameters including cursor and any future fields - """ - return await self.send_request( - types.ListResourceTemplatesRequest(params=params), - types.ListResourceTemplatesResult, - ) - - @overload - async def read_resource( - self, - uri: str, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: Literal[False] = False, - ) -> types.ReadResourceResult: ... - - @overload - async def read_resource( - self, - uri: str, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: bool, - ) -> types.ReadResourceResult | types.InputRequiredResult: ... - - async def read_resource( - self, - uri: str, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: bool = False, - ) -> types.ReadResourceResult | types.InputRequiredResult: - """Send a resources/read request. - - Args: - input_responses: Responses to a prior `InputRequiredResult.input_requests`. - request_state: Opaque state echoed from a prior `InputRequiredResult`. - allow_input_required: When `False` (default), an `InputRequiredResult` - from the server raises `RuntimeError`; when `True`, it is returned - so the caller can resolve the requests and retry. - - Raises: - RuntimeError: If the server returns an `InputRequiredResult` and - `allow_input_required` is `False`. - """ - result = await self.send_request( - types.ReadResourceRequest( - params=types.ReadResourceRequestParams( - uri=uri, - input_responses=input_responses, - request_state=request_state, - _meta=meta, - ), - ), - _ReadResourceResultAdapter, - ) - if isinstance(result, types.InputRequiredResult) and not allow_input_required: - raise _input_required_unexpected("read_resource") - return result - - @deprecated( - "resources/subscribe is removed as of 2026-07-28; use Client.listen() instead.", - category=MCPDeprecationWarning, - ) - async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: - """Send a resources/subscribe request (2025-era servers only).""" - return await self.send_request( - types.SubscribeRequest(params=types.SubscribeRequestParams(uri=uri, _meta=meta)), - types.EmptyResult, - ) - - @deprecated( - "resources/unsubscribe is removed as of 2026-07-28; use Client.listen() instead.", - category=MCPDeprecationWarning, - ) - async def unsubscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: - """Send a resources/unsubscribe request (2025-era servers only).""" - return await self.send_request( - types.UnsubscribeRequest(params=types.UnsubscribeRequestParams(uri=uri, _meta=meta)), - types.EmptyResult, - ) - - @overload - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: Literal[False] = False, - allow_claimed: Literal[False] = False, - ) -> types.CallToolResult: ... - - @overload - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: bool, - allow_claimed: Literal[False] = False, - ) -> types.CallToolResult | types.InputRequiredResult: ... - - @overload - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: Literal[False] = False, - allow_claimed: bool, - ) -> types.CallToolResult | types.Result: ... - - @overload - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: bool, - allow_claimed: bool, - ) -> types.CallToolResult | types.InputRequiredResult | types.Result: ... - - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: bool = False, - allow_claimed: bool = False, - ) -> types.CallToolResult | types.InputRequiredResult | types.Result: - """Send a tools/call request with optional progress callback support. - - On a modern (2026-07-28) connection, arguments annotated with `x-mcp-header` - in the tool's input schema are mirrored into `Mcp-Param-*` request headers. - The annotations are read from the tool's last `list_tools` entry, so list - the tool before calling it to enable header emission. - - Args: - input_responses: Responses to a prior `InputRequiredResult.input_requests`. - request_state: Opaque state echoed from a prior `InputRequiredResult`. - allow_input_required: When ``False`` (default), an `InputRequiredResult` - from the server raises `RuntimeError`; when ``True``, it is returned - so the caller can resolve the requests and retry. - allow_claimed: When `False` (default), a claimed extension result raises - `UnexpectedClaimedResult`; when `True`, the parsed claim model is returned. - - Raises: - RuntimeError: If the server returns an `InputRequiredResult` and - ``allow_input_required`` is ``False``. - UnexpectedClaimedResult: Claimed result with `allow_claimed` False; carries the parsed value. - """ - result = await self.send_request( - types.CallToolRequest( - params=types.CallToolRequestParams( - name=name, - arguments=arguments, - input_responses=input_responses, - request_state=request_state, - _meta=meta, - ), - ), - self._call_tool_adapter, - request_read_timeout_seconds=read_timeout_seconds, - progress_callback=progress_callback, - ) - - if isinstance(result, types.CallToolResult) and not result.is_error: - await self.validate_tool_result(name, result) - - # The input_required arm stays first; a claimed shape is terminal for the multi-round-trip driver. - if isinstance(result, types.InputRequiredResult) and not allow_input_required: - raise _input_required_unexpected("call_tool") - if not isinstance(result, types.CallToolResult | types.InputRequiredResult) and not allow_claimed: - raise UnexpectedClaimedResult(result) - return result - - def _resolve_param_headers(self, name: str, arguments: Mapping[str, Any]) -> dict[str, str]: - """`Mcp-Param-*` headers for a `tools/call`, or empty when the tool was never listed.""" - header_map = self._x_mcp_header_maps.get(name) - if header_map is None: - return {} - return mcp_param_headers(header_map, arguments) - - async def validate_tool_result(self, name: str, result: types.CallToolResult) -> None: - """Revalidate a `CallToolResult` against the tool's declared output schema. - - Raises: - RuntimeError: Structured content is missing or does not conform to the schema, or the - schema is invalid or has a `$ref` that does not resolve within the schema document. - """ - if name not in self._tool_output_schemas: - # refresh output schema cache - await self.list_tools() - - output_schema = None - if name in self._tool_output_schemas: - output_schema = self._tool_output_schemas.get(name) - else: - logger.warning(f"Tool {name} not listed by server, cannot validate any structured content") - - if output_schema is not None: - from jsonschema import exceptions as jsonschema_exceptions - from referencing.exceptions import Unresolvable - - if result.structured_content is None: - raise RuntimeError(f"Tool {name} has an output schema but did not return structured content") - validator = self._output_schema_validator(name, output_schema) - # `best_match` picks the same error the previous `jsonschema.validate()` call raised, - # so the message a caller sees is unchanged. It is untyped upstream. - errors = validator.iter_errors(result.structured_content) - try: - error = cast( - "Exception | None", - jsonschema_exceptions.best_match(errors), # pyright: ignore[reportUnknownMemberType] - ) - except Unresolvable as e: - # A `$ref` did not resolve within the schema document. - raise RuntimeError(f"Invalid schema for tool {name}: {e}") from e - if error is not None: - raise RuntimeError(f"Invalid structured content returned by tool {name}: {error}") from error - - def _output_schema_validator(self, name: str, output_schema: dict[str, Any]) -> Validator: - """Compiled validator for the tool's cached output schema, built once per schema value. - - Compiling is ~60x the cost of validating, so a one-shot `jsonschema.validate()` per - result dominates `call_tool`; the compiled validator is cached instead. It stays valid - because `_absorb_tool_listing` evicts a tool's validator whenever it absorbs a different - schema for that tool, so a cached entry always matches `output_schema`. - - Raises: - RuntimeError: The schema is not a valid JSON Schema. Raised on every call, since a - failed compile is never cached. - """ - from jsonschema import SchemaError - from jsonschema.validators import validator_for - from referencing import Registry - - if (validator := self._tool_output_validators.get(name)) is not None: - return validator - - validator_cls = validator_for(output_schema) - try: - validator_cls.check_schema(output_schema) - except SchemaError as e: - raise RuntimeError(f"Invalid schema for tool {name}: {e}") - # An explicit empty registry: `$ref`s resolve within the schema document and the bundled metaschemas. - validator = validator_cls(output_schema, registry=Registry()) - self._tool_output_validators[name] = validator - return validator - - async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult: - """Send a prompts/list request. - - Args: - params: Full pagination parameters including cursor and any future fields - """ - return await self.send_request(types.ListPromptsRequest(params=params), types.ListPromptsResult) - - @overload - async def get_prompt( - self, - name: str, - arguments: dict[str, str] | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: Literal[False] = False, - ) -> types.GetPromptResult: ... - - @overload - async def get_prompt( - self, - name: str, - arguments: dict[str, str] | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: bool, - ) -> types.GetPromptResult | types.InputRequiredResult: ... - - async def get_prompt( - self, - name: str, - arguments: dict[str, str] | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: RequestParamsMeta | None = None, - allow_input_required: bool = False, - ) -> types.GetPromptResult | types.InputRequiredResult: - """Send a prompts/get request. - - Args: - input_responses: Responses to a prior `InputRequiredResult.input_requests`. - request_state: Opaque state echoed from a prior `InputRequiredResult`. - allow_input_required: When `False` (default), an `InputRequiredResult` - from the server raises `RuntimeError`; when `True`, it is returned - so the caller can resolve the requests and retry. - - Raises: - RuntimeError: If the server returns an `InputRequiredResult` and - `allow_input_required` is `False`. - """ - result = await self.send_request( - types.GetPromptRequest( - params=types.GetPromptRequestParams( - name=name, - arguments=arguments, - input_responses=input_responses, - request_state=request_state, - _meta=meta, - ), - ), - _GetPromptResultAdapter, - ) - if isinstance(result, types.InputRequiredResult) and not allow_input_required: - raise _input_required_unexpected("get_prompt") - return result - - async def complete( - self, - ref: types.ResourceTemplateReference | types.PromptReference, - argument: dict[str, str], - context_arguments: dict[str, str] | None = None, - ) -> types.CompleteResult: - """Send a completion/complete request.""" - context = None - if context_arguments is not None: - context = types.CompletionContext(arguments=context_arguments) - - return await self.send_request( - types.CompleteRequest( - params=types.CompleteRequestParams( - ref=ref, - argument=types.CompletionArgument(**argument), - context=context, - ), - ), - types.CompleteResult, - ) - - async def list_tools(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListToolsResult: - """Send a tools/list request. - - Args: - params: Full pagination parameters including cursor and any future fields - """ - result = await self.send_request( - types.ListToolsRequest(params=params), - types.ListToolsResult, - ) - complete = (params is None or params.cursor is None) and result.next_cursor is None - return self._absorb_tool_listing(result, complete=complete) - - def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool) -> types.ListToolsResult: - """Filter the listing per the 2026 x-mcp-header MUST and rebuild derived per-tool state, in place. - - Idempotent: cached values are already post-filter, so the response cache can re-absorb a served listing. - `complete` (an uncursored single-page listing) prunes per-tool state down to the listing's tools. - """ - if self._negotiated_version in MODERN_PROTOCOL_VERSIONS: - # 2026-07-28: clients MUST drop tools whose x-mcp-header annotations are invalid. - kept: list[types.Tool] = [] - for tool in result.tools: - if (reason := find_invalid_x_mcp_header(tool.input_schema)) is not None: - logger.warning("dropping tool %r: invalid x-mcp-header (%s)", tool.name, reason) - # Evict any map cached from a prior valid listing so a stale entry can't - # mirror headers for a tool this listing dropped. - self._x_mcp_header_maps.pop(tool.name, None) - continue - # Cache the arg→header map so a later tools/call mirrors it into Mcp-Param-* headers. - self._x_mcp_header_maps[tool.name] = x_mcp_header_map(tool.input_schema) - kept.append(tool) - result.tools = kept - - # Cache tool output schemas for future validation; cursor pages only ever add. A - # changed schema evicts its compiled validator; an unchanged one (a re-listing, or the - # response cache re-absorbing a served hit) keeps it. Only validated tools pay the check. - for tool in result.tools: - if tool.name in self._tool_output_validators and not _same_schema( - self._tool_output_schemas.get(tool.name), tool.output_schema - ): - del self._tool_output_validators[tool.name] - self._tool_output_schemas[tool.name] = tool.output_schema - - if complete: - # The listing is the full tool universe, so state for unlisted tools is stale - # (the server dropped them, or a shared-cache writer's filter did). - names = {tool.name for tool in result.tools} - self._x_mcp_header_maps = {k: v for k, v in self._x_mcp_header_maps.items() if k in names} - self._tool_output_schemas = {k: v for k, v in self._tool_output_schemas.items() if k in names} - self._tool_output_validators = {k: v for k, v in self._tool_output_validators.items() if k in names} - - return result - - @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def send_roots_list_changed(self) -> None: - """Send a roots/list_changed notification.""" - await self.send_notification(types.RootsListChangedNotification()) - - async def _on_request( - self, dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None - ) -> dict[str, Any]: - """Answer a server-initiated request via the registered callbacks.""" - # Literal, not LATEST_PROTOCOL_VERSION: the fallback covers the initialize - # handshake (which only exists at <=2025) and stateless until the header - # is plumbed; its meaning is fixed regardless of LATEST bumps. - version = self._negotiated_version or "2025-11-25" - try: - request = cast(types.ServerRequest, _methods.parse_server_request(method, version, params)) - except KeyError: - raise MCPError(code=METHOD_NOT_FOUND, message="Method not found", data=method) from None - - response: types.ClientResult | types.ErrorData - if isinstance(request, types.PingRequest): - # Answered without a context: ping has no callback that would need one. - response = types.EmptyResult() - else: - assert dctx.request_id is not None # the callback-driving dispatchers always assign ids - ctx = ClientRequestContext( - session=self, request_id=dctx.request_id, meta=request.params.meta if request.params else None - ) - response = await self.dispatch_input_request(ctx, request) - client_response = ClientResponse.validate_python(response) - if isinstance(client_response, types.ErrorData): - raise MCPError.from_error_data(client_response) - dumped = client_response.model_dump(by_alias=True, mode="json", exclude_none=True) - try: - _methods.validate_client_result(method, version, dumped) - except ValidationError: - logger.exception("client callback for %r returned an invalid result", method) - raise MCPError(code=INTERNAL_ERROR, message="Client callback returned an invalid result") from None - return dumped - - async def dispatch_input_request( - self, ctx: ClientRequestContext, request: types.InputRequest - ) -> types.InputResponse | types.ErrorData: - """Route an input request through the client's callback table. - - Shared by the legacy server→client RPC path (`_on_request`) and the - 2026-07-28 multi-round-trip driver, which dispatches the embedded - `InputRequiredResult.input_requests` through the same callbacks. - - Returns the callback's `InputResponse`, or `ErrorData` when the callback declines. - """ - match request: - case types.CreateMessageRequest(params=p): - return await self._sampling_callback(ctx, p) - case types.ElicitRequest(params=p): - return await self._elicitation_callback(ctx, p) - case types.ListRootsRequest(): # pragma: no branch - return await self._list_roots_callback(ctx) - - def _register_listen_route(self, request_id: RequestId) -> ListenRoute: - """Create the demux route for a listen request id; the caller registers BEFORE sending.""" - route = ListenRoute() - self._listen_routes[request_id] = route - return route - - def _unregister_listen_route(self, request_id: RequestId) -> None: - """Drop a listen route; the handle owns membership, so a missing key is a no-op.""" - self._listen_routes.pop(request_id, None) - - def _settle_listen_routes_closed(self) -> None: - """Settle all open listen routes as lost on session exit; cancelled driver tasks cannot.""" - closed = MCPError(code=CONNECTION_CLOSED, message="Connection closed") - for route in self._listen_routes.values(): - route.settle("lost", error=closed) - self._listen_routes.clear() - - def _intercept_notification(self, method: str, params: Mapping[str, Any] | None) -> bool: - """Wire-order listen demux, run synchronously on the dispatcher's receive path. - - Bookkeeping must advance in receive order with the listen result (resolved on - this same path); the spawned `_on_notify` path would race it and drop events. - Returns True to consume the frame: a live route's ack is driver state, never surfaced. - """ - if not self._listen_routes: - return False - if method == "notifications/cancelled": - request_id = cancelled_request_id_from_params(params) - if request_id is not None and (listen_route := self._listen_routes.get(request_id)) is not None: - # a server-sent cancel naming a listen request is that stream's teardown signal - listen_route.settle("lost") - return False # _on_notify swallows every cancelled either way (v1 parity) - if params is None: - return False - meta = params.get("_meta") - if not isinstance(meta, Mapping): - return False - # as_request_id is not a tripwire: raw wire _meta can carry a non-id (even unhashable) value - subscription_id = as_request_id(cast("Mapping[str, Any]", meta).get(SUBSCRIPTION_ID_META_KEY)) - if subscription_id is None or (listen_route := self._listen_routes.get(subscription_id)) is None: - return False - if method == "notifications/subscriptions/acknowledged": - raw_filter = params.get("notifications") - if raw_filter is None: - # malformed, not an empty filter: leave it to the spawned path's validation warning - return False - try: - honored = types.SubscriptionFilter.model_validate(raw_filter) - except ValidationError: - return False - listen_route.set_acked(honored) - return True - if (event := event_from_wire(method, params)) is not None: - listen_route.deliver(event) - return False # events (and any other stamped frame) still tee as usual - - async def _on_notify( - self, dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None - ) -> None: - """Route a server notification: validate, run the typed callback, tee to message_handler.""" - # Same fallback as `_on_request`: covers pre-handshake and stateless. - version = self._negotiated_version or "2025-11-25" - try: - notification = cast(types.ServerNotification, _methods.parse_server_notification(method, version, params)) - except KeyError: - # Only methods unknown to the negotiated version's core tables reach the bindings. - binding = self._notification_bindings.get(method) - if binding is None: - logger.debug("dropped %r: not defined at %s", method, version) - return - try: - bound_params = binding.params_type.model_validate(params or {}) - except ValidationError: - logger.warning("Failed to validate notification: %s", method, exc_info=True) - return - send, receive = self._binding_queues[method] - try: - # Must not await: DirectDispatcher calls _on_notify inline; blocking deadlocks in-process servers. - send.send_nowait(bound_params) - except anyio.WouldBlock: - # Evict the oldest event; no checkpoint since the failed send, - # so the buffer is still full and the retry cannot block. - receive.receive_nowait() - logger.warning("notification queue for %r is full; dropped the oldest event", method) - send.send_nowait(bound_params) - return - except ValidationError: - logger.warning("Failed to validate notification: %s", method, exc_info=True) - return - if isinstance(notification, types.CancelledNotification): - # Never surfaced (v1 parity): the dispatcher already applied it; listen cancels settled by the intercept. - return - try: - if isinstance(notification, types.LoggingMessageNotification): - await self._logging_callback(notification.params) - await self._message_handler(notification) - except Exception: - # Contain here, not in the dispatcher: DirectDispatcher awaits this - # handler inline in the peer's notify() call, so a raising callback - # would otherwise fail the peer's send. A raising logging_callback - # skips the message_handler tee for that notification (v1 parity). - logger.exception("notification callback for %r raised", method) - - async def _on_stream_exception(self, exc: Exception) -> None: - """Deliver a transport-level fault to message_handler via a spawned task. - - Running the handler inline would park the dispatcher's read loop and - deadlock handlers that await session I/O. - """ - assert self._task_group is not None - self._task_group.start_soon(self._deliver_stream_exception, exc) - - async def _deliver_stream_exception(self, exc: Exception) -> None: - try: - await self._message_handler(exc) - except Exception: - logger.exception("message_handler raised on transport exception") +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index a544cecbe8..30045bb00e 100644 --- a/src/mcp/client/session_group.py +++ b/src/mcp/client/session_group.py @@ -1,449 +1,20 @@ -"""SessionGroup concurrently manages multiple MCP session connections. - -Tools, resources, and prompts are aggregated across servers. Servers may -be connected to or disconnected from at any point after initialization. - -This abstraction can handle naming collisions using a custom user-provided hook. -""" - -import contextlib -import logging -from collections.abc import Callable -from dataclasses import dataclass -from types import TracebackType -from typing import Any, Literal, TypeAlias, overload - -import anyio -import httpx2 -import mcp_types as types -from pydantic import BaseModel, Field -from typing_extensions import Self - -import mcp -from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT -from mcp.client.sse import sse_client -from mcp.client.stdio import StdioServerParameters -from mcp.client.streamable_http import streamable_http_client -from mcp.shared._httpx_utils import create_mcp_http_client -from mcp.shared.dispatcher import ProgressFnT -from mcp.shared.exceptions import MCPError - - -class SseServerParameters(BaseModel): - """Parameters for initializing an sse_client.""" - - # The endpoint URL. - url: str - - # Optional headers to include in requests. - headers: dict[str, Any] | None = None - - # HTTP timeout for regular operations (in seconds). - timeout: float = 5.0 - - # Timeout for SSE read operations (in seconds). - sse_read_timeout: float = 300.0 - - -class StreamableHttpParameters(BaseModel): - """Parameters for initializing a streamable_http_client.""" - - # The endpoint URL. - url: str - - # Optional headers to include in requests. - headers: dict[str, Any] | None = None - - # HTTP timeout for regular operations (in seconds). - timeout: float = 30.0 - - # Timeout for SSE read operations (in seconds). - sse_read_timeout: float = 300.0 - - # Close the client session when the transport closes. - terminate_on_close: bool = True - - -ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters - - -# Use dataclass instead of Pydantic BaseModel -# because Pydantic BaseModel cannot handle Protocol fields. -@dataclass -class ClientSessionParameters: - """Parameters for establishing a client session to an MCP server.""" - - read_timeout_seconds: float | None = None - sampling_callback: SamplingFnT | None = None - elicitation_callback: ElicitationFnT | None = None - list_roots_callback: ListRootsFnT | None = None - logging_callback: LoggingFnT | None = None - message_handler: MessageHandlerFnT | None = None - client_info: types.Implementation | None = None - - -class ClientSessionGroup: - """Client for managing connections to multiple MCP servers. - - This class is responsible for encapsulating management of server connections. - It aggregates tools, resources, and prompts from all connected servers. - - For auxiliary handlers, such as resource subscription, this is delegated to - the client and can be accessed via the session. - - Example: - ```python - name_fn = lambda name, server_info: f"{(server_info.name)}_{name}" - async with ClientSessionGroup(component_name_hook=name_fn) as group: - for server_param in server_params: - await group.connect_to_server(server_param) - ... - ``` - """ - - class _ComponentNames(BaseModel): - """Used for reverse index to find components.""" - - prompts: set[str] = Field(default_factory=set) - resources: set[str] = Field(default_factory=set) - tools: set[str] = Field(default_factory=set) - - # Standard MCP components. - _prompts: dict[str, types.Prompt] - _resources: dict[str, types.Resource] - _tools: dict[str, types.Tool] - - # Client-server connection management. - _sessions: dict[mcp.ClientSession, _ComponentNames] - _tool_to_session: dict[str, mcp.ClientSession] - _exit_stack: contextlib.AsyncExitStack - _session_exit_stacks: dict[mcp.ClientSession, contextlib.AsyncExitStack] - - # Optional fn consuming (component_name, server_info) for custom names. - # This is to provide a means to mitigate naming conflicts across servers. - # Example: (tool_name, server_info) => "{result.server_info.name}.{tool_name}" - _ComponentNameHook: TypeAlias = Callable[[str, types.Implementation], str] - _component_name_hook: _ComponentNameHook | None - - def __init__( - self, - exit_stack: contextlib.AsyncExitStack | None = None, - component_name_hook: _ComponentNameHook | None = None, - ) -> None: - """Initializes the MCP client.""" - - self._tools = {} - self._resources = {} - self._prompts = {} - - self._sessions = {} - self._tool_to_session = {} - if exit_stack is None: - self._exit_stack = contextlib.AsyncExitStack() - self._owns_exit_stack = True - else: - self._exit_stack = exit_stack - self._owns_exit_stack = False - self._session_exit_stacks = {} - self._component_name_hook = component_name_hook - - async def __aenter__(self) -> Self: # pragma: no cover - # Enter the exit stack only if we created it ourselves - if self._owns_exit_stack: - await self._exit_stack.__aenter__() - return self - - async def __aexit__( - self, - _exc_type: type[BaseException] | None, - _exc_val: BaseException | None, - _exc_tb: TracebackType | None, - ) -> bool | None: # pragma: no cover - """Closes session exit stacks and main exit stack upon completion.""" - - # Only close the main exit stack if we created it - if self._owns_exit_stack: - await self._exit_stack.aclose() - - # Concurrently close session stacks. - async with anyio.create_task_group() as tg: - for exit_stack in self._session_exit_stacks.values(): - tg.start_soon(exit_stack.aclose) - - @property - def sessions(self) -> list[mcp.ClientSession]: - """Returns the list of sessions being managed.""" - return list(self._sessions.keys()) # pragma: no cover - - @property - def prompts(self) -> dict[str, types.Prompt]: - """Returns the prompts as a dictionary of names to prompts.""" - return self._prompts - - @property - def resources(self) -> dict[str, types.Resource]: - """Returns the resources as a dictionary of names to resources.""" - return self._resources - - @property - def tools(self) -> dict[str, types.Tool]: - """Returns the tools as a dictionary of names to tools.""" - return self._tools - - @overload - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: types.RequestParamsMeta | None = None, - allow_input_required: Literal[False] = False, - ) -> types.CallToolResult: ... - - @overload - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: types.RequestParamsMeta | None = None, - allow_input_required: bool, - ) -> types.CallToolResult | types.InputRequiredResult: ... - - async def call_tool( - self, - name: str, - arguments: dict[str, Any] | None = None, - read_timeout_seconds: float | None = None, - progress_callback: ProgressFnT | None = None, - *, - input_responses: types.InputResponses | None = None, - request_state: str | None = None, - meta: types.RequestParamsMeta | None = None, - allow_input_required: bool = False, - ) -> types.CallToolResult | types.InputRequiredResult: - """Executes a tool given its name and arguments. - - Raises: - RuntimeError: If the server returns an `InputRequiredResult` and - ``allow_input_required`` is ``False``. - """ - session = self._tool_to_session[name] - session_tool_name = self.tools[name].name - return await session.call_tool( - session_tool_name, - arguments=arguments, - read_timeout_seconds=read_timeout_seconds, - progress_callback=progress_callback, - input_responses=input_responses, - request_state=request_state, - meta=meta, - allow_input_required=allow_input_required, - ) - - async def disconnect_from_server(self, session: mcp.ClientSession) -> None: - """Disconnects from a single MCP server.""" - - session_known_for_components = session in self._sessions - session_known_for_stack = session in self._session_exit_stacks - - if not session_known_for_components and not session_known_for_stack: - raise MCPError( - code=types.INVALID_PARAMS, - message="Provided session is not managed or already disconnected.", - ) - - if session_known_for_components: # pragma: no branch - component_names = self._sessions.pop(session) # Pop from _sessions tracking - - # Remove prompts associated with the session. - for name in component_names.prompts: - if name in self._prompts: # pragma: no branch - del self._prompts[name] - # Remove resources associated with the session. - for name in component_names.resources: - if name in self._resources: # pragma: no branch - del self._resources[name] - # Remove tools associated with the session. - for name in component_names.tools: - if name in self._tools: # pragma: no branch - del self._tools[name] - if name in self._tool_to_session: # pragma: no branch - del self._tool_to_session[name] - - # Clean up the session's resources via its dedicated exit stack - if session_known_for_stack: - session_stack_to_close = self._session_exit_stacks.pop(session) # pragma: no cover - await session_stack_to_close.aclose() # pragma: no cover - - async def connect_with_session( - self, server_info: types.Implementation, session: mcp.ClientSession - ) -> mcp.ClientSession: - """Connects to a single MCP server.""" - await self._aggregate_components(server_info, session) - return session - - async def connect_to_server( - self, - server_params: ServerParameters, - session_params: ClientSessionParameters | None = None, - ) -> mcp.ClientSession: - """Connects to a single MCP server.""" - server_info, session = await self._establish_session(server_params, session_params or ClientSessionParameters()) - return await self.connect_with_session(server_info, session) - - async def _establish_session( - self, - server_params: ServerParameters, - session_params: ClientSessionParameters, - ) -> tuple[types.Implementation, mcp.ClientSession]: - """Establish a client session to an MCP server.""" - - session_stack = contextlib.AsyncExitStack() - try: - # Create read and write streams that facilitate io with the server. - if isinstance(server_params, StdioServerParameters): - client = mcp.stdio_client(server_params) - read, write = await session_stack.enter_async_context(client) - elif isinstance(server_params, SseServerParameters): - client = sse_client( - url=server_params.url, - headers=server_params.headers, - timeout=server_params.timeout, - sse_read_timeout=server_params.sse_read_timeout, - ) - read, write = await session_stack.enter_async_context(client) - else: - httpx_client = create_mcp_http_client( - headers=server_params.headers, - timeout=httpx2.Timeout( - server_params.timeout, - read=server_params.sse_read_timeout, - ), - ) - await session_stack.enter_async_context(httpx_client) - - client = streamable_http_client( - url=server_params.url, - http_client=httpx_client, - terminate_on_close=server_params.terminate_on_close, - ) - read, write = await session_stack.enter_async_context(client) - - session = await session_stack.enter_async_context( - mcp.ClientSession( - read, - write, - read_timeout_seconds=session_params.read_timeout_seconds, - sampling_callback=session_params.sampling_callback, - elicitation_callback=session_params.elicitation_callback, - list_roots_callback=session_params.list_roots_callback, - logging_callback=session_params.logging_callback, - message_handler=session_params.message_handler, - client_info=session_params.client_info, - ) - ) - - result = await session.initialize() - - # Session successfully initialized. - # Store its stack and register the stack with the main group stack. - self._session_exit_stacks[session] = session_stack - # session_stack itself becomes a resource managed by the - # main _exit_stack. - await self._exit_stack.enter_async_context(session_stack) - - return result.server_info, session - except Exception: # pragma: no cover - # If anything during this setup fails, ensure the session-specific - # stack is closed. - await session_stack.aclose() - raise - - async def _aggregate_components(self, server_info: types.Implementation, session: mcp.ClientSession) -> None: - """Aggregates prompts, resources, and tools from a given session.""" - - # Create a reverse index so we can find all prompts, resources, and - # tools belonging to this session. Used for removing components from - # the session group via self.disconnect_from_server. - component_names = self._ComponentNames() - - # Temporary components dicts. We do not want to modify the aggregate - # lists in case of an intermediate failure. - prompts_temp: dict[str, types.Prompt] = {} - resources_temp: dict[str, types.Resource] = {} - tools_temp: dict[str, types.Tool] = {} - tool_to_session_temp: dict[str, mcp.ClientSession] = {} - - # Query the server for its prompts and aggregate to list. - try: - prompts = (await session.list_prompts()).prompts - for prompt in prompts: - name = self._component_name(prompt.name, server_info) - prompts_temp[name] = prompt - component_names.prompts.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch prompts: {err}") - - # Query the server for its resources and aggregate to list. - try: - resources = (await session.list_resources()).resources - for resource in resources: - name = self._component_name(resource.name, server_info) - resources_temp[name] = resource - component_names.resources.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch resources: {err}") - - # Query the server for its tools and aggregate to list. - try: - tools = (await session.list_tools()).tools - for tool in tools: - name = self._component_name(tool.name, server_info) - tools_temp[name] = tool - tool_to_session_temp[name] = session - component_names.tools.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch tools: {err}") - - # Clean up exit stack for session if we couldn't retrieve anything - # from the server. - if not any((prompts_temp, resources_temp, tools_temp)): - del self._session_exit_stacks[session] # pragma: no cover - - # Check for duplicates. - matching_prompts = prompts_temp.keys() & self._prompts.keys() - if matching_prompts: - raise MCPError( # pragma: no cover - code=types.INVALID_PARAMS, - message=f"{matching_prompts} already exist in group prompts.", - ) - matching_resources = resources_temp.keys() & self._resources.keys() - if matching_resources: - raise MCPError( # pragma: no cover - code=types.INVALID_PARAMS, - message=f"{matching_resources} already exist in group resources.", - ) - matching_tools = tools_temp.keys() & self._tools.keys() - if matching_tools: - raise MCPError(code=types.INVALID_PARAMS, message=f"{matching_tools} already exist in group tools.") - - # Aggregate components. - self._sessions[session] = component_names - self._prompts.update(prompts_temp) - self._resources.update(resources_temp) - self._tools.update(tools_temp) - self._tool_to_session.update(tool_to_session_temp) - - def _component_name(self, name: str, server_info: types.Implementation) -> str: - if self._component_name_hook: - return self._component_name_hook(name, server_info) - return name +import sys + +import mcp_client.client.session_group as _implementation +from mcp_client.client.session_group import ( + ClientSessionGroup as ClientSessionGroup, +) +from mcp_client.client.session_group import ( + ClientSessionParameters as ClientSessionParameters, +) +from mcp_client.client.session_group import ( + ServerParameters as ServerParameters, +) +from mcp_client.client.session_group import ( + SseServerParameters as SseServerParameters, +) +from mcp_client.client.session_group import ( + StreamableHttpParameters as StreamableHttpParameters, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py index f72ff273a9..0abeb8d991 100644 --- a/src/mcp/client/sse.py +++ b/src/mcp/client/sse.py @@ -1,171 +1,17 @@ -import logging -from collections.abc import Callable -from contextlib import asynccontextmanager -from typing import Any -from urllib.parse import parse_qs, urljoin, urlparse +import sys -import anyio -import httpx2 -import mcp_types as types -from anyio.abc import TaskStatus -from httpx2 import SSEError - -from mcp.shared._compat import resync_tracer -from mcp.shared._context_streams import create_context_streams -from mcp.shared._httpx_utils import ( - McpHttpClientFactory, - create_mcp_http_client, - request_within_origin, - sse_within_origin, +import mcp_client.client.sse as _implementation +from mcp_client.client.sse import ( + _extract_session_id_from_endpoint as _extract_session_id_from_endpoint, +) +from mcp_client.client.sse import ( + logger as logger, +) +from mcp_client.client.sse import ( + remove_request_params as remove_request_params, +) +from mcp_client.client.sse import ( + sse_client as sse_client, ) -from mcp.shared.message import SessionMessage - -logger = logging.getLogger(__name__) - - -def remove_request_params(url: str) -> str: - return urljoin(url, urlparse(url).path) - - -def _extract_session_id_from_endpoint(endpoint_url: str) -> str | None: - query_params = parse_qs(urlparse(endpoint_url).query) - return query_params.get("sessionId", [None])[0] or query_params.get("session_id", [None])[0] - - -@asynccontextmanager -async def sse_client( - url: str, - headers: dict[str, Any] | None = None, - timeout: float = 5.0, - sse_read_timeout: float = 300.0, - httpx_client_factory: McpHttpClientFactory = create_mcp_http_client, - auth: httpx2.Auth | None = None, - on_session_created: Callable[[str], None] | None = None, -): - """Client transport for SSE. - - `sse_read_timeout` determines how long (in seconds) the client will wait for a new - event before disconnecting. All other HTTP operations are controlled by `timeout`. - - Args: - url: The SSE endpoint URL. - headers: Optional headers to include in requests. - timeout: HTTP timeout for regular operations (in seconds). - sse_read_timeout: Timeout for SSE read operations (in seconds). - httpx_client_factory: Factory function for creating the httpx2 client. Whichever client it - returns, MCP requests follow a redirect only when it stays on the endpoint's origin - (same scheme, host and port, or http to https on the same host with default ports) and - keeps the request method (any status for the SSE GET, 307/308 for a message POST); any - other redirect is not followed, so connecting fails with - `httpx2.HTTPStatusError` for the redirect response. The client's `follow_redirects` - setting is not consulted; the SDK's OAuth providers apply the same rule to the requests - they make. - auth: Optional httpx2 authentication handler. - on_session_created: Optional callback invoked with the session ID when received. - """ - logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}") - async with httpx_client_factory( - headers=headers, auth=auth, timeout=httpx2.Timeout(timeout, read=sse_read_timeout) - ) as client: - async with sse_within_origin(client, url) as event_source: - event_source.response.raise_for_status() - logger.debug("SSE connection established") - - read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0) - write_stream, write_stream_reader = create_context_streams[SessionMessage](0) - - async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED): - try: - async for sse in event_source: # pragma: no branch - logger.debug(f"Received SSE event: {sse.event}") - match sse.event: - case "endpoint": - endpoint_url = urljoin(url, sse.data) - logger.debug(f"Received endpoint URL: {endpoint_url}") - - url_parsed = urlparse(url) - endpoint_parsed = urlparse(endpoint_url) - if ( # pragma: no cover - url_parsed.netloc != endpoint_parsed.netloc - or url_parsed.scheme != endpoint_parsed.scheme - ): - error_msg = ( # pragma: no cover - f"Endpoint origin does not match connection origin: {endpoint_url}" - ) - logger.error(error_msg) # pragma: no cover - raise ValueError(error_msg) # pragma: no cover - - if on_session_created: - session_id = _extract_session_id_from_endpoint(endpoint_url) - if session_id: - on_session_created(session_id) - - task_status.started(endpoint_url) - - case "message": - # Skip empty data (keep-alive pings) - if not sse.data: - continue - try: - message = types.jsonrpc_message_adapter.validate_json(sse.data, by_name=False) - logger.debug(f"Received server message: {message}") - except Exception as exc: # pragma: no cover - logger.exception("Error parsing server message") # pragma: no cover - await read_stream_writer.send(exc) # pragma: no cover - continue # pragma: no cover - - session_message = SessionMessage(message) - await read_stream_writer.send(session_message) - case _: # pragma: no cover - logger.warning(f"Unknown SSE event: {sse.event}") # pragma: no cover - except SSEError as sse_exc: # pragma: lax no cover - logger.exception("Encountered SSE exception") - raise sse_exc - except Exception as exc: # pragma: lax no cover - logger.exception("Error in sse_reader") - await read_stream_writer.send(exc) - finally: - await read_stream_writer.aclose() - - async def post_writer(endpoint_url: str): - try: - async with write_stream_reader, write_stream: - - async def _send_message(session_message: SessionMessage) -> None: - logger.debug(f"Sending client message: {session_message}") - response = await request_within_origin( - client, - "POST", - endpoint_url, - json=session_message.message.model_dump(by_alias=True, mode="json", exclude_unset=True), - ) - response.raise_for_status() - logger.debug(f"Client message sent successfully: {response.status_code}") - - async for session_message in write_stream_reader: - sender_ctx = write_stream_reader.last_context - if sender_ctx is not None: - async with anyio.create_task_group() as tg: - sender_ctx.run(tg.start_soon, _send_message, session_message) - else: - await _send_message(session_message) # pragma: no cover - except Exception: # pragma: lax no cover - logger.exception("Error in post_writer") - - # On Python 3.14, coverage.py reports a phantom branch arc on this - # line (->yield) when nested two async-with levels deep. The branch - # is the unreachable "did __aexit__ suppress?" arm for memory streams. - async with ( # pragma: no branch - read_stream_writer, - read_stream, - write_stream, - write_stream_reader, - anyio.create_task_group() as tg, - ): - endpoint_url = await tg.start(sse_reader) - logger.debug(f"Starting post writer with endpoint URL: {endpoint_url}") - tg.start_soon(post_writer, endpoint_url) - yield read_stream, write_stream - tg.cancel_scope.cancel() - await resync_tracer() +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/stdio.py b/src/mcp/client/stdio.py index 6a3ad12111..2cf67b2860 100644 --- a/src/mcp/client/stdio.py +++ b/src/mcp/client/stdio.py @@ -1,355 +1,65 @@ -"""stdio client transport. - -Runs an MCP server as a subprocess and exchanges newline-delimited JSON-RPC -messages with it over stdin/stdout. Two pipe tasks bridge the server's pipes -to the session's in-memory streams; shutdown follows the MCP spec sequence -(close stdin, wait, then kill the process tree) inside a cancellation shield -with every wait bounded, so a cancelled caller can neither leak a live server -process nor hang on one. -""" - -import logging -import os import sys -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager, suppress -from pathlib import Path -from typing import Literal, TextIO - -import anyio -import anyio.lowlevel -import anyio.to_thread -import mcp_types as types -from anyio.abc import AsyncResource, Process -from anyio.streams.text import TextReceiveStream -from pydantic import BaseModel, Field -from mcp.client._transport import TransportStreams -from mcp.os.posix.utilities import terminate_posix_process_tree -from mcp.os.win32.utilities import ( - ServerProcess, - close_process_job, - create_windows_process, - get_windows_executable_command, - terminate_windows_process_tree, +import mcp_client.client.stdio as _implementation +from mcp_client.client.stdio import ( + _EXIT_POLL_INTERVAL as _EXIT_POLL_INTERVAL, ) -from mcp.shared.message import SessionMessage - -logger = logging.getLogger(__name__) - -# Environment variables to inherit by default -DEFAULT_INHERITED_ENV_VARS = ( - [ - "APPDATA", - "HOMEDRIVE", - "HOMEPATH", - "LOCALAPPDATA", - "PATH", - "PATHEXT", - "PROCESSOR_ARCHITECTURE", - "SYSTEMDRIVE", - "SYSTEMROOT", - "TEMP", - "USERNAME", - "USERPROFILE", - ] - if sys.platform == "win32" - else ["HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER"] +from mcp_client.client.stdio import ( + _KILL_REAP_TIMEOUT as _KILL_REAP_TIMEOUT, +) +from mcp_client.client.stdio import ( + _WRITER_FLUSH_TIMEOUT as _WRITER_FLUSH_TIMEOUT, +) +from mcp_client.client.stdio import ( + DEFAULT_INHERITED_ENV_VARS as DEFAULT_INHERITED_ENV_VARS, +) +from mcp_client.client.stdio import ( + FORCE_KILL_TIMEOUT as FORCE_KILL_TIMEOUT, +) +from mcp_client.client.stdio import ( + PROCESS_TERMINATION_TIMEOUT as PROCESS_TERMINATION_TIMEOUT, +) +from mcp_client.client.stdio import ( + StdioServerParameters as StdioServerParameters, +) +from mcp_client.client.stdio import ( + _aclose_all as _aclose_all, +) +from mcp_client.client.stdio import ( + _close_pipe as _close_pipe, +) +from mcp_client.client.stdio import ( + _close_subprocess_transport as _close_subprocess_transport, +) +from mcp_client.client.stdio import ( + _create_platform_compatible_process as _create_platform_compatible_process, +) +from mcp_client.client.stdio import ( + _drain_stdout as _drain_stdout, +) +from mcp_client.client.stdio import ( + _get_executable_command as _get_executable_command, +) +from mcp_client.client.stdio import ( + _parse_line as _parse_line, +) +from mcp_client.client.stdio import ( + _stop_server_process as _stop_server_process, +) +from mcp_client.client.stdio import ( + _terminate_process_tree as _terminate_process_tree, +) +from mcp_client.client.stdio import ( + _wait_for_process_exit as _wait_for_process_exit, +) +from mcp_client.client.stdio import ( + get_default_environment as get_default_environment, +) +from mcp_client.client.stdio import ( + logger as logger, +) +from mcp_client.client.stdio import ( + stdio_client as stdio_client, ) -# Grace period for the server to exit on its own after its stdin closes. -PROCESS_TERMINATION_TIMEOUT = 2.0 - -# Extra time after SIGTERM before SIGKILL; POSIX only (Windows kills hard). -FORCE_KILL_TIMEOUT = 2.0 - -# Time for the event loop to observe a kill; only an unkillable process runs this out. -_KILL_REAP_TIMEOUT = 2.0 - -# Time for the writer to flush accepted messages before stdin closes. -_WRITER_FLUSH_TIMEOUT = 0.5 - -# How often to poll returncode while waiting for the process to die. -_EXIT_POLL_INTERVAL = 0.01 - - -def get_default_environment() -> dict[str, str]: - """Returns only the environment variables that are safe to inherit.""" - env: dict[str, str] = {} - - for key in DEFAULT_INHERITED_ENV_VARS: - value = os.environ.get(key) - if value is None: # pragma: lax no cover - continue - - if value.startswith("()"): # pragma: no cover - # Skip functions, which are a security risk - continue # pragma: no cover - - env[key] = value - - return env - - -class StdioServerParameters(BaseModel): - command: str - """The executable to run to start the server.""" - - args: list[str] = Field(default_factory=list) - """Command line arguments to pass to the executable.""" - - env: dict[str, str] | None = None - """Extra environment variables, merged over get_default_environment().""" - - cwd: str | Path | None = None - """The working directory to use when spawning the process.""" - - encoding: str = "utf-8" - """Text encoding for messages to and from the server.""" - - encoding_error_handler: Literal["strict", "ignore", "replace"] = "strict" - """Encoding error handler; see https://docs.python.org/3/library/codecs.html#error-handlers.""" - - -@asynccontextmanager -async def stdio_client( - server: StdioServerParameters, errlog: TextIO = sys.stderr -) -> AsyncGenerator[TransportStreams, None]: - """Spawns an MCP server subprocess and connects to it over stdin/stdout. - - Raises: - OSError: If the server process cannot be spawned. - ValueError: If the spawn parameters are invalid (embedded NUL bytes). - """ - command = await _get_executable_command(server.command) - - process = await _create_platform_compatible_process( - command=command, - args=server.args, - env=get_default_environment() | (server.env or {}), - errlog=errlog, - cwd=server.cwd, - ) - - # The spawn succeeded; no awaits until the task group is entered, or a - # cancellation delivered in the gap would leak the live process. - read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0) - write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0) - - shutting_down = False - writer_done = anyio.Event() - - async def stdout_reader() -> None: - assert process.stdout, "Opened process is missing stdout" - - stdout = TextReceiveStream(process.stdout, encoding=server.encoding, errors=server.encoding_error_handler) - try: - async with read_stream_writer: - try: - # One line at a time; no read-ahead while a delivery is blocked. - buffer = "" - async for chunk in stdout: - lines = (buffer + chunk).split("\n") - buffer = lines.pop() - for line in lines: - try: - await read_stream_writer.send(_parse_line(line)) - except (anyio.ClosedResourceError, anyio.BrokenResourceError): - return # the session is gone; only the drain below remains - finally: - await _drain_stdout(process) - except anyio.ClosedResourceError: - pass # our own shutdown closed the stdout stream under the read - except (anyio.BrokenResourceError, ConnectionError): - # Teardown noise during shutdown, a real failure otherwise; either way - # the session sees clean closure when the read stream closes. - if not shutting_down: - logger.exception("Reading from the MCP server's stdout failed mid-session") - - async def stdin_writer() -> None: - assert process.stdin, "Opened process is missing stdin" - - try: - async with write_stream_reader: - async for session_message in write_stream_reader: - json = session_message.message.model_dump_json(by_alias=True, exclude_unset=True) - data = (json + "\n").encode(encoding=server.encoding, errors=server.encoding_error_handler) - await process.stdin.send(data) - except (anyio.ClosedResourceError, anyio.BrokenResourceError, OSError): - # The server may still be alive: close the read stream so the session - # sees the connection end instead of a request hanging forever. - await read_stream_writer.aclose() - finally: - writer_done.set() - - async def shutdown() -> None: - """Winds the transport down: stop traffic, flush, stop the server, release the streams.""" - # Unblock the reader into its drain: a server stuck writing stdout cannot - # read its stdin, so draining is what lets the flush below complete. - read_stream.close() - # Bounded window for the writer to flush already-accepted messages. - write_stream.close() - with anyio.move_on_after(_WRITER_FLUSH_TIMEOUT) as flush_scope: - await writer_done.wait() - if flush_scope.cancelled_caught: - await anyio.lowlevel.cancel_shielded_checkpoint() # resync coverage on 3.11 (gh-106749) - await _stop_server_process(process) - await _aclose_all(read_stream, write_stream, read_stream_writer, write_stream_reader) - # One pass so unblocked tasks exit via their except paths before the cancel. - await anyio.lowlevel.checkpoint() - - async with anyio.create_task_group() as tg: - tg.start_soon(stdout_reader) - tg.start_soon(stdin_writer) - try: - yield read_stream, write_stream - finally: - shutting_down = True - # Shutdown must finish even under caller cancellation, or the server - # process would leak; every wait inside is bounded. (Native - # task.cancel() and the fallback's worker threads can still defeat it.) - with anyio.CancelScope(shield=True): - await shutdown() - # Unstick pipe tasks a kill survivor's open pipe end could still block. - tg.cancel_scope.cancel() - # The cancel lands via throw(); one yield resyncs 3.11 coverage (gh-106749). - await anyio.lowlevel.cancel_shielded_checkpoint() - - -def _parse_line(line: str) -> SessionMessage | Exception: - """Parses one stdout line, returning parse errors as values for the session to surface.""" - try: - message = types.jsonrpc_message_adapter.validate_json(line, by_name=False) - except ValueError as exc: - logger.exception("Failed to parse JSONRPC message from server") - return exc - return SessionMessage(message) - - -async def _drain_stdout(process: ServerProcess) -> None: - """Consumes and discards the server's remaining stdout. - - Keeps a server flushing buffered output from blocking on a full pipe and - missing its chance to exit; shielded, raw bytes, ends when shutdown closes - the pipe. - """ - assert process.stdout - with anyio.CancelScope(shield=True): - with suppress( - anyio.EndOfStream, - anyio.ClosedResourceError, - anyio.BrokenResourceError, - ConnectionError, - OSError, - ): - while True: - await process.stdout.receive() - - -async def _stop_server_process(process: ServerProcess) -> None: - """Closes stdin, waits out the grace period, then kills the whole tree. - - The escalation order is spec text; timeouts and tree-wide scope are SDK policy: - https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle#shutdown - """ - assert process.stdin and process.stdout, "server process is spawned with pipes" - - await _close_pipe(process.stdin) - if not await _wait_for_process_exit(process, PROCESS_TERMINATION_TIMEOUT): - await _terminate_process_tree(process) - # Until the event loop observes the death, the transport cannot close. - if not await _wait_for_process_exit(process, _KILL_REAP_TIMEOUT): - logger.warning("MCP server process %d is still alive after the kill escalation; abandoning it", process.pid) - - # Reaps surviving Windows job members now, not at GC; no-op on POSIX. - close_process_job(process) - # A kill survivor can hold the stdout pipe open; poison the reader anyway. - await _close_pipe(process.stdout) - _close_subprocess_transport(process) - - -async def _close_pipe(stream: AsyncResource) -> None: - """Closes a pipe stream, tolerating one already closed, broken, or contended.""" - with suppress(OSError, anyio.BrokenResourceError, anyio.ClosedResourceError): - await stream.aclose() - - -async def _wait_for_process_exit(process: ServerProcess, timeout: float) -> bool: - """Returns whether the process died within the timeout, by polling returncode. - - Not process.wait(): on asyncio 3.11+ it also waits for pipe EOF, and a - child that inherited the pipes makes an exited server look hung. - """ - deadline = anyio.current_time() + timeout - while process.returncode is None: - if anyio.current_time() >= deadline: - return False - await anyio.sleep(_EXIT_POLL_INTERVAL) - return True - - -async def _terminate_process_tree(process: ServerProcess) -> None: - """Kills the process and all its descendants. - - POSIX: SIGTERM to the process group, SIGKILL after FORCE_KILL_TIMEOUT. - Windows: immediate Job Object termination (already a hard kill). - """ - if sys.platform == "win32": # pragma: no cover - await terminate_windows_process_tree(process) - else: # pragma: lax no cover - # The Windows-only FallbackProcess never reaches the POSIX path. - assert isinstance(process, Process) - await terminate_posix_process_tree(process, FORCE_KILL_TIMEOUT) - - -def _close_subprocess_transport(process: ServerProcess) -> None: - """Closes the asyncio subprocess transport, if there is one. - - The transport otherwise stays open (and warns at GC) while a surviving - descendant holds a pipe end; nothing public exposes it, hence the attribute - walk. No-op on trio and the Windows fallback. - """ - transport = getattr(getattr(process, "_process", None), "_transport", None) - # Duck-typed: uvloop's UVProcessTransport is not an asyncio.SubprocessTransport. - close = getattr(transport, "close", None) - if callable(close): - # close() on <=3.12 can raise PermissionError re-killing a setuid child. - with suppress(PermissionError): - close() - - -async def _get_executable_command(command: str) -> str: - """Normalizes the command for the current platform.""" - if sys.platform == "win32": - return await anyio.to_thread.run_sync(get_windows_executable_command, command, abandon_on_cancel=True) - else: # pragma: lax no cover - return command - - -async def _create_platform_compatible_process( - command: str, - args: list[str], - env: dict[str, str] | None = None, - errlog: TextIO = sys.stderr, - cwd: Path | str | None = None, -) -> ServerProcess: - """Spawns the server in its own kill scope. - - A new session/process group on POSIX, a Job Object on Windows. - """ - if sys.platform == "win32": # pragma: no cover - return await create_windows_process(command, args, env, errlog, cwd) - else: # pragma: lax no cover - return await anyio.open_process( - [command, *args], - env=env, - stderr=errlog, - cwd=cwd, - start_new_session=True, - ) - - -async def _aclose_all(*streams: AsyncResource) -> None: - """Closes every given stream.""" - for stream in streams: - await stream.aclose() +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py index 82de50fd05..0795e123d6 100644 --- a/src/mcp/client/streamable_http.py +++ b/src/mcp/client/streamable_http.py @@ -1,758 +1,50 @@ -"""Implements StreamableHTTP transport for MCP clients.""" +import sys -from __future__ import annotations as _annotations - -import contextlib -import logging -from collections.abc import AsyncGenerator, Awaitable, Callable -from contextlib import asynccontextmanager -from dataclasses import dataclass - -import anyio -import httpx2 -from anyio.abc import TaskGroup -from httpx2 import EventSource, ServerSentEvent -from mcp_types import ( - CONNECTION_CLOSED, - INTERNAL_ERROR, - INVALID_REQUEST, - METHOD_NOT_FOUND, - PARSE_ERROR, - ErrorData, - JSONRPCError, - JSONRPCMessage, - JSONRPCNotification, - JSONRPCRequest, - JSONRPCResponse, - RequestId, - jsonrpc_message_adapter, +import mcp_client.client.streamable_http as _implementation +from mcp_client.client.streamable_http import ( + DEFAULT_RECONNECTION_DELAY_MS as DEFAULT_RECONNECTION_DELAY_MS, ) -from mcp_types.version import MODERN_PROTOCOL_VERSIONS -from pydantic import ValidationError - -from mcp.client._transport import TransportStreams -from mcp.shared._compat import resync_tracer -from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams -from mcp.shared._httpx_utils import ( - create_mcp_http_client, - redirect_location, - request_within_origin, - sse_within_origin, - stream_within_origin, +from mcp_client.client.streamable_http import ( + LAST_EVENT_ID as LAST_EVENT_ID, +) +from mcp_client.client.streamable_http import ( + MAX_RECONNECTION_ATTEMPTS as MAX_RECONNECTION_ATTEMPTS, +) +from mcp_client.client.streamable_http import ( + MCP_SESSION_ID as MCP_SESSION_ID, +) +from mcp_client.client.streamable_http import ( + RequestContext as RequestContext, +) +from mcp_client.client.streamable_http import ( + ResumptionError as ResumptionError, +) +from mcp_client.client.streamable_http import ( + SessionMessageOrError as SessionMessageOrError, +) +from mcp_client.client.streamable_http import ( + StreamableHTTPError as StreamableHTTPError, +) +from mcp_client.client.streamable_http import ( + StreamableHTTPTransport as StreamableHTTPTransport, +) +from mcp_client.client.streamable_http import ( + StreamReader as StreamReader, +) +from mcp_client.client.streamable_http import ( + StreamWriter as StreamWriter, +) +from mcp_client.client.streamable_http import ( + _InFlightPost as _InFlightPost, +) +from mcp_client.client.streamable_http import ( + _unfollowed_redirect as _unfollowed_redirect, +) +from mcp_client.client.streamable_http import ( + logger as logger, +) +from mcp_client.client.streamable_http import ( + streamable_http_client as streamable_http_client, ) -from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER -from mcp.shared.jsonrpc_dispatcher import cancelled_request_id_from_params -from mcp.shared.message import ClientMessageMetadata, SessionMessage - -logger = logging.getLogger(__name__) - - -# TODO(Marcelo): Put the TransportStreams in a module under shared, so we can import here. -SessionMessageOrError = SessionMessage | Exception -StreamWriter = ContextSendStream[SessionMessageOrError] -StreamReader = ContextReceiveStream[SessionMessage] - -MCP_SESSION_ID = "mcp-session-id" -LAST_EVENT_ID = "last-event-id" - -# Reconnection defaults -DEFAULT_RECONNECTION_DELAY_MS = 1000 # 1 second fallback when server doesn't provide retry -MAX_RECONNECTION_ATTEMPTS = 2 # Max retry attempts before giving up - - -class StreamableHTTPError(Exception): - """Base exception for StreamableHTTP transport errors.""" - - -class ResumptionError(StreamableHTTPError): - """Raised when resumption request is invalid.""" - - -def _unfollowed_redirect(response: httpx2.Response) -> str | None: - """Describe a redirect `stream_within_origin` left unfollowed, or None if `response` is not one.""" - location = redirect_location(response) - if location is None: - return None - if response.request.url.scheme == "https" and location.scheme == "http": - return ( - f"Redirect to {location} not followed: it would downgrade this HTTPS endpoint to plain HTTP.\n" - "The server is likely behind a TLS-terminating proxy whose forwarded headers it does not trust,\n" - f"often combined with a trailing-slash difference. Try {location.copy_with(scheme='https')} instead, " - "or fix the proxy settings." - ) - return f"Redirect to {location} not followed; use that URL as the endpoint if it is the intended server" - - -@dataclass -class RequestContext: - """Context for a request operation.""" - - client: httpx2.AsyncClient - session_id: str | None - session_message: SessionMessage - metadata: ClientMessageMetadata | None - read_stream_writer: StreamWriter - - -@dataclass(slots=True) -class _InFlightPost: - """A request POST in flight: its abort scope and the era it was sent under. - - `modern` is the negotiated-version cache as of this request's dequeue, so a - later cancel frame is interpreted under the era the request actually ran - with, not whatever the cache says by then. - """ - - scope: anyio.CancelScope - modern: bool - - -class StreamableHTTPTransport: - """StreamableHTTP client transport implementation.""" - - def __init__(self, url: str) -> None: - """Initialize the StreamableHTTP transport. - - Args: - url: The endpoint URL. - """ - self.url = url - self.session_id: str | None = None - # Captured from each stamped message's metadata, synchronously in the - # post_writer loop so the cache always reflects wire order (a POST task's - # scheduling is arbitrary). Reused on outbound HTTP that carries no - # per-message header (transport-internal GET/DELETE, and dispatcher-written - # response/error POSTs that bypass the session's stamp), and consulted by - # `_consume_modern_cancellation`. Cleared when an `initialize` message is - # dequeued so a probe-stamped value cannot leak onto the handshake. - self._protocol_version_header: str | None = None - # Every request's POST runs inside one of these so an outbound - # `notifications/cancelled` at 2026 can abort it; see - # `_consume_modern_cancellation`. Keys are verbatim-typed ("1" is not 1). - self._in_flight_posts: dict[RequestId, _InFlightPost] = {} - - def _prepare_headers(self) -> dict[str, str]: - """Build MCP-specific request headers for any outbound HTTP request. - - These are merged with the ``httpx2.AsyncClient`` defaults (these take - precedence). The cached ``MCP-Protocol-Version`` is included whenever - present so messages that don't pass through the session's stamp — - response/error POSTs, legacy cancel frames, transport-internal - GET/DELETE — still carry the negotiated version. Per-message headers - are layered on top by the caller. - """ - headers: dict[str, str] = { - "accept": "application/json, text/event-stream", - "content-type": "application/json", - } - if self.session_id: - headers[MCP_SESSION_ID] = self.session_id - if self._protocol_version_header: - headers[MCP_PROTOCOL_VERSION_HEADER] = self._protocol_version_header - return headers - - def _is_initialization_request(self, message: JSONRPCMessage) -> bool: - """Check if the message is an initialization request.""" - return isinstance(message, JSONRPCRequest) and message.method == "initialize" - - def _is_initialized_notification(self, message: JSONRPCMessage) -> bool: - """Check if the message is an initialized notification.""" - return isinstance(message, JSONRPCNotification) and message.method == "notifications/initialized" - - def _maybe_extract_session_id_from_response(self, response: httpx2.Response) -> None: - """Extract and store session ID from response headers.""" - new_session_id = response.headers.get(MCP_SESSION_ID) - if new_session_id: - self.session_id = new_session_id - logger.info(f"Received session ID: {self.session_id}") - - async def _handle_sse_event( - self, - sse: ServerSentEvent, - read_stream_writer: StreamWriter, - original_request_id: RequestId | None = None, - resumption_callback: Callable[[str], Awaitable[None]] | None = None, - ) -> bool: - """Handle an SSE event, returning True if the response is complete.""" - if sse.event == "message": - # Handle priming events (empty data with ID) for resumability - if not sse.data: - # Call resumption callback for priming events that have an ID - if sse.id and resumption_callback: - await resumption_callback(sse.id) - return False - try: - message = jsonrpc_message_adapter.validate_json(sse.data, by_name=False) - logger.debug(f"SSE message: {message}") - - # If this is a response and we have original_request_id, replace it - if original_request_id is not None and isinstance(message, JSONRPCResponse | JSONRPCError): - message.id = original_request_id - - session_message = SessionMessage(message) - await read_stream_writer.send(session_message) - - # Call resumption token callback if we have an ID - if sse.id and resumption_callback: - await resumption_callback(sse.id) - - # If this is a response or error return True indicating completion - # Otherwise, return False to continue listening - return isinstance(message, JSONRPCResponse | JSONRPCError) - - # Forwarding to a closed read stream lands here when the caller cancels mid-SSE - # (BrokenResourceError, not a parse failure); coverage is timing-dependent in the - # streaming story's modern HTTP cancellation leg. - except Exception as exc: # pragma: lax no cover - logger.exception("Error parsing SSE message") - if original_request_id is not None: - error_data = ErrorData(code=PARSE_ERROR, message=f"Failed to parse SSE message: {exc}") - error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=original_request_id, error=error_data)) - await read_stream_writer.send(error_msg) - return True - await read_stream_writer.send(exc) - return False - else: # pragma: no cover - logger.warning(f"Unknown SSE event: {sse.event}") - return False - - async def handle_get_stream(self, client: httpx2.AsyncClient, read_stream_writer: StreamWriter) -> None: - """Handle GET stream for server-initiated messages with auto-reconnect.""" - last_event_id: str | None = None - retry_interval_ms: int | None = None - attempt: int = 0 - - while attempt < MAX_RECONNECTION_ATTEMPTS: # pragma: no branch - try: - if not self.session_id: - return - - headers = self._prepare_headers() - if last_event_id: - headers[LAST_EVENT_ID] = last_event_id - - async with sse_within_origin(client, self.url, headers=headers) as event_source: - if (redirect := _unfollowed_redirect(event_source.response)) is not None: - # The same GET would be redirected again, so retrying cannot help. - logger.warning(f"GET stream not opened: {redirect}") - return - event_source.response.raise_for_status() - logger.debug("GET SSE connection established") - - async for sse in event_source: - # Track last event ID for reconnection - if sse.id: - last_event_id = sse.id - # Track retry interval from server - if sse.retry is not None: - retry_interval_ms = sse.retry - - await self._handle_sse_event(sse, read_stream_writer) - - # Stream ended normally (server closed) - reset attempt counter - attempt = 0 - - except Exception: - logger.debug("GET stream error", exc_info=True) - attempt += 1 - - if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover - logger.debug(f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded") - return - - # Wait before reconnecting - delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS - logger.info(f"GET stream disconnected, reconnecting in {delay_ms}ms...") - await anyio.sleep(delay_ms / 1000.0) - - async def _handle_resumption_request(self, ctx: RequestContext) -> None: - """Handle a resumption request using GET with SSE.""" - headers = self._prepare_headers() - if ctx.metadata and ctx.metadata.resumption_token: - headers[LAST_EVENT_ID] = ctx.metadata.resumption_token - else: - raise ResumptionError("Resumption request requires a resumption token") # pragma: no cover - - # Extract original request ID to map responses - original_request_id = None - if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch - original_request_id = ctx.session_message.message.id - - async with sse_within_origin(ctx.client, self.url, headers=headers) as event_source: - if (redirect := _unfollowed_redirect(event_source.response)) is not None: - logger.warning(redirect) - assert original_request_id is not None - await self._resolve_abandoned_request( - ctx.read_stream_writer, original_request_id, redirect, code=INVALID_REQUEST - ) - return - event_source.response.raise_for_status() - logger.debug("Resumption GET SSE connection established") - - async for sse in event_source: # pragma: no branch - is_complete = await self._handle_sse_event( - sse, - ctx.read_stream_writer, - original_request_id, - ctx.metadata.on_resumption_token_update if ctx.metadata else None, - ) - if is_complete: - await event_source.response.aclose() - break - - def _consume_modern_cancellation(self, session_message: SessionMessage) -> bool: - """Translate an outbound `notifications/cancelled` at 2026; True means "do not POST". - - The 2026 wire defines no client-to-server notifications over streamable - HTTP: closing a request's response stream IS its cancellation signal. - The dispatcher still emits the courtesy frame as its abandon signal - (every outbound cancel names one of our own request ids - the spec - forbids cancelling a request the sender did not issue), so this - transport translates it: when the named request's POST is in flight, - that POST's own recorded era decides - abort-and-swallow at 2026, POST - the frame below it (where the frame is the signal and a disconnect - explicitly is not). With no POST to consult, the cached negotiated - version decides; at 2026 the frame is swallowed even unmatched, so a - late cancel racing the response cannot leak onto the wire. - """ - message = session_message.message - if not (isinstance(message, JSONRPCNotification) and message.method == "notifications/cancelled"): - return False - request_id = cancelled_request_id_from_params(message.params) - post = self._in_flight_posts.get(request_id) if request_id is not None else None - if post is not None: - if not post.modern: - return False - logger.debug("aborting in-flight POST for cancelled request %r", request_id) - post.scope.cancel() - return True - return self._protocol_version_header in MODERN_PROTOCOL_VERSIONS - - async def _run_request_post( - self, - post_fn: Callable[[], Awaitable[None]], - post: _InFlightPost, - request_id: RequestId, - ) -> None: - """Run one request's POST inside its abort scope (see `_consume_modern_cancellation`).""" - try: - with post.scope: - await post_fn() - finally: - # Identity-guarded: a reused id may already have a successor - # registered while this task unwinds - popping by key alone would - # evict the live entry and leave the new POST unabortable. - if self._in_flight_posts.get(request_id) is post: - del self._in_flight_posts[request_id] - - async def _handle_post_request(self, ctx: RequestContext) -> None: - """Handle a POST request with response processing.""" - message = ctx.session_message.message - headers = self._prepare_headers() - if ctx.metadata is not None and ctx.metadata.headers is not None: - headers.update(ctx.metadata.headers) - - async with stream_within_origin( - ctx.client, - "POST", - self.url, - json=message.model_dump(by_alias=True, mode="json", exclude_unset=True), - headers=headers, - ) as response: - if response.status_code == 202: - logger.debug("Received 202 Accepted") - if isinstance(message, JSONRPCRequest): - # A request's response arrives on this POST's body; 202 says - # none will follow. Resolve rather than park the caller forever. - await self._resolve_abandoned_request( - ctx.read_stream_writer, - message.id, - "server answered a request with 202 Accepted", - code=INVALID_REQUEST, - ) - return - - if (redirect := _unfollowed_redirect(response)) is not None: - logger.warning(redirect) - if isinstance(message, JSONRPCRequest): - await self._resolve_abandoned_request( - ctx.read_stream_writer, message.id, redirect, code=INVALID_REQUEST - ) - return - - if response.status_code >= 400: - if isinstance(message, JSONRPCRequest): - # A spec-correct server may return the JSON-RPC error in the - # body at a non-2xx status (e.g. 400 for INVALID_PARAMS, 404 - # for METHOD_NOT_FOUND). Surface that error rather than the - # status-derived stand-in below. - if response.headers.get("content-type", "").lower().startswith("application/json"): - try: - body = await response.aread() - parsed = jsonrpc_message_adapter.validate_json(body, by_name=False) - if isinstance(parsed, JSONRPCError): - # The server may have set `id: null` (request rejected before its - # id was parsed); use this request's id so correlation works. - reply = JSONRPCError(jsonrpc="2.0", id=message.id, error=parsed.error) - await ctx.read_stream_writer.send(SessionMessage(reply)) - return - except (httpx2.StreamError, ValidationError): - pass - logger.debug("Non-2xx body was not a JSON-RPC error; using fallback") - if response.status_code == 404: - if self.session_id is None: - # No session yet → 404 is the HTTP-level spelling of - # METHOD_NOT_FOUND (gateway / legacy server doesn't know - # this method); "Session terminated" would be a lie here. - error_data = ErrorData(code=METHOD_NOT_FOUND, message="Not Found") - else: - error_data = ErrorData(code=INVALID_REQUEST, message="Session terminated") - else: - error_data = ErrorData(code=INTERNAL_ERROR, message="Server returned an error response") - session_message = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data)) - await ctx.read_stream_writer.send(session_message) - return - - if self._is_initialization_request(message): - self._maybe_extract_session_id_from_response(response) - - # Per https://modelcontextprotocol.io/specification/2025-06-18/basic#notifications: - # The server MUST NOT send a response to notifications. - if isinstance(message, JSONRPCRequest): - content_type = response.headers.get("content-type", "").lower() - if content_type.startswith("application/json"): - await self._handle_json_response(response, ctx.read_stream_writer, request_id=message.id) - elif content_type.startswith("text/event-stream"): - await self._handle_sse_response(response, ctx) - else: - logger.error(f"Unexpected content type: {content_type}") - error_data = ErrorData(code=INVALID_REQUEST, message=f"Unexpected content type: {content_type}") - error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data)) - await ctx.read_stream_writer.send(error_msg) - - async def _handle_json_response( - self, - response: httpx2.Response, - read_stream_writer: StreamWriter, - *, - request_id: RequestId, - ) -> None: - """Handle JSON response from the server.""" - try: - content = await response.aread() - message = jsonrpc_message_adapter.validate_json(content, by_name=False) - session_message = SessionMessage(message) - await read_stream_writer.send(session_message) - except (httpx2.StreamError, ValidationError) as exc: - logger.exception("Error parsing JSON response") - error_data = ErrorData(code=PARSE_ERROR, message=f"Failed to parse JSON response: {exc}") - error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=request_id, error=error_data)) - await read_stream_writer.send(error_msg) - - async def _handle_sse_response( - self, - response: httpx2.Response, - ctx: RequestContext, - ) -> None: - """Handle SSE response from the server.""" - last_event_id: str | None = None - retry_interval_ms: int | None = None - - # The caller (_handle_post_request) only reaches here inside - # isinstance(message, JSONRPCRequest), so this is always a JSONRPCRequest. - assert isinstance(ctx.session_message.message, JSONRPCRequest) - original_request_id = ctx.session_message.message.id - - try: - event_source = EventSource(response) - async for sse in event_source: # pragma: no branch - # Track last event ID for potential reconnection - if sse.id: - last_event_id = sse.id - - # Track retry interval from server - if sse.retry is not None: - retry_interval_ms = sse.retry - - is_complete = await self._handle_sse_event( - sse, - ctx.read_stream_writer, - original_request_id=original_request_id, - resumption_callback=(ctx.metadata.on_resumption_token_update if ctx.metadata else None), - ) - # If the SSE event indicates completion, like returning response/error - # break the loop - if is_complete: - await response.aclose() - return # Normal completion, no reconnect needed - except Exception: - logger.debug("SSE stream ended", exc_info=True) # pragma: lax no cover - - # Stream ended without response - reconnect if we received an event with ID - if last_event_id is not None: - logger.info("SSE stream disconnected, reconnecting...") - await self._handle_reconnection(ctx, last_event_id, retry_interval_ms) - else: - # Not resumable: resolve the waiter, else a listen stream's consumer - # would hang forever instead of learning the subscription is lost. - await self._resolve_abandoned_request( - ctx.read_stream_writer, original_request_id, "SSE stream ended without a response" - ) - - async def _resolve_abandoned_request( - self, read_stream_writer: StreamWriter, request_id: RequestId, message: str, *, code: int = CONNECTION_CLOSED - ) -> None: - """Resolve a request whose response can never arrive with a synthesized error. - - Best-effort: a closed read stream means the session is tearing down. - """ - error_data = ErrorData(code=code, message=message) - error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=request_id, error=error_data)) - try: - await read_stream_writer.send(error_msg) - except (anyio.BrokenResourceError, anyio.ClosedResourceError): - logger.debug("read stream closed before request %r could be resolved", request_id) - - async def _handle_reconnection( - self, - ctx: RequestContext, - last_event_id: str, - retry_interval_ms: int | None = None, - attempt: int = 0, - ) -> None: - """Reconnect with Last-Event-ID to resume stream after server disconnect.""" - # Only requests reconnect: every caller arrives from a request's response stream. - assert isinstance(ctx.session_message.message, JSONRPCRequest) - original_request_id = ctx.session_message.message.id - - if attempt >= MAX_RECONNECTION_ATTEMPTS: - # Resolve on give-up: a request with no read timeout (a listen - # stream) would otherwise hang its caller forever. - logger.debug(f"Max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded") - await self._resolve_abandoned_request( - ctx.read_stream_writer, original_request_id, "SSE stream ended and reconnection attempts were exhausted" - ) - return - - # Always wait - use server value or default - delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS - await anyio.sleep(delay_ms / 1000.0) - - headers = self._prepare_headers() - headers[LAST_EVENT_ID] = last_event_id - - try: - async with sse_within_origin(ctx.client, self.url, headers=headers) as event_source: - event_source.response.raise_for_status() - logger.info("Reconnected to SSE stream") - - # Track for potential further reconnection - reconnect_last_event_id: str = last_event_id - reconnect_retry_ms = retry_interval_ms - - async for sse in event_source: - if sse.id: # pragma: no branch - reconnect_last_event_id = sse.id - if sse.retry is not None: - reconnect_retry_ms = sse.retry - - is_complete = await self._handle_sse_event( - sse, - ctx.read_stream_writer, - original_request_id, - ctx.metadata.on_resumption_token_update if ctx.metadata else None, - ) - if is_complete: - await event_source.response.aclose() - return - - # Stream ended again without response - reconnect again (reset attempt counter) - logger.info("SSE stream disconnected, reconnecting...") - await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, 0) - except Exception as e: # pragma: no cover - logger.debug(f"Reconnection failed: {e}") - # Try to reconnect again if we still have an event ID - await self._handle_reconnection(ctx, last_event_id, retry_interval_ms, attempt + 1) - - async def post_writer( - self, - client: httpx2.AsyncClient, - write_stream_reader: StreamReader, - read_stream_writer: StreamWriter, - write_stream: ContextSendStream[SessionMessage], - start_get_stream: Callable[[], None], - tg: TaskGroup, - ) -> None: - """Handle writing requests to the server.""" - try: - async with write_stream_reader, read_stream_writer, write_stream: - - async def _handle_message(session_message: SessionMessage) -> None: - message = session_message.message - if self._consume_modern_cancellation(session_message): - return - metadata = ( - session_message.metadata - if isinstance(session_message.metadata, ClientMessageMetadata) - else None - ) - - # Check if this is a resumption request - is_resumption = bool(metadata and metadata.resumption_token) - - logger.debug(f"Sending client message: {message}") - - # Handle initialized notification - if self._is_initialized_notification(message): - start_get_stream() - - if self._is_initialization_request(message): - # `initialize` is the negotiation, not a "subsequent request" — discard any - # probe-stamped value so the discover→fallback path can't leak it onto the handshake. - self._protocol_version_header = None - elif metadata is not None and metadata.headers is not None: - stamped_version = metadata.headers.get(MCP_PROTOCOL_VERSION_HEADER) - if stamped_version is not None: - self._protocol_version_header = stamped_version - - ctx = RequestContext( - client=client, - session_id=self.session_id, - session_message=session_message, - metadata=metadata, - read_stream_writer=read_stream_writer, - ) - - async def handle_request_async(): - if is_resumption: - await self._handle_resumption_request(ctx) - else: - await self._handle_post_request(ctx) - - # If this is a request, start a new task to handle it - if isinstance(message, JSONRPCRequest): - # Register the abort scope before the spawn: the next - # message through this loop can already be the abandon - # signal for this id, ahead of the task ever running. - post = _InFlightPost( - scope=anyio.CancelScope(), - modern=self._protocol_version_header in MODERN_PROTOCOL_VERSIONS, - ) - superseded = self._in_flight_posts.get(message.id) - if superseded is not None: - # A reused id means the waiter belongs to this attempt now: - # sever the old POST so its zombie stream cannot answer, - # fail, or resolve the successor's request. - superseded.scope.cancel() - self._in_flight_posts[message.id] = post - tg.start_soon(self._run_request_post, handle_request_async, post, message.id) - else: - await handle_request_async() - - async for session_message in write_stream_reader: - sender_ctx = write_stream_reader.last_context - if sender_ctx is not None: - async with anyio.create_task_group() as tg_local: - sender_ctx.run(tg_local.start_soon, _handle_message, session_message) - else: - await _handle_message(session_message) # pragma: no cover - - except Exception: # pragma: lax no cover - logger.exception("Error in post_writer") - - async def terminate_session(self, client: httpx2.AsyncClient) -> None: - """Terminate the session by sending a DELETE request.""" - if not self.session_id: - return # pragma: no cover - - try: - headers = self._prepare_headers() - response = await request_within_origin(client, "DELETE", self.url, headers=headers) - - if response.status_code == 405: - logger.debug("Server does not allow session termination") - elif response.status_code not in (200, 204): - logger.warning(f"Session termination failed: {response.status_code}") # pragma: no cover - except Exception as exc: # pragma: no cover - logger.warning(f"Session termination failed: {exc}") - - -@asynccontextmanager -async def streamable_http_client( - url: str, - *, - http_client: httpx2.AsyncClient | None = None, - terminate_on_close: bool = True, -) -> AsyncGenerator[TransportStreams, None]: - """Client transport for StreamableHTTP. - - Args: - url: The MCP server endpoint URL. - http_client: Optional pre-configured httpx2.AsyncClient. If None, a default - client with recommended MCP timeouts will be created. To configure headers, - authentication, or other HTTP settings, create an httpx2.AsyncClient and pass it here. - Whichever client is used, MCP requests follow a redirect only when it stays on the - endpoint's origin (same scheme, host and port, or http to https on the same host with - default ports) and keeps the request method (307/308 for a POST; any status for the GET - stream); any other redirect is not followed and the message it answered fails with an - error naming the location. The - client's `follow_redirects` setting is not consulted; the SDK's OAuth providers apply the - same rule to the requests they make. - terminate_on_close: If True, send a DELETE request to terminate the session when the context exits. - - Yields: - Tuple containing: - - read_stream: Stream for reading messages from the server - - write_stream: Stream for sending messages to the server - - Example: - See examples/snippets/clients/ for usage patterns. - """ - # Determine if we need to create and manage the client - client_provided = http_client is not None - client = http_client - - if client is None: - # Create default client with recommended MCP timeouts - client = create_mcp_http_client() - - transport = StreamableHTTPTransport(url) - - logger.debug(f"Connecting to StreamableHTTP endpoint: {url}") - - async with contextlib.AsyncExitStack() as stack: - # Only manage client lifecycle if we created it - if not client_provided: - await stack.enter_async_context(client) - - read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0) - write_stream, write_stream_reader = create_context_streams[SessionMessage](0) - - async with ( - read_stream_writer, - read_stream, - write_stream, - write_stream_reader, - anyio.create_task_group() as tg, - ): - - def start_get_stream() -> None: - tg.start_soon(transport.handle_get_stream, client, read_stream_writer) - - tg.start_soon( - transport.post_writer, - client, - write_stream_reader, - read_stream_writer, - write_stream, - start_get_stream, - tg, - ) - try: - yield read_stream, write_stream - finally: - if transport.session_id and terminate_on_close: - await transport.terminate_session(client) - tg.cancel_scope.cancel() - await resync_tracer() +sys.modules[__name__] = _implementation diff --git a/src/mcp/client/subscriptions.py b/src/mcp/client/subscriptions.py index 27283909be..e169b7a4ae 100644 --- a/src/mcp/client/subscriptions.py +++ b/src/mcp/client/subscriptions.py @@ -1,282 +1,47 @@ -"""Client-side `subscriptions/listen` driver (2026-07-28, SEP-2575). +import sys -`listen()` opens the stream as an async context manager: entering waits for -the server's acknowledgment, iteration yields typed change events, a graceful -server close ends the loop, and an abrupt drop raises `SubscriptionLost`. -There is no replay and no automatic re-listen: a client that re-opens a -subscription refetches what it depends on. -""" - -from __future__ import annotations - -from collections.abc import AsyncIterator, Awaitable, Callable, Sequence -from contextlib import asynccontextmanager -from itertools import count -from typing import TYPE_CHECKING, Literal - -import anyio -import mcp_types as types -from mcp_types.version import MODERN_PROTOCOL_VERSIONS - -from mcp.shared.dispatcher import CallOptions -from mcp.shared.exceptions import MCPError -from mcp.shared.subscriptions import ( - PromptsListChanged, - ResourcesListChanged, - ResourceUpdated, - ServerEvent, - ToolsListChanged, - event_matches, +import mcp_client.client.subscriptions as _implementation +from mcp_client.client.subscriptions import ( + _MAX_PENDING_EVENTS as _MAX_PENDING_EVENTS, +) +from mcp_client.client.subscriptions import ( + ListenNotSupportedError as ListenNotSupportedError, +) +from mcp_client.client.subscriptions import ( + ListenRoute as ListenRoute, +) +from mcp_client.client.subscriptions import ( + OnEvent as OnEvent, +) +from mcp_client.client.subscriptions import ( + PromptsListChanged as PromptsListChanged, +) +from mcp_client.client.subscriptions import ( + ResourcesListChanged as ResourcesListChanged, +) +from mcp_client.client.subscriptions import ( + ResourceUpdated as ResourceUpdated, +) +from mcp_client.client.subscriptions import ( + ServerEvent as ServerEvent, +) +from mcp_client.client.subscriptions import ( + Subscription as Subscription, +) +from mcp_client.client.subscriptions import ( + SubscriptionLost as SubscriptionLost, +) +from mcp_client.client.subscriptions import ( + ToolsListChanged as ToolsListChanged, +) +from mcp_client.client.subscriptions import ( + _listen_ids as _listen_ids, +) +from mcp_client.client.subscriptions import ( + _SubscriptionEnd as _SubscriptionEnd, +) +from mcp_client.client.subscriptions import ( + listen as listen, ) -if TYPE_CHECKING: - from mcp.client.session import ClientSession - -__all__ = [ - "ListenNotSupportedError", - "OnEvent", - "PromptsListChanged", - "ResourceUpdated", - "ResourcesListChanged", - "ServerEvent", - "Subscription", - "SubscriptionLost", - "ToolsListChanged", - "listen", -] - -_listen_ids = count(1) -"""Process-wide `listen-N` sequence: string ids can never collide with a dispatcher's minted ints.""" - -_MAX_PENDING_EVENTS = 1024 -"""Backlog backstop: the spec allows sub-resource URIs, so distinct pending -`ResourceUpdated` events are unbounded; overflowing this cap settles the -subscription lost rather than growing client memory.""" - -_SubscriptionEnd = Literal["graceful", "lost", "local"] - - -class ListenNotSupportedError(RuntimeError): - """`subscriptions/listen` requires a 2026-07-28 connection.""" - - def __init__(self, negotiated_version: str | None) -> None: - self.negotiated_version = negotiated_version - super().__init__( - f"subscriptions/listen is not available at protocol version {negotiated_version!r}; it requires " - "2026-07-28. On earlier versions use subscribe_resource() and the change notifications delivered " - "through message_handler." - ) - - -class SubscriptionLost(RuntimeError): - """The stream ended without the server's graceful close; re-listen and refetch.""" - - -class ListenRoute: - """Package-internal demux state for one listen stream, fed synchronously in receive order by the session.""" - - def __init__(self) -> None: - self.honored: types.SubscriptionFilter | None = None - self.acked = anyio.Event() - self.error: MCPError | None = None - self.end: _SubscriptionEnd | None = None - self._honored_uris: frozenset[str] = frozenset() - self._pending: dict[ServerEvent, None] = {} - self._wake = anyio.Event() - - def set_acked(self, honored: types.SubscriptionFilter) -> None: - """Record the acknowledged filter; the first ack wins.""" - if not self.acked.is_set(): - self.honored = honored - self._honored_uris = frozenset(honored.resource_subscriptions or ()) - self.acked.set() - - def deliver(self, event: ServerEvent) -> None: - """Queue an event within the honored filter, deduplicated against the backlog. - - Any `ResourceUpdated` is admitted once URI subscriptions were honored at - all: the spec allows the stamped URI to be a sub-resource of a subscribed one. - """ - if self.end is not None or self.honored is None: - return - if isinstance(event, ResourceUpdated): - admitted = bool(self._honored_uris) - else: - admitted = event_matches(self.honored, self._honored_uris, event) - if not admitted or event in self._pending: - return - if len(self._pending) >= _MAX_PENDING_EVENTS: - self.settle( - "lost", - error=MCPError( - types.INTERNAL_ERROR, - f"subscription backlog exceeded {_MAX_PENDING_EVENTS} unconsumed events; re-listen and refetch", - ), - ) - return - self._pending[event] = None - self._wake.set() - - def settle(self, end: _SubscriptionEnd, error: MCPError | None = None) -> None: - """Record the stream's end; the first reason wins and wakes both waiters.""" - if self.end is None: - self.end = end - self.error = error - self.acked.set() - self._wake.set() - - async def next_event(self) -> ServerEvent | _SubscriptionEnd: - """Peek the next pending event, or the stream's end once the backlog drains. - - A "local" end short-circuits the backlog; the other endings drain it first, - so a graceful close never swallows events that preceded it. - """ - while True: - # Snapshot the wake event before checking state so a deliver landing after the checks cannot be missed. - wake = self._wake - if self.end == "local": - return self.end - if self._pending: - return next(iter(self._pending)) - if self.end is not None: - return self.end - await wake.wait() - self._wake = anyio.Event() - - def consume(self, event: ServerEvent) -> None: - """Remove a peeked event from the backlog.""" - self._pending.pop(event, None) - - -OnEvent = Callable[[ServerEvent], Awaitable[None]] -"""Per-event barrier awaited before a `Subscription` returns each event to its consumer.""" - - -class Subscription: - """One open `subscriptions/listen` stream: an async iterator of typed events. - - Produced by `listen()` / `Client.listen()`, not constructed directly. - """ - - def __init__( - self, - route: ListenRoute, - subscription_id: types.RequestId, - honored: types.SubscriptionFilter, - on_event: OnEvent | None = None, - ): - self._route = route - self._on_event = on_event - self.subscription_id = subscription_id - """The listen request's JSON-RPC id, stamped into every frame's `_meta`.""" - self.honored = honored - """The subset of the requested filter the server agreed to deliver.""" - - def __aiter__(self) -> Subscription: - return self - - async def __anext__(self) -> ServerEvent: - """Yield the next change event; the loop ends when the stream does. - - Raises: - SubscriptionLost: the stream dropped without the server's graceful close. - """ - outcome = await self._route.next_event() - if isinstance(outcome, str): - if outcome == "lost": - raise SubscriptionLost( - f"subscription {self.subscription_id!r} ended without the server's graceful close;" - " re-listen and refetch" - ) from self._route.error - raise StopAsyncIteration - if self._on_event is not None: - # The event stays pending while the barrier runs: a cancellation or a - # raising barrier leaves it for the next anext instead of dropping it. - await self._on_event(outcome) - self._route.consume(outcome) - return outcome - - -@asynccontextmanager -async def listen( - session: ClientSession, - *, - tools_list_changed: bool = False, - prompts_list_changed: bool = False, - resources_list_changed: bool = False, - resource_subscriptions: Sequence[str] = (), - on_event: OnEvent | None = None, -) -> AsyncIterator[Subscription]: - """Open one `subscriptions/listen` stream on `session` (2026-07-28 only). - - Entering sends the request and returns once the server's acknowledgment - arrives; exiting ends the subscription. `on_event` is awaited before each - event is returned - the seam `Client.listen` uses to finish cache eviction - before the consumer can refetch. - - Raises: - ListenNotSupportedError: negotiated version predates 2026-07-28. - MCPError: the server rejected the request, or the connection failed pre-ack. - SubscriptionLost: the stream ended before it was acknowledged. - TimeoutError: the session's read timeout elapsed before the acknowledgment. - """ - if session.protocol_version not in MODERN_PROTOCOL_VERSIONS: - raise ListenNotSupportedError(session.protocol_version) - if isinstance(resource_subscriptions, str): - raise TypeError("resource_subscriptions takes a sequence of URIs, not a bare string") - request = types.SubscriptionsListenRequest( - params=types.SubscriptionsListenRequestParams( - notifications=types.SubscriptionFilter( - tools_list_changed=tools_list_changed or None, - prompts_list_changed=prompts_list_changed or None, - resources_list_changed=resources_list_changed or None, - resource_subscriptions=list(resource_subscriptions) or None, - ) - ) - ) - task_group = session._task_group # pyright: ignore[reportPrivateUsage] - if task_group is None: - raise RuntimeError("listen() requires an entered session") - request_id: types.RequestId = f"listen-{next(_listen_ids)}" - data = request.model_dump(by_alias=True, mode="json", exclude_none=True) - opts: CallOptions = {"request_id": request_id} - session._stamp(data, opts) # pyright: ignore[reportPrivateUsage] - driver_scope = anyio.CancelScope() - - async def drive() -> None: - # Deliberately no result timeout: the response arrives when the stream ends. - with driver_scope: - try: - await session._dispatcher.send_raw_request( # pyright: ignore[reportPrivateUsage] - data["method"], data.get("params"), opts - ) - except MCPError as error: - route.settle("lost", error=error) - return - except ValueError as error: - # A raw request id collided with our minted listen id: fail this subscription - # and release the route in this same slice, so it cannot consume the raw caller's ack. - session._unregister_listen_route(request_id) # pyright: ignore[reportPrivateUsage] - route.settle("lost", error=MCPError(types.INTERNAL_ERROR, str(error))) - return - # A result, whatever its body, is the spec's graceful close; with no prior ack - # it opens the subscription already closed. - route.set_acked(types.SubscriptionFilter()) - route.settle("graceful") - - # Register the demux route before the request is written so the ack cannot race it. - route = session._register_listen_route(request_id) # pyright: ignore[reportPrivateUsage] - try: - task_group.start_soon(drive) - with anyio.fail_after(session._session_read_timeout_seconds): # pyright: ignore[reportPrivateUsage] - await route.acked.wait() - if route.honored is None: - # Only reachable on failure paths: a graceful no-ack result acked an empty filter in drive(). - if route.error is not None: - raise route.error - raise SubscriptionLost(f"subscription {request_id!r} ended before it was acknowledged") - yield Subscription(route, request_id, route.honored, on_event) - finally: - route.settle("local") - driver_scope.cancel() - session._unregister_listen_route(request_id) # pyright: ignore[reportPrivateUsage] +sys.modules[__name__] = _implementation diff --git a/src/mcp/os/__init__.py b/src/mcp/os/__init__.py index fa5dbc809c..a9a2c5b3bb 100644 --- a/src/mcp/os/__init__.py +++ b/src/mcp/os/__init__.py @@ -1 +1 @@ -"""Platform-specific utilities for MCP.""" +__all__ = [] diff --git a/src/mcp/os/posix/__init__.py b/src/mcp/os/posix/__init__.py index 23aff8bb02..33d48b46e0 100644 --- a/src/mcp/os/posix/__init__.py +++ b/src/mcp/os/posix/__init__.py @@ -1 +1,5 @@ -"""POSIX-specific utilities for MCP.""" +from mcp_client.os.posix import ( + utilities as utilities, +) + +__all__ = [] diff --git a/src/mcp/os/posix/utilities.py b/src/mcp/os/posix/utilities.py index d15be17194..606d68e426 100644 --- a/src/mcp/os/posix/utilities.py +++ b/src/mcp/os/posix/utilities.py @@ -1,63 +1,17 @@ -"""POSIX-specific functionality for stdio client operations.""" - -import logging -import os -import signal -from contextlib import suppress - -import anyio -from anyio.abc import Process - -logger = logging.getLogger(__name__) - -# How often to probe for surviving group members between SIGTERM and SIGKILL. -_GROUP_POLL_INTERVAL = 0.01 - - -async def terminate_posix_process_tree(process: Process, timeout_seconds: float = 2.0) -> None: - """Terminates a process and all its descendants on POSIX. - - SIGTERMs the process group, waits up to timeout_seconds for it to - disappear, then SIGKILLs whatever remains. killpg reaches every descendant - atomically, even ones whose parent already exited; daemonizers that left - the group escape by design. A group only disappears once every member is - dead and reaped, so a client running as PID 1 should reap orphans (e.g. - docker run --init) or the wait below runs its full timeout. - """ - # The leader's pid is the pgid (start_new_session). Never use getpgid(): - # it fails once the leader is reaped, even with live members left. - pgid = process.pid - - try: - os.killpg(pgid, signal.SIGTERM) - except ProcessLookupError: - return # the whole group is already gone - except PermissionError: - # EPERM never proves the group is gone (macOS raises it for zombie or - # foreign-euid members), so keep waiting and escalating. - logger.warning( - "No permission to signal some of process group %d; waiting for it to exit anyway", pgid, exc_info=True - ) - - with anyio.move_on_after(timeout_seconds): - while _group_alive(pgid): - # Reading returncode reaps the leader on trio; a zombie leader would - # otherwise keep the group alive for the full timeout. - _ = process.returncode - await anyio.sleep(_GROUP_POLL_INTERVAL) - return - - # ESRCH: died since the last probe. EPERM: we killed what we were allowed to. - with suppress(ProcessLookupError, PermissionError): - os.killpg(pgid, signal.SIGKILL) - - -def _group_alive(pgid: int) -> bool: - """Probes the group with signal 0; only ESRCH proves it is gone.""" - try: - os.killpg(pgid, 0) - except ProcessLookupError: - return False - except PermissionError: - pass # unsignalable survivors or unreaped zombies; EPERM is ambiguous - return True +import sys + +import mcp_client.os.posix.utilities as _implementation +from mcp_client.os.posix.utilities import ( + _GROUP_POLL_INTERVAL as _GROUP_POLL_INTERVAL, +) +from mcp_client.os.posix.utilities import ( + _group_alive as _group_alive, +) +from mcp_client.os.posix.utilities import ( + logger as logger, +) +from mcp_client.os.posix.utilities import ( + terminate_posix_process_tree as terminate_posix_process_tree, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/os/win32/__init__.py b/src/mcp/os/win32/__init__.py index f1ebab98df..a9a2c5b3bb 100644 --- a/src/mcp/os/win32/__init__.py +++ b/src/mcp/os/win32/__init__.py @@ -1 +1 @@ -"""Windows-specific utilities for MCP.""" +__all__ = [] diff --git a/src/mcp/os/win32/utilities.py b/src/mcp/os/win32/utilities.py index 321fda8a66..50c0c1bf62 100644 --- a/src/mcp/os/win32/utilities.py +++ b/src/mcp/os/win32/utilities.py @@ -1,292 +1,47 @@ -"""Windows-specific functionality for stdio transport operations.""" - -import logging -import shutil -import subprocess import sys -import weakref -from contextlib import suppress -from pathlib import Path -from typing import BinaryIO, TextIO, TypeAlias, cast - -import anyio -from anyio.abc import Process -from anyio.streams.file import FileReadStream, FileWriteStream - -logger = logging.getLogger(__name__) - -# Windows-specific imports for Job Objects -if sys.platform == "win32": - import msvcrt - - import pywintypes - import win32api - import win32con - import win32job -else: - # Type stubs for non-Windows platforms - win32api = None - win32con = None - msvcrt = None - win32job = None - pywintypes = None - - -def rebind_std_handle_to_fd(fd: int) -> None: - """Points the Win32 standard-handle slot for fd 0, 1, or 2 at fd's current OS handle. - - os.dup2 updates only the CRT descriptor table; subprocess handle inheritance - reads the Win32 slot, so it must be repointed too. - - Raises: - OSError: The slot could not be set. - """ - if sys.platform != "win32" or not win32api or not msvcrt or not pywintypes: - return - std_ids = {0: win32api.STD_INPUT_HANDLE, 1: win32api.STD_OUTPUT_HANDLE, 2: win32api.STD_ERROR_HANDLE} - try: - win32api.SetStdHandle(std_ids[fd], msvcrt.get_osfhandle(fd)) - except pywintypes.error as exc: - # Normalized so callers' OSError-based best-effort handling covers it. - raise OSError(f"SetStdHandle failed for fd {fd}") from exc - - -# How often FallbackProcess polls the underlying Popen for exit. -_EXIT_POLL_INTERVAL = 0.01 - -# Job Object handle per spawned process, for tree termination at shutdown. -# Values stay pywin32 PyHANDLEs: if no pop site ever runs, the dying weak entry -# drops the last reference and the PyHANDLE destructor closes the handle, which -# is what makes KILL_ON_JOB_CLOSE reap an abandoned tree. -_process_jobs: "weakref.WeakKeyDictionary[Process | FallbackProcess, object]" = weakref.WeakKeyDictionary() - - -def get_windows_executable_command(command: str) -> str: - """Resolves the command to a Windows executable path. - - Tries the bare name first, then the common script extensions (.cmd, .bat, - .exe, .ps1). - """ - try: - if command_path := shutil.which(command): - return command_path - - for ext in [".cmd", ".bat", ".exe", ".ps1"]: - ext_version = f"{command}{ext}" - if ext_path := shutil.which(ext_version): - return ext_path - - return command - except OSError: - return command # path probing failed (permissions, broken symlinks) - - -class FallbackProcess: - """Async wrapper around subprocess.Popen for SelectorEventLoop. - - Windows event loops without async subprocess support get this Popen-backed - fallback, with anyio file streams wrapping the pipes. - """ - - def __init__(self, popen_obj: subprocess.Popen[bytes]) -> None: - self.popen: subprocess.Popen[bytes] = popen_obj - stdin = popen_obj.stdin - stdout = popen_obj.stdout - - self.stdin = FileWriteStream(cast(BinaryIO, stdin)) if stdin else None - self.stdout = FileReadStream(cast(BinaryIO, stdout)) if stdout else None - - async def wait(self) -> int: - """Waits for exit by polling the Popen. - - A thread blocked in Popen.wait() cannot be cancelled by anyio, which - would defeat every timeout placed around this call. - """ - while (returncode := self.popen.poll()) is None: - await anyio.sleep(_EXIT_POLL_INTERVAL) - return returncode - - def terminate(self) -> None: - """Terminates the subprocess.""" - self.popen.terminate() - - def kill(self) -> None: - """Kills the subprocess (on Windows the same hard kill as terminate).""" - self.popen.kill() - - @property - def pid(self) -> int: - """Returns the process ID.""" - return self.popen.pid - - @property - def returncode(self) -> int | None: - """The exit code, or None while the process is still running. - - Polls the Popen so death is observable without anyone calling wait(). - """ - return self.popen.poll() - - -# The process handle stdio_client drives: anyio's Process, or the Popen-backed -# fallback used on Windows event loops without async subprocess support. -ServerProcess: TypeAlias = Process | FallbackProcess - - -async def create_windows_process( - command: str, - args: list[str], - env: dict[str, str] | None = None, - errlog: TextIO | None = sys.stderr, - cwd: Path | str | None = None, -) -> Process | FallbackProcess: - """Creates a subprocess with Job Object support for tree termination. - - Spawns via anyio's open_process; event loops without async subprocess - support (notably the SelectorEventLoop) raise NotImplementedError, in which - case the spawn falls back to a Popen-backed FallbackProcess. Either way the - process is then assigned to a Job Object so its children can be terminated - with it; children spawned before the assignment completes are not captured - (see the inline note below). - - Returns: - Process | FallbackProcess: The spawned process with async stdin/stdout streams. - """ - try: - process = await anyio.open_process( - [command, *args], - env=env, - # Ensure we don't create console windows for each process - creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0), - stderr=errlog, - cwd=cwd, - ) - except NotImplementedError: - # Windows event loops without async subprocess support (SelectorEventLoop) - process = await _create_windows_fallback_process(command, args, env, errlog, cwd) - - # Children spawned before the assignment completes land outside the job - # (membership is inherited at CreateProcess, never acquired retroactively); - # if that ever bites, the fix is a CREATE_SUSPENDED spawn -> assign -> resume. - job = _create_job_object() - _maybe_assign_process_to_job(process, job) - return process - - -async def _create_windows_fallback_process( - command: str, - args: list[str], - env: dict[str, str] | None = None, - errlog: TextIO | None = sys.stderr, - cwd: Path | str | None = None, -) -> FallbackProcess: - """Spawns via subprocess.Popen and wraps it in FallbackProcess.""" - popen_obj = subprocess.Popen( - [command, *args], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=errlog, - env=env, - cwd=cwd, - bufsize=0, # Unbuffered output - creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0), - ) - return FallbackProcess(popen_obj) - - -def _create_job_object() -> object | None: - """Creates a Windows Job Object configured to terminate all its processes when closed.""" - if sys.platform != "win32" or not win32api or not win32job: - return None - - job = None - try: - job = win32job.CreateJobObject(None, "") - extended_info = win32job.QueryInformationJobObject(job, win32job.JobObjectExtendedLimitInformation) - - extended_info["BasicLimitInformation"]["LimitFlags"] |= win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE - win32job.SetInformationJobObject(job, win32job.JobObjectExtendedLimitInformation, extended_info) - return job - except pywintypes.error: - logger.warning("Failed to create Job Object for process tree management", exc_info=True) - # If creation succeeded but configuration failed, close the handle now. - if job is not None: - _close_job_handle(job) - return None - - -def _maybe_assign_process_to_job(process: Process | FallbackProcess, job: object | None) -> None: - """Assigns the process to the job and records it for tree termination. - - On any failure the job handle is closed instead. - """ - if job is None: - return - - if sys.platform != "win32" or not win32api or not win32con or not win32job: - return - - try: - process_handle = win32api.OpenProcess( - win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, process.pid - ) - if not process_handle: - raise pywintypes.error(0, "OpenProcess", "Failed to open process handle") - - try: - win32job.AssignProcessToJobObject(job, process_handle) - finally: - win32api.CloseHandle(process_handle) - # Record only after the CloseHandle above succeeded: had it failed, the - # except below would close the job and KILL_ON_JOB_CLOSE takes the server. - _process_jobs[process] = job - except pywintypes.error: - logger.warning("Failed to assign process %d to Job Object", process.pid, exc_info=True) - _close_job_handle(job) - - -def close_process_job(process: Process | FallbackProcess) -> None: - """Closes the process's Job Object handle, if it still has one. - - KILL_ON_JOB_CLOSE makes the close also kill any members still alive, - deterministically rather than at GC time; a deliberate divergence from - POSIX, where a graceful server's children are left alive. - """ - if sys.platform != "win32": - return - - job = _process_jobs.pop(process, None) - if job is not None: - _close_job_handle(job) - - -async def terminate_windows_process_tree(process: Process | FallbackProcess) -> None: - """Terminates the process's job, or just the process if it has no job. - - Job termination is an immediate hard kill of every member. Windows has no - tree-wide SIGTERM; the stdin-close grace period is the server's chance to - exit cleanly. - """ - if sys.platform != "win32": - return - - job = _process_jobs.pop(process, None) - if job is not None and win32job: - try: - with suppress(pywintypes.error): # the job might already be terminated - win32job.TerminateJobObject(job, 1) - finally: - _close_job_handle(job) - - # The process may have no job (creation or assignment failed); kill it directly too. - try: - process.terminate() - except OSError: - pass - -def _close_job_handle(job: object) -> None: - """Closes a Job Object handle, tolerating one that is already closed.""" - if win32api and pywintypes: - with suppress(pywintypes.error): - win32api.CloseHandle(job) +import mcp_client.os.win32.utilities as _implementation +from mcp_client.os.win32.utilities import ( + _EXIT_POLL_INTERVAL as _EXIT_POLL_INTERVAL, +) +from mcp_client.os.win32.utilities import ( + FallbackProcess as FallbackProcess, +) +from mcp_client.os.win32.utilities import ( + ServerProcess as ServerProcess, +) +from mcp_client.os.win32.utilities import ( + _close_job_handle as _close_job_handle, +) +from mcp_client.os.win32.utilities import ( + _create_job_object as _create_job_object, +) +from mcp_client.os.win32.utilities import ( + _create_windows_fallback_process as _create_windows_fallback_process, +) +from mcp_client.os.win32.utilities import ( + _maybe_assign_process_to_job as _maybe_assign_process_to_job, +) +from mcp_client.os.win32.utilities import ( + _process_jobs as _process_jobs, +) +from mcp_client.os.win32.utilities import ( + close_process_job as close_process_job, +) +from mcp_client.os.win32.utilities import ( + create_windows_process as create_windows_process, +) +from mcp_client.os.win32.utilities import ( + get_windows_executable_command as get_windows_executable_command, +) +from mcp_client.os.win32.utilities import ( + logger as logger, +) +from mcp_client.os.win32.utilities import ( + rebind_std_handle_to_fd as rebind_std_handle_to_fd, +) +from mcp_client.os.win32.utilities import ( + terminate_windows_process_tree as terminate_windows_process_tree, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 8a886dcc24..11615f2c3a 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -40,11 +40,12 @@ async def main(): import logging import warnings from collections.abc import AsyncIterator, Awaitable, Callable, Mapping -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager from dataclasses import dataclass from functools import cached_property from typing import Any, Generic, overload +import anyio import mcp_types as types from mcp_types.version import MODERN_PROTOCOL_VERSIONS from pydantic import BaseModel @@ -63,7 +64,7 @@ async def main(): from mcp.server.caching import CacheableMethod, CacheHint, validate_cache_hints from mcp.server.context import HandlerResult, ServerMiddleware, ServerRequestContext from mcp.server.models import InitializationOptions -from mcp.server.runner import serve_dual_era_loop +from mcp.server.runner import modern_on_request, serve_dual_era_loop from mcp.server.streamable_http import EventStore from mcp.server.streamable_http_manager import ( DEFAULT_MAX_SESSIONS, @@ -73,7 +74,10 @@ async def main(): ) from mcp.server.transport_security import DEFAULT_MAX_REQUEST_BODY_SIZE, TransportSecuritySettings from mcp.shared._stream_protocols import ReadStream, WriteStream +from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair +from mcp.shared.dispatcher import Dispatcher from mcp.shared.exceptions import MCPDeprecationWarning +from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher from mcp.shared.message import SessionMessage logger = logging.getLogger(__name__) @@ -126,6 +130,10 @@ async def _ping_handler(ctx: ServerRequestContext[Any], params: types.RequestPar return types.EmptyResult() +async def _no_inbound_client_notifications(_dctx: Any, _method: str, _params: Mapping[str, Any] | None) -> None: + """Modern in-process clients send no notifications; cancellation uses the caller's cancel scope.""" + + class Server(Generic[LifespanResultT]): @overload def __init__( @@ -689,6 +697,24 @@ def session_manager(self) -> StreamableHTTPSessionManager: ) return self._session_manager + async def __mcp_client_connect__( + self, exit_stack: AsyncExitStack, mode: str, raise_exceptions: bool + ) -> Dispatcher[Any]: + """Connect a client to this server without a network transport.""" + if mode == "legacy": + from mcp.client._memory import InMemoryTransport + + transport = InMemoryTransport(self, raise_exceptions=raise_exceptions) + read_stream, write_stream = await exit_stack.enter_async_context(transport) + return JSONRPCDispatcher(read_stream, write_stream) + lifespan_state = await exit_stack.enter_async_context(self.lifespan(self)) + client_disp, server_disp = create_direct_dispatcher_pair(raise_handler_exceptions=raise_exceptions) + tg = await exit_stack.enter_async_context(anyio.create_task_group()) + exit_stack.callback(server_disp.close) + on_request = modern_on_request(self, lifespan_state) + await tg.start(server_disp.run, on_request, _no_inbound_client_notifications) + return client_disp + async def run( self, read_stream: ReadStream[SessionMessage | Exception], diff --git a/src/mcp/server/mcpserver/server.py b/src/mcp/server/mcpserver/server.py index fbd2c26dd8..19fc612c89 100644 --- a/src/mcp/server/mcpserver/server.py +++ b/src/mcp/server/mcpserver/server.py @@ -5,7 +5,7 @@ import base64 import inspect from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping, Sequence -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager from typing import Any, Generic, Literal, TypeVar, cast, overload import anyio @@ -100,6 +100,7 @@ ) from mcp.server.subscriptions import InMemorySubscriptionBus, ListenHandler, SubscriptionBus from mcp.server.transport_security import DEFAULT_MAX_REQUEST_BODY_SIZE, TransportSecuritySettings +from mcp.shared.dispatcher import Dispatcher from mcp.shared.exceptions import MCPError from mcp.shared.uri_template import UriTemplate @@ -267,6 +268,12 @@ def __init__( self._apply_extension(extension) self._install_extension_interceptor() + async def __mcp_client_connect__( + self, exit_stack: AsyncExitStack, mode: str, raise_exceptions: bool + ) -> Dispatcher[Any]: + """Connect a client to this server without a network transport.""" + return await self._lowlevel_server.__mcp_client_connect__(exit_stack, mode, raise_exceptions) + @property def name(self) -> str: return self._lowlevel_server.name diff --git a/src/mcp/shared/__init__.py b/src/mcp/shared/__init__.py index e69de29bb2..a9a2c5b3bb 100644 --- a/src/mcp/shared/__init__.py +++ b/src/mcp/shared/__init__.py @@ -0,0 +1 @@ +__all__ = [] diff --git a/src/mcp/shared/_callable_inspection.py b/src/mcp/shared/_callable_inspection.py index 0e89e446f8..1b42eecf04 100644 --- a/src/mcp/shared/_callable_inspection.py +++ b/src/mcp/shared/_callable_inspection.py @@ -1,33 +1,14 @@ -"""Callable inspection utilities. - -Adapted from Starlette's `is_async_callable` implementation. -https://github.com/encode/starlette/blob/main/starlette/_utils.py -""" - -from __future__ import annotations - -import functools -import inspect -from collections.abc import Awaitable, Callable -from typing import Any, TypeGuard, TypeVar, overload - -T = TypeVar("T") - -AwaitableCallable = Callable[..., Awaitable[T]] - - -@overload -def is_async_callable(obj: AwaitableCallable[T]) -> TypeGuard[AwaitableCallable[T]]: ... - - -@overload -def is_async_callable(obj: Any) -> TypeGuard[AwaitableCallable[Any]]: ... - - -def is_async_callable(obj: Any) -> Any: - while isinstance(obj, functools.partial): # pragma: lax no cover - obj = obj.func - - return inspect.iscoroutinefunction(obj) or ( - callable(obj) and inspect.iscoroutinefunction(getattr(obj, "__call__", None)) - ) +import sys + +import mcp_client.shared._callable_inspection as _implementation +from mcp_client.shared._callable_inspection import ( + AwaitableCallable as AwaitableCallable, +) +from mcp_client.shared._callable_inspection import ( + T as T, +) +from mcp_client.shared._callable_inspection import ( + is_async_callable as is_async_callable, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/_compat.py b/src/mcp/shared/_compat.py index 88d50ba20a..9793237872 100644 --- a/src/mcp/shared/_compat.py +++ b/src/mcp/shared/_compat.py @@ -1,19 +1,8 @@ -"""Workarounds for CPython interpreter bugs the SDK papers over.""" +import sys -import anyio.lowlevel +import mcp_client.shared._compat as _implementation +from mcp_client.shared._compat import ( + resync_tracer as resync_tracer, +) -__all__ = ["resync_tracer"] - - -async def resync_tracer() -> None: - """Resync coverage tracing after a cancelled task-group join. - - A cancel delivered at a join resumes the awaiting coroutine chain via - `coro.throw()`; on CPython 3.11 (python/cpython#106749) that drops the - `'call'` trace events for the outer frames and desyncs coverage's CTracer - until the chain next suspends and resumes normally. Yielding once here - resumes via `.send()`, which re-stamps the missing events. Shielded so a - pending outer cancel is not re-delivered at this point; behaviorally a - no-op. Delete this module when Python 3.11 support ends (EOL 2027-10). - """ - await anyio.lowlevel.cancel_shielded_checkpoint() +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/_context_streams.py b/src/mcp/shared/_context_streams.py index 04c33306d9..477e9c10e6 100644 --- a/src/mcp/shared/_context_streams.py +++ b/src/mcp/shared/_context_streams.py @@ -1,119 +1,20 @@ -"""Context-aware memory stream wrappers. - -anyio memory streams do not propagate ``contextvars.Context`` across task -boundaries. These thin wrappers capture the sender's context at ``send()`` -time and expose it on the receive side via ``last_context``, so consumers -can restore it with ``ctx.run(handler, item)``. - -The iteration interface is unchanged (yields ``T``, not tuples), keeping -these wrappers duck-type compatible with plain ``MemoryObjectSendStream`` -and ``MemoryObjectReceiveStream``. -""" - -from __future__ import annotations - -import contextvars -from types import TracebackType -from typing import Any, Generic, TypeVar - -import anyio -from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream - -T = TypeVar("T") - -# Internal payload carried through the underlying raw stream. -_Envelope = tuple[contextvars.Context, T] - - -class ContextSendStream(Generic[T]): - """Send-side wrapper that snapshots ``contextvars.copy_context()`` on every ``send()``.""" - - __slots__ = ("_inner",) - - def __init__(self, inner: MemoryObjectSendStream[_Envelope[T]]) -> None: - self._inner = inner - - async def send(self, item: T) -> None: - await self._inner.send((contextvars.copy_context(), item)) - - def close(self) -> None: - self._inner.close() - - async def aclose(self) -> None: - await self._inner.aclose() - - def clone(self) -> ContextSendStream[T]: # pragma: no cover - return ContextSendStream(self._inner.clone()) - - async def __aenter__(self) -> ContextSendStream[T]: - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: TracebackType | None, - ) -> bool | None: - await self.aclose() - return None - - -class ContextReceiveStream(Generic[T]): - """Receive-side wrapper that yields ``T`` and stores the sender's context in ``last_context``.""" - - __slots__ = ("_inner", "last_context") - - def __init__(self, inner: MemoryObjectReceiveStream[_Envelope[T]]) -> None: - self._inner = inner - self.last_context: contextvars.Context | None = None - - async def receive(self) -> T: - ctx, item = await self._inner.receive() - self.last_context = ctx - return item - - def close(self) -> None: - self._inner.close() - - async def aclose(self) -> None: - await self._inner.aclose() - - def clone(self) -> ContextReceiveStream[T]: # pragma: no cover - return ContextReceiveStream(self._inner.clone()) - - def __aiter__(self) -> ContextReceiveStream[T]: - return self - - async def __anext__(self) -> T: - try: - return await self.receive() - except anyio.EndOfStream: - raise StopAsyncIteration - - async def __aenter__(self) -> ContextReceiveStream[T]: - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: TracebackType | None, - ) -> bool | None: - await self.aclose() - return None - - -class create_context_streams( - tuple[ContextSendStream[T], ContextReceiveStream[T]], -): - """Create context-aware memory object streams. - - Supports ``create_context_streams[T](n)`` bracket syntax, - matching anyio's ``create_memory_object_stream`` API style. - """ - - def __new__(cls, max_buffer_size: float = 0) -> tuple[ContextSendStream[T], ContextReceiveStream[T]]: # type: ignore[type-var] - raw_send: MemoryObjectSendStream[Any] - raw_receive: MemoryObjectReceiveStream[Any] - raw_send, raw_receive = anyio.create_memory_object_stream(max_buffer_size) - return (ContextSendStream(raw_send), ContextReceiveStream(raw_receive)) +import sys + +import mcp_client.shared._context_streams as _implementation +from mcp_client.shared._context_streams import ( + ContextReceiveStream as ContextReceiveStream, +) +from mcp_client.shared._context_streams import ( + ContextSendStream as ContextSendStream, +) +from mcp_client.shared._context_streams import ( + T as T, +) +from mcp_client.shared._context_streams import ( + _Envelope as _Envelope, +) +from mcp_client.shared._context_streams import ( + create_context_streams as create_context_streams, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/_httpx_utils.py b/src/mcp/shared/_httpx_utils.py index 940b9f08cc..74e2e0e602 100644 --- a/src/mcp/shared/_httpx_utils.py +++ b/src/mcp/shared/_httpx_utils.py @@ -1,218 +1,47 @@ -"""Utilities for creating and using httpx2 AsyncClient instances in the MCP transports.""" - -from abc import ABC, abstractmethod -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager -from typing import Any, Protocol - -import httpx2 - -__all__ = ["create_mcp_http_client", "MCP_DEFAULT_TIMEOUT", "MCP_DEFAULT_SSE_READ_TIMEOUT"] - -# Default MCP timeout configuration -MCP_DEFAULT_TIMEOUT = 30.0 # General operations (seconds) -MCP_DEFAULT_SSE_READ_TIMEOUT = 300.0 # SSE streams - 5 minutes (seconds) - -# The headers httpx2.AsyncClient.sse() adds to an event-stream request. -_SSE_HEADERS = {"Accept": "text/event-stream", "Cache-Control": "no-store"} - -# How many redirects one auth-flow request may follow within its origin (see RedirectAwareAuth). -_AUTH_REDIRECT_LIMIT = 5 - - -class McpHttpClientFactory(Protocol): # pragma: no branch - def __call__( # pragma: no branch - self, - headers: dict[str, str] | None = None, - timeout: httpx2.Timeout | None = None, - auth: httpx2.Auth | None = None, - ) -> httpx2.AsyncClient: ... - - -def create_mcp_http_client( - headers: dict[str, str] | None = None, - timeout: httpx2.Timeout | None = None, - auth: httpx2.Auth | None = None, -) -> httpx2.AsyncClient: - """Create an httpx2 AsyncClient with the MCP transports' default timeouts. - - The client uses a 30-second timeout for connect/write/pool and a 300-second - read timeout, because a server may hold a response stream open. Redirect - following is left at the httpx2 default (off): the MCP transports follow - redirects within the endpoint's origin themselves, see `stream_within_origin`. - - Args: - headers: Optional headers to include with all requests. - timeout: Request timeout as httpx2.Timeout object. Defaults to 30s for - connect/write/pool and 300s for read (for long-lived SSE streams). - auth: Optional authentication handler. - - Returns: - Configured httpx2.AsyncClient instance. - - Note: - The returned AsyncClient must be used as a context manager to ensure - proper cleanup of connections. - """ - if timeout is None: - timeout = httpx2.Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT) - kwargs: dict[str, Any] = {"timeout": timeout} - if headers is not None: - kwargs["headers"] = headers - if auth is not None: # pragma: no cover - kwargs["auth"] = auth - return httpx2.AsyncClient(**kwargs) - - -def _within_origin(url: httpx2.URL, location: httpx2.URL) -> bool: - """Whether `location` is on `url`'s origin, or is its https upgrade on the default ports. - - httpx2 normalises a scheme's default port to None and lower-cases hosts, so - plain tuple comparison is exact. The upgrade rule is the one httpx2 itself - uses to decide a redirect has not left the origin (`_is_https_redirect`). - """ - if (url.scheme, url.host, url.port) == (location.scheme, location.host, location.port): - return True - return ( - url.host == location.host - and url.scheme == "http" - and url.port is None - and location.scheme == "https" - and location.port is None - ) - - -def next_request_within_origin(response: httpx2.Response) -> httpx2.Request | None: - """The request that follows `response`'s redirect, if it is one the MCP transports follow. - - That is when httpx2 built a next request for it (a redirect status with a - Location), the next request keeps the method (307/308, or any redirect of a - GET: httpx2 turns a POST into a body-less GET for 301/302/303, which would - drop the message), its URL stays within the origin of the request just sent - (same scheme, host and port, or http to https on the same host with default - ports), and the Location does not bring userinfo of its own (which httpx2 - would otherwise send as Basic auth; userinfo the configured URL already had - is kept by a relative Location and is fine). None for anything else, - including a non-redirect. - """ - next_request = response.next_request - if next_request is None: - return None - sent = response.request - if ( - next_request.method != sent.method - or (next_request.url.userinfo and next_request.url.userinfo != sent.url.userinfo) - or not _within_origin(sent.url, next_request.url) - ): - return None - return next_request - - -@asynccontextmanager -async def stream_within_origin( - client: httpx2.AsyncClient, method: str, url: httpx2.URL | str, **kwargs: Any -) -> AsyncGenerator[httpx2.Response]: - """`client.stream(...)`, following redirects only while they stay within the request's origin. - - An MCP transport talks to one configured endpoint, and everything on a request - (headers, auth, body) was configured for that endpoint. A redirect that - `next_request_within_origin` accepts, such as a 307/308 trailing-slash - normalisation, is followed, at most `client.max_redirects` times. Any other - redirect (or one past that budget) is not followed: the redirect response - itself is yielded, the way httpx2 hands one back when `follow_redirects` is - off, and the caller treats it as the non-success it is. The client's own - `follow_redirects` setting is not consulted. Requests an `httpx2.Auth` flow - makes during the call are sent without following either; the SDK's OAuth - providers apply the same rule to their own requests. - """ - request = client.build_request(method, url, **kwargs) - followed = 0 - while True: - response = await client.send(request, stream=True, follow_redirects=False) - next_request = next_request_within_origin(response) - if next_request is None or followed == client.max_redirects: - break - try: - # Drain the redirect body so the connection returns to the pool, as httpx2 does when it follows. - await response.aread() - finally: - await response.aclose() - request = next_request - followed += 1 - try: - yield response - finally: - await response.aclose() - - -async def request_within_origin( - client: httpx2.AsyncClient, method: str, url: httpx2.URL | str, **kwargs: Any -) -> httpx2.Response: - """`client.request(...)` with the redirect handling of `stream_within_origin`.""" - async with stream_within_origin(client, method, url, **kwargs) as response: - await response.aread() - return response - - -@asynccontextmanager -async def sse_within_origin( - client: httpx2.AsyncClient, url: httpx2.URL | str, *, headers: dict[str, str] | None = None -) -> AsyncGenerator[httpx2.EventSource]: - """`client.sse(url)` with the redirect handling of `stream_within_origin`.""" - merged = httpx2.Headers(_SSE_HEADERS) - merged.update(headers or {}) - async with stream_within_origin(client, "GET", url, headers=merged) as response: - yield httpx2.EventSource(response) - - -def redirect_location(response: httpx2.Response) -> httpx2.URL | None: - """Where `response` redirects to, for use in a message: without userinfo, query or fragment, - which can carry state that does not belong in an error or a log line. None if not a redirect.""" - if response.next_request is None: - return None - return response.next_request.url.copy_with(userinfo=b"", query=None, fragment=None) - - -def redirect_note(response: httpx2.Response) -> str: - """A suffix naming the location of a redirect response that was not followed, else empty.""" - location = redirect_location(response) - if location is None: - return "" - return f" (redirected to {location}; not followed)" - - -class RedirectAwareAuth(ABC, httpx2.Auth): - """An `httpx2.Auth` whose own requests follow redirects the way MCP transport requests do. - - The transports send every request with redirect following off and follow a - redirect themselves only within the endpoint's origin (`stream_within_origin`). - httpx2 applies that per-request setting to the requests an auth flow makes - too (metadata discovery, registration, token), so on their own those would - follow nothing. Subclasses write their flow as `_auth_flow`; this class - drives it and, for each request the flow makes other than the one being - authenticated, follows a redirect that `next_request_within_origin` accepts, - up to `_AUTH_REDIRECT_LIMIT` times. Any other redirect response is handed - to the flow as it is. - """ - - @abstractmethod - def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: - """The subclass's flow, written as `httpx2.Auth.async_auth_flow` otherwise would be.""" - - async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: - flow = self._auth_flow(request) - try: - outgoing = await flow.__anext__() - while True: - response = yield outgoing - if outgoing is not request: - for _ in range(_AUTH_REDIRECT_LIMIT): - follow = next_request_within_origin(response) - if follow is None: - break - response = yield follow - outgoing = await flow.asend(response) - except StopAsyncIteration: - return - finally: - await flow.aclose() +import sys + +import mcp_client.shared._httpx_utils as _implementation +from mcp_client.shared._httpx_utils import ( + _AUTH_REDIRECT_LIMIT as _AUTH_REDIRECT_LIMIT, +) +from mcp_client.shared._httpx_utils import ( + _SSE_HEADERS as _SSE_HEADERS, +) +from mcp_client.shared._httpx_utils import ( + MCP_DEFAULT_SSE_READ_TIMEOUT as MCP_DEFAULT_SSE_READ_TIMEOUT, +) +from mcp_client.shared._httpx_utils import ( + MCP_DEFAULT_TIMEOUT as MCP_DEFAULT_TIMEOUT, +) +from mcp_client.shared._httpx_utils import ( + McpHttpClientFactory as McpHttpClientFactory, +) +from mcp_client.shared._httpx_utils import ( + RedirectAwareAuth as RedirectAwareAuth, +) +from mcp_client.shared._httpx_utils import ( + _within_origin as _within_origin, +) +from mcp_client.shared._httpx_utils import ( + create_mcp_http_client as create_mcp_http_client, +) +from mcp_client.shared._httpx_utils import ( + next_request_within_origin as next_request_within_origin, +) +from mcp_client.shared._httpx_utils import ( + redirect_location as redirect_location, +) +from mcp_client.shared._httpx_utils import ( + redirect_note as redirect_note, +) +from mcp_client.shared._httpx_utils import ( + request_within_origin as request_within_origin, +) +from mcp_client.shared._httpx_utils import ( + sse_within_origin as sse_within_origin, +) +from mcp_client.shared._httpx_utils import ( + stream_within_origin as stream_within_origin, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/_otel.py b/src/mcp/shared/_otel.py index b7b05b11ab..b708895d84 100644 --- a/src/mcp/shared/_otel.py +++ b/src/mcp/shared/_otel.py @@ -1,60 +1,17 @@ -"""OpenTelemetry helpers for MCP.""" - -from __future__ import annotations - -from collections.abc import Generator, Mapping -from contextlib import contextmanager -from typing import Any - -from opentelemetry.context import Context -from opentelemetry.propagate import extract, inject -from opentelemetry.trace import SpanKind, get_current_span, get_tracer -from opentelemetry.trace.span import Span - -_tracer = get_tracer("mcp-python-sdk") - - -@contextmanager -def otel_span( - name: str, - *, - kind: SpanKind, - attributes: dict[str, Any] | None = None, - context: Context | None = None, - record_exception: bool = True, - set_status_on_exception: bool = True, -) -> Generator[Span]: - """Create an OTel span.""" - with _tracer.start_as_current_span( - name, - kind=kind, - attributes=attributes, - context=context, - record_exception=record_exception, - set_status_on_exception=set_status_on_exception, - ) as span: - yield span - - -def inject_trace_context(meta: dict[str, Any]) -> None: - """Inject W3C trace context (traceparent/tracestate) into a `_meta` dict.""" - inject(meta) - - -def extract_trace_context(meta: Mapping[str, Any] | None) -> Context | None: - """Extract W3C trace context from a `_meta` dict. - - Returns `None` when the carrier is absent, malformed, or carries no - valid `traceparent`, so callers fall through to ambient parenting; an - explicit empty `Context` would orphan the span instead of nesting under - the current one. - """ - if not meta: - return None - try: - ctx = extract(meta) - except (ValueError, TypeError): - return None - if not get_current_span(ctx).get_span_context().is_valid: - return None - return ctx +import sys + +import mcp_client.shared._otel as _implementation +from mcp_client.shared._otel import ( + _tracer as _tracer, +) +from mcp_client.shared._otel import ( + extract_trace_context as extract_trace_context, +) +from mcp_client.shared._otel import ( + inject_trace_context as inject_trace_context, +) +from mcp_client.shared._otel import ( + otel_span as otel_span, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/_stream_protocols.py b/src/mcp/shared/_stream_protocols.py index b799751329..d5c29b77a5 100644 --- a/src/mcp/shared/_stream_protocols.py +++ b/src/mcp/shared/_stream_protocols.py @@ -1,49 +1,17 @@ -"""Stream protocols for MCP transports. - -These are general-purpose protocols satisfied by both ``MemoryObjectSendStream``/ -``MemoryObjectReceiveStream`` and the context-aware wrappers in ``_context_streams``. -""" - -from __future__ import annotations - -from types import TracebackType -from typing import Protocol, TypeVar - -from typing_extensions import Self - -T_co = TypeVar("T_co", covariant=True) -T_contra = TypeVar("T_contra", contravariant=True) - - -class ReadStream(Protocol[T_co]): - """Protocol for reading items from a stream. - - Consumers that need the sender's context should use - ``getattr(stream, 'last_context', None)``. - """ - - async def receive(self) -> T_co: ... - async def aclose(self) -> None: ... - def __aiter__(self) -> ReadStream[T_co]: ... - async def __anext__(self) -> T_co: ... - async def __aenter__(self) -> Self: ... - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: TracebackType | None, - ) -> bool | None: ... - - -class WriteStream(Protocol[T_contra]): - """Protocol for writing items to a stream.""" - - async def send(self, item: T_contra, /) -> None: ... - async def aclose(self) -> None: ... - async def __aenter__(self) -> Self: ... - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: TracebackType | None, - ) -> bool | None: ... +import sys + +import mcp_client.shared._stream_protocols as _implementation +from mcp_client.shared._stream_protocols import ( + ReadStream as ReadStream, +) +from mcp_client.shared._stream_protocols import ( + T_co as T_co, +) +from mcp_client.shared._stream_protocols import ( + T_contra as T_contra, +) +from mcp_client.shared._stream_protocols import ( + WriteStream as WriteStream, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/auth.py b/src/mcp/shared/auth.py index 881379d381..491632b78c 100644 --- a/src/mcp/shared/auth.py +++ b/src/mcp/shared/auth.py @@ -1,258 +1,44 @@ -from typing import Any, Literal, cast - -from pydantic import AnyHttpUrl, AnyUrl, BaseModel, ConfigDict, Field, field_validator, model_validator - -# RFC 7523 JWT bearer grant; SEP-990 leg 2 uses this to present the ID-JAG. -JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer" - -# Token-endpoint client authentication methods this SDK's clients request, and the set -# `OAuthContext.prepare_token_auth` recognizes on a registered client (`private_key_jwt` is -# applied by `PrivateKeyJWTOAuthProvider`; the rest send a client secret or nothing). -TokenEndpointAuthMethod = Literal["none", "client_secret_post", "client_secret_basic", "private_key_jwt"] - -# grant_types a client requests when it does not specify its own (RFC 7591 §2). -DEFAULT_GRANT_TYPES = ["authorization_code", "refresh_token"] - - -def _empty_str_to_none(v: object) -> object: - # RFC 7591 §2 marks these URL fields OPTIONAL; a "" placeholder means absent, so it - # must not fail AnyHttpUrl validation. (The registered-client record applies the same - # rule to every member; this coercion serves the request model.) - if v == "": - return None - return v - - -class OAuthToken(BaseModel): - """See https://datatracker.ietf.org/doc/html/rfc6749#section-5.1""" - - access_token: str - token_type: Literal["Bearer"] = "Bearer" - expires_in: int | None = None - scope: str | None = None - refresh_token: str | None = None - - @field_validator("token_type", mode="before") - @classmethod - def normalize_token_type(cls, v: str | None) -> str | None: - if isinstance(v, str): - # Bearer is title-cased in the spec, so we normalize it - # https://datatracker.ietf.org/doc/html/rfc6750#section-4 - return v.title() - return v # pragma: no cover - - -class AuthorizationCodeResult(BaseModel): - """Authorization-code-grant redirect parameters returned by a callback handler. - - `iss` carries the RFC 9207 authorization-response issuer when the authorization server - includes it in the redirect; the client validates it against the expected issuer. - """ - - code: str - state: str | None = None - iss: str | None = None - - -class InvalidScopeError(Exception): - def __init__(self, message: str): - self.message = message - - -class InvalidRedirectUriError(Exception): - def __init__(self, message: str): - self.message = message - - -class OAuthClientMetadataBase(BaseModel): - """RFC 7591 OAuth 2.0 Dynamic Client Registration metadata shared verbatim by the - registration request (`OAuthClientMetadata`) and the authorization server's record of a - registered client (`OAuthClientInformationFull`). Fields whose acceptable values differ - between the two - what this SDK sends versus what a third-party server may echo - are - declared on each model rather than here. - See https://datatracker.ietf.org/doc/html/rfc7591#section-2 - """ - - model_config = ConfigDict(url_preserve_empty_path=True) - - # The MCP spec requires the "code" response type, but OAuth - # servers may also return additional types they support - response_types: list[str] = ["code"] - scope: str | None = None - - # these fields are currently unused, but we support & store them for potential - # future use - client_name: str | None = None - client_uri: AnyHttpUrl | None = None - logo_uri: AnyHttpUrl | None = None - contacts: list[str] | None = None - tos_uri: AnyHttpUrl | None = None - policy_uri: AnyHttpUrl | None = None - jwks_uri: AnyHttpUrl | None = None - jwks: Any | None = None - software_id: str | None = None - software_version: str | None = None - - @field_validator( - "client_uri", - "logo_uri", - "tos_uri", - "policy_uri", - "jwks_uri", - mode="before", - ) - @classmethod - def _empty_string_optional_url_to_none(cls, v: object) -> object: - # These URL fields are OPTIONAL; an echoed "" would otherwise fail AnyHttpUrl - # and throw away an otherwise valid registration response. - return _empty_str_to_none(v) - - -class OAuthClientMetadata(OAuthClientMetadataBase): - """RFC 7591 OAuth 2.0 Dynamic Client Registration request metadata: what an MCP - client sends when it registers. Field values are narrowed to what this SDK will put - on the wire; parsing the authorization server's response is `OAuthClientInformationFull`'s - job. See https://datatracker.ietf.org/doc/html/rfc7591#section-2 - """ - - redirect_uris: list[AnyUrl] | None = Field(..., min_length=1) - # supported auth methods for the token endpoint - token_endpoint_auth_method: TokenEndpointAuthMethod | None = None - # supported grant_types of this implementation - grant_types: list[ - Literal["authorization_code", "refresh_token", "urn:ietf:params:oauth:grant-type:jwt-bearer"] | str - ] = list(DEFAULT_GRANT_TYPES) - # SEP-837: OIDC application_type. Defaults to "native" since MCP clients typically use - # loopback redirect URIs; set "web" for remote browser-based clients on a non-local host. - application_type: Literal["web", "native"] = "native" - - -class OAuthClientInformationFull(OAuthClientMetadataBase): - """RFC 7591 OAuth 2.0 Dynamic Client Registration client information response - (client information plus metadata) - the authorization server's record of a - registered client. See https://datatracker.ietf.org/doc/html/rfc7591#section-3.2.1 - - A third-party authorization server "MAY reject or replace any of the client's - requested metadata values submitted during the registration and substitute them with - suitable values", so `application_type`, `token_endpoint_auth_method`, and `grant_types` - are typed to accept any string the server echoes, and `redirect_uris` may be absent or - empty. A member the server serializes as a placeholder - an explicit `null`, or `""` - - is read as an omitted key, so the field's default applies rather than the parse failing. - Whether a substituted value is usable is decided where the value is used, not at parse. - `redirect_uris` elements are still parsed as URLs, as the authorization server compares - them against a client's requested `redirect_uri`. - """ - - redirect_uris: list[AnyUrl] | None = None - # RFC 7591 §3.2.1: the server may assign an auth method other than the one requested, - # including methods this SDK does not implement, or omit it. - token_endpoint_auth_method: str | None = None - grant_types: list[str] = list(DEFAULT_GRANT_TYPES) - # SEP-837: OIDC application_type. OIDC Registration §2 defines "web" and "native", but - # servers echo other strings or an explicit null; the value is informational here. - application_type: str | None = None - - # RFC 7591 §3.2.1: client_id is REQUIRED in a client information response - a body - # without one is not a registration, whatever else it echoes. - client_id: str - client_secret: str | None = None - client_id_issued_at: int | None = None - client_secret_expires_at: int | None = None - # SEP-2352: the issuer these credentials were registered with, recorded by the SDK (not an - # RFC 7591 field) to detect authorization-server migration and avoid cross-AS credential reuse. - issuer: str | None = None - - @model_validator(mode="before") - @classmethod - def _placeholder_members_read_as_omitted(cls, data: object) -> object: - # Servers dump unset members of their client record as null, or echo them as "", - # instead of omitting the keys. Either placeholder would otherwise fail the parse of a - # list field (or read "" as an unrecognized method) and discard an already-provisioned - # registration; a placeholder and an absent key mean the same thing. - if isinstance(data, dict): - members = cast(dict[str, Any], data) - return {key: value for key, value in members.items() if value is not None and value != ""} - return data - - def validate_scope(self, requested_scope: str | None) -> list[str] | None: - if requested_scope is None: - return None - requested_scopes = requested_scope.split(" ") - allowed_scopes = [] if self.scope is None else self.scope.split(" ") - for scope in requested_scopes: - if scope not in allowed_scopes: - raise InvalidScopeError(f"Client was not registered with scope {scope}") - return requested_scopes - - def validate_redirect_uri(self, redirect_uri: AnyUrl | None) -> AnyUrl: - if redirect_uri is not None: - # Validate redirect_uri against client's registered redirect URIs - if not self.redirect_uris or redirect_uri not in self.redirect_uris: - raise InvalidRedirectUriError(f"Redirect URI '{redirect_uri}' not registered for client") - return redirect_uri - elif self.redirect_uris and len(self.redirect_uris) == 1: - return self.redirect_uris[0] - else: - raise InvalidRedirectUriError( - "redirect_uri must be specified unless the client has exactly one registered URI" - ) - - -class OAuthMetadata(BaseModel): - """RFC 8414 OAuth 2.0 Authorization Server Metadata. - See https://datatracker.ietf.org/doc/html/rfc8414#section-2 - """ - - model_config = ConfigDict(url_preserve_empty_path=True) - - issuer: AnyHttpUrl - authorization_endpoint: AnyHttpUrl - token_endpoint: AnyHttpUrl - registration_endpoint: AnyHttpUrl | None = None - scopes_supported: list[str] | None = None - response_types_supported: list[str] = ["code"] - response_modes_supported: list[str] | None = None - grant_types_supported: list[str] | None = None - token_endpoint_auth_methods_supported: list[str] | None = None - token_endpoint_auth_signing_alg_values_supported: list[str] | None = None - service_documentation: AnyHttpUrl | None = None - ui_locales_supported: list[str] | None = None - op_policy_uri: AnyHttpUrl | None = None - op_tos_uri: AnyHttpUrl | None = None - revocation_endpoint: AnyHttpUrl | None = None - revocation_endpoint_auth_methods_supported: list[str] | None = None - revocation_endpoint_auth_signing_alg_values_supported: list[str] | None = None - introspection_endpoint: AnyHttpUrl | None = None - introspection_endpoint_auth_methods_supported: list[str] | None = None - introspection_endpoint_auth_signing_alg_values_supported: list[str] | None = None - code_challenge_methods_supported: list[str] | None = None - client_id_metadata_document_supported: bool | None = None - authorization_response_iss_parameter_supported: bool | None = None - # SEP-990 / draft-ietf-oauth-identity-assertion-authz-grant §7.2: profiles whose grants the - # authorization server supports, e.g. `urn:ietf:params:oauth:grant-profile:id-jag`. - authorization_grant_profiles_supported: list[str] | None = None - - -class ProtectedResourceMetadata(BaseModel): - """RFC 9728 OAuth 2.0 Protected Resource Metadata. - See https://datatracker.ietf.org/doc/html/rfc9728#section-2 - """ - - model_config = ConfigDict(url_preserve_empty_path=True) - - resource: AnyHttpUrl - authorization_servers: list[AnyHttpUrl] = Field(..., min_length=1) - jwks_uri: AnyHttpUrl | None = None - scopes_supported: list[str] | None = None - bearer_methods_supported: list[str] | None = Field(default=["header"]) # MCP only supports header method - resource_signing_alg_values_supported: list[str] | None = None - resource_name: str | None = None - resource_documentation: AnyHttpUrl | None = None - resource_policy_uri: AnyHttpUrl | None = None - resource_tos_uri: AnyHttpUrl | None = None - # tls_client_certificate_bound_access_tokens default is False, but omitted here for clarity - tls_client_certificate_bound_access_tokens: bool | None = None - authorization_details_types_supported: list[str] | None = None - dpop_signing_alg_values_supported: list[str] | None = None - # dpop_bound_access_tokens_required default is False, but omitted here for clarity - dpop_bound_access_tokens_required: bool | None = None +import sys + +import mcp_client.shared.auth as _implementation +from mcp_client.shared.auth import ( + DEFAULT_GRANT_TYPES as DEFAULT_GRANT_TYPES, +) +from mcp_client.shared.auth import ( + JWT_BEARER_GRANT_TYPE as JWT_BEARER_GRANT_TYPE, +) +from mcp_client.shared.auth import ( + AuthorizationCodeResult as AuthorizationCodeResult, +) +from mcp_client.shared.auth import ( + InvalidRedirectUriError as InvalidRedirectUriError, +) +from mcp_client.shared.auth import ( + InvalidScopeError as InvalidScopeError, +) +from mcp_client.shared.auth import ( + OAuthClientInformationFull as OAuthClientInformationFull, +) +from mcp_client.shared.auth import ( + OAuthClientMetadata as OAuthClientMetadata, +) +from mcp_client.shared.auth import ( + OAuthClientMetadataBase as OAuthClientMetadataBase, +) +from mcp_client.shared.auth import ( + OAuthMetadata as OAuthMetadata, +) +from mcp_client.shared.auth import ( + OAuthToken as OAuthToken, +) +from mcp_client.shared.auth import ( + ProtectedResourceMetadata as ProtectedResourceMetadata, +) +from mcp_client.shared.auth import ( + TokenEndpointAuthMethod as TokenEndpointAuthMethod, +) +from mcp_client.shared.auth import ( + _empty_str_to_none as _empty_str_to_none, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/auth_utils.py b/src/mcp/shared/auth_utils.py index 3ba880f40d..c96123bdb5 100644 --- a/src/mcp/shared/auth_utils.py +++ b/src/mcp/shared/auth_utils.py @@ -1,80 +1,14 @@ -"""Utilities for OAuth 2.0 Resource Indicators (RFC 8707) and PKCE (RFC 7636).""" - -import time -from urllib.parse import urlparse, urlsplit, urlunsplit - -from pydantic import AnyUrl, HttpUrl - - -def resource_url_from_server_url(url: str | HttpUrl | AnyUrl) -> str: - """Convert server URL to canonical resource URL per RFC 8707. - - RFC 8707 section 2 states that resource URIs "MUST NOT include a fragment component". - Returns absolute URI with lowercase scheme/host for canonical form. - - Args: - url: Server URL to convert - - Returns: - Canonical resource URL string - """ - # Convert to string if needed - url_str = str(url) - - # Parse the URL and remove fragment, create canonical form - parsed = urlsplit(url_str) - canonical = urlunsplit(parsed._replace(scheme=parsed.scheme.lower(), netloc=parsed.netloc.lower(), fragment="")) - - return canonical - - -def check_resource_allowed(requested_resource: str, configured_resource: str) -> bool: - """Check if a requested resource URL matches a configured resource URL. - - A requested resource matches if it has the same scheme, domain, port, - and its path starts with the configured resource's path. This allows - hierarchical matching where a token for a parent resource can be used - for child resources. - - Args: - requested_resource: The resource URL being requested - configured_resource: The resource URL that has been configured - - Returns: - True if the requested resource matches the configured resource - """ - # Parse both URLs - requested = urlparse(requested_resource) - configured = urlparse(configured_resource) - - # Compare scheme, host, and port (origin) - if requested.scheme.lower() != configured.scheme.lower() or requested.netloc.lower() != configured.netloc.lower(): - return False - - # Normalize trailing slashes before comparison so that - # "/foo" and "/foo/" are treated as equivalent. - requested_path = requested.path - configured_path = configured.path - if not requested_path.endswith("/"): - requested_path += "/" - if not configured_path.endswith("/"): - configured_path += "/" - - # Check hierarchical match: requested must start with configured path. - # The trailing-slash normalization ensures "/api123/" won't match "/api/". - return requested_path.startswith(configured_path) - - -def calculate_token_expiry(expires_in: int | str | None) -> float | None: - """Calculate token expiry timestamp from expires_in seconds. - - Args: - expires_in: Seconds until token expiration (may be string from some servers) - - Returns: - Unix timestamp when token expires, or None if no expiry specified - """ - if expires_in is None: - return None # pragma: no cover - # Defensive: handle servers that return expires_in as string - return time.time() + int(expires_in) +import sys + +import mcp_client.shared.auth_utils as _implementation +from mcp_client.shared.auth_utils import ( + calculate_token_expiry as calculate_token_expiry, +) +from mcp_client.shared.auth_utils import ( + check_resource_allowed as check_resource_allowed, +) +from mcp_client.shared.auth_utils import ( + resource_url_from_server_url as resource_url_from_server_url, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/context.py b/src/mcp/shared/context.py index 13c145be5b..298fb6d03f 100644 --- a/src/mcp/shared/context.py +++ b/src/mcp/shared/context.py @@ -1,85 +1,11 @@ -"""`BaseContext` - the user-facing per-request context. +import sys -Composition over a `DispatchContext`: forwards the transport metadata, the -back-channel (`send_raw_request`/`notify`), progress reporting, and the cancel -event. Adds `meta` (the inbound request's `_meta` field). +import mcp_client.shared.context as _implementation +from mcp_client.shared.context import ( + BaseContext as BaseContext, +) +from mcp_client.shared.context import ( + TransportT as TransportT, +) -Satisfies `Outbound`, so `ClientPeer` can wrap it. Shared between client and -server: the server's `Context` extends this with `lifespan`/`connection`; -`ClientContext` is just an alias. -""" - -from collections.abc import Mapping -from typing import Any, Generic - -import anyio -from mcp_types import RequestParamsMeta -from typing_extensions import TypeVar - -from mcp.shared.dispatcher import CallOptions, DispatchContext -from mcp.shared.transport_context import TransportContext - -__all__ = ["BaseContext"] - -TransportT = TypeVar("TransportT", bound=TransportContext, default=TransportContext, covariant=True) - - -class BaseContext(Generic[TransportT]): - """Per-request context wrapping a `DispatchContext`. - - `ServerRunner` constructs one per inbound request and passes it to the - user's handler. - """ - - def __init__(self, dctx: DispatchContext[TransportT], meta: RequestParamsMeta | None = None) -> None: - self._dctx = dctx - self._meta = meta - - @property - def transport(self) -> TransportT: - """Transport-specific metadata for this inbound request.""" - return self._dctx.transport - - @property - def cancel_requested(self) -> anyio.Event: - """Set when the peer sends `notifications/cancelled` for this request.""" - return self._dctx.cancel_requested - - @property - def can_send_request(self) -> bool: - """Whether the back-channel can currently deliver server-initiated requests. - - `False` when the transport has no back-channel, or when the underlying - dispatch context has been closed because the inbound request finished. - """ - return self._dctx.can_send_request - - @property - def meta(self) -> RequestParamsMeta | None: - """The inbound request's `_meta` field, if present.""" - return self._meta - - async def send_raw_request( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None = None, - ) -> dict[str, Any]: - """Send a request to the peer on the back-channel. - - Raises: - MCPError: The peer responded with an error. - NoBackChannelError: `can_send_request` is `False`. - """ - return await self._dctx.send_raw_request(method, params, opts) - - async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: - """Send a notification to the peer on the back-channel.""" - await self._dctx.notify(method, params, opts) - - async def report_progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: - """Report progress for this request, if the peer supplied a progress token. - - A no-op when no token was supplied. - """ - await self._dctx.progress(progress, total, message) +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/direct_dispatcher.py b/src/mcp/shared/direct_dispatcher.py index e17283afa2..df5993fd5e 100644 --- a/src/mcp/shared/direct_dispatcher.py +++ b/src/mcp/shared/direct_dispatcher.py @@ -1,334 +1,26 @@ -"""In-memory `Dispatcher` that wires two peers together with no transport. +import sys -`DirectDispatcher` is the simplest possible `Dispatcher` implementation: a -request on one side directly invokes the other side's `on_request`. There is no -serialization, no JSON-RPC framing, and no streams. It exists to: - -* prove the `Dispatcher` Protocol is implementable without JSON-RPC -* provide a fast substrate for testing the layers above the dispatcher - (`ServerRunner`, `Context`, `Connection`) without wire-level moving parts -* embed a server in-process when the JSON-RPC overhead is unnecessary - -Like `JSONRPCDispatcher`, this is an exception-to-error boundary: a handler -exception surfaces to the caller as `MCPError`. The `raise_handler_exceptions` -knob controls whether unmapped exceptions are sanitized (matching the wire -path) or chained as ``__cause__`` for in-process debugging. -""" - -from __future__ import annotations - -import logging -from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, field -from typing import Any - -import anyio -import anyio.abc -from mcp_types import CONNECTION_CLOSED, INTERNAL_ERROR, INVALID_PARAMS, REQUEST_TIMEOUT, RequestId -from pydantic import ValidationError - -from mcp.shared._compat import resync_tracer -from mcp.shared.dispatcher import ( - CallOptions, - OnNotify, - OnNotifyIntercept, - OnRequest, - ProgressFnT, - coerce_request_id, - run_notify_intercept, +import mcp_client.shared.direct_dispatcher as _implementation +from mcp_client.shared.direct_dispatcher import ( + DIRECT_TRANSPORT_KIND as DIRECT_TRANSPORT_KIND, +) +from mcp_client.shared.direct_dispatcher import ( + DirectDispatcher as DirectDispatcher, +) +from mcp_client.shared.direct_dispatcher import ( + _DirectDispatchContext as _DirectDispatchContext, +) +from mcp_client.shared.direct_dispatcher import ( + _Notify as _Notify, +) +from mcp_client.shared.direct_dispatcher import ( + _Request as _Request, +) +from mcp_client.shared.direct_dispatcher import ( + create_direct_dispatcher_pair as create_direct_dispatcher_pair, +) +from mcp_client.shared.direct_dispatcher import ( + logger as logger, ) -from mcp.shared.exceptions import MCPError, NoBackChannelError -from mcp.shared.message import MessageMetadata -from mcp.shared.transport_context import TransportContext - -logger = logging.getLogger(__name__) - -__all__ = ["DirectDispatcher", "create_direct_dispatcher_pair"] - -DIRECT_TRANSPORT_KIND = "direct" - - -_Request = Callable[[str, Mapping[str, Any] | None, CallOptions | None], Awaitable[dict[str, Any]]] -_Notify = Callable[[str, Mapping[str, Any] | None], Awaitable[None]] - - -@dataclass -class _DirectDispatchContext: - """`DispatchContext` for an inbound request on a `DirectDispatcher`. - - The back-channel callables target the *originating* side, so a handler's - `send_raw_request` reaches the peer that made the inbound request. - """ - - transport: TransportContext - _back_request: _Request - _back_notify: _Notify - request_id: RequestId | None = None - """The caller-supplied `CallOptions["request_id"]`, else a dispatcher-synthesized - id for requests; `None` for notifications.""" - message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework - """Always `None`: in-memory dispatch attaches no transport metadata.""" - _on_progress: ProgressFnT | None = None - cancel_requested: anyio.Event = field(default_factory=anyio.Event) - - @property - def can_send_request(self) -> bool: - return self.transport.can_send_request - - async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: - await self._back_notify(method, params) - - async def send_raw_request( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None = None, - ) -> dict[str, Any]: - if not self.can_send_request: - raise NoBackChannelError(method) - return await self._back_request(method, params, opts) - - async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: - if self._on_progress is not None: - await self._on_progress(progress, total, message) - - -class DirectDispatcher: - """A `Dispatcher` that calls a peer's handlers directly, in-process. - - Two instances are wired together with `create_direct_dispatcher_pair`; each - holds a reference to the other. `send_raw_request` on one awaits the peer's - `on_request`. `run` parks until `close` is called. - - Lifecycle mirrors `JSONRPCDispatcher`: `send_raw_request` requires `run()` - to have started, and once a side has closed - via `close()` or `run()` - ending - `send_raw_request` raises `MCPError` (`CONNECTION_CLOSED`) and - inbound requests fail the peer's call the same way instead of invoking the - handler. Notifications are fire-and-forget in both directions: after close - they are silently dropped. - """ - - def __init__(self, transport_ctx: TransportContext, *, raise_handler_exceptions: bool = True): - self._transport_ctx = transport_ctx - self._raise_handler_exceptions = raise_handler_exceptions - self._peer: DirectDispatcher | None = None - self._on_request: OnRequest | None = None - self._on_notify: OnNotify | None = None - self._on_notify_intercept: OnNotifyIntercept | None = None - self._next_id = 0 - self._in_flight_ids: set[RequestId] = set() - self._ready = anyio.Event() - self._close_event = anyio.Event() - self._running = False - self._closed = False - - def connect_to(self, peer: DirectDispatcher) -> None: - self._peer = peer - - async def send_raw_request( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None = None, - ) -> dict[str, Any]: - """Send a request by invoking the peer's `on_request` directly. - - Raises: - MCPError: The peer's handler raised; `REQUEST_TIMEOUT` if - `opts["timeout"]` elapsed; `CONNECTION_CLOSED` if either - side has closed. - RuntimeError: Called before `run()`. - """ - if self._peer is None: - raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()") - # Post-close sends get the same CONNECTION_CLOSED contract as JSONRPCDispatcher. - if self._closed: - raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") - if not self._running: - raise RuntimeError("DirectDispatcher.send_raw_request called before run()") - return await self._peer._dispatch_request(method, params, opts) - - async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: - """Send a notification by invoking the peer's `on_notify` directly. - - Fire-and-forget: usable before `run()` (delivery waits for the peer to - start), and after close it is silently dropped, matching - `JSONRPCDispatcher.notify`. `opts` is accepted for `Dispatcher` - conformance; there is no HTTP layer here so `headers` is ignored. - """ - if self._peer is None: - raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()") - if self._closed: - logger.debug("dropped notification %r on closed DirectDispatcher", method) - return - await self._peer._dispatch_notify(method, params) - - async def run( - self, - on_request: OnRequest, - on_notify: OnNotify, - on_notify_intercept: OnNotifyIntercept | None = None, - *, - task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, - ) -> None: - """Mark this side ready and park until `close()` is called. - - Single-shot, like `JSONRPCDispatcher.run`: once it returns the - dispatcher stays closed and cannot be restarted. - """ - try: - self._on_request = on_request - self._on_notify = on_notify - self._on_notify_intercept = on_notify_intercept - self._running = True - self._ready.set() - task_status.started() - await self._close_event.wait() - finally: - self._running = False - self._closed = True - # run() may end via cancellation without close() ever being - # called; setting the event wakes `_wait_ready` waiters so they - # observe the closed state instead of parking forever. - self._close_event.set() - - def close(self) -> None: - self._closed = True - self._close_event.set() - - def _make_context( - self, on_progress: ProgressFnT | None = None, request_id: RequestId | None = None - ) -> _DirectDispatchContext: - assert self._peer is not None - peer = self._peer - return _DirectDispatchContext( - transport=self._transport_ctx, - _back_request=lambda m, p, o: peer._dispatch_request(m, p, o), - _back_notify=lambda m, p: peer._dispatch_notify(m, p), - request_id=request_id, - _on_progress=on_progress, - ) - - async def _wait_ready(self) -> None: - """Park until `run()` has started, waking early if this side closes. - - Raises: - MCPError: `CONNECTION_CLOSED` if this side has closed. - """ - if not self._ready.is_set() and not self._close_event.is_set(): - async with anyio.create_task_group() as tg: - - async def wake_on(event: anyio.Event) -> None: - await event.wait() - tg.cancel_scope.cancel() - - tg.start_soon(wake_on, self._ready) - tg.start_soon(wake_on, self._close_event) - if self._closed: - raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") - - async def _dispatch_request( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None, - ) -> dict[str, Any]: - opts = opts or {} - try: - with anyio.fail_after(opts.get("timeout")): - # Inside the timeout scope, so a configured timeout also bounds - # waiting on a peer whose run() has not started yet. - await self._wait_ready() - assert self._on_request is not None - supplied_id = opts.get("request_id") - if supplied_id is not None: - request_id: RequestId = supplied_id - # Collisions use the same coerced domain as JSONRPCDispatcher's - # pending keys, so this in-memory stand-in raises for exactly - # the ids the wire dispatcher would; the context still sees - # the verbatim value. - in_flight_key = coerce_request_id(request_id) - if in_flight_key in self._in_flight_ids: - raise ValueError(f"request id {request_id!r} is already in flight") - else: - # Synthesize an id (the DispatchContext contract reserves None - # for notifications), minting past any key a supplied id - # occupies: the collision error is reserved for the caller - # who actually chose the id. - self._next_id += 1 - while self._next_id in self._in_flight_ids: - self._next_id += 1 - request_id = self._next_id - in_flight_key = request_id - self._in_flight_ids.add(in_flight_key) - dctx = self._make_context(on_progress=opts.get("on_progress"), request_id=request_id) - try: - return await self._on_request(dctx, method, params) - except MCPError: - raise - except ValidationError as e: - # Same shape JSONRPCDispatcher writes, so runner-over-direct - # tests see what runner-over-JSONRPC would. - raise MCPError(code=INVALID_PARAMS, message="Invalid request parameters", data="") from e - except Exception as e: - # Single owner of the in-proc exception-to-error policy (mirrors - # JSONRPCDispatcher / `_streamable_http_modern._to_jsonrpc_response` - # for the wire paths). True chains the original for in-process - # debugging; False sanitizes to match the wire path's leak guard. - if self._raise_handler_exceptions: - raise MCPError(code=INTERNAL_ERROR, message=str(e)) from e - logger.exception("request handler raised") - raise MCPError(code=INTERNAL_ERROR, message="Internal server error") from None - finally: - self._in_flight_ids.discard(in_flight_key) - except TimeoutError: - raise MCPError( - code=REQUEST_TIMEOUT, - message=f"Timed out after {opts.get('timeout')}s waiting for {method!r}", - ) from None - finally: - await resync_tracer() - - async def _dispatch_notify(self, method: str, params: Mapping[str, Any] | None) -> None: - try: - await self._wait_ready() - except MCPError: - # Notifications are fire-and-forget: a notify to a closed peer is - # dropped, not raised back into the sender's call. - logger.debug("dropped notification %r to closed DirectDispatcher", method) - return - if run_notify_intercept(self._on_notify_intercept, method, params): - return - assert self._on_notify is not None - dctx = self._make_context() - await self._on_notify(dctx, method, params) - - -def create_direct_dispatcher_pair( - *, - can_send_request: bool = True, - headers: Mapping[str, str] | None = None, - raise_handler_exceptions: bool = True, -) -> tuple[DirectDispatcher, DirectDispatcher]: - """Create two `DirectDispatcher` instances wired to each other. - - Args: - can_send_request: Sets `TransportContext.can_send_request` on both - sides. Pass `False` to simulate a transport with no back-channel. - headers: Sets `TransportContext.headers` on both sides. - raise_handler_exceptions: When `True` (the default - this is an - in-process debugging substrate), an unmapped handler exception - reaches the caller as `MCPError` with the original chained as - ``__cause__``. When `False` it is sanitized to an opaque - `INTERNAL_ERROR` so the in-process path matches the wire. - Returns: - A `(client, server)` pair. The wiring is symmetric, so the roles - are conventional only. - """ - ctx = TransportContext(kind=DIRECT_TRANSPORT_KIND, can_send_request=can_send_request, headers=headers) - client = DirectDispatcher(ctx, raise_handler_exceptions=raise_handler_exceptions) - server = DirectDispatcher(ctx, raise_handler_exceptions=raise_handler_exceptions) - client.connect_to(server) - server.connect_to(client) - return client, server +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/dispatcher.py b/src/mcp/shared/dispatcher.py index f2ff96e7d5..7e0387444d 100644 --- a/src/mcp/shared/dispatcher.py +++ b/src/mcp/shared/dispatcher.py @@ -1,277 +1,44 @@ -"""Dispatcher Protocol - the call/return boundary between transports and handlers. - -A Dispatcher turns a duplex message channel into two things: - -* an outbound API: `send_raw_request(method, params)` and `notify(method, params)` -* an inbound pump: `run(on_request, on_notify)` that drives the receive loop - and invokes the supplied handlers for each incoming request/notification - -It is deliberately *not* MCP-aware. Method names are strings, params and -results are `dict[str, Any]`. The MCP type layer (request/result models, -capability negotiation, `Context`) sits above this; the wire encoding -(JSON-RPC, gRPC, in-process direct calls) sits below it. - -See `JSONRPCDispatcher` for the production implementation and -`DirectDispatcher` for an in-memory implementation used in tests and for -embedding a server in-process. -""" - -import logging -from collections.abc import Awaitable, Callable, Mapping -from typing import Any, Protocol, TypedDict, TypeVar, runtime_checkable - -import anyio -import anyio.abc -from mcp_types import RequestId - -from mcp.shared.message import MessageMetadata -from mcp.shared.transport_context import TransportContext - -logger = logging.getLogger(__name__) - -__all__ = [ - "CallOptions", - "DispatchContext", - "Dispatcher", - "OnNotify", - "OnNotifyIntercept", - "OnRequest", - "Outbound", - "ProgressFnT", - "as_request_id", - "coerce_request_id", - "run_notify_intercept", -] - -TransportT_co = TypeVar("TransportT_co", bound=TransportContext, covariant=True) - - -def as_request_id(value: object) -> RequestId | None: - """Narrow an untyped wire value to a `RequestId`, or None; rejects bool (True would alias request id 1).""" - if isinstance(value, str | int) and not isinstance(value, bool): - return value - return None - - -def coerce_request_id(request_id: RequestId) -> RequestId: - """Coerce a stringified int request id back to int so a peer-echoed id still correlates (matches the TS SDK). - - This is the collision/correlation domain dispatchers share: "7" and 7 are one - id for correlation purposes, even where the wire carries the verbatim value. - """ - if isinstance(request_id, str): - try: - return int(request_id) - except ValueError: - pass - return request_id - - -class ProgressFnT(Protocol): - """Callback invoked when a progress notification arrives for a pending request.""" - - async def __call__(self, progress: float, total: float | None, message: str | None) -> None: ... - - -class CallOptions(TypedDict, total=False): - """Per-call options for `Outbound.send_raw_request`. - - All keys are optional. Dispatchers ignore keys they do not understand. - """ - - request_id: RequestId - """Send the request under this caller-supplied id instead of a dispatcher-minted one. - - The peer sees the value verbatim ("7" stays a string). A value that collides - with one of the sender's own in-flight request ids raises `ValueError`. - Callers that need to know a request's id before its result arrives (a - `subscriptions/listen` stream is demultiplexed by it) mint their own ids - here; string ids that don't parse as integers can never collide with the - dispatcher's minted sequence. Per the class contract, dispatchers that - predate this key ignore it and mint as usual. - """ - - timeout: float - """Seconds to wait for a result before raising and sending `notifications/cancelled`.""" - - cancel_on_abandon: bool - """Whether abandoning this request (timeout or caller cancellation) sends `notifications/cancelled`. - - Defaults to `True`. Set `False` for requests the protocol forbids cancelling, such as `initialize`. - Also suppressed when resumption hints reach the transport, or when the request was never written. - """ - - on_progress: ProgressFnT - """Receive `notifications/progress` updates for this request.""" - - resumption_token: str - """Opaque token to resume a previously interrupted request. - - Client-side, streamable-HTTP only. Ignored by server dispatchers and other - transports, and also ignored (with a debug log) for requests sent from a - `DispatchContext`, where routing onto the inbound request's stream takes - precedence. Supports protocol version 2025-11-25 and earlier; SSE-stream - resumption is removed in the next protocol revision. - """ - - on_resumption_token: Callable[[str], Awaitable[None]] - """Receive a resumption token when the transport issues one for this request. - - Client-side, streamable-HTTP only. Ignored by server dispatchers and other - transports, and also ignored (with a debug log) for requests sent from a - `DispatchContext`, where routing onto the inbound request's stream takes - precedence. Supports protocol version 2025-11-25 and earlier; SSE-stream - resumption is removed in the next protocol revision. - """ - - headers: dict[str, str] - """Transport-layer hint: HTTP transports merge these onto the outgoing request; non-HTTP transports ignore.""" - - -@runtime_checkable -class Outbound(Protocol): - """Anything that can send requests and notifications to the peer. - - Both `Dispatcher` (top-level outbound) and `DispatchContext` (back-channel - during an inbound request) extend this. The MCP type layer (`ClientPeer`, - `Connection`) builds typed `send_request` / convenience methods on top of - this raw channel. - """ - - async def send_raw_request( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None = None, - ) -> dict[str, Any]: - """Send a request and await its raw result dict. - - Raises: - MCPError: If the peer responded with an error, or the handler - raised. Implementations normalize all handler exceptions to - `MCPError` so callers see a single exception type. - """ - ... - - async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: - """Send a fire-and-forget notification.""" - ... - - -class DispatchContext(Outbound, Protocol[TransportT_co]): - """Per-request context handed to `on_request` / `on_notify`. - - Carries the transport metadata for the inbound message and provides the - back-channel for sending requests/notifications to the peer while handling - it. `send_raw_request` raises `NoBackChannelError` if `can_send_request` - is `False`. - """ - - @property - def transport(self) -> TransportT_co: - """Transport-specific metadata for this inbound message.""" - ... - - @property - def can_send_request(self) -> bool: - """Whether the back-channel can currently deliver server-initiated requests. - - `False` when the transport has no back-channel, or when this context has - been closed (the inbound request finished). `send_raw_request` raises - `NoBackChannelError` exactly when this is `False`. - """ - ... - - @property - def request_id(self) -> RequestId | None: - """The id of the inbound request, or `None` for a notification. - - For JSON-RPC this is the wire `id` field. Handlers thread it through - as `related_request_id` on outbound notifications so HTTP transports - can route them onto the originating request's response stream. - """ - ... - - @property - def message_metadata(self) -> MessageMetadata: - """The metadata the transport attached to this inbound message, if any. - - This is `SessionMessage.metadata` passed through verbatim: HTTP - transports attach `ServerMessageMetadata` (the HTTP request, SSE - stream-close callbacks); stdio and in-memory dispatch attach nothing. - Tied to the `SessionMessage` wire format - goes away when transports - stop delivering messages that way. - """ - # TODO(maxisbey): remove for context rework - ... - - @property - def cancel_requested(self) -> anyio.Event: - """Set when the peer sends `notifications/cancelled` for this request.""" - ... - - async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: - """Report progress for the inbound request, if the peer supplied a progress token. - - A no-op when no token was supplied. - """ - ... - - -OnRequest = Callable[[DispatchContext[TransportContext], str, Mapping[str, Any] | None], Awaitable[dict[str, Any]]] -"""Handler for inbound requests: `(ctx, method, params) -> result`. Raise `MCPError` to send an error response.""" - -OnNotify = Callable[[DispatchContext[TransportContext], str, Mapping[str, Any] | None], Awaitable[None]] -"""Handler for inbound notifications: `(ctx, method, params)`.""" - -OnNotifyIntercept = Callable[[str, Mapping[str, Any] | None], bool] -"""Synchronous receive-order intercept for inbound notifications: `(method, params) -> consumed`. - -Runs before `on_notify` is scheduled so correlation state advances in wire order -relative to response resolution (the client's listen demux depends on this). -Returning True consumes the notification. Must not block the receive path. -""" - - -def run_notify_intercept(intercept: OnNotifyIntercept | None, method: str, params: Mapping[str, Any] | None) -> bool: - """Invoke `intercept`, containing a raise to that one notification (never the receive loop).""" - if intercept is None: - return False - try: - return intercept(method, params) - except Exception: - logger.exception("notification intercept raised; passing %r through", method) - return False - - -class Dispatcher(Outbound, Protocol[TransportT_co]): - """A duplex request/notification channel with call-return semantics. - - Implementations own correlation of outbound requests to inbound results, the - receive loop, per-request concurrency, and cancellation/progress wiring. - - The lifecycle surface is provisional; `run()` may change in a 2.x minor - release. - """ - - async def run( - self, - on_request: OnRequest, - on_notify: OnNotify, - on_notify_intercept: OnNotifyIntercept | None = None, - *, - task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, - ) -> None: - """Drive the receive loop until the underlying channel closes. - - Each inbound request is dispatched to `on_request` in its own task; - the returned dict (or raised `MCPError`) is sent back as the response. - Implementations MUST offer every inbound notification to - `on_notify_intercept` synchronously in receive order (via - `run_notify_intercept`), handing only unconsumed ones to `on_notify`. - - `task_status.started()` is called once the dispatcher is ready to - accept `send_request`/`notify` calls, so callers can use - `await tg.start(dispatcher.run, on_request, on_notify)`. - """ - ... +import sys + +import mcp_client.shared.dispatcher as _implementation +from mcp_client.shared.dispatcher import ( + CallOptions as CallOptions, +) +from mcp_client.shared.dispatcher import ( + DispatchContext as DispatchContext, +) +from mcp_client.shared.dispatcher import ( + Dispatcher as Dispatcher, +) +from mcp_client.shared.dispatcher import ( + OnNotify as OnNotify, +) +from mcp_client.shared.dispatcher import ( + OnNotifyIntercept as OnNotifyIntercept, +) +from mcp_client.shared.dispatcher import ( + OnRequest as OnRequest, +) +from mcp_client.shared.dispatcher import ( + Outbound as Outbound, +) +from mcp_client.shared.dispatcher import ( + ProgressFnT as ProgressFnT, +) +from mcp_client.shared.dispatcher import ( + TransportT_co as TransportT_co, +) +from mcp_client.shared.dispatcher import ( + as_request_id as as_request_id, +) +from mcp_client.shared.dispatcher import ( + coerce_request_id as coerce_request_id, +) +from mcp_client.shared.dispatcher import ( + logger as logger, +) +from mcp_client.shared.dispatcher import ( + run_notify_intercept as run_notify_intercept, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/exceptions.py b/src/mcp/shared/exceptions.py index c2a7fd44e7..d94589541d 100644 --- a/src/mcp/shared/exceptions.py +++ b/src/mcp/shared/exceptions.py @@ -1,119 +1,20 @@ -from __future__ import annotations - -from typing import Any, cast - -from mcp_types import INVALID_REQUEST, URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError - - -class MCPDeprecationWarning(UserWarning): - """A custom deprecation warning for the MCP SDK. - - Unlike the built-in `DeprecationWarning`, this inherits from `UserWarning` so - it is shown by default, helping users discover deprecated features without - enabling warnings explicitly. - - Reference: https://sethmlarson.dev/deprecations-via-warnings-dont-work-for-python-libraries - """ - - -class MCPError(Exception): - """Exception type raised when an error arrives over an MCP connection.""" - - error: ErrorData - - def __init__(self, code: int, message: str, data: Any = None): - super().__init__(code, message, data) - if data is not None: - self.error = ErrorData(code=code, message=message, data=data) - else: - self.error = ErrorData(code=code, message=message) - - @property - def code(self) -> int: - return self.error.code - - @property - def message(self) -> str: - return self.error.message - - @property - def data(self) -> Any: - return self.error.data - - @classmethod - def from_jsonrpc_error(cls, error: JSONRPCError) -> MCPError: - return cls.from_error_data(error.error) - - @classmethod - def from_error_data(cls, error: ErrorData) -> MCPError: - return cls(code=error.code, message=error.message, data=error.data) - - def __str__(self) -> str: - return self.message - - -class NoBackChannelError(MCPError): - """Raised when a server-initiated request has no channel that can deliver it. - - Raised by `DispatchContext.send_raw_request` when its request-scoped channel - reports `TransportContext.can_send_request` as `False` (the cases are - documented on that field), and by a connection's standalone channel when it - has none; serializes to an `INVALID_REQUEST` error response. - """ - - def __init__(self, method: str): - super().__init__( - code=INVALID_REQUEST, - message=( - f"Cannot send {method!r}: this transport context has no back-channel for server-initiated requests." - ), - ) - self.method = method - - -class UrlElicitationRequiredError(MCPError): - """Specialized error for when a tool requires URL mode elicitation(s) before proceeding. - - Servers can raise this error from tool handlers to indicate that the client - must complete one or more URL elicitations before the request can be processed. - - Example: - ```python - raise UrlElicitationRequiredError([ - ElicitRequestURLParams( - message="Authorization required for your files", - url="https://example.com/oauth/authorize", - elicitation_id="auth-001" - ) - ]) - ``` - """ - - def __init__(self, elicitations: list[ElicitRequestURLParams], message: str | None = None): - """Initialize UrlElicitationRequiredError.""" - if message is None: - message = f"URL elicitation{'s' if len(elicitations) > 1 else ''} required" - - self._elicitations = elicitations - - super().__init__( - code=URL_ELICITATION_REQUIRED, - message=message, - data={"elicitations": [e.model_dump(by_alias=True, exclude_none=True) for e in elicitations]}, - ) - - @property - def elicitations(self) -> list[ElicitRequestURLParams]: - """The list of URL elicitations required before the request can proceed.""" - return self._elicitations - - @classmethod - def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError: - """Reconstruct from an ErrorData received over the wire.""" - if error.code != URL_ELICITATION_REQUIRED: - raise ValueError(f"Expected error code {URL_ELICITATION_REQUIRED}, got {error.code}") - - data = cast(dict[str, Any], error.data or {}) - raw_elicitations = cast(list[dict[str, Any]], data.get("elicitations", [])) - elicitations = [ElicitRequestURLParams.model_validate(e) for e in raw_elicitations] - return cls(elicitations, error.message) +import sys + +import mcp_client.shared.exceptions as _implementation +from mcp_client.shared.exceptions import ( + MCPDeprecationWarning as MCPDeprecationWarning, +) +from mcp_client.shared.exceptions import ( + MCPError as MCPError, +) +from mcp_client.shared.exceptions import ( + NoBackChannelError as NoBackChannelError, +) +from mcp_client.shared.exceptions import ( + UrlElicitationRequiredError as UrlElicitationRequiredError, +) + +for _exception in (MCPDeprecationWarning, MCPError, NoBackChannelError, UrlElicitationRequiredError): + _exception.__module__ = __name__ + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/extension.py b/src/mcp/shared/extension.py index 283e9ba89b..972dc4fa44 100644 --- a/src/mcp/shared/extension.py +++ b/src/mcp/shared/extension.py @@ -1,28 +1,17 @@ -"""Extension-identifier grammar shared by the server and client extension surfaces.""" +import sys -from __future__ import annotations +import mcp_client.shared.extension as _implementation +from mcp_client.shared.extension import ( + _IDENTIFIER_RE as _IDENTIFIER_RE, +) +from mcp_client.shared.extension import ( + _LABEL as _LABEL, +) +from mcp_client.shared.extension import ( + _NAME as _NAME, +) +from mcp_client.shared.extension import ( + validate_extension_identifier as validate_extension_identifier, +) -import re -from typing import Any - -__all__ = ["validate_extension_identifier"] - -# Extension identifiers follow the `_meta` key grammar with a mandatory prefix -# (SEP-2133 / basic/index.mdx): dot-separated labels, each starting with a -# letter and ending with a letter or digit (hyphens interior), then `/`, then a -# name that starts and ends alphanumeric (`.`/`_`/`-` interior). -_LABEL = r"[A-Za-z](?:[A-Za-z0-9-]*[A-Za-z0-9])?" -_NAME = r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?" -_IDENTIFIER_RE = re.compile(rf"{_LABEL}(?:\.{_LABEL})*/{_NAME}") - - -def validate_extension_identifier(identifier: Any, *, owner: str) -> None: - """Raise `TypeError` unless `identifier` is a `vendor-prefix/name` string. - - SEP-2133 requires extension identifiers to carry a reverse-DNS prefix. - """ - if not isinstance(identifier, str) or not _IDENTIFIER_RE.fullmatch(identifier): - raise TypeError( - f"{owner}.identifier must be a `vendor-prefix/name` string " - f"(reverse-DNS prefix required), got {identifier!r}" - ) +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/inbound.py b/src/mcp/shared/inbound.py index e33d6f502f..eddade29d7 100644 --- a/src/mcp/shared/inbound.py +++ b/src/mcp/shared/inbound.py @@ -1,595 +1,101 @@ -"""Inbound request classification for the modern per-request-envelope path. +import sys -Pure module: no I/O, no transport, no `mcp.server` imports. Runs the -validation ladder against a decoded JSON-RPC body and returns either an -:class:`InboundModernRoute` (every rung passed) or an -:class:`InboundLadderRejection` (the first rung that failed). Callers map a -rejection's `code` through :data:`ERROR_CODE_HTTP_STATUS` to pick the HTTP -status. - -Also hosts the shared header-value codec and the `x-mcp-header` schema -validator so client emit and server validate read the same source of truth. -""" - -import base64 -import binascii -import re -from collections.abc import Iterable, Iterator, Mapping, Sequence -from dataclasses import dataclass -from types import MappingProxyType -from typing import Any, Final, cast - -from mcp_types import ( - CLIENT_CAPABILITIES_META_KEY, - CLIENT_INFO_META_KEY, - PROTOCOL_VERSION_META_KEY, - UnsupportedProtocolVersionErrorData, -) -from mcp_types.jsonrpc import ( - HEADER_MISMATCH, - INVALID_PARAMS, - INVALID_REQUEST, - METHOD_NOT_FOUND, - MISSING_REQUIRED_CLIENT_CAPABILITY, - PARSE_ERROR, - UNSUPPORTED_PROTOCOL_VERSION, -) -from mcp_types.version import MODERN_PROTOCOL_VERSIONS - -__all__ = [ - "ERROR_CODE_HTTP_STATUS", - "InboundLadderRejection", - "InboundModernRoute", - "MCP_METHOD_HEADER", - "MCP_NAME_HEADER", - "MCP_PARAM_HEADER_PREFIX", - "MCP_PROTOCOL_VERSION_HEADER", - "NAME_BEARING_METHODS", - "X_MCP_HEADER_KEY", - "classify_inbound_request", - "decode_header_value", - "encode_header_value", - "find_duplicated_routing_header", - "find_invalid_x_mcp_header", - "mcp_param_headers", - "unsupported_protocol_version_rejection", - "validate_mcp_param_headers", - "x_mcp_header_map", -] - -MCP_PROTOCOL_VERSION_HEADER: Final = "mcp-protocol-version" -"""Canonical lowercase name of the HTTP header carrying the MCP protocol version.""" - -MCP_METHOD_HEADER: Final = "mcp-method" -"""Canonical lowercase name of the HTTP header carrying the JSON-RPC method.""" - -MCP_NAME_HEADER: Final = "mcp-name" -"""Canonical lowercase name of the HTTP header carrying the resource name (tool/prompt/resource URI).""" - -X_MCP_HEADER_KEY: Final = "x-mcp-header" -"""JSON-Schema property annotation that designates an `Mcp-Param-*` HTTP header.""" - -NAME_BEARING_METHODS: Final[Mapping[str, str]] = MappingProxyType( - { - "tools/call": "name", - "prompts/get": "name", - "resources/read": "uri", - } -) -"""Method → params key whose value is mirrored as the `Mcp-Name` HTTP header. - -Shared by client emit (which header to send) and server validate (which body -field to compare against), so both ends agree on the field by construction. -""" - -_B64_SENTINEL = re.compile(r"^=\?base64\?(?P.*)\?=$") -# RFC 7230 token chars minus DEL; visible ASCII 0x20-0x7E is the practical bound for a header value. -_HEADER_SAFE = re.compile(r"^[\x20-\x7E]*$") -# RFC 9110 §5.6.2 token: the only characters permitted in an HTTP field name. -_RFC9110_TOKEN = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") -# JSON-Schema types the spec permits to carry `x-mcp-header` (transports.mdx -# §Custom Headers). `number` is explicitly forbidden — float→str is not -# portable across implementations. -_X_MCP_HEADER_PRIMITIVE_TYPES: Final = frozenset({"string", "integer", "boolean"}) - -# JSON Schema 2020-12 applicator keywords whose values are themselves schema -# positions, grouped by value shape. `properties` is handled separately as the -# only keyword that preserves the statically-reachable chain; every keyword -# here drops the chain to None. Instance-data keywords (`default`, `examples`, -# `const`, `enum`) and `$ref`/`$dynamicRef` are deliberately absent so the -# walk never mistakes data for an annotation and never dereferences. -_SUBSCHEMA_SINGLE: Final = frozenset( - { - "items", - "contains", - "unevaluatedItems", - "additionalProperties", - "propertyNames", - "unevaluatedProperties", - "not", - "if", - "then", - "else", - "contentSchema", - } -) -_SUBSCHEMA_LIST: Final = frozenset({"allOf", "anyOf", "oneOf", "prefixItems"}) -_SUBSCHEMA_MAP: Final = frozenset({"patternProperties", "dependentSchemas", "$defs", "definitions"}) - - -def _walk_schema_positions(root: Any) -> Iterator[tuple[tuple[str, ...] | None, dict[str, Any]]]: - """Yield `(properties_path, schema)` for every schema position in `root`. - - `properties_path` is the chain of `properties` keys from the root to the - position, or `None` once any other applicator keyword has been crossed. - The root itself yields `()`. Only the JSON Schema 2020-12 applicators - listed above are entered; instance-data keywords are not, and `$ref` is - not dereferenced, so the walk terminates on any finite JSON value. An - explicit stack keeps the function total even on pathologically deep input. - """ - stack: list[tuple[tuple[str, ...] | None, Any]] = [((), root)] - while stack: - path, node = stack.pop() - if not isinstance(node, dict): - continue - schema = cast(dict[str, Any], node) - yield path, schema - for kw, val in schema.items(): - if kw == "properties" and isinstance(val, dict): - for name, sub in cast(dict[str, Any], val).items(): - stack.append(((*path, name) if path is not None else None, sub)) - elif kw in _SUBSCHEMA_SINGLE: - stack.append((None, val)) - elif kw in _SUBSCHEMA_LIST and isinstance(val, list): - stack.extend((None, sub) for sub in cast(list[Any], val)) - elif kw in _SUBSCHEMA_MAP and isinstance(val, dict): - stack.extend((None, sub) for sub in cast(dict[str, Any], val).values()) - - -def encode_header_value(value: str) -> str: - """Wrap `value` in the `=?base64?...?=` sentinel when it would not survive an HTTP field round-trip. - - Plain printable ASCII without leading/trailing whitespace passes verbatim; - anything else (control chars, non-ASCII, edge whitespace, or a value that - already looks like the sentinel) is base64-wrapped so the receiver can - recover the exact bytes. - """ - if _HEADER_SAFE.fullmatch(value) and value == value.strip() and not _B64_SENTINEL.fullmatch(value): - return value - return f"=?base64?{base64.b64encode(value.encode('utf-8')).decode('ascii')}?=" - - -def decode_header_value(value: str | None) -> str | None: - """Inverse of :func:`encode_header_value`. - - Returns the value verbatim unless it carries the `=?base64?...?=` sentinel, - in which case the payload is decoded as UTF-8. A malformed sentinel (bad - base64, non-canonical base64, or bad UTF-8) yields `None` so a corrupt - header never matches a body value by accident. `None` in → `None` out so - callers can pass `headers.get(...)` directly. - """ - if value is None: - return None - m = _B64_SENTINEL.fullmatch(value) - if m is None: - return value - payload = m.group("payload") - try: - decoded = base64.b64decode(payload, validate=True) - except binascii.Error: - return None - # Reject non-canonical base64 (e.g. non-zero trailing bits), which - # `validate=True` tolerates; the encoder only ever emits canonical form. - if base64.b64encode(decoded).decode("ascii") != payload: - return None - try: - return decoded.decode("utf-8") - except UnicodeDecodeError: - return None - - -def find_invalid_x_mcp_header(input_schema: Any) -> str | None: - """Return a reason string if any `x-mcp-header` annotation in `input_schema` is invalid; else `None`. - - Walks every JSON Schema 2020-12 schema position. An annotation is valid - only when it sits on a property statically reachable from the root via a - chain of pure `properties` keys, names a non-empty RFC 9110 token, is on - an integer/string/boolean property, and is case-insensitively unique - across the whole schema. A `None` / non-mapping schema has no schema - positions and returns `None`. - """ - seen: dict[str, str] = {} - for path, schema in _walk_schema_positions(input_schema): - if X_MCP_HEADER_KEY not in schema: - continue - if not path: # None (off the pure-properties chain) or () (the root itself) - return f"{X_MCP_HEADER_KEY} found at a schema position not reachable via a pure `properties` chain" - where = ".".join(path) - header = schema[X_MCP_HEADER_KEY] - # Wrong type and malformed value are distinct failures with distinct messages: the - # non-str arm returns before any interpolation, because `repr` of an arbitrary - # schema value is not total (a large `int` exceeds `sys.get_int_max_str_digits`). - if not isinstance(header, str): - return f"property {where!r}: {X_MCP_HEADER_KEY} must be a string, not {type(header).__name__}" - if not _RFC9110_TOKEN.fullmatch(header): - return f"property {where!r}: {X_MCP_HEADER_KEY} {header!r} is not an RFC 9110 token" - prop_type = schema.get("type") - if not isinstance(prop_type, str): - return ( - f"property {where!r}: {X_MCP_HEADER_KEY} is only permitted on " - f"integer/string/boolean properties (the type keyword is {type(prop_type).__name__}, not a string)" - ) - if prop_type not in _X_MCP_HEADER_PRIMITIVE_TYPES: - return ( - f"property {where!r}: {X_MCP_HEADER_KEY} is only permitted on " - f"integer/string/boolean properties (got {prop_type!r})" - ) - lower = header.lower() - if lower in seen: - return f"{X_MCP_HEADER_KEY} {header!r} on property {where!r} duplicates property {seen[lower]!r}" - seen[lower] = where - return None - - -MCP_PARAM_HEADER_PREFIX: Final = "Mcp-Param-" -"""Prefix the `x-mcp-header` token is joined to, forming the per-parameter HTTP header name.""" - - -def x_mcp_header_map(input_schema: Any) -> dict[tuple[str, ...], str]: - """Map each property carrying a valid `x-mcp-header` to its annotation token, keyed by property path. - - The key is the chain of `properties` keys from the schema root to the - annotated property; a top-level property has a one-element path, a nested - one a longer path. Call only on a schema that - :func:`find_invalid_x_mcp_header` accepts; an invalid schema yields an - undefined subset. - """ - return {path: token for path, token, _ in _annotated_positions(input_schema)} - - -def _annotated_positions(input_schema: Any) -> Iterator[tuple[tuple[str, ...], str, dict[str, Any]]]: - """Yield `(path, token, schema)` for every statically-reachable `x-mcp-header` annotation. - - Shared by client emit and server validate so both ends agree on what counts as a declared header. - """ - for path, schema in _walk_schema_positions(input_schema): - if path and isinstance(token := schema.get(X_MCP_HEADER_KEY), str): - yield path, token, schema - - -def _render_header_scalar(value: Any) -> str | None: - """Render `value` the way the client mirrors it into a header, or `None` when no rendering exists. - - Shared by emit and validate so both sides agree on what is mirrorable: - non-primitives and ints beyond CPython's int-to-str digit limit are not. - """ - if isinstance(value, bool): - return "true" if value else "false" - if not isinstance(value, str | int | float): - return None - try: - return str(value) - except ValueError: - return None - - -def mcp_param_headers(header_map: Mapping[tuple[str, ...], str], arguments: Mapping[str, Any]) -> dict[str, str]: - """Build the `Mcp-Param-*` headers a `tools/call` mirrors from its arguments. - - For each `(path, token)` in `header_map`, read the value at that property - path in `arguments` and, when it is present and not `None`, emit - `Mcp-Param-` carrying it: `bool` as `true`/`false`, other scalars via - `str`, each passed through :func:`encode_header_value` so a non-token value - is base64-wrapped. A path that hits a missing key or a non-mapping node is - skipped, matching the spec's "omit the header when no value is present", - as is a value with no header rendering. - """ - headers: dict[str, str] = {} - for path, token in header_map.items(): - value = _value_at_path(arguments, path) - if value is None or (rendered := _render_header_scalar(value)) is None: - continue - headers[f"{MCP_PARAM_HEADER_PREFIX}{token}"] = encode_header_value(rendered) - return headers - - -def _value_at_path(arguments: Mapping[str, Any], path: tuple[str, ...]) -> Any: - """Read the value at a `properties`-key path in `arguments`, or `None` if any step is missing or non-mapping.""" - node: Any = arguments - for key in path: - if not isinstance(node, Mapping): - return None - node = cast("Mapping[str, Any]", node).get(key) - return node - - -# INTERNAL_ERROR is deliberately unmapped (→ HTTP 200): the spec assigns no status to -# -32603, and whether handler-origin errors get 5xx is an open S4 question — see TODO(L66). -ERROR_CODE_HTTP_STATUS: Final[Mapping[int, int]] = MappingProxyType( - { - PARSE_ERROR: 400, - INVALID_REQUEST: 400, - INVALID_PARAMS: 400, - HEADER_MISMATCH: 400, - MISSING_REQUIRED_CLIENT_CAPABILITY: 400, - UNSUPPORTED_PROTOCOL_VERSION: 400, - METHOD_NOT_FOUND: 404, - } -) -"""HTTP status to send for a JSON-RPC `error.code`. - -Consulted for classifier-origin *and* handler-origin errors, so one table -decides the wire status regardless of where the error was produced. Unmapped -codes fall back to the caller's default (typically 200). -""" - - -@dataclass(frozen=True) -class InboundModernRoute: - """A modern-protocol request whose envelope passed every ladder rung. - - `client_info` and `client_capabilities` are the raw envelope values; the - classifier checks presence only, not shape, and `client_info` is `None` - when the (optional, SHOULD-include) key is absent. Method existence is not - a ladder rung — kernel dispatch is the single source of truth for that. - """ - - protocol_version: str - client_info: Any - client_capabilities: Any - - -@dataclass(frozen=True) -class InboundLadderRejection: - """The first ladder rung that failed, as JSON-RPC error fields.""" - - code: int - message: str - data: Any = None - - -_ROUTING_HEADER_NAMES: Final = frozenset({MCP_PROTOCOL_VERSION_HEADER, MCP_METHOD_HEADER, MCP_NAME_HEADER}) - - -def find_duplicated_routing_header(headers: Iterable[tuple[str, str]]) -> str | None: - """Name of a routing header supplied more than once in raw header lines, or `None`. - - Takes raw `(name, value)` pairs — a folded mapping hides duplicates. A - duplicate is rejected because first-copy and last-copy readers would - disagree. `Mcp-Param-*` duplicates are :func:`validate_mcp_param_headers`'s job. - """ - seen: set[str] = set() - for name, _ in headers: - key = name.lower() - if key in _ROUTING_HEADER_NAMES: - if key in seen: - return key - seen.add(key) - return None - - -def unsupported_protocol_version_rejection( - requested: str, supported_modern_versions: Sequence[str] = MODERN_PROTOCOL_VERSIONS -) -> InboundLadderRejection | None: - """The `UNSUPPORTED_PROTOCOL_VERSION` rejection for `requested`, or `None` if it is served. - - The request ladder's last rung, shared with the transport's notification arm - so both message kinds name the same `supported` list in the same words. - """ - if requested in supported_modern_versions: - return None - return InboundLadderRejection( - code=UNSUPPORTED_PROTOCOL_VERSION, - message="Unsupported protocol version", - data=UnsupportedProtocolVersionErrorData( - supported=list(supported_modern_versions), requested=requested - ).model_dump(mode="json"), - ) - - -def classify_inbound_request( - body: Mapping[str, Any], - *, - headers: Mapping[str, str] | None = None, - supported_modern_versions: Sequence[str] = MODERN_PROTOCOL_VERSIONS, -) -> InboundModernRoute | InboundLadderRejection: - """Run the modern-protocol validation ladder over a decoded JSON-RPC body. - - Rungs, in order — first failure wins: - - 1. `params._meta` is a mapping carrying the required envelope pair - (protocol version, client capabilities) → else - :data:`~mcp_types.jsonrpc.INVALID_PARAMS` naming the missing key(s) - (basic/index.mdx "Per-request protocol fields"). Client info is - optional (SHOULD-include, spec PR #3002); absent reads as `None`. - 2. When `headers` is given, `MCP-Protocol-Version` equals the envelope's - protocol version, `Mcp-Method` equals `body.method`, and — for the - methods in :data:`NAME_BEARING_METHODS` — `Mcp-Name` equals the named - body param → else :data:`~mcp_types.jsonrpc.HEADER_MISMATCH`. Runs - before the supported-version rung so a client that disagrees with itself - is told so, rather than told the body's version is unsupported. - 3. The envelope's protocol version is a string in - `supported_modern_versions` → non-string values are - :data:`~mcp_types.jsonrpc.INVALID_PARAMS` (a shape defect, not a - negotiation outcome), else - :data:`~mcp_types.jsonrpc.UNSUPPORTED_PROTOCOL_VERSION` with - `data = {"supported": [...], "requested": }`. - - Method existence is *not* a rung: kernel dispatch owns that decision so - custom-registered methods route and the answer lives in one place. - - Args: - body: The decoded JSON-RPC request mapping. Envelope shape - (`jsonrpc` / `id`) is not checked here. - headers: Transport headers keyed by lowercase name, or `None` to - skip the header rung (non-HTTP callers). - supported_modern_versions: Modern protocol revisions this server - accepts on the per-request-envelope path. - """ - try: - meta_value = body["params"]["_meta"] - except (KeyError, TypeError): - meta_value = None - if not isinstance(meta_value, Mapping): - return InboundLadderRejection( - code=INVALID_PARAMS, - message="params._meta must be an object carrying the required " - f"{PROTOCOL_VERSION_META_KEY!r} and {CLIENT_CAPABILITIES_META_KEY!r} envelope keys", - ) - meta = cast("Mapping[str, Any]", meta_value) - if missing := [key for key in (PROTOCOL_VERSION_META_KEY, CLIENT_CAPABILITIES_META_KEY) if key not in meta]: - return InboundLadderRejection( - code=INVALID_PARAMS, - message=f"params._meta is missing the required envelope key(s): {', '.join(missing)}", - ) - protocol_version: Any = meta[PROTOCOL_VERSION_META_KEY] - client_info: Any = meta.get(CLIENT_INFO_META_KEY) - client_capabilities: Any = meta[CLIENT_CAPABILITIES_META_KEY] - if headers is not None: - version_header = headers.get(MCP_PROTOCOL_VERSION_HEADER) - # Presence is checked explicitly: a null body version would otherwise - # slip the equality check (None == None) and mask the absent header. - if version_header is None or version_header != protocol_version: - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{MCP_PROTOCOL_VERSION_HEADER} header does not match the request envelope's protocol version", - ) - method: Any = body.get("method") - if headers.get(MCP_METHOD_HEADER) != method: - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{MCP_METHOD_HEADER} header does not match the request body's method", - ) - name_key = NAME_BEARING_METHODS.get(method) - if name_key is not None: - # Rung 1 already proved body["params"] is a mapping (its `_meta` is one). - body_value = cast("Mapping[str, Any]", body["params"]).get(name_key) - if body_value is not None and decode_header_value(headers.get(MCP_NAME_HEADER)) != body_value: - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{MCP_NAME_HEADER} header does not match the request body's {name_key!r} parameter", - ) - - if not isinstance(protocol_version, str): - # Rung 3's precondition: a shape defect, not a version-negotiation - # outcome - -32022 is the one code auto-negotiating clients do NOT - # fall back from, and the typed rung-3 payload itself requires a - # string `requested`. Sits after the header rung, which fires first - # for every header-bearing entry (an absent version header is a - # mismatch, and a present one is a string that can never equal a - # non-string body value) - so this rejection is reachable only on - # header-less transports. - return InboundLadderRejection( - code=INVALID_PARAMS, - message="the protocol-version envelope value must be a string", - ) - - if (unsupported := unsupported_protocol_version_rejection(protocol_version, supported_modern_versions)) is not None: - return unsupported - - return InboundModernRoute( - protocol_version=protocol_version, - client_info=client_info, - client_capabilities=client_capabilities, - ) - - -# Header values eligible for the spec's numeric-comparison SHOULD; scientific -# notation never compares numerically (matching the typescript-sdk's gate). -_CANONICAL_DECIMAL = re.compile(r"^-?[0-9]+(\.[0-9]+)?$") - - -def _mcp_param_value_matches(prop_type: Any, value: Any, rendered: str, decoded: str) -> bool: - """True when a decoded `Mcp-Param-*` header value agrees with the body argument. - - Integer-typed declarations with an integral body value compare numerically - (`42` matches `42.0`, the spec's SHOULD) for canonical-decimal headers — - exact, no float round-trip, so values beyond the IEEE754 safe range still - compare. Anything else compares against `rendered`, the emit-side rendering. - """ - if ( - prop_type == "integer" - and not isinstance(value, bool) - and (isinstance(value, int) or (isinstance(value, float) and value.is_integer())) - and _CANONICAL_DECIMAL.fullmatch(decoded) is not None - ): - whole, _, fraction = decoded.partition(".") - if fraction and set(fraction) != {"0"}: - return False - try: - return int(whole) == int(value) - except ValueError: - return False - return decoded == rendered - - -def validate_mcp_param_headers( - input_schema: Any, - arguments: Mapping[str, Any], - headers: Mapping[str, str], -) -> InboundLadderRejection | None: - """Compare a `tools/call` request's `Mcp-Param-*` headers against its body arguments. - - Each annotated property's header and argument must agree: present together - and equal after sentinel decoding, or absent together (`null` counts as - absent). Returns the first failure as a `HEADER_MISMATCH` rejection, else `None`. +import mcp_client.shared.inbound as _implementation +from mcp_client.shared.inbound import ( + _B64_SENTINEL as _B64_SENTINEL, +) +from mcp_client.shared.inbound import ( + _CANONICAL_DECIMAL as _CANONICAL_DECIMAL, +) +from mcp_client.shared.inbound import ( + _HEADER_SAFE as _HEADER_SAFE, +) +from mcp_client.shared.inbound import ( + _RFC9110_TOKEN as _RFC9110_TOKEN, +) +from mcp_client.shared.inbound import ( + _ROUTING_HEADER_NAMES as _ROUTING_HEADER_NAMES, +) +from mcp_client.shared.inbound import ( + _SUBSCHEMA_LIST as _SUBSCHEMA_LIST, +) +from mcp_client.shared.inbound import ( + _SUBSCHEMA_MAP as _SUBSCHEMA_MAP, +) +from mcp_client.shared.inbound import ( + _SUBSCHEMA_SINGLE as _SUBSCHEMA_SINGLE, +) +from mcp_client.shared.inbound import ( + _X_MCP_HEADER_PRIMITIVE_TYPES as _X_MCP_HEADER_PRIMITIVE_TYPES, +) +from mcp_client.shared.inbound import ( + ERROR_CODE_HTTP_STATUS as ERROR_CODE_HTTP_STATUS, +) +from mcp_client.shared.inbound import ( + MCP_METHOD_HEADER as MCP_METHOD_HEADER, +) +from mcp_client.shared.inbound import ( + MCP_NAME_HEADER as MCP_NAME_HEADER, +) +from mcp_client.shared.inbound import ( + MCP_PARAM_HEADER_PREFIX as MCP_PARAM_HEADER_PREFIX, +) +from mcp_client.shared.inbound import ( + MCP_PROTOCOL_VERSION_HEADER as MCP_PROTOCOL_VERSION_HEADER, +) +from mcp_client.shared.inbound import ( + NAME_BEARING_METHODS as NAME_BEARING_METHODS, +) +from mcp_client.shared.inbound import ( + X_MCP_HEADER_KEY as X_MCP_HEADER_KEY, +) +from mcp_client.shared.inbound import ( + InboundLadderRejection as InboundLadderRejection, +) +from mcp_client.shared.inbound import ( + InboundModernRoute as InboundModernRoute, +) +from mcp_client.shared.inbound import ( + _annotated_positions as _annotated_positions, +) +from mcp_client.shared.inbound import ( + _mcp_param_value_matches as _mcp_param_value_matches, +) +from mcp_client.shared.inbound import ( + _render_header_scalar as _render_header_scalar, +) +from mcp_client.shared.inbound import ( + _value_at_path as _value_at_path, +) +from mcp_client.shared.inbound import ( + _walk_schema_positions as _walk_schema_positions, +) +from mcp_client.shared.inbound import ( + classify_inbound_request as classify_inbound_request, +) +from mcp_client.shared.inbound import ( + decode_header_value as decode_header_value, +) +from mcp_client.shared.inbound import ( + encode_header_value as encode_header_value, +) +from mcp_client.shared.inbound import ( + find_duplicated_routing_header as find_duplicated_routing_header, +) +from mcp_client.shared.inbound import ( + find_invalid_x_mcp_header as find_invalid_x_mcp_header, +) +from mcp_client.shared.inbound import ( + mcp_param_headers as mcp_param_headers, +) +from mcp_client.shared.inbound import ( + unsupported_protocol_version_rejection as unsupported_protocol_version_rejection, +) +from mcp_client.shared.inbound import ( + validate_mcp_param_headers as validate_mcp_param_headers, +) +from mcp_client.shared.inbound import ( + x_mcp_header_map as x_mcp_header_map, +) - A header whose argument is absent or unrenderable is deliberately rejected: - the spec's purpose clause is exactly an intermediary routing on a value the - body never carried. A duplicated recognized header is rejected — first-copy - and last-copy readers would disagree. A schema :func:`find_invalid_x_mcp_header` - rejects validates nothing: conforming clients drop the tool and emit no headers. - """ - if find_invalid_x_mcp_header(input_schema) is not None: - return None - folded: dict[str, str] = {} - duplicated: set[str] = set() - for name, value in headers.items(): - key = name.lower() - if key in folded: - duplicated.add(key) - folded[key] = value - for path, token, schema in _annotated_positions(input_schema): - header_name = f"{MCP_PARAM_HEADER_PREFIX}{token}" - key = header_name.lower() - raw = folded.get(key) - value = _value_at_path(arguments, path) - argument = ".".join(path) - if raw is not None and key in duplicated: - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{header_name} header appears more than once", - ) - if value is None: - if raw is not None: - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{header_name} header is present but the request body's {argument!r} argument is absent", - ) - continue - rendered = _render_header_scalar(value) - if rendered is None: - # Unrenderable value: a conforming client omitted the header, so one claiming it can never match. - if raw is not None: - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{header_name} header does not match the request body's {argument!r} argument", - ) - continue - if raw is None: - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{header_name} header is missing but the request body's {argument!r} argument is present", - ) - decoded = decode_header_value(raw) - if decoded is None: - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{header_name} header carries a malformed base64 sentinel value", - ) - if not _mcp_param_value_matches(schema.get("type"), value, rendered, decoded): - return InboundLadderRejection( - code=HEADER_MISMATCH, - message=f"{header_name} header does not match the request body's {argument!r} argument", - ) - return None +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 87bdf31ceb..f1f69269d4 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -1,836 +1,56 @@ -"""JSON-RPC `Dispatcher` over the `SessionMessage` stream contract all transports speak. +import sys -Owns request-id correlation, the receive loop, per-request task isolation, -cancellation/progress wiring, and the single exception-to-wire boundary; -methods and params are otherwise opaque strings and dicts. -""" - -from __future__ import annotations - -import contextvars -import logging -from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, field -from functools import partial -from typing import Any, Generic, Literal, cast - -import anyio -import anyio.abc -import anyio.lowlevel -from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream -from mcp_types import ( - CONNECTION_CLOSED, - INTERNAL_ERROR, - INVALID_PARAMS, - REQUEST_TIMEOUT, - ErrorData, - JSONRPCError, - JSONRPCMessage, - JSONRPCNotification, - JSONRPCRequest, - JSONRPCResponse, - ProgressToken, - RequestId, +import mcp_client.shared.jsonrpc_dispatcher as _implementation +from mcp_client.shared.jsonrpc_dispatcher import ( + _ABANDON_WRITE_TIMEOUT as _ABANDON_WRITE_TIMEOUT, ) -from opentelemetry.trace import SpanKind -from pydantic import ValidationError -from typing_extensions import TypeVar - -from mcp.shared._compat import resync_tracer -from mcp.shared._otel import inject_trace_context, otel_span -from mcp.shared._stream_protocols import ReadStream, WriteStream -from mcp.shared.dispatcher import ( - CallOptions, - DispatchContext, - Dispatcher, - OnNotify, - OnNotifyIntercept, - OnRequest, - ProgressFnT, - as_request_id, - coerce_request_id, - run_notify_intercept, +from mcp_client.shared.jsonrpc_dispatcher import ( + _SHUTDOWN_WRITE_TIMEOUT as _SHUTDOWN_WRITE_TIMEOUT, ) -from mcp.shared.exceptions import MCPError, NoBackChannelError -from mcp.shared.message import ( - ClientMessageMetadata, - MessageMetadata, - ServerMessageMetadata, - SessionMessage, +from mcp_client.shared.jsonrpc_dispatcher import ( + JSONRPCDispatcher as JSONRPCDispatcher, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + PeerCancelMode as PeerCancelMode, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + TransportT as TransportT, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + _contained_notify as _contained_notify, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + _default_transport_builder as _default_transport_builder, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + _InFlight as _InFlight, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + _JSONRPCDispatchContext as _JSONRPCDispatchContext, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + _OutboundPlan as _OutboundPlan, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + _Pending as _Pending, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + _plan_outbound as _plan_outbound, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + _shielded_progress as _shielded_progress, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + cancelled_request_id_from_params as cancelled_request_id_from_params, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + handler_exception_to_error_data as handler_exception_to_error_data, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + logger as logger, +) +from mcp_client.shared.jsonrpc_dispatcher import ( + progress_token_from_params as progress_token_from_params, ) -from mcp.shared.transport_context import TransportContext - -__all__ = [ - "JSONRPCDispatcher", - "cancelled_request_id_from_params", - "handler_exception_to_error_data", - "progress_token_from_params", -] - -logger = logging.getLogger(__name__) - -_ABANDON_WRITE_TIMEOUT: float = 5 -"""Bound for courtesy-cancel writes on the abandon paths; the caller-cancel -arm shields its write, so a wedged transport would otherwise hang it uncancellably.""" - -_SHUTDOWN_WRITE_TIMEOUT: float = 1 -"""Tighter bound for the shutdown-arm error write so a wedged transport can't hold session close.""" - -TransportT = TypeVar("TransportT", bound=TransportContext, default=TransportContext) - -PeerCancelMode = Literal["interrupt", "signal"] -"""How `notifications/cancelled` is applied: `"interrupt"` (default) cancels -the handler's scope; `"signal"` only sets `ctx.cancel_requested` and lets the -handler run to completion. Either way the cancelled request is never -answered - the handler's eventual result or error is dropped, not written.""" - - -def handler_exception_to_error_data(exc: BaseException) -> ErrorData | None: - """Map a handler-raised exception to its wire `ErrorData`. - - The two rungs every dispatcher shares: an `MCPError` carries its own - `ErrorData`; a pydantic `ValidationError` is the spec's INVALID_PARAMS - with empty ``data`` (no pydantic text on the wire). Returns ``None`` for - any other exception so each caller applies its own catch-all - - `JSONRPCDispatcher` currently pins ``code=0`` for v1 compat, - the modern HTTP entry uses `INTERNAL_ERROR`. - """ - if isinstance(exc, MCPError): - return exc.error - if isinstance(exc, ValidationError): - return ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data="") - return None - - -def progress_token_from_params(params: Mapping[str, Any] | None) -> ProgressToken | None: - """Read `params._meta.progressToken`; reject bool (bool subclasses int, so True would alias 1).""" - match params: - case {"_meta": {"progressToken": str() | int() as token}} if not isinstance(token, bool): - return token - case _: - return None - - -def cancelled_request_id_from_params(params: Mapping[str, Any] | None) -> RequestId | None: - """Read `params.requestId` from a `notifications/cancelled` (`as_request_id` shape rules).""" - return as_request_id((params or {}).get("requestId")) - - -@dataclass(slots=True) -class _Pending: - """An outbound request awaiting its response.""" - - send: MemoryObjectSendStream[dict[str, Any] | ErrorData] - receive: MemoryObjectReceiveStream[dict[str, Any] | ErrorData] - on_progress: ProgressFnT | None = None - - -@dataclass(slots=True) -class _InFlight(Generic[TransportT]): - """An inbound request currently being handled.""" - - scope: anyio.CancelScope - dctx: _JSONRPCDispatchContext[TransportT] - - -@dataclass -class _JSONRPCDispatchContext(Generic[TransportT]): - """Concrete `DispatchContext` produced for each inbound JSON-RPC message.""" - - transport: TransportT - _dispatcher: JSONRPCDispatcher[TransportT] - _request_id: RequestId | None - message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework - """Transport-attached `SessionMessage.metadata` that the server lifts onto its request context.""" - _progress_token: ProgressToken | None = None - _closed: bool = False - cancel_requested: anyio.Event = field(default_factory=anyio.Event) - - @property - def request_id(self) -> RequestId | None: - return self._request_id - - @property - def can_send_request(self) -> bool: - return self.transport.can_send_request and not self._closed - - async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: - if self._closed: - logger.debug("dropped %s: dispatch context closed", method) - return - await self._dispatcher.notify(method, params, opts, _related_request_id=self._request_id) - - async def send_raw_request( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None = None, - ) -> dict[str, Any]: - if not self.can_send_request: - raise NoBackChannelError(method) - return await self._dispatcher.send_raw_request(method, params, opts, _related_request_id=self._request_id) - - async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: - if self._progress_token is None: - return - params: dict[str, Any] = {"progressToken": self._progress_token, "progress": progress} - if total is not None: - params["total"] = total - if message is not None: - params["message"] = message - await self.notify("notifications/progress", params) - - def close(self) -> None: - self._closed = True - - -def _default_transport_builder(metadata: MessageMetadata) -> TransportContext: - """The `TransportContext` for a message, honoring the transport's own verdict when it stamps one. - - A message reads as riding a full duplex pipe (`can_send_request=True`) - unless the transport that framed it says otherwise on the metadata it - attached, so a transport whose response has no room for a server request - (streamable HTTP in JSON-response mode) needs no wiring from whoever drives - its streams. - """ - can_send_request = metadata.can_send_request if isinstance(metadata, ServerMessageMetadata) else True - return TransportContext(kind="jsonrpc", can_send_request=can_send_request) - - -def _shielded_progress(fn: ProgressFnT) -> ProgressFnT: - """Wrap a user progress callback so an exception can't cancel the dispatcher's task group.""" - - async def _wrapped(progress: float, total: float | None, message: str | None) -> None: - try: - await fn(progress, total, message) - except Exception: - logger.exception("progress callback raised") - - return _wrapped - - -def _contained_notify(fn: OnNotify) -> OnNotify: - """Wrap a notification handler so it can't crash the dispatcher (same boundary as `_shielded_progress`).""" - - async def _wrapped(dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None: - try: - await fn(dctx, method, params) - except Exception: - logger.exception("notification handler for %r raised", method) - - return _wrapped - - -@dataclass(slots=True, frozen=True) -class _OutboundPlan: - """Outbound metadata plus whether abandoning the request sends a courtesy `notifications/cancelled`.""" - - metadata: MessageMetadata - cancel_on_abandon: bool - - -def _plan_outbound(related_request_id: RequestId | None, opts: CallOptions | None) -> _OutboundPlan: - """Choose the outbound `SessionMessage.metadata` and the abandon-cancellation policy. - - `related_request_id` wins over resumption hints (they are dropped). Only - hints that actually reach the transport suppress the courtesy cancel - a - request that is neither resumable nor cancelled would leak the peer's work. - """ - opts = opts or {} - cancel_on_abandon = opts.get("cancel_on_abandon", True) - token = opts.get("resumption_token") - on_token = opts.get("on_resumption_token") - headers = opts.get("headers") - if related_request_id is not None: - if token is not None or on_token is not None: - logger.debug( - "dropping resumption hints: related_request_id %r takes precedence on metadata", related_request_id - ) - return _OutboundPlan(ServerMessageMetadata(related_request_id=related_request_id), cancel_on_abandon) - if token is not None or on_token is not None: - return _OutboundPlan( - ClientMessageMetadata(resumption_token=token, on_resumption_token_update=on_token, headers=headers), - cancel_on_abandon=False, - ) - if headers: - return _OutboundPlan(ClientMessageMetadata(headers=headers), cancel_on_abandon) - return _OutboundPlan(None, cancel_on_abandon) - - -class JSONRPCDispatcher(Dispatcher[TransportT]): - """`Dispatcher` over the `SessionMessage` stream contract. - - Explicit Protocol base so pyright checks conformance at the class definition. - """ - - def __init__( - self, - read_stream: ReadStream[SessionMessage | Exception], - write_stream: WriteStream[SessionMessage], - *, - transport_builder: Callable[[MessageMetadata], TransportT] | None = None, - peer_cancel_mode: PeerCancelMode = "interrupt", - raise_handler_exceptions: bool = False, - inline_methods: frozenset[str] = frozenset(), - on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None, - ) -> None: - """Wire a dispatcher over a transport's `SessionMessage` stream pair. - - Args: - transport_builder: Builds each message's `TransportContext` from - its `SessionMessage.metadata`. - raise_handler_exceptions: Re-raise handler exceptions out of - `run()` after the error response is written. - inline_methods: Methods awaited in the read loop before the next - message is dequeued (e.g. `initialize`); an inline handler - that awaits the peer deadlocks the parked loop. - on_stream_exception: Observer for `Exception` items on the read - stream; without it they are debug-logged and dropped. Awaited - inline in the read loop, so a slow observer stalls dispatch. - """ - self._read_stream = read_stream - self._write_stream = write_stream - # With transport_builder omitted, TransportT defaults to - # TransportContext; pyright can't connect the two, hence the cast. - self._transport_builder = cast( - "Callable[[MessageMetadata], TransportT]", - transport_builder or _default_transport_builder, - ) - self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode - self._raise_handler_exceptions = raise_handler_exceptions - self._inline_methods = inline_methods - self.on_stream_exception = on_stream_exception - """Observer for ``Exception`` items on the read stream. Mutable so a session can - bind it after the dispatcher is built (e.g. ``ClientSession`` routing into - ``message_handler``); only consulted inside ``run()`` so pre-enter assignment is safe.""" - - self._next_id = 0 - self._pending: dict[RequestId, _Pending] = {} - self._in_flight: dict[RequestId, _InFlight[TransportT]] = {} - self._on_notify_intercept: OnNotifyIntercept | None = None - self._tg: anyio.abc.TaskGroup | None = None - self._running = False - self._closed = False - - async def send_raw_request( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None = None, - *, - _related_request_id: RequestId | None = None, - ) -> dict[str, Any]: - """Send a JSON-RPC request and await its response. - - `_related_request_id` is set only by `_JSONRPCDispatchContext` so that - mid-handler requests route onto the inbound request's SSE stream. - - Raises: - MCPError: Peer error response; `REQUEST_TIMEOUT` if - `opts["timeout"]` elapsed; `CONNECTION_CLOSED` if the - transport closed or the dispatcher shut down. - RuntimeError: Called before `run()`. - """ - # Post-close sends get the same CONNECTION_CLOSED contract as in-flight waiters. - if self._closed: - raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") - if not self._running: - raise RuntimeError("JSONRPCDispatcher.send_raw_request called before run()") - opts = opts or {} - supplied_id = opts.get("request_id") - if supplied_id is not None: - request_id: RequestId = supplied_id - # The pending key gets the same coercion `_resolve_pending` applies - # to inbound response ids, so a supplied "7" still correlates - # whether the peer echoes "7" or 7. The wire id stays verbatim. - pending_key = coerce_request_id(request_id) - if pending_key in self._pending: - raise ValueError(f"request id {request_id!r} is already in flight") - else: - # Mint past any key a supplied id occupies: the collision error is - # reserved for the caller who actually chose the id. - request_id = self._allocate_id() - while request_id in self._pending: - request_id = self._allocate_id() - pending_key = request_id - out_params = dict(params) if params is not None else {} - out_meta = dict(out_params.get("_meta") or {}) - on_progress = opts.get("on_progress") - if on_progress is not None: - # The request id doubles as the progress token, so `_pending[token]` finds `on_progress` directly. - out_meta["progressToken"] = request_id - out_params["_meta"] = out_meta - - # buffer=1: a close signal can arrive before the waiter parks in receive(); - # a WouldBlock later just means the waiter already has its one outcome. - send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1) - pending = _Pending(send=send, receive=receive, on_progress=on_progress) - self._pending[pending_key] = pending - - plan = _plan_outbound(_related_request_id, opts) - # Spec MUST: only previously-issued requests may be cancelled. A write - # interrupted by cancellation may still have delivered (a memory-stream - # send can hand its item to the receiver and still raise), so a started - # write counts as issued: the peer ignores a cancel for an id it never - # saw, while skipping it would leak a delivered request's handler. - request_write_started = False - timeout_armed = False - - target = out_params.get("name") - span_name = f"MCP send {method}{f' {target}' if isinstance(target, str) else ''}" - # TODO(maxisbey): move the otel span + inject into an outbound - # middleware once that seam exists; the dispatcher should not own otel. - try: - with otel_span( - span_name, - kind=SpanKind.CLIENT, - attributes={"mcp.method.name": method, "jsonrpc.request.id": str(request_id)}, - ): - # SEP-414: inject W3C trace context; `_meta` stays on the wire even with a no-op tracer. - inject_trace_context(out_meta) - msg = JSONRPCRequest(jsonrpc="2.0", id=request_id, method=method, params=out_params) - # Surface a pre-existing cancellation while the request provably - # never started; past this point a cancelled write counts as issued. - await anyio.lowlevel.checkpoint_if_cancelled() - request_write_started = True - try: - await self._write(msg, plan.metadata) - except (anyio.BrokenResourceError, anyio.ClosedResourceError): - # Transport tore down before run() noticed EOF; surface the documented contract. - raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") from None - with anyio.fail_after(opts.get("timeout")): - timeout_armed = True - outcome = await receive.receive() - except TimeoutError: - if not timeout_armed: - # `fail_after` arms only after the write, so this TimeoutError is the - # transport's own bounded send() failing - a transport error, not - # `opts["timeout"]` elapsing. Propagate it raw (v1 kept the write - # outside the timeout-catching try and did the same). - raise - # Courtesy cancel (spec-recommended, new vs v1) so the peer stops work; - # unshielded so an outer caller cancellation can still interrupt the write. - if plan.cancel_on_abandon: - await self._final_write( - partial( - self._cancel_outbound, - request_id, - f"timed out after {opts.get('timeout')}s", - _related_request_id, - ), - shield=False, - timeout=_ABANDON_WRITE_TIMEOUT, - describe=f"courtesy cancel for timed-out request {request_id!r}", - ) - raise MCPError(code=REQUEST_TIMEOUT, message=f"Request {method!r} timed out") from None - except anyio.get_cancelled_exc_class(): - # Caller cancelled: bare awaits re-raise here, so the shielded helper - # lets the courtesy cancel go out before we propagate. - if plan.cancel_on_abandon and request_write_started: - await self._final_write( - partial(self._cancel_outbound, request_id, "caller cancelled", _related_request_id), - shield=True, - timeout=_ABANDON_WRITE_TIMEOUT, - describe=f"courtesy cancel for caller-cancelled request {request_id!r}", - ) - raise - finally: - # Remove the waiter on every path so a late response is dropped, not leaked. - self._pending.pop(pending_key, None) - send.close() - receive.close() - - if isinstance(outcome, ErrorData): - raise MCPError(code=outcome.code, message=outcome.message, data=outcome.data) - return outcome - - async def notify( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None = None, - *, - _related_request_id: RequestId | None = None, - ) -> None: - """Send a fire-and-forget notification. - - Fire-and-forget all the way: a post-close send or a write onto a - torn-down transport drops the notification with a debug log instead - of raising (same policy as the response writes and `ctx.notify`). - """ - if self._closed: - logger.debug("dropped %s: dispatcher closed", method) - return - # Leave `params` unset when None: with `exclude_unset=True` an explicit - # None would serialize as `"params": null`, which JSON-RPC 2.0 forbids. - if params is not None: - msg = JSONRPCNotification(jsonrpc="2.0", method=method, params=dict(params)) - else: - msg = JSONRPCNotification(jsonrpc="2.0", method=method) - try: - await self._write(msg, _plan_outbound(_related_request_id, opts).metadata) - except (anyio.BrokenResourceError, anyio.ClosedResourceError): - # Transport tore down before run() noticed EOF. - logger.debug("dropped %s: write stream closed", method) - - async def run( - self, - on_request: OnRequest, - on_notify: OnNotify, - on_notify_intercept: OnNotifyIntercept | None = None, - *, - task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, - ) -> None: - """Drive the receive loop until the read stream closes. - - `task_status.started()` fires once `send_raw_request` is usable. - Single-shot: once the loop ends the dispatcher stays closed and cannot be restarted. - """ - self._on_notify_intercept = on_notify_intercept - try: - # LIFO exits: the write stream closes only after the task-group join, so teardown writes still land. - async with self._write_stream: - async with anyio.create_task_group() as tg: - self._tg = tg - self._running = True - task_status.started() - try: - async with self._read_stream: - try: - async for item in self._read_stream: - # Duck-typed: only `ContextReceiveStream` carries the - # sender's per-message contextvars snapshot. - sender_ctx: contextvars.Context | None = getattr( - self._read_stream, "last_context", None - ) - await self._dispatch(item, on_request, on_notify, sender_ctx) - except anyio.ClosedResourceError: - # Receive end closed under us (stateless SHTTP teardown); same as EOF. - logger.debug("read stream closed by transport; treating as EOF") - # EOF: wake blocked `send_raw_request` waiters with CONNECTION_CLOSED. - self._running = False - self._closed = True - self._fan_out_closed() - finally: - # Cancel in-flight handlers; otherwise the task-group join - # waits on handlers whose callers are already gone. - tg.cancel_scope.cancel() - finally: - # Covers cancel/crash paths that skip the inline fan-out; idempotent. - self._running = False - self._closed = True - self._tg = None - self._fan_out_closed() - await resync_tracer() - - async def _dispatch( - self, - item: SessionMessage | Exception, - on_request: OnRequest, - on_notify: OnNotify, - sender_ctx: contextvars.Context | None, - ) -> None: - """Route one inbound item. - - Only `inline_methods` requests and the `on_stream_exception` observer - are awaited; any other `await` would head-of-line block the read loop. - """ - if isinstance(item, Exception): - if self.on_stream_exception is None: - logger.debug("transport yielded exception: %r", item) - return - try: - await self.on_stream_exception(item) - except Exception: - logger.exception("on_stream_exception observer raised") - return - metadata = item.metadata - msg = item.message - match msg: - case JSONRPCRequest(): - await self._dispatch_request(msg, metadata, on_request, sender_ctx) - case JSONRPCNotification(): - self._dispatch_notification(msg, metadata, on_notify, sender_ctx) - case JSONRPCResponse(): - self._resolve_pending(msg.id, msg.result) - case JSONRPCError(): # pragma: no branch - # Exhaustive over JSONRPCMessage, so the no-match arc is unreachable. - self._resolve_pending(msg.id, msg.error) - - async def _dispatch_request( - self, - req: JSONRPCRequest, - metadata: MessageMetadata, - on_request: OnRequest, - sender_ctx: contextvars.Context | None, - ) -> None: - progress_token = progress_token_from_params(req.params) - try: - transport_ctx = self._transport_builder(metadata) - except Exception: - # A raising builder must cost only this message, not the connection. - logger.exception("transport_builder raised; rejecting request %r", req.id) - self._spawn( - self._write_error, - req.id, - ErrorData(code=INTERNAL_ERROR, message="transport context unavailable"), - sender_ctx=sender_ctx, - ) - return - dctx = _JSONRPCDispatchContext( - transport=transport_ctx, - _dispatcher=self, - _request_id=req.id, - message_metadata=metadata, - _progress_token=progress_token, - ) - scope = anyio.CancelScope() - # TODO(maxisbey): duplicate ids blind-overwrite (v1/TS parity); revisit - # rejecting with INVALID_REQUEST. Key coerced so a stringified - # `notifications/cancelled` id still correlates. - self._in_flight[coerce_request_id(req.id)] = _InFlight(scope=scope, dctx=dctx) - if req.method in self._inline_methods: - # Spawn so `sender_ctx` applies, but park the read loop until the - # handler returns - that's the inline ordering guarantee. - done = anyio.Event() - - async def _run_inline() -> None: - try: - await self._handle_request(req, dctx, scope, on_request) - finally: - done.set() - - self._spawn(_run_inline, sender_ctx=sender_ctx) - await done.wait() - else: - self._spawn(self._handle_request, req, dctx, scope, on_request, sender_ctx=sender_ctx) - - def _dispatch_notification( - self, - msg: JSONRPCNotification, - metadata: MessageMetadata, - on_notify: OnNotify, - sender_ctx: contextvars.Context | None, - ) -> None: - """Route one inbound notification. - - `notifications/cancelled` and `notifications/progress` are intercepted - here (they correlate against the `_in_flight`/`_pending` tables this - layer owns) and still teed to `on_notify` afterwards. The caller's - `on_notify_intercept` then runs in receive order; only unconsumed - notifications reach the spawned `on_notify`. - """ - if msg.method == "notifications/cancelled": - rid = cancelled_request_id_from_params(msg.params) - if rid is not None and (in_flight := self._in_flight.get(coerce_request_id(rid))) is not None: - in_flight.dctx.cancel_requested.set() - if self._peer_cancel_mode == "interrupt": - in_flight.scope.cancel() - elif msg.method == "notifications/progress": - match msg.params: - case {"progressToken": str() | int() as token, "progress": int() | float() as progress} if ( - not isinstance(token, bool) - and not isinstance(progress, bool) - and (pending := self._pending.get(coerce_request_id(token))) is not None - and pending.on_progress is not None - ): - total = msg.params.get("total") - message = msg.params.get("message") - self._spawn( - _shielded_progress(pending.on_progress), - float(progress), - float(total) if isinstance(total, int | float) else None, - message if isinstance(message, str) else None, - sender_ctx=sender_ctx, - ) - case _: - pass - if run_notify_intercept(self._on_notify_intercept, msg.method, msg.params): - return - try: - transport_ctx = self._transport_builder(metadata) - except Exception: - # Same containment as `_dispatch_request`: drop the notification, keep the loop. - logger.exception("transport_builder raised; dropping notification %r", msg.method) - return - dctx = _JSONRPCDispatchContext( - transport=transport_ctx, _dispatcher=self, _request_id=None, message_metadata=metadata - ) - self._spawn(_contained_notify(on_notify), dctx, msg.method, msg.params, sender_ctx=sender_ctx) - - def _resolve_pending(self, request_id: RequestId | None, outcome: dict[str, Any] | ErrorData) -> None: - pending = self._pending.get(coerce_request_id(request_id)) if request_id is not None else None - if pending is None: - logger.debug("dropping response for unknown/late request id %r", request_id) - return - try: - pending.send.send_nowait(outcome) - except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError): - logger.debug("waiter for request id %r already gone", request_id) - - def _spawn( - self, - fn: Callable[..., Awaitable[Any]], - *args: object, - sender_ctx: contextvars.Context | None, - ) -> None: - """Schedule `fn(*args)` in the run() task group, propagating the sender's contextvars. - - ASGI middleware (auth, OTel) sets contextvars on the task that wrote the - message; `Context.run` makes the spawned handler inherit that context. - """ - assert self._tg is not None - if sender_ctx is not None: - sender_ctx.run(self._tg.start_soon, fn, *args) - else: - self._tg.start_soon(fn, *args) - - def _fan_out_closed(self) -> None: - """Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED`. - - Synchronous: callers may be inside a cancelled scope. Idempotent. - """ - closed = ErrorData(code=CONNECTION_CLOSED, message="Connection closed") - for pending in self._pending.values(): - try: - pending.send.send_nowait(closed) - except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError): - pass - self._pending.clear() - - async def _handle_request( - self, - req: JSONRPCRequest, - dctx: _JSONRPCDispatchContext[TransportT], - scope: anyio.CancelScope, - on_request: OnRequest, - ) -> None: - """Run `on_request` for one inbound request and write its response. - - The single exception-to-wire boundary: handler exceptions become - `JSONRPCError` here. A request the peer cancelled is never answered - (spec: MUST NOT send further messages for it) - it settles unanswered - instead, and `_settle_unanswered` tells the transport. - """ - answer_write_started = False - handler_failure: BaseException | None = None # re-raised once the request settles - try: - with scope: - try: - result = await on_request(dctx, req.method, req.params) - finally: - # Close the back-channel and drop from `_in_flight`; no checkpoint - # since handler return, so a peer cancel can't interleave. - # Identity guard: don't evict a duplicate id's newer entry. - dctx.close() - key = coerce_request_id(req.id) - if (entry := self._in_flight.get(key)) is not None and entry.dctx is dctx: - del self._in_flight[key] - if not dctx.cancel_requested.is_set(): - # A write interrupted by cancellation may still have delivered - # (a memory-stream send can hand its item to the receiver and - # still raise), so a started answer write counts as sent below: - # peers drop late responses, while a second answer for one id - # would break JSON-RPC. - answer_write_started = True - await self._write_result(req.id, result) - except anyio.get_cancelled_exc_class(): - # Shutdown: answer the request so the peer isn't left waiting - unless - # an answer write already started (it may have reached the transport; - # prefer possibly-zero answers over possibly-two), or the peer already - # cancelled it and stopped waiting. The shielded helper is needed - # because bare awaits re-raise here. - if not answer_write_started and not dctx.cancel_requested.is_set(): - await self._final_write( - partial(self._write_error, req.id, ErrorData(code=CONNECTION_CLOSED, message="Connection closed")), - shield=True, - timeout=_SHUTDOWN_WRITE_TIMEOUT, - describe=f"shutdown error response for request {req.id!r}", - ) - raise - except Exception as e: - error = handler_exception_to_error_data(e) - if error is None: - logger.exception("handler for %r raised", req.method) - # TODO(L58): code=0 pins existing-server compat; JSON-RPC says - # INTERNAL_ERROR. Revisit per the suite's divergence entry. - error = ErrorData(code=0, message=str(e)) - if self._raise_handler_exceptions: - handler_failure = e - # A cancel silences only the wire; the failure stays as visible as before. - if not dctx.cancel_requested.is_set(): - answer_write_started = True - await self._write_error(req.id, error) - # The one place a cancelled request settles: the handler is done (any - # mode) with nothing written. A peer-interrupt cancel is absorbed at - # scope __exit__ and lands here too. - if not answer_write_started: - await self._settle_unanswered(dctx) - if handler_failure is not None: - raise handler_failure - # No `_in_flight` pop here: the inner finally covers every path, and a late pop could evict a reused id. - - def _allocate_id(self) -> int: - self._next_id += 1 - return self._next_id - - async def _write(self, message: JSONRPCMessage, metadata: MessageMetadata = None) -> None: - await self._write_stream.send(SessionMessage(message=message, metadata=metadata)) - - async def _write_result(self, request_id: RequestId, result: dict[str, Any]) -> None: - try: - await self._write(JSONRPCResponse(jsonrpc="2.0", id=request_id, result=result)) - except (anyio.BrokenResourceError, anyio.ClosedResourceError): - logger.debug("dropped result for %r: write stream closed", request_id) - - async def _write_error(self, request_id: RequestId, error: ErrorData) -> None: - try: - await self._write(JSONRPCError(jsonrpc="2.0", id=request_id, error=error)) - except (anyio.BrokenResourceError, anyio.ClosedResourceError): - logger.debug("dropped error for %r: write stream closed", request_id) - - async def _settle_unanswered(self, dctx: _JSONRPCDispatchContext[TransportT]) -> None: - """Run the transport's `on_request_unanswered` hook: this request settled with no response. - - The dispatcher writes nothing for it; a transport whose wire must still - end the request (2025-era streamable HTTP) does so from this hook. A - raising hook is contained here, like the other callback boundaries. - """ - metadata = dctx.message_metadata - if not isinstance(metadata, ServerMessageMetadata) or metadata.on_request_unanswered is None: - return - try: - await metadata.on_request_unanswered() - except (anyio.BrokenResourceError, anyio.ClosedResourceError): - logger.debug("on_request_unanswered dropped: connection closing") - except Exception: - logger.exception("on_request_unanswered hook raised") - - async def _final_write( - self, - write: Callable[[], Awaitable[None]], - *, - shield: bool, - timeout: float, - describe: str, - ) -> None: - """Attempt one last write under the shared abandon/teardown policy. - - `shield=True` is for arms already inside a cancelled scope (a bare - `await` would re-raise); the bound keeps a wedged transport write - from becoming an uncancellable hang. - """ - with anyio.move_on_after(timeout, shield=shield) as scope: - await write() - if scope.cancelled_caught: - logger.warning("%s gave up: transport write blocked", describe) - async def _cancel_outbound(self, request_id: RequestId, reason: str, related_request_id: RequestId | None) -> None: - # Thread `related_request_id` so streamable HTTP routes the cancel onto - # the request's own SSE stream instead of a possibly-absent GET stream. - # `notify` swallows connection-state errors itself, so no guard here. - await self.notify( - "notifications/cancelled", - {"requestId": request_id, "reason": reason}, - _related_request_id=related_request_id, - ) +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/memory.py b/src/mcp/shared/memory.py index 01cab77c85..72da9db734 100644 --- a/src/mcp/shared/memory.py +++ b/src/mcp/shared/memory.py @@ -1,33 +1,11 @@ -"""In-memory transports""" +import sys -from __future__ import annotations +import mcp_client.shared.memory as _implementation +from mcp_client.shared.memory import ( + MessageStream as MessageStream, +) +from mcp_client.shared.memory import ( + create_client_server_memory_streams as create_client_server_memory_streams, +) -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager - -from mcp.shared._compat import resync_tracer -from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams -from mcp.shared.message import SessionMessage - -MessageStream = tuple[ContextReceiveStream[SessionMessage | Exception], ContextSendStream[SessionMessage | Exception]] - - -@asynccontextmanager -async def create_client_server_memory_streams() -> AsyncGenerator[tuple[MessageStream, MessageStream], None]: - """Creates a pair of bidirectional memory streams for client-server communication. - - Yields: - A tuple of (client_streams, server_streams) where each is a tuple of - (read_stream, write_stream) - """ - # Create streams for both directions - server_to_client_send, server_to_client_receive = create_context_streams[SessionMessage | Exception](1) - client_to_server_send, client_to_server_receive = create_context_streams[SessionMessage | Exception](1) - - client_streams = (server_to_client_receive, client_to_server_send) - server_streams = (client_to_server_receive, server_to_client_send) - - async with server_to_client_receive, client_to_server_send, client_to_server_receive, server_to_client_send: - yield client_streams, server_streams - # Heals caller-driven cancels; closing memory streams never suspends. - await resync_tracer() +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/message.py b/src/mcp/shared/message.py index 31e51e7128..9d1c249d25 100644 --- a/src/mcp/shared/message.py +++ b/src/mcp/shared/message.py @@ -1,63 +1,26 @@ -"""Message wrapper with metadata support. - -This module defines a wrapper type that combines JSONRPCMessage with metadata -to support transport-specific features like resumability. -""" - -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from typing import Any - -from mcp_types import JSONRPCMessage, RequestId - -ResumptionToken = str - -ResumptionTokenUpdateCallback = Callable[[ResumptionToken], Awaitable[None]] - -# Callback type for closing SSE streams without terminating -CloseSSEStreamCallback = Callable[[], Awaitable[None]] - - -@dataclass -class ClientMessageMetadata: - """Metadata specific to client messages.""" - - resumption_token: ResumptionToken | None = None - on_resumption_token_update: Callable[[ResumptionToken], Awaitable[None]] | None = None - # Per-message HTTP headers (e.g. MCP-Protocol-Version, Mcp-Method) the transport should set. - headers: dict[str, str] | None = None - - -@dataclass -class ServerMessageMetadata: - """Metadata specific to server messages.""" - - related_request_id: RequestId | None = None - # Transport-specific request context (e.g. starlette Request for HTTP - # transports, None for stdio). Typed as Any because the server layer is - # transport-agnostic. - request_context: Any = None - # Callback to close SSE stream for the current request without terminating - close_sse_stream: CloseSSEStreamCallback | None = None - # Callback to close the standalone GET SSE stream (for unsolicited notifications) - close_standalone_sse_stream: CloseSSEStreamCallback | None = None - # Callback the dispatcher runs when this request settles without a response - # (e.g. it was cancelled), for a transport whose wire must still end the - # request even though no response is written. - on_request_unanswered: Callable[[], Awaitable[None]] | None = None - # The transport's verdict on whether this message's request-scoped channel - # can deliver a server-initiated request (see - # `TransportContext.can_send_request`); a transport that says nothing leaves - # it True. - can_send_request: bool = True - - -MessageMetadata = ClientMessageMetadata | ServerMessageMetadata | None - - -@dataclass -class SessionMessage: - """A message with specific metadata for transport-specific features.""" - - message: JSONRPCMessage - metadata: MessageMetadata = None +import sys + +import mcp_client.shared.message as _implementation +from mcp_client.shared.message import ( + ClientMessageMetadata as ClientMessageMetadata, +) +from mcp_client.shared.message import ( + CloseSSEStreamCallback as CloseSSEStreamCallback, +) +from mcp_client.shared.message import ( + MessageMetadata as MessageMetadata, +) +from mcp_client.shared.message import ( + ResumptionToken as ResumptionToken, +) +from mcp_client.shared.message import ( + ResumptionTokenUpdateCallback as ResumptionTokenUpdateCallback, +) +from mcp_client.shared.message import ( + ServerMessageMetadata as ServerMessageMetadata, +) +from mcp_client.shared.message import ( + SessionMessage as SessionMessage, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/metadata_utils.py b/src/mcp/shared/metadata_utils.py index b646133477..1967c1cc18 100644 --- a/src/mcp/shared/metadata_utils.py +++ b/src/mcp/shared/metadata_utils.py @@ -1,46 +1,8 @@ -"""Utility functions for working with metadata in MCP types. +import sys -These utilities are primarily intended for client-side usage to properly display -human-readable names in user interfaces in a spec-compliant way. -""" +import mcp_client.shared.metadata_utils as _implementation +from mcp_client.shared.metadata_utils import ( + get_display_name as get_display_name, +) -from mcp_types import Implementation, Prompt, Resource, ResourceTemplate, Tool - - -def get_display_name(obj: Tool | Resource | Prompt | ResourceTemplate | Implementation) -> str: - """Get the display name for an MCP object with proper precedence. - - This is a client-side utility function designed to help MCP clients display - human-readable names in their user interfaces. When servers provide a 'title' - field, it should be preferred over the programmatic 'name' field for display. - - For tools: title > annotations.title > name - For other objects: title > name - - Example: - ```python - # In a client displaying available tools - tools = await session.list_tools() - for tool in tools.tools: - display_name = get_display_name(tool) - print(f"Available tool: {display_name}") - ``` - - Args: - obj: An MCP object with name and optional title fields - - Returns: - The display name to use for UI presentation - """ - if isinstance(obj, Tool): - # Tools have special precedence: title > annotations.title > name - if hasattr(obj, "title") and obj.title is not None: - return obj.title - if obj.annotations and hasattr(obj.annotations, "title") and obj.annotations.title is not None: - return obj.annotations.title - return obj.name - else: - # All other objects: title > name - if hasattr(obj, "title") and obj.title is not None: - return obj.title - return obj.name +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/path_security.py b/src/mcp/shared/path_security.py index abe810dcfa..de170e24f9 100644 --- a/src/mcp/shared/path_security.py +++ b/src/mcp/shared/path_security.py @@ -1,176 +1,17 @@ -"""Filesystem path safety primitives for resource handlers. - -These functions help MCP servers reject paths that would resolve -outside the served root when extracted URI template parameters are -used in filesystem operations. They are standalone utilities usable from both the -high-level :class:`~mcp.server.mcpserver.MCPServer` and lowlevel server -implementations. - -The canonical safe pattern:: - - from mcp.shared.path_security import safe_join - - @mcp.resource("file://docs/{+path}") - def read_doc(path: str) -> str: - return safe_join("/data/docs", path).read_text(encoding="utf-8") -""" - -import string -from pathlib import Path - -__all__ = ["PathEscapeError", "contains_path_traversal", "is_absolute_path", "safe_join"] - - -class PathEscapeError(ValueError): - """Raised by :func:`safe_join` when the resolved path escapes the base.""" - - -def contains_path_traversal(value: str) -> bool: - r"""Check whether a value, treated as a relative path, escapes its origin. - - This is a **base-free** check: it does not know the sandbox root, so - it detects only whether ``..`` components would move above the - starting point. Use :func:`safe_join` when you know the root — it - additionally catches symlink escapes and absolute-path injection. - - Note: - This is a string-level check on the value as supplied. It does - not model platform-specific filesystem normalisation (e.g. Win32 - stripping of trailing dots and spaces from the final path - component). For filesystem access, use :func:`safe_join`, which - resolves through the OS and verifies containment. - - The check is component-based: ``..`` is dangerous only as a - standalone path segment, not as a substring. Both ``/`` and ``\`` - are treated as separators. - - Example:: - - >>> contains_path_traversal("a/b/c") - False - >>> contains_path_traversal("../etc") - True - >>> contains_path_traversal("a/../../b") - True - >>> contains_path_traversal("a/../b") - False - >>> contains_path_traversal("1.0..2.0") - False - >>> contains_path_traversal("..") - True - - Args: - value: A string that may be used as a filesystem path. - - Returns: - ``True`` if the path would escape its starting directory. - """ - depth = 0 - for part in value.replace("\\", "/").split("/"): - if part == "..": - depth -= 1 - if depth < 0: - return True - elif part and part != ".": - depth += 1 - return False - - -def is_absolute_path(value: str) -> bool: - r"""Check whether a value is an absolute filesystem path. - - Absolute paths are dangerous when joined onto a base: in Python, - ``Path("/data") / "/etc/passwd"`` yields ``/etc/passwd`` — the - absolute right-hand side silently discards the base. - - Detects POSIX absolute (``/foo``), Windows drive-absolute - (``C:\foo``) and drive-relative (``C:foo``), and Windows - UNC/root-relative (``\\server\share``, ``\foo``). - - Example:: - - >>> is_absolute_path("relative/path") - False - >>> is_absolute_path("/etc/passwd") - True - >>> is_absolute_path("C:\\Windows") - True - >>> is_absolute_path("") - False - - Args: - value: A string that may be used as a filesystem path. - - Returns: - ``True`` if the path is absolute on any common platform. - """ - if not value: - return False - if value[0] in ("/", "\\"): - return True - # Windows drive form: C:, C:\, C:foo (drive-relative). A drive- - # relative right-hand side discards the join base when drives - # differ, so flag it even though PureWindowsPath.is_absolute() - # is False. This means single-letter-prefixed identifiers like - # "x:y" also match — opt out via ResourceSecurity(exempt_params=). - if len(value) >= 2 and value[1] == ":" and value[0] in string.ascii_letters: - return True - return False - - -def safe_join(base: str | Path, *parts: str) -> Path: - """Join path components onto a base, rejecting escapes. - - Resolves the joined path and verifies it remains within ``base``. - This is the **gold-standard** check: it catches ``..`` traversal, - absolute-path injection, and symlink escapes that the base-free - checks cannot. - - The symlink check is point-in-time: a directory swapped for a - symlink between this call and the caller's subsequent open would not - be re-checked. Handlers serving a tree that may be modified - concurrently should additionally open with ``O_NOFOLLOW`` or use - platform path-confinement primitives. - - Example:: - - >>> safe_join("/data/docs", "readme.txt") - PosixPath('/data/docs/readme.txt') - >>> safe_join("/data/docs", "../../../etc/passwd") - Traceback (most recent call last): - ... - PathEscapeError: ... - - Args: - base: The sandbox root. May be relative; it will be resolved. - parts: Path components to join. Each is checked for null bytes - and absolute form before joining. - - Returns: - The resolved path, verified to be within ``base`` at resolution - time. - - Raises: - PathEscapeError: If any part contains a null byte, any part is - absolute, or the resolved path is not contained within the - resolved base. - """ - base_resolved = Path(base).resolve() - - for part in parts: - # Null bytes pass through Path construction but fail at the - # syscall boundary with a cryptic error. Reject here so callers - # get a clear PathEscapeError instead. - if "\0" in part: - raise PathEscapeError(f"Path component contains a null byte; refusing to join onto {base_resolved}") - # Absolute parts would silently discard everything to the left - # in Path's / operator. - if is_absolute_path(part): - raise PathEscapeError(f"Path component {part!r} is absolute; refusing to join onto {base_resolved}") - - target = base_resolved.joinpath(*parts).resolve() - - if not target.is_relative_to(base_resolved): - raise PathEscapeError(f"Path {target} escapes base {base_resolved}") - - return target +import sys + +import mcp_client.shared.path_security as _implementation +from mcp_client.shared.path_security import ( + PathEscapeError as PathEscapeError, +) +from mcp_client.shared.path_security import ( + contains_path_traversal as contains_path_traversal, +) +from mcp_client.shared.path_security import ( + is_absolute_path as is_absolute_path, +) +from mcp_client.shared.path_security import ( + safe_join as safe_join, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/peer.py b/src/mcp/shared/peer.py index 14e8fe1c29..6653f3e036 100644 --- a/src/mcp/shared/peer.py +++ b/src/mcp/shared/peer.py @@ -1,240 +1,14 @@ -"""Typed MCP request sugar over an `Outbound`. +import sys -`ClientPeer` wraps any `Outbound` (anything with `send_raw_request` and -`notify`) and exposes the server-to-client request methods (sampling, -elicitation, roots, ping) as typed methods. - -`ClientPeer` does no capability gating: it builds the params, calls -`send_raw_request(method, params)`, and parses the result into the typed -model. Gating (and `NoBackChannelError`) is the wrapped `Outbound`'s job. -""" - -from collections.abc import Mapping -from typing import Any, cast, overload - -from mcp_types import ( - CreateMessageRequestParams, - CreateMessageResult, - CreateMessageResultWithTools, - ElicitRequestedSchema, - ElicitRequestFormParams, - ElicitRequestURLParams, - ElicitResult, - IncludeContext, - ListRootsResult, - ModelPreferences, - RequestParams, - RequestParamsMeta, - SamplingMessage, - Tool, - ToolChoice, +import mcp_client.shared.peer as _implementation +from mcp_client.shared.peer import ( + ClientPeer as ClientPeer, +) +from mcp_client.shared.peer import ( + Meta as Meta, +) +from mcp_client.shared.peer import ( + dump_params as dump_params, ) -from pydantic import BaseModel -from typing_extensions import deprecated - -from mcp.shared.dispatcher import CallOptions, Outbound -from mcp.shared.exceptions import MCPDeprecationWarning - -__all__ = ["ClientPeer", "Meta"] - -Meta = dict[str, Any] -"""Type alias for the `_meta` field carried on request/notification params.""" - - -def dump_params(model: BaseModel | None, meta: Meta | None = None) -> dict[str, Any] | None: - """Serialize a params model to a wire dict, merging `meta` into `_meta`. - - Shared by `ClientPeer` and `Connection` so every typed convenience method - gets the same `_meta` handling. `meta` keys take precedence over any - `_meta` already present on the model. - - `meta` is serialized through `RequestParams` so Python field names emit - their wire aliases: an inbound `ctx.meta` carries `progress_token` (the - key `_extract_meta` validation produces), and forwarding it outbound via - `meta=ctx.meta` must put `progressToken` back on the wire. Keys not - declared on `RequestParamsMeta` pass through unchanged. - """ - out = model.model_dump(by_alias=True, mode="json", exclude_none=True) if model is not None else None - if meta: - wire_meta = RequestParams(_meta=cast(RequestParamsMeta, meta)).model_dump(by_alias=True, mode="json")["_meta"] - out = dict(out or {}) - out["_meta"] = {**out.get("_meta", {}), **wire_meta} - return out - - -class ClientPeer: - """Typed server-to-client request methods over a wrapped `Outbound`. - - Use this when you have a bare dispatcher (or any `Outbound`) and want the - typed methods (`sample`, `elicit_form`, `elicit_url`, `list_roots`, - `ping`) without writing your own host class. - """ - - def __init__(self, outbound: Outbound) -> None: - self._outbound = outbound - - async def send_raw_request( - self, - method: str, - params: Mapping[str, Any] | None, - opts: CallOptions | None = None, - ) -> dict[str, Any]: - return await self._outbound.send_raw_request(method, params, opts) - - async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: - await self._outbound.notify(method, params, opts) - - @overload - @deprecated("The sampling capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def sample( - self, - messages: list[SamplingMessage], - *, - max_tokens: int, - system_prompt: str | None = None, - include_context: IncludeContext | None = None, - temperature: float | None = None, - stop_sequences: list[str] | None = None, - metadata: dict[str, Any] | None = None, - model_preferences: ModelPreferences | None = None, - tools: None = None, - tool_choice: None = None, - meta: Meta | None = None, - opts: CallOptions | None = None, - ) -> CreateMessageResult: ... - @overload - @deprecated("The sampling capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def sample( - self, - messages: list[SamplingMessage], - *, - max_tokens: int, - system_prompt: str | None = None, - include_context: IncludeContext | None = None, - temperature: float | None = None, - stop_sequences: list[str] | None = None, - metadata: dict[str, Any] | None = None, - model_preferences: ModelPreferences | None = None, - tools: list[Tool], - tool_choice: ToolChoice | None = None, - meta: Meta | None = None, - opts: CallOptions | None = None, - ) -> CreateMessageResultWithTools: ... - @overload - @deprecated("The sampling capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def sample( - self, - messages: list[SamplingMessage], - *, - max_tokens: int, - system_prompt: str | None = None, - include_context: IncludeContext | None = None, - temperature: float | None = None, - stop_sequences: list[str] | None = None, - metadata: dict[str, Any] | None = None, - model_preferences: ModelPreferences | None = None, - tools: list[Tool] | None = None, - tool_choice: ToolChoice, - meta: Meta | None = None, - opts: CallOptions | None = None, - ) -> CreateMessageResultWithTools: ... - @deprecated("The sampling capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def sample( - self, - messages: list[SamplingMessage], - *, - max_tokens: int, - system_prompt: str | None = None, - include_context: IncludeContext | None = None, - temperature: float | None = None, - stop_sequences: list[str] | None = None, - metadata: dict[str, Any] | None = None, - model_preferences: ModelPreferences | None = None, - tools: list[Tool] | None = None, - tool_choice: ToolChoice | None = None, - meta: Meta | None = None, - opts: CallOptions | None = None, - ) -> CreateMessageResult | CreateMessageResultWithTools: - """Send a `sampling/createMessage` request to the peer. - - Raises: - MCPError: The peer responded with an error. - NoBackChannelError: No back-channel for server-initiated requests. - pydantic.ValidationError: The peer's result does not match the expected result type. - """ - params = CreateMessageRequestParams( - messages=messages, - system_prompt=system_prompt, - include_context=include_context, - temperature=temperature, - max_tokens=max_tokens, - stop_sequences=stop_sequences, - metadata=metadata, - model_preferences=model_preferences, - tools=tools, - tool_choice=tool_choice, - ) - result = await self.send_raw_request("sampling/createMessage", dump_params(params, meta), opts) - if tools is not None or tool_choice is not None: - return CreateMessageResultWithTools.model_validate(result, by_name=False) - return CreateMessageResult.model_validate(result, by_name=False) - - async def elicit_form( - self, - message: str, - requested_schema: ElicitRequestedSchema, - *, - meta: Meta | None = None, - opts: CallOptions | None = None, - ) -> ElicitResult: - """Send a form-mode `elicitation/create` request. - - Raises: - MCPError: The peer responded with an error. - NoBackChannelError: No back-channel for server-initiated requests. - pydantic.ValidationError: The peer's result does not match the expected result type. - """ - params = ElicitRequestFormParams(message=message, requested_schema=requested_schema) - result = await self.send_raw_request("elicitation/create", dump_params(params, meta), opts) - return ElicitResult.model_validate(result, by_name=False) - - async def elicit_url( - self, - message: str, - url: str, - elicitation_id: str, - *, - meta: Meta | None = None, - opts: CallOptions | None = None, - ) -> ElicitResult: - """Send a URL-mode `elicitation/create` request. - - Raises: - MCPError: The peer responded with an error. - NoBackChannelError: No back-channel for server-initiated requests. - pydantic.ValidationError: The peer's result does not match the expected result type. - """ - params = ElicitRequestURLParams(message=message, url=url, elicitation_id=elicitation_id) - result = await self.send_raw_request("elicitation/create", dump_params(params, meta), opts) - return ElicitResult.model_validate(result, by_name=False) - - @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) - async def list_roots(self, *, meta: Meta | None = None, opts: CallOptions | None = None) -> ListRootsResult: - """Send a `roots/list` request. - - Raises: - MCPError: The peer responded with an error. - NoBackChannelError: No back-channel for server-initiated requests. - pydantic.ValidationError: The peer's result does not match the expected result type. - """ - result = await self.send_raw_request("roots/list", dump_params(None, meta), opts) - return ListRootsResult.model_validate(result, by_name=False) - - async def ping(self, *, meta: Meta | None = None, opts: CallOptions | None = None) -> None: - """Send a `ping` request and ignore the result. - Raises: - MCPError: The peer responded with an error. - NoBackChannelError: No back-channel for server-initiated requests. - """ - await self.send_raw_request("ping", dump_params(None, meta), opts) +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/subscriptions.py b/src/mcp/shared/subscriptions.py index 30449a82ff..478097e30f 100644 --- a/src/mcp/shared/subscriptions.py +++ b/src/mcp/shared/subscriptions.py @@ -1,111 +1,38 @@ -"""Typed event vocabulary for `subscriptions/listen` (2026-07-28, SEP-2575), shared by server and client. +import sys -Every event is a level trigger ("this changed, refetch if you care"), so both sides bound buffers by dedupe. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Any - -from mcp_types import ( - NotificationParams, - PromptListChangedNotification, - ResourceListChangedNotification, - ResourceUpdatedNotification, - ResourceUpdatedNotificationParams, - ServerNotification, - SubscriptionFilter, - ToolListChangedNotification, +import mcp_client.shared.subscriptions as _implementation +from mcp_client.shared.subscriptions import ( + _LIST_CHANGED_EVENTS as _LIST_CHANGED_EVENTS, +) +from mcp_client.shared.subscriptions import ( + LISTEN_STREAM_METHODS as LISTEN_STREAM_METHODS, +) +from mcp_client.shared.subscriptions import ( + SUBSCRIPTION_ID_META_KEY as SUBSCRIPTION_ID_META_KEY, +) +from mcp_client.shared.subscriptions import ( + PromptsListChanged as PromptsListChanged, +) +from mcp_client.shared.subscriptions import ( + ResourcesListChanged as ResourcesListChanged, +) +from mcp_client.shared.subscriptions import ( + ResourceUpdated as ResourceUpdated, +) +from mcp_client.shared.subscriptions import ( + ServerEvent as ServerEvent, +) +from mcp_client.shared.subscriptions import ( + ToolsListChanged as ToolsListChanged, +) +from mcp_client.shared.subscriptions import ( + event_from_wire as event_from_wire, +) +from mcp_client.shared.subscriptions import ( + event_matches as event_matches, +) +from mcp_client.shared.subscriptions import ( + event_to_notification as event_to_notification, ) -__all__ = [ - "LISTEN_STREAM_METHODS", - "SUBSCRIPTION_ID_META_KEY", - "PromptsListChanged", - "ResourceUpdated", - "ResourcesListChanged", - "ServerEvent", - "ToolsListChanged", - "event_from_wire", - "event_matches", - "event_to_notification", -] - -SUBSCRIPTION_ID_META_KEY = "io.modelcontextprotocol/subscriptionId" -"""The `_meta` key on every listen-stream frame; the value is the `subscriptions/listen` request's JSON-RPC id.""" - - -@dataclass(frozen=True) -class ToolsListChanged: - """The server's tool list changed.""" - - -@dataclass(frozen=True) -class PromptsListChanged: - """The server's prompt list changed.""" - - -@dataclass(frozen=True) -class ResourcesListChanged: - """The server's resource list changed.""" - - -@dataclass(frozen=True) -class ResourceUpdated: - """The resource at `uri` changed and may need to be read again.""" - - uri: str - - -ServerEvent = ToolsListChanged | PromptsListChanged | ResourcesListChanged | ResourceUpdated -"""An event a server publishes for delivery to listen subscribers.""" - - -def event_to_notification(event: ServerEvent, meta: dict[str, Any]) -> ServerNotification: - """Build the stamped wire notification for `event` (the server's direction).""" - if isinstance(event, ToolsListChanged): - return ToolListChangedNotification(params=NotificationParams(_meta=meta)) - if isinstance(event, PromptsListChanged): - return PromptListChangedNotification(params=NotificationParams(_meta=meta)) - if isinstance(event, ResourcesListChanged): - return ResourceListChangedNotification(params=NotificationParams(_meta=meta)) - return ResourceUpdatedNotification(params=ResourceUpdatedNotificationParams(uri=event.uri, _meta=meta)) - - -_LIST_CHANGED_EVENTS: dict[str, ServerEvent] = { - "notifications/tools/list_changed": ToolsListChanged(), - "notifications/prompts/list_changed": PromptsListChanged(), - "notifications/resources/list_changed": ResourcesListChanged(), -} - -LISTEN_STREAM_METHODS: frozenset[str] = frozenset({*_LIST_CHANGED_EVENTS, "notifications/resources/updated"}) -"""The notification methods that ride `subscriptions/listen` streams at 2026-07-28 -(and, at that era, nowhere else): the change-notification vocabulary.""" - - -def event_from_wire(method: str, params: Mapping[str, Any] | None) -> ServerEvent | None: - """The event a raw listen-stream frame announces, or None if it carries none. - - Takes the raw wire dict: the client demultiplexes before the typed notification parse.""" - if (event := _LIST_CHANGED_EVENTS.get(method)) is not None: - return event - if method == "notifications/resources/updated": - uri = (params or {}).get("uri") - if isinstance(uri, str): - return ResourceUpdated(uri=uri) - return None - - -def event_matches(honored: SubscriptionFilter, uris: frozenset[str], event: ServerEvent) -> bool: - """Whether `event` is within the stream's honored filter (`uris`: the honored resource subscriptions as a set). - - The admission predicate both sides share: server delivery and client intake honor only what was acknowledged.""" - if isinstance(event, ToolsListChanged): - return honored.tools_list_changed is True - if isinstance(event, PromptsListChanged): - return honored.prompts_list_changed is True - if isinstance(event, ResourcesListChanged): - return honored.resources_list_changed is True - return event.uri in uris +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/tool_name_validation.py b/src/mcp/shared/tool_name_validation.py index 96c34f7826..8e6991a080 100644 --- a/src/mcp/shared/tool_name_validation.py +++ b/src/mcp/shared/tool_name_validation.py @@ -1,129 +1,26 @@ -"""Tool name validation utilities according to SEP-986. - -Tool names SHOULD be between 1 and 128 characters in length (inclusive). -Tool names are case-sensitive. -Allowed characters: uppercase and lowercase ASCII letters (A-Z, a-z), -digits (0-9), underscore (_), dash (-), and dot (.). -Tool names SHOULD NOT contain spaces, commas, or other special characters. - -See: https://modelcontextprotocol.io/specification/2025-11-25/server/tools#tool-names -""" - -from __future__ import annotations - -import logging -import re -from dataclasses import dataclass, field - -logger = logging.getLogger(__name__) - -# Regular expression for valid tool names according to SEP-986 specification -TOOL_NAME_REGEX = re.compile(r"^[A-Za-z0-9._-]{1,128}$") - -# SEP reference URL for warning messages -SEP_986_URL = "https://modelcontextprotocol.io/specification/2025-11-25/server/tools#tool-names" - - -@dataclass -class ToolNameValidationResult: - """Result of tool name validation. - - Attributes: - is_valid: Whether the tool name conforms to SEP-986 requirements. - warnings: List of warning messages for non-conforming aspects. - """ - - is_valid: bool - warnings: list[str] = field(default_factory=lambda: []) - - -def validate_tool_name(name: str) -> ToolNameValidationResult: - """Validate a tool name according to the SEP-986 specification. - - Args: - name: The tool name to validate. - - Returns: - ToolNameValidationResult containing validation status and any warnings. - """ - warnings: list[str] = [] - - # Check for empty name - if not name: - return ToolNameValidationResult( - is_valid=False, - warnings=["Tool name cannot be empty"], - ) - - # Check length - if len(name) > 128: - return ToolNameValidationResult( - is_valid=False, - warnings=[f"Tool name exceeds maximum length of 128 characters (current: {len(name)})"], - ) - - # Check for problematic patterns (warnings, not validation failures) - if " " in name: - warnings.append("Tool name contains spaces, which may cause parsing issues") - - if "," in name: - warnings.append("Tool name contains commas, which may cause parsing issues") - - # Check for potentially confusing leading/trailing characters - if name.startswith("-") or name.endswith("-"): - warnings.append("Tool name starts or ends with a dash, which may cause parsing issues in some contexts") - - if name.startswith(".") or name.endswith("."): - warnings.append("Tool name starts or ends with a dot, which may cause parsing issues in some contexts") - - # Check for invalid characters - if not TOOL_NAME_REGEX.fullmatch(name): - # Find all invalid characters (unique, preserving order) - invalid_chars: list[str] = [] - seen: set[str] = set() - for char in name: - if not re.match(r"[A-Za-z0-9._-]", char) and char not in seen: - invalid_chars.append(char) - seen.add(char) - - warnings.append(f"Tool name contains invalid characters: {', '.join(repr(c) for c in invalid_chars)}") - warnings.append("Allowed characters are: A-Z, a-z, 0-9, underscore (_), dash (-), and dot (.)") - - return ToolNameValidationResult(is_valid=False, warnings=warnings) - - return ToolNameValidationResult(is_valid=True, warnings=warnings) - - -def issue_tool_name_warning(name: str, warnings: list[str]) -> None: - """Log warnings for non-conforming tool names. - - Args: - name: The tool name that triggered the warnings. - warnings: List of warning messages to log. - """ - if not warnings: - return - - logger.warning(f'Tool name validation warning for "{name}":') - for warning in warnings: - logger.warning(f" - {warning}") - logger.warning("Tool registration will proceed, but this may cause compatibility issues.") - logger.warning("Consider updating the tool name to conform to the MCP tool naming standard.") - logger.warning(f"See SEP-986 ({SEP_986_URL}) for more details.") - - -def validate_and_warn_tool_name(name: str) -> bool: - """Validate a tool name and issue warnings for non-conforming names. - - This is the primary entry point for tool name validation. It validates - the name and logs any warnings via the logging module. - - Args: - name: The tool name to validate. - - Returns: - True if the name is valid, False otherwise. - """ - result = validate_tool_name(name) - issue_tool_name_warning(name, result.warnings) - return result.is_valid +import sys + +import mcp_client.shared.tool_name_validation as _implementation +from mcp_client.shared.tool_name_validation import ( + SEP_986_URL as SEP_986_URL, +) +from mcp_client.shared.tool_name_validation import ( + TOOL_NAME_REGEX as TOOL_NAME_REGEX, +) +from mcp_client.shared.tool_name_validation import ( + ToolNameValidationResult as ToolNameValidationResult, +) +from mcp_client.shared.tool_name_validation import ( + issue_tool_name_warning as issue_tool_name_warning, +) +from mcp_client.shared.tool_name_validation import ( + logger as logger, +) +from mcp_client.shared.tool_name_validation import ( + validate_and_warn_tool_name as validate_and_warn_tool_name, +) +from mcp_client.shared.tool_name_validation import ( + validate_tool_name as validate_tool_name, +) + +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/transport_context.py b/src/mcp/shared/transport_context.py index 8d15a2eaa2..eb56ac4127 100644 --- a/src/mcp/shared/transport_context.py +++ b/src/mcp/shared/transport_context.py @@ -1,45 +1,8 @@ -"""Transport-specific metadata attached to each inbound message. +import sys -`TransportContext` is the base; each transport defines its own subclass with -whatever fields make sense (HTTP request id, ASGI scope, stdio process handle, -etc.). The dispatcher passes it through opaquely; only the layers above the -dispatcher (`ServerRunner`, `Context`, user handlers) read its concrete fields. -""" +import mcp_client.shared.transport_context as _implementation +from mcp_client.shared.transport_context import ( + TransportContext as TransportContext, +) -from collections.abc import Mapping -from dataclasses import dataclass - -__all__ = ["TransportContext"] - - -@dataclass(kw_only=True, frozen=True) -class TransportContext: - """Base transport metadata for an inbound message. - - Subclass per transport and add fields as needed. Instances are immutable. - """ - - kind: str - """Short identifier for the transport (e.g. `"stdio"`, `"streamable-http"`).""" - - can_send_request: bool - """Whether this message's request-scoped channel can deliver a server-initiated request. - - `False` for any of three reasons: the response has no room (streamable - HTTP in JSON-response mode and the 2026-07-28 single-exchange entry answer - with one JSON-RPC reply), the client's reply has nowhere to land (stateless - HTTP, no session), or the protocol forbids server-initiated requests (any - 2026-07-28 connection, whose dispatch masks the flag off). `True` for a - plain duplex pipe (stdio, SSE) and stateful streamable HTTP with SSE - responses, all pre-2026-07-28. When `False`, - `DispatchContext.send_raw_request` raises `NoBackChannelError` instead of - parking a waiter no reply can reach. Says nothing about the connection's - standalone channel, which refuses separately. - """ - - headers: Mapping[str, str] | None = None - """Request headers carried by this message, when the transport has them. - - Populated by HTTP-based transports; `None` on stdio. Handlers should - None-check before use. - """ +sys.modules[__name__] = _implementation diff --git a/src/mcp/shared/uri_template.py b/src/mcp/shared/uri_template.py index 20d6fa9c2e..3da4b0158d 100644 --- a/src/mcp/shared/uri_template.py +++ b/src/mcp/shared/uri_template.py @@ -1,1116 +1,107 @@ -"""RFC 6570 URI Templates with bidirectional support. - -Provides both expansion (template + variables → URI) and matching -(URI → variables). RFC 6570 only specifies expansion; matching is the -inverse operation needed by MCP servers to route ``resources/read`` -requests to handlers. - -Supports Levels 1-3 fully, plus Level 4 explode modifier for path-like -operators (``{/var*}``, ``{.var*}``, ``{;var*}``). The Level 4 prefix -modifier (``{var:N}``) and query-explode (``{?var*}``) are not supported. - -Matching semantics ------------------- - -Matching is not specified by RFC 6570 (§1.4 explicitly defers to regex -languages). This implementation uses a two-ended scan that never -backtracks: match time is O(n·v) where n is URI length and v is the -number of template variables. Realistic templates have v < 10, making -this effectively linear; there is no input that produces -superpolynomial time. - -A template may contain **at most one multi-segment variable** — -``{+var}``, ``{#var}``, or an explode-modified variable (``{/var*}``, -``{.var*}``, ``{;var*}``). This variable greedily consumes whatever the -surrounding bounded variables and literals do not. Two such variables -in one template are inherently ambiguous (which one gets the extra -segment?) and are rejected at parse time. So are any two variables -adjacent with no literal between them — including a variable adjacent -to the multi-segment variable: the scan has nothing to anchor the -boundary on. Operators that emit their own lead character supply that -literal themselves, so ``{+path}{.ext}`` and ``{a}{.b}`` are fine -while ``{+path}{ext}`` and ``{a}{b}`` are not. - -Bounded variables before the multi-segment variable match **lazily** -(first occurrence of the following literal); those after match -**greedily** (last occurrence of the preceding literal). Templates -without a multi-segment variable match greedily throughout, identical -to regex semantics. - -Reserved expansion ``{+var}`` leaves ``?`` and ``#`` unencoded, but -the scan stops at those characters so ``{+path}{?q}`` can separate path -from query. A value containing a literal ``?`` or ``#`` expands fine -but will not round-trip through ``match()``. -""" - -from __future__ import annotations - -import re -from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field -from typing import Literal, TypeAlias, cast -from urllib.parse import quote, unquote - -__all__ = [ - "DEFAULT_MAX_TEMPLATE_LENGTH", - "DEFAULT_MAX_VARIABLES", - "DEFAULT_MAX_URI_LENGTH", - "InvalidUriTemplate", - "Operator", - "UriTemplate", - "Variable", -] - -Operator = Literal["", "+", "#", ".", "/", ";", "?", "&"] - -_OPERATORS: frozenset[str] = frozenset({"+", "#", ".", "/", ";", "?", "&"}) - -# RFC 6570 §2.3: varname = varchar *(["."] varchar), varchar = ALPHA / DIGIT / "_" -# Dots appear only between varchar groups — not consecutive, not trailing. -# (Percent-encoded varchars are technically allowed but unseen in practice.) -_VARNAME_RE = re.compile(r"^[A-Za-z0-9_]+(?:\.[A-Za-z0-9_]+)*$") - -DEFAULT_MAX_TEMPLATE_LENGTH = 8_192 -DEFAULT_MAX_VARIABLES = 256 -DEFAULT_MAX_URI_LENGTH = 65_536 - -# RFC 3986 reserved characters, kept unencoded by {+var} and {#var}. -_RESERVED = ":/?#[]@!$&'()*+,;=" - - -@dataclass(frozen=True) -class _OperatorSpec: - """Expansion behavior for a single operator (RFC 6570 §3.2, Table in §A).""" - - prefix: str - """Leading character emitted before the first variable.""" - separator: str - """Character between variables (and between exploded list items).""" - named: bool - """Emit ``name=value`` pairs (query/path-param style) rather than bare values.""" - allow_reserved: bool - """Keep reserved characters unencoded ({+var}, {#var}).""" - ifemp: str - """Suffix after a named variable whose expanded value is empty (RFC §A): '' for ;, '=' for ?/&.""" - - -_OPERATOR_SPECS: dict[Operator, _OperatorSpec] = { - "": _OperatorSpec(prefix="", separator=",", named=False, allow_reserved=False, ifemp=""), - "+": _OperatorSpec(prefix="", separator=",", named=False, allow_reserved=True, ifemp=""), - "#": _OperatorSpec(prefix="#", separator=",", named=False, allow_reserved=True, ifemp=""), - ".": _OperatorSpec(prefix=".", separator=".", named=False, allow_reserved=False, ifemp=""), - "/": _OperatorSpec(prefix="/", separator="/", named=False, allow_reserved=False, ifemp=""), - ";": _OperatorSpec(prefix=";", separator=";", named=True, allow_reserved=False, ifemp=""), - "?": _OperatorSpec(prefix="?", separator="&", named=True, allow_reserved=False, ifemp="="), - "&": _OperatorSpec(prefix="&", separator="&", named=True, allow_reserved=False, ifemp="="), -} - -# Per-operator stop characters for the linear scan. A bounded variable's -# value ends at the first occurrence of any character in its stop set, -# mirroring the character-class boundaries a regex would use but without -# the backtracking. -_STOP_CHARS: dict[Operator, str] = { - "": "/?#&,", # simple: everything structural is pct-encoded - "+": "?#", # reserved: / allowed, stop at query/fragment - "#": "", # fragment: tail of URI, nothing stops it - ".": "./?#", # label: stop at next . - "/": "/?#", # path segment: stop at next / - ";": ";/?#", # path-param value (may be empty: ;name) - "?": "&#", # query value (may be empty: ?name=) - "&": "&#", # query-cont value -} - - -class InvalidUriTemplate(ValueError): - """Raised when a URI template string is malformed or unsupported. - - Attributes: - template: The template string that failed to parse. - position: Character offset where the error was detected, or None - if the error is not tied to a specific position. - """ - - def __init__(self, message: str, *, template: str, position: int | None = None) -> None: - super().__init__(message) - self.template = template - self.position = position - - -@dataclass(frozen=True) -class Variable: - """A single variable within a URI template expression.""" - - name: str - operator: Operator - explode: bool = False - - -@dataclass -class _Expression: - """A parsed ``{...}`` expression: one operator, one or more variables.""" - - operator: Operator - variables: list[Variable] - - -_Part = str | _Expression - - -@dataclass(frozen=True) -class _Lit: - """A literal run in the flattened match-atom sequence.""" - - text: str - - -@dataclass(frozen=True) -class _Cap: - """A single-variable capture in the flattened match-atom sequence. - - ``ifemp`` marks the ``;`` operator's optional-equals quirk: ``{;id}`` - expands to ``;id=value`` or bare ``;id`` when the value is empty, so - the scan must accept both forms. - """ - - var: Variable - ifemp: bool = False - - -_Atom: TypeAlias = _Lit | _Cap - - -def _is_greedy(var: Variable) -> bool: - """Return True if this variable can span multiple path segments. - - Reserved/fragment expansion and explode variables are the only - constructs whose match range is not bounded by a single structural - delimiter. A template may contain at most one such variable. - """ - return var.explode or var.operator in ("+", "#") - - -def _is_str_sequence(value: object) -> bool: - """Check if value is a non-string sequence whose items are all strings.""" - if isinstance(value, str) or not isinstance(value, Sequence): - return False - seq = cast(Sequence[object], value) - return all(isinstance(item, str) for item in seq) - - -_PCT_TRIPLET_RE = re.compile(r"%[0-9A-Fa-f]{2}") - - -def _encode(value: str, *, allow_reserved: bool) -> str: - """Percent-encode a value per RFC 6570 §3.2.1. - - Simple expansion encodes everything except unreserved characters. - Reserved expansion (``{+var}``, ``{#var}``) additionally keeps - RFC 3986 reserved characters intact and passes through existing - ``%XX`` pct-triplets unchanged (RFC 6570 §3.2.3). A bare ``%`` not - followed by two hex digits is still encoded to ``%25``. - """ - if not allow_reserved: - return quote(value, safe="") - - # Reserved expansion: walk the string, pass through triplets as-is, - # quote the gaps between them. A bare % with no triplet lands in a - # gap and gets encoded normally. - out: list[str] = [] - last = 0 - for m in _PCT_TRIPLET_RE.finditer(value): - out.append(quote(value[last : m.start()], safe=_RESERVED)) - out.append(m.group()) - last = m.end() - out.append(quote(value[last:], safe=_RESERVED)) - return "".join(out) - - -def _expand_expression(expr: _Expression, variables: Mapping[str, str | Sequence[str]]) -> str: - """Expand a single ``{...}`` expression into its URI fragment. - - Walks the expression's variables, encoding and joining defined ones - according to the operator's spec. Undefined variables are skipped - (RFC 6570 §2.3); if all are undefined, the expression contributes - nothing (no prefix is emitted). - """ - spec = _OPERATOR_SPECS[expr.operator] - rendered: list[str] = [] - - for var in expr.variables: - if var.name not in variables: - # Undefined: skip entirely, no placeholder. - continue - - value = variables[var.name] - - # Explicit type guard: reject non-str scalars with a clear message - # rather than a confusing "not iterable" from the sequence branch. - if not isinstance(value, str) and not _is_str_sequence(value): - raise TypeError(f"Variable {var.name!r} must be str or a sequence of str, got {type(value).__name__}") - - if isinstance(value, str): - encoded = _encode(value, allow_reserved=spec.allow_reserved) - if spec.named: - rendered.append(f"{var.name}{spec.ifemp}" if value == "" else f"{var.name}={encoded}") - else: - rendered.append(encoded) - else: - # Sequence value. - items = [_encode(v, allow_reserved=spec.allow_reserved) for v in value] - if not items: - continue - if var.explode: - # Each item gets the operator's separator; named ops repeat the key. - if spec.named: - rendered.append( - spec.separator.join(f"{var.name}{spec.ifemp}" if v == "" else f"{var.name}={v}" for v in items) - ) - else: - rendered.append(spec.separator.join(items)) - else: - # Non-explode: comma-join into a single value, then apply - # ifemp to the joined result (RFC §3.2.1: behaves as if the - # value were the joined string). - joined = ",".join(items) - if spec.named: - rendered.append(f"{var.name}{spec.ifemp}" if joined == "" else f"{var.name}={joined}") - else: - rendered.append(joined) - - if not rendered: - return "" - return spec.prefix + spec.separator.join(rendered) - - -@dataclass(frozen=True) -class UriTemplate: - """A parsed RFC 6570 URI template. - - Construct via :meth:`parse`. Instances are immutable and hashable; - equality is based on the template string alone. - """ - - template: str - _parts: list[_Part] = field(repr=False, compare=False) - _variables: list[Variable] = field(repr=False, compare=False) - _prefix: list[_Atom] = field(repr=False, compare=False) - _greedy: Variable | None = field(repr=False, compare=False) - _suffix: list[_Atom] = field(repr=False, compare=False) - _query_variables: list[Variable] = field(repr=False, compare=False) - - @staticmethod - def is_template(value: str) -> bool: - """Check whether a string contains URI template expressions. - - A cheap heuristic for distinguishing concrete URIs from templates - without the cost of full parsing. Returns ``True`` if the string - contains at least one ``{...}`` pair. - - Example:: - - >>> UriTemplate.is_template("file://docs/{name}") - True - >>> UriTemplate.is_template("file://docs/readme.txt") - False - - Note: - This does not validate the template. A ``True`` result does - not guarantee :meth:`parse` will succeed. - """ - open_i = value.find("{") - return open_i != -1 and value.find("}", open_i) != -1 - - @classmethod - def parse( - cls, - template: str, - *, - max_length: int = DEFAULT_MAX_TEMPLATE_LENGTH, - max_variables: int = DEFAULT_MAX_VARIABLES, - ) -> UriTemplate: - """Parse a URI template string. - - Args: - template: An RFC 6570 URI template. - max_length: Maximum permitted length of the template string. - Guards against resource exhaustion. - max_variables: Maximum number of variables permitted across - all expressions. Counting variables rather than - ``{...}`` expressions closes the gap where a single - ``{v0,v1,...,vN}`` expression packs arbitrarily many - variables under one expression count. - - Raises: - InvalidUriTemplate: If the template is malformed, exceeds the - size limits, or uses unsupported RFC 6570 features. - """ - if len(template) > max_length: - raise InvalidUriTemplate( - f"Template exceeds maximum length of {max_length}", - template=template, - ) - - parts, variables = _parse(template, max_variables=max_variables) - - # Trailing {?...}/{&...} expressions are split off and matched as - # a query string (order-agnostic, partial, extras ignored) rather - # than via the linear scan. - path_parts, query_vars = _split_query_tail(parts) - atoms = _flatten(path_parts) - prefix, greedy, suffix = _partition_greedy(atoms, template) - - return cls( - template=template, - _parts=parts, - _variables=variables, - _prefix=prefix, - _greedy=greedy, - _suffix=suffix, - _query_variables=query_vars, - ) - - @property - def variables(self) -> list[Variable]: - """All variables in the template, in order of appearance.""" - return list(self._variables) - - @property - def variable_names(self) -> list[str]: - """All variable names in the template, in order of appearance.""" - return [v.name for v in self._variables] - - @property - def query_variable_names(self) -> frozenset[str]: - """Names of variables that :meth:`match` treats as optional query parameters. - - These are the variables in a trailing run of ``{?...}``/``{&...}`` - expressions, which are matched leniently: a URI that omits some - (or all) of them still matches, and the omitted names are simply - absent from the result. Any value bound to such a name therefore - needs a fallback for the omitted case. - - Every other variable is bound on every successful :meth:`match` - (possibly to an empty string) and is *not* in this set. That - includes a ``{&...}`` expression with no preceding ``{?...}``: it - never emits the ``?`` the lenient query split keys on, so it is - matched strictly. - """ - return frozenset(v.name for v in self._query_variables) - - def expand(self, variables: Mapping[str, str | Sequence[str]]) -> str: - """Expand the template by substituting variable values. - - String values are percent-encoded according to their operator: - simple ``{var}`` encodes reserved characters; ``{+var}`` and - ``{#var}`` leave them intact. Sequence values are joined with - commas for non-explode variables, or with the operator's - separator for explode variables. - - Example:: - - >>> t = UriTemplate.parse("file://docs/{name}") - >>> t.expand({"name": "hello world.txt"}) - 'file://docs/hello%20world.txt' - - >>> t = UriTemplate.parse("file://docs/{+path}") - >>> t.expand({"path": "src/main.py"}) - 'file://docs/src/main.py' - - >>> t = UriTemplate.parse("/search{?q,lang}") - >>> t.expand({"q": "mcp", "lang": "en"}) - '/search?q=mcp&lang=en' - - >>> t = UriTemplate.parse("/files{/path*}") - >>> t.expand({"path": ["a", "b", "c"]}) - '/files/a/b/c' - - Args: - variables: Values for each template variable. Keys must be - strings; values must be ``str`` or a sequence of ``str``. - - Returns: - The expanded URI string. - - Note: - Per RFC 6570, variables absent from the mapping are - **silently omitted**. This is the correct behavior for - optional query parameters (``{?page}`` with no page yields - no ``?page=``), but for required path segments it produces - a structurally incomplete URI. If you need all variables - present, validate before calling:: - - missing = set(t.variable_names) - variables.keys() - if missing: - raise ValueError(f"Missing: {missing}") - - Raises: - TypeError: If a value is neither ``str`` nor an iterable of - ``str``. Non-string scalars (``int``, ``None``) are not - coerced. - """ - out: list[str] = [] - for part in self._parts: - if isinstance(part, str): - out.append(part) - else: - out.append(_expand_expression(part, variables)) - return "".join(out) - - def match(self, uri: str, *, max_uri_length: int = DEFAULT_MAX_URI_LENGTH) -> dict[str, str | list[str]] | None: - """Match a concrete URI against this template and extract variables. - - This is the inverse of :meth:`expand`. The URI is matched via a - linear scan of the template and captured values are - percent-decoded. The round-trip ``match(expand({k: v})) == {k: v}`` - holds when ``v`` does not contain its operator's separator - unencoded: ``{.ext}`` with ``ext="tar.gz"`` expands to - ``.tar.gz`` but does not match — the scan stops ``ext`` at the - first ``.`` and the trailing ``.gz`` has nothing to consume it. - RFC 6570 §1.4 notes this is an inherent reversal limitation. - - Matching is structural at the URI level only: a simple ``{name}`` - will not match across a literal ``/`` in the URI (the scan stops - there), but a percent-encoded ``%2F`` that decodes to ``/`` is - accepted as part of the value. Path-safety validation belongs at - a higher layer; see :mod:`mcp.shared.path_security`. - - Example:: - - >>> t = UriTemplate.parse("file://docs/{name}") - >>> t.match("file://docs/readme.txt") - {'name': 'readme.txt'} - >>> t.match("file://docs/hello%20world.txt") - {'name': 'hello world.txt'} - - >>> t = UriTemplate.parse("file://docs/{+path}") - >>> t.match("file://docs/src/main.py") - {'path': 'src/main.py'} - - >>> t = UriTemplate.parse("/files{/path*}") - >>> t.match("/files/a/b/c") - {'path': ['a', 'b', 'c']} - - **Query parameters** (``{?q,lang}`` at the end of a template) - are matched leniently: order-agnostic, partial, and unrecognized - params are ignored. Absent params are omitted from the result so - downstream function defaults can apply:: - - >>> t = UriTemplate.parse("logs://{service}{?since,level}") - >>> t.match("logs://api") - {'service': 'api'} - >>> t.match("logs://api?level=error") - {'service': 'api', 'level': 'error'} - >>> t.match("logs://api?level=error&since=5m&utm=x") - {'service': 'api', 'since': '5m', 'level': 'error'} - - Args: - uri: A concrete URI string. - max_uri_length: Maximum permitted length of the input URI. - Oversized inputs return ``None`` without scanning, - guarding against resource exhaustion. - - Returns: - A mapping from variable names to decoded values (``str`` for - scalar variables, ``list[str]`` for explode variables), or - ``None`` if the URI does not match the template or exceeds - ``max_uri_length``. - """ - if len(uri) > max_uri_length: - return None - - if self._query_variables: - # Two-phase: scan matches the path, the query is split and - # decoded manually. Query params may be partial, reordered, - # or include extras; absent params stay absent so downstream - # defaults can apply. Fragment is stripped first since the - # template's {?...} tail never describes a fragment. - before_fragment, _, _ = uri.partition("#") - path, _, query = before_fragment.partition("?") - result = self._scan(path) - if result is None: - return None - if query: - parsed = _parse_query(query) - for var in self._query_variables: - if var.name in parsed: - result[var.name] = parsed[var.name] - return result - - return self._scan(uri) - - def _scan(self, uri: str) -> dict[str, str | list[str]] | None: - """Run the two-ended linear scan against the path portion of a URI.""" - n = len(uri) - - if self._greedy is None: - # No greedy var: the suffix IS the whole template, scanned - # right-to-left and anchored so atoms[0] matches at position 0. - suffix = _scan_suffix(self._suffix, uri, n, anchored=True) - if suffix is None: - return None - suffix_result, suffix_start = suffix - return suffix_result if suffix_start == 0 else None - - # Greedy var present. The parser rejects a capture adjacent to - # the greedy slot, so a non-empty suffix begins with a _Lit whose - # rfind-derived anchor does not depend on how far the prefix - # scans. Scan the suffix first, then give the prefix that exact - # position as its ceiling so it cannot consume past the anchor. - suffix = _scan_suffix(self._suffix, uri, n, anchored=False) - if suffix is None: - return None - suffix_result, suffix_start = suffix - prefix = _scan_prefix(self._prefix, uri, 0, suffix_start) - if prefix is None: - return None - prefix_result, prefix_end = prefix - - # Prefix consumed [0, prefix_end); suffix consumed [suffix_start, n); - # the greedy var takes the gap. The prefix scan is bounded by - # suffix_start, so this holds by construction; guard explicitly - # rather than asserting so a future regression surfaces as a - # non-match, not an exception. - if suffix_start < prefix_end: - return None # pragma: no cover - unreachable while bounds hold - middle = uri[prefix_end:suffix_start] - greedy_value = _extract_greedy(self._greedy, middle) - if greedy_value is None: - return None - - return {**prefix_result, self._greedy.name: greedy_value, **suffix_result} - - def __str__(self) -> str: - return self.template - - -def _parse_query(query: str) -> dict[str, str]: - """Parse a query string into a name→value mapping. - - Unlike ``urllib.parse.parse_qs``, this follows RFC 3986 semantics: - ``+`` is a literal sub-delim, not a space. Form-urlencoding treats - ``+`` as space for HTML form submissions, but RFC 6570 and MCP - resource URIs follow RFC 3986 where only ``%20`` encodes a space. - - Parameter names are **not** percent-decoded. RFC 6570 expansion - never encodes variable names, so a legitimate match will always - have the name in literal form. Decoding names would let - ``%74oken=evil&token=real`` shadow the real ``token`` parameter - via first-wins. - - Duplicate keys keep the first value. Pairs without ``=`` are - treated as empty-valued. - """ - result: dict[str, str] = {} - for pair in query.split("&"): - name, _, value = pair.partition("=") - if name and name not in result: - result[name] = unquote(value) - return result - - -def _extract_greedy(var: Variable, raw: str) -> str | list[str] | None: - """Decode the greedy variable's isolated middle span. - - For scalar greedy (``{+var}``, ``{#var}``) this is a stop-char - validation and a single ``unquote``. For explode variables the span - is a run of separator-delimited segments (``/a/b/c`` or - ``;keys=a;keys=b``) that is split, validated, and decoded per item. - """ - spec = _OPERATOR_SPECS[var.operator] - stops = _STOP_CHARS[var.operator] - - if not var.explode: - if any(c in stops for c in raw): - return None - return unquote(raw) - - sep = spec.separator - if not raw: - return [] - # A non-empty explode span must begin with the separator: {/a*} - # expands to "/x/y", never "x/y". The scan does not consume the - # separator itself, so it must be the first character here. - if raw[0] != sep: - return None - # Segments must not contain the operator's non-separator stop - # characters (e.g. {/path*} segments may contain neither ? nor #). - body_stops = set(stops) - {sep} - if any(c in body_stops for c in raw): - return None - - segments: list[str] = [] - prefix = f"{var.name}=" - # split()[0] is always "" because raw starts with the separator; - # subsequent empties are legitimate values ({/path*} with - # ["a","","c"] expands to /a//c). - for seg in raw.split(sep)[1:]: - if spec.named: - # Named explode emits name=value per item (or bare name - # under ; with empty value). Validate the name and strip - # the prefix before decoding. - if seg.startswith(prefix): - seg = seg[len(prefix) :] - elif seg == var.name: - seg = "" - else: - return None - segments.append(unquote(seg)) - return segments - - -def _split_query_tail(parts: list[_Part]) -> tuple[list[_Part], list[Variable]]: - """Separate trailing ``?``/``&`` expressions from the path portion. - - Lenient query matching (order-agnostic, partial, ignores extras) - applies when a template ends with one or more consecutive ``?``/``&`` - expressions and the preceding path portion contains no literal - ``?``. If the path has a literal ``?`` (e.g., ``?fixed=1{&page}``), - the URI's ``?`` split won't align with the template's expression - boundary, so the strict scan is used instead. - - Returns: - A pair ``(path_parts, query_vars)``. If lenient matching does - not apply, ``query_vars`` is empty and ``path_parts`` is the - full input. - """ - split = len(parts) - for i in range(len(parts) - 1, -1, -1): - part = parts[i] - if isinstance(part, _Expression) and part.operator in ("?", "&"): - split = i - else: - break - - if split == len(parts): - return parts, [] - - # The tail must start with a {?...} expression so that expand() - # emits a ? the URI can split on. A standalone {&page} expands - # with an & prefix, which partition("?") won't find. - first = parts[split] - assert isinstance(first, _Expression) - if first.operator != "?": - return parts, [] - - # If the path portion contains a literal ?/# or a {?...}/{#...} - # expression, lenient matching's partition("#") then partition("?") - # would strip content the path scan expects to see. Fall back to - # the strict scan. - for part in parts[:split]: - if isinstance(part, str): - if "?" in part or "#" in part: - return parts, [] - elif part.operator in ("?", "#"): - return parts, [] - - query_vars: list[Variable] = [] - for part in parts[split:]: - assert isinstance(part, _Expression) - query_vars.extend(part.variables) - - return parts[:split], query_vars - - -def _parse(template: str, *, max_variables: int) -> tuple[list[_Part], list[Variable]]: - """Split a template into an ordered sequence of literals and expressions. - - Walks the string, alternating between collecting literal runs and - parsing ``{...}`` expressions. The resulting ``parts`` sequence - preserves positional interleaving so ``match()`` and ``expand()`` can - walk it in order. - - Raises: - InvalidUriTemplate: On unclosed braces, too many expressions, or - any error surfaced by :func:`_parse_expression`. - """ - parts: list[_Part] = [] - variables: list[Variable] = [] - i = 0 - n = len(template) - - while i < n: - # Find the next expression opener from the current cursor. - brace = template.find("{", i) - - if brace == -1: - # No more expressions; everything left is a trailing literal. - parts.append(template[i:]) - break - - if brace > i: - # Literal text between cursor and the brace. - parts.append(template[i:brace]) - - end = template.find("}", brace) - if end == -1: - raise InvalidUriTemplate( - f"Unclosed expression at position {brace}", - template=template, - position=brace, - ) - - # Delegate body (between braces, exclusive) to the expression parser. - expr = _parse_expression(template, template[brace + 1 : end], brace) - parts.append(expr) - variables.extend(expr.variables) - - if len(variables) > max_variables: - raise InvalidUriTemplate( - f"Template exceeds maximum of {max_variables} variables", - template=template, - ) - - # Advance past the closing brace. - i = end + 1 - - _check_duplicate_variables(template, variables) - _check_single_query_expression(template, parts) - return parts, variables - - -def _parse_expression(template: str, body: str, pos: int) -> _Expression: - """Parse the body of a single ``{...}`` expression. - - The body is everything between the braces. It consists of an optional - leading operator character followed by one or more comma-separated - variable specifiers. Each specifier is a name with an optional - trailing ``*`` (explode modifier). - - Args: - template: The full template string, for error reporting. - body: The expression body, braces excluded. - pos: Character offset of the opening brace, for error reporting. - - Raises: - InvalidUriTemplate: On empty body, invalid variable names, or - unsupported modifiers. - """ - if not body: - raise InvalidUriTemplate(f"Empty expression at position {pos}", template=template, position=pos) - - # Peel off the operator, if any. Membership check justifies the cast. - operator: Operator = "" - if body[0] in _OPERATORS: - operator = cast(Operator, body[0]) - body = body[1:] - if not body: - raise InvalidUriTemplate( - f"Expression has operator but no variables at position {pos}", - template=template, - position=pos, - ) - - # Remaining body is comma-separated variable specs: name[*] - variables: list[Variable] = [] - for spec in body.split(","): - if ":" in spec: - raise InvalidUriTemplate( - f"Prefix modifier {{var:N}} is not supported (in {spec!r} at position {pos})", - template=template, - position=pos, - ) - - explode = spec.endswith("*") - name = spec[:-1] if explode else spec - - if not _VARNAME_RE.fullmatch(name): - raise InvalidUriTemplate( - f"Invalid variable name {name!r} at position {pos}", - template=template, - position=pos, - ) - - # Explode only makes sense for operators that repeat a separator. - # Simple/reserved/fragment have no per-item separator; query-explode - # needs order-agnostic dict matching which we don't support yet. - if explode and operator in ("", "+", "#", "?", "&"): - raise InvalidUriTemplate( - f"Explode modifier on {{{operator}{name}*}} is not supported for matching", - template=template, - position=pos, - ) - - variables.append(Variable(name=name, operator=operator, explode=explode)) - - return _Expression(operator=operator, variables=variables) - - -def _check_duplicate_variables(template: str, variables: list[Variable]) -> None: - """Reject templates that use the same variable name more than once. - - RFC 6570 requires repeated variables to expand to the same value, - which would require backreference matching with potentially - exponential cost. Rather than silently returning only the last - captured value, we reject at parse time. - - Raises: - InvalidUriTemplate: If any variable name appears more than once. - """ - seen: set[str] = set() - for var in variables: - if var.name in seen: - raise InvalidUriTemplate( - f"Variable {var.name!r} appears more than once; repeated variables are not supported", - template=template, - ) - seen.add(var.name) - - -def _check_single_query_expression(template: str, parts: list[_Part]) -> None: - """Reject templates with more than one ``{?...}`` expression. - - The ``?`` operator emits a leading ``?``, so two such expressions - expand to a URI with two ``?`` characters — malformed per RFC 3986 - §3.4. Use ``{?a,b}`` or ``{?a}{&b}`` for multiple query parameters. - """ - seen = False - for part in parts: - if isinstance(part, _Expression) and part.operator == "?": - if seen: - raise InvalidUriTemplate( - "Template contains more than one {?...} expression; " - "use {?a,b} or {?a}{&b} for multiple query parameters", - template=template, - ) - seen = True - - -def _flatten(parts: list[_Part]) -> list[_Atom]: - """Lower expressions into a flat sequence of literals and single-variable captures. - - Operator prefixes and separators become explicit ``_Lit`` atoms so - the scan only ever sees two atom kinds. Adjacent literals are - coalesced so that anchor-finding (``find``/``rfind``) operates on - the longest possible literal, reducing false matches. - - Explode variables emit no lead literal: the explode capture - includes its own separator-prefixed repetitions (``{/a*}`` → - ``/x/y/z``, not ``/`` then ``x/y/z``). - """ - atoms: list[_Atom] = [] - - def push_lit(text: str) -> None: - if not text: - return - if atoms and isinstance(atoms[-1], _Lit): - atoms[-1] = _Lit(atoms[-1].text + text) - else: - atoms.append(_Lit(text)) - - for part in parts: - if isinstance(part, str): - push_lit(part) - continue - spec = _OPERATOR_SPECS[part.operator] - for i, var in enumerate(part.variables): - lead = spec.prefix if i == 0 else spec.separator - if var.explode: - atoms.append(_Cap(var)) - elif spec.named: - # ; uses ifemp (bare name when empty); ? and & always - # emit name= so the equals is part of the literal. - if part.operator == ";": - push_lit(f"{lead}{var.name}") - atoms.append(_Cap(var, ifemp=True)) - else: - push_lit(f"{lead}{var.name}=") - atoms.append(_Cap(var)) - else: - push_lit(lead) - atoms.append(_Cap(var)) - return atoms - - -def _partition_greedy(atoms: list[_Atom], template: str) -> tuple[list[_Atom], Variable | None, list[_Atom]]: - """Split atoms at the single greedy variable, if any. - - Returns ``(prefix, greedy_var, suffix)``. If there is no greedy - variable the entire atom list is returned as the suffix so that - the right-to-left scan (which matches regex-greedy semantics) - handles it. - - Raises: - InvalidUriTemplate: If two variables are adjacent with no - literal between them — whether or not one is the - multi-segment variable, the scan has nothing to anchor the - boundary on — or if more than one multi-segment variable - is present (two are inherently ambiguous: there is no - principled way to decide which one absorbs an extra - segment). - """ - greedy_idx: int | None = None - prev: _Atom | None = None - for i, atom in enumerate(atoms): - if isinstance(atom, _Cap): - if isinstance(prev, _Cap): - raise InvalidUriTemplate( - f"Variables {prev.var.name!r} and {atom.var.name!r} are adjacent " - "with no literal separator; matching cannot determine where one " - "ends and the other begins. Add a literal between them or use a " - "single variable.", - template=template, - ) - if _is_greedy(atom.var): - if greedy_idx is not None: - raise InvalidUriTemplate( - "Template contains more than one multi-segment variable " - "({+var}, {#var}, or explode modifier); matching would be ambiguous", - template=template, - ) - greedy_idx = i - prev = atom - if greedy_idx is None: - return [], None, atoms - greedy = atoms[greedy_idx] - assert isinstance(greedy, _Cap) - return atoms[:greedy_idx], greedy.var, atoms[greedy_idx + 1 :] - - -def _scan_suffix( - atoms: Sequence[_Atom], uri: str, end: int, *, anchored: bool -) -> tuple[dict[str, str | list[str]], int] | None: - """Scan atoms right-to-left from ``end``, returning captures and start position. - - Each bounded variable takes the minimum span that lets its - preceding literal match (found via ``rfind``), which makes the - *first* variable in template order greedy — identical to Python - regex semantics for a sequence of greedy groups. - - When ``anchored`` is true the atom sequence is the entire template - (no greedy variable), so ``atoms[0]`` must match at URI position 0 - rather than at its rightmost occurrence. - """ - result: dict[str, str | list[str]] = {} - pos = end - i = len(atoms) - 1 - while i >= 0: - atom = atoms[i] - if isinstance(atom, _Lit): - n = len(atom.text) - if pos < n or uri[pos - n : pos] != atom.text: - return None - pos -= n - i -= 1 - continue - - var = atom.var - stops = _STOP_CHARS[var.operator] - prev = atoms[i - 1] if i > 0 else None - - if atom.ifemp: - # ;name or ;name=value. The preceding _Lit is ";name". - # Try empty first: if the lit ends at pos the value is - # absent (RFC ifemp). Otherwise require =value. - assert isinstance(prev, _Lit) - if uri.endswith(prev.text, 0, pos): - result[var.name] = "" - i -= 1 - continue - earliest = pos - while earliest > 0 and uri[earliest - 1] not in stops: - earliest -= 1 - eq = uri.find("=", earliest, pos) - if eq == -1: - return None - result[var.name] = unquote(uri[eq + 1 : pos]) - pos = eq - i -= 1 - continue - - # Earliest valid start: the var cannot extend left past any - # stop-char, so scan backward to find that boundary. - earliest = pos - while earliest > 0 and uri[earliest - 1] not in stops: - earliest -= 1 - - if prev is None: - start = earliest - else: - # prev is a _Lit: the parser rejects two adjacent captures, - # so the only possible neighbour kind is a literal. - assert isinstance(prev, _Lit) - if anchored and i - 1 == 0: - # First atom of the whole template: positionally fixed at - # 0, not rightmost occurrence. rfind would land inside the - # value when the literal repeats there (e.g. "prefix-{id}" - # against "prefix-prefix-123"). - start = len(prev.text) - if start < earliest or start > pos: - return None - else: - # Rightmost occurrence of the preceding literal whose end - # falls within the var's valid range. - idx = uri.rfind(prev.text, 0, pos) - if idx == -1 or idx + len(prev.text) < earliest: - return None - start = idx + len(prev.text) - - result[var.name] = unquote(uri[start:pos]) - pos = start - i -= 1 - return result, pos - - -def _scan_prefix( - atoms: Sequence[_Atom], uri: str, start: int, limit: int -) -> tuple[dict[str, str | list[str]], int] | None: - """Scan atoms left-to-right from ``start``, not exceeding ``limit``. - - Each bounded variable takes the minimum span that lets its - following literal match (found via ``find``), leaving the - greedy variable as much of the URI as possible. - """ - result: dict[str, str | list[str]] = {} - pos = start - for i, atom in enumerate(atoms): - if isinstance(atom, _Lit): - end = pos + len(atom.text) - if end > limit or uri[pos:end] != atom.text: - return None - pos = end - continue - - var = atom.var - stops = _STOP_CHARS[var.operator] - # Every capture here is followed by a literal: the parser rejects - # two adjacent captures, and a capture at the END of the prefix - # would be adjacent to the greedy variable. - nxt = atoms[i + 1] - assert isinstance(nxt, _Lit) - - if atom.ifemp: - # RFC §3.2.7 ifemp: ;name=val for non-empty, bare ;name for - # empty. Decide which form is present without falling through - # to the stop-char scan when the value is empty. - if uri.startswith(nxt.text, pos): - # Following literal begins immediately: value is empty. - # Checked before '=' so a literal that itself starts - # with '=' is not mistaken for the ifemp separator. - result[var.name] = "" - continue - if pos < limit and uri[pos] == "=": - pos += 1 # value follows; fall through to the scan - else: - # The following literal does not start here and there is - # no '=': the URI's name continued past the template's - # (e.g. ;keys vs ;key) — no parse. - return None - - # Latest valid end: the var stops at the first stop-char or - # the scan limit, whichever comes first. - latest = pos - while latest < limit and uri[latest] not in stops: - latest += 1 - - # First occurrence of the following literal: the capture takes - # the minimum span, leaving the greedy variable as much of the - # URI as possible. The search window's upper bound already - # forces any hit to start at or before ``latest``, so the var - # never extends past a stop-char. - end = uri.find(nxt.text, pos, latest + len(nxt.text)) - if end == -1: - return None - - result[var.name] = unquote(uri[pos:end]) - pos = end - return result, pos +import sys + +import mcp_client.shared.uri_template as _implementation +from mcp_client.shared.uri_template import ( + _OPERATOR_SPECS as _OPERATOR_SPECS, +) +from mcp_client.shared.uri_template import ( + _OPERATORS as _OPERATORS, +) +from mcp_client.shared.uri_template import ( + _PCT_TRIPLET_RE as _PCT_TRIPLET_RE, +) +from mcp_client.shared.uri_template import ( + _RESERVED as _RESERVED, +) +from mcp_client.shared.uri_template import ( + _STOP_CHARS as _STOP_CHARS, +) +from mcp_client.shared.uri_template import ( + _VARNAME_RE as _VARNAME_RE, +) +from mcp_client.shared.uri_template import ( + DEFAULT_MAX_TEMPLATE_LENGTH as DEFAULT_MAX_TEMPLATE_LENGTH, +) +from mcp_client.shared.uri_template import ( + DEFAULT_MAX_URI_LENGTH as DEFAULT_MAX_URI_LENGTH, +) +from mcp_client.shared.uri_template import ( + DEFAULT_MAX_VARIABLES as DEFAULT_MAX_VARIABLES, +) +from mcp_client.shared.uri_template import ( + InvalidUriTemplate as InvalidUriTemplate, +) +from mcp_client.shared.uri_template import ( + Operator as Operator, +) +from mcp_client.shared.uri_template import ( + UriTemplate as UriTemplate, +) +from mcp_client.shared.uri_template import ( + Variable as Variable, +) +from mcp_client.shared.uri_template import ( + _Atom as _Atom, +) +from mcp_client.shared.uri_template import ( + _Cap as _Cap, +) +from mcp_client.shared.uri_template import ( + _check_duplicate_variables as _check_duplicate_variables, +) +from mcp_client.shared.uri_template import ( + _check_single_query_expression as _check_single_query_expression, +) +from mcp_client.shared.uri_template import ( + _encode as _encode, +) +from mcp_client.shared.uri_template import ( + _expand_expression as _expand_expression, +) +from mcp_client.shared.uri_template import ( + _Expression as _Expression, +) +from mcp_client.shared.uri_template import ( + _extract_greedy as _extract_greedy, +) +from mcp_client.shared.uri_template import ( + _flatten as _flatten, +) +from mcp_client.shared.uri_template import ( + _is_greedy as _is_greedy, +) +from mcp_client.shared.uri_template import ( + _is_str_sequence as _is_str_sequence, +) +from mcp_client.shared.uri_template import ( + _Lit as _Lit, +) +from mcp_client.shared.uri_template import ( + _OperatorSpec as _OperatorSpec, +) +from mcp_client.shared.uri_template import ( + _parse as _parse, +) +from mcp_client.shared.uri_template import ( + _parse_expression as _parse_expression, +) +from mcp_client.shared.uri_template import ( + _parse_query as _parse_query, +) +from mcp_client.shared.uri_template import ( + _Part as _Part, +) +from mcp_client.shared.uri_template import ( + _partition_greedy as _partition_greedy, +) +from mcp_client.shared.uri_template import ( + _scan_prefix as _scan_prefix, +) +from mcp_client.shared.uri_template import ( + _scan_suffix as _scan_suffix, +) +from mcp_client.shared.uri_template import ( + _split_query_tail as _split_query_tail, +) + +sys.modules[__name__] = _implementation diff --git a/tests/client/test_package.py b/tests/client/test_package.py new file mode 100644 index 0000000000..175c810bfd --- /dev/null +++ b/tests/client/test_package.py @@ -0,0 +1,141 @@ +import importlib +import pkgutil +import subprocess +import sys +from contextlib import AsyncExitStack +from typing import Any, get_type_hints + +import anyio +import mcp_client +import mcp_types +import pytest +from typing_extensions import Self + +import mcp +from mcp.client import Transport +from mcp.server import Server, ServerRequestContext +from mcp.server.mcpserver import MCPServer + + +class ContextManagedServer(Server): + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *args: object) -> None: + return None + + +class ContextManagedMCPServer(MCPServer): + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *args: object) -> None: + return None + + +def test_client_exports_share_identity_with_the_full_sdk() -> None: + """The package split preserves the SDK's existing client classes and exceptions.""" + for name in set(mcp_client.__all__) & set(mcp.__all__): + assert vars(mcp_client)[name] is vars(mcp)[name] + + +def test_legacy_modules_export_the_extracted_implementations() -> None: + """Public legacy modules retain their objects rather than loading a second implementation.""" + for info in pkgutil.walk_packages(mcp_client.__path__, prefix="mcp_client."): + if any(part.startswith("_") for part in info.name.split(".")): + continue + implementation = importlib.import_module(info.name) + legacy = importlib.import_module(info.name.replace("mcp_client.", "mcp.", 1)) + if info.ispkg: + for name in vars(implementation).get("__all__", []): + assert vars(legacy)[name] is vars(implementation)[name] + else: + assert legacy is implementation + + +def test_full_sdk_client_annotations_remain_resolvable() -> None: + """The split keeps runtime annotation inspection available through the full SDK.""" + expected = Server[Any] | MCPServer | Transport | mcp.StdioServerParameters | str + assert get_type_hints(mcp.Client.__init__)["server"] == expected + + +@pytest.mark.anyio +@pytest.mark.parametrize("mode", ["auto", "legacy"]) +async def test_extracted_client_calls_an_in_process_sdk_server(mode: str) -> None: + """Both distributions use the same protocol machinery for modern and legacy connections.""" + result = mcp_types.ListToolsResult(tools=[mcp_types.Tool(name="example", input_schema={"type": "object"})]) + + async def list_tools( + ctx: ServerRequestContext, params: mcp_types.PaginatedRequestParams | None + ) -> mcp_types.ListToolsResult: + return result + + server = Server("example", on_list_tools=list_tools) + with anyio.fail_after(5): + async with mcp_client.Client(server, mode=mode) as client: + received = await client.list_tools() + assert isinstance(received, mcp_types.ListToolsResult) + assert received.tools == result.tools + + +@pytest.mark.anyio +@pytest.mark.parametrize("mode", ["auto", "legacy"]) +@pytest.mark.parametrize("server_type", [ContextManagedServer, ContextManagedMCPServer]) +async def test_servers_take_precedence_over_context_manager_transports( + server_type: type[ContextManagedServer] | type[ContextManagedMCPServer], mode: str +) -> None: + """The SDK connects server subclasses in-process even when they also manage an async context.""" + name = "context-managed-server" + server = server_type(name) + with anyio.fail_after(5): + async with AsyncExitStack() as stack: + await stack.enter_async_context(server) + client = await stack.enter_async_context(mcp_client.Client(server, mode=mode)) + assert client.server_info is not None + assert client.server_info.name == name + + +@pytest.mark.parametrize("first_import", ["mcp", "mcp_client"]) +def test_full_sdk_imports_preserve_module_attributes_and_exception_names(first_import: str) -> None: + """The SDK keeps its namespace in either import order. + + A fresh interpreter prevents other tests from populating missing module attributes. + """ + result = subprocess.run( + [ + sys.executable, + "-c", + f"import {first_import}\n" + + """ +import pickle + +import mcp +import mcp_client + +for name in ( + "caching", "client", "context", "extension", "session", "session_group", + "sse", "stdio", "streamable_http", "subscriptions", +): + assert vars(mcp.client)[name] is vars(mcp_client.client)[name], name + +import mcp.client.auth +import mcp_client.client.auth + +for name in ("exceptions", "oauth2", "utils"): + assert vars(mcp.client.auth)[name] is vars(mcp_client.client.auth)[name], name + +for name in ("MCPError", "MCPDeprecationWarning", "NoBackChannelError", "UrlElicitationRequiredError"): + error_type = vars(mcp.shared.exceptions)[name] + assert error_type is vars(mcp_client.shared.exceptions)[name] + assert error_type.__module__ == "mcp.shared.exceptions", name + assert pickle.loads(pickle.dumps(error_type)) is error_type +""", + ], + capture_output=True, + text=True, + encoding="utf-8", + check=False, + # Cold interpreter imports include coverage startup under xdist. + timeout=20, + ) + assert result.returncode == 0, result.stderr diff --git a/tests/client/test_session_group.py b/tests/client/test_session_group.py index b75d22b7a0..805a143e6e 100644 --- a/tests/client/test_session_group.py +++ b/tests/client/test_session_group.py @@ -303,7 +303,7 @@ async def test_client_session_group_disconnect_non_existent_server(): ( StdioServerParameters(command="test_stdio_cmd"), "stdio", - "mcp.client.session_group.mcp.stdio_client", + "mcp.client.session_group.stdio_client", ), ( SseServerParameters(url="http://test.com/sse", timeout=10.0), @@ -322,7 +322,7 @@ async def test_client_session_group_establish_session_parameterized( client_type_name: str, # Just for clarity or conditional logic if needed patch_target_for_client_func: str, ): - with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_ClientSession_class: + with mock.patch("mcp.client.session_group.ClientSession") as mock_ClientSession_class: with mock.patch(patch_target_for_client_func) as mock_specific_client_func: mock_client_cm_instance = mock.AsyncMock(name=f"{client_type_name}ClientCM") mock_read_stream = mock.AsyncMock(name=f"{client_type_name}Read") diff --git a/uv.lock b/uv.lock index 0d93802eb4..25b7762b40 100644 --- a/uv.lock +++ b/uv.lock @@ -10,6 +10,7 @@ resolution-markers = [ [manifest] members = [ "mcp", + "mcp-client", "mcp-everything-server", "mcp-example-stories", "mcp-simple-auth", @@ -1025,12 +1026,12 @@ dependencies = [ { name = "anyio" }, { name = "httpx2" }, { name = "jsonschema" }, + { name = "mcp-client" }, { name = "mcp-types" }, { name = "opentelemetry-api" }, { name = "pydantic" }, { name = "pyjwt", extra = ["crypto"] }, { name = "python-multipart" }, - { name = "pywin32", marker = "sys_platform == 'win32'" }, { name = "sse-starlette" }, { name = "starlette" }, { name = "typing-extensions" }, @@ -1090,13 +1091,13 @@ requires-dist = [ { name = "anyio", marker = "python_full_version >= '3.14'", specifier = ">=4.10" }, { name = "httpx2", specifier = ">=2.5.0" }, { name = "jsonschema", specifier = ">=4.20.0" }, + { name = "mcp-client", editable = "src/mcp-client" }, { name = "mcp-types", editable = "src/mcp-types" }, { name = "opentelemetry-api", specifier = ">=1.28.0" }, { name = "pydantic", specifier = ">=2.12.0" }, { name = "pyjwt", extras = ["crypto"], specifier = ">=2.10.1" }, { name = "python-dotenv", marker = "extra == 'cli'", specifier = ">=1.0.0" }, { name = "python-multipart", specifier = ">=0.0.9" }, - { name = "pywin32", marker = "sys_platform == 'win32'", specifier = ">=311" }, { name = "rich", marker = "extra == 'rich'", specifier = ">=13.9.4" }, { name = "sse-starlette", specifier = ">=3.0.0" }, { name = "starlette", marker = "python_full_version < '3.14'", specifier = ">=0.27" }, @@ -1141,6 +1142,35 @@ docs = [ ] translate = [{ name = "anthropic", specifier = ">=0.121.0" }] +[[package]] +name = "mcp-client" +source = { editable = "src/mcp-client" } +dependencies = [ + { name = "anyio" }, + { name = "httpx2" }, + { name = "jsonschema" }, + { name = "mcp-types" }, + { name = "opentelemetry-api" }, + { name = "pydantic" }, + { name = "pyjwt", extra = ["crypto"] }, + { name = "pywin32", marker = "sys_platform == 'win32'" }, + { name = "typing-extensions" }, +] + +[package.metadata] +requires-dist = [ + { name = "anyio", marker = "python_full_version < '3.14'", specifier = ">=4.9" }, + { name = "anyio", marker = "python_full_version >= '3.14'", specifier = ">=4.10" }, + { name = "httpx2", specifier = ">=2.5.0" }, + { name = "jsonschema", specifier = ">=4.20.0" }, + { name = "mcp-types", editable = "src/mcp-types" }, + { name = "opentelemetry-api", specifier = ">=1.28.0" }, + { name = "pydantic", specifier = ">=2.12.0" }, + { name = "pyjwt", extras = ["crypto"], specifier = ">=2.10.1" }, + { name = "pywin32", marker = "sys_platform == 'win32'", specifier = ">=311" }, + { name = "typing-extensions", specifier = ">=4.13.0" }, +] + [[package]] name = "mcp-everything-server" version = "0.1.0"