feat: support MCP SDK 2.x alongside 1.x

.github/workflows/ci.yml

@@ -42,3 +42,36 @@ jobs:
 
       - name: Build package
         run: uv build
+
+  # The proxy hands its own HTTP client to the SDK's transport, and SDK 2.0
+  # both renamed that library (httpx -> httpx2) and changed the transport and
+  # message shapes. The locked resolution only ever proves whichever major the
+  # lockfile happens to hold, so each supported major is resolved and tested
+  # explicitly here.
+  test-sdk-majors:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        include:
+          - mcp-version: "mcp>=1.24,<2"
+            http-version: "httpx>=0.28,<0.29"
+          - mcp-version: "mcp>=2,<3"
+            http-version: "httpx2>=2.5,<3"
+    steps:
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5  # v4
+
+      - uses: astral-sh/setup-uv@e4db8464a088ece1b920f60402e813ea4de65b8f  # v4
+        with:
+          python-version: "3.13"
+
+      - name: Install dependencies
+        run: uv sync
+
+      - name: Resolve latest supported SDK and HTTP client
+        run: uv pip install --upgrade "${{ matrix.mcp-version }}" "${{ matrix.http-version }}"
+
+      - name: Run unit tests
+        run: uv run --no-sync pytest -m unit
+
+      - name: Typecheck
+        run: uv run --no-sync mypy

CLAUDE.md

@@ -39,10 +39,42 @@ Package in `src/uc_mcp_proxy/`:
 - `errors.py` — HTTP error diagnosis and reporting from the remote server
 - `token_exchange.py` — RFC 8693 exchange of a PAT for an app-scoped OAuth token
 - `app_discovery.py` — App-host / classic-PAT detection and lookup of an app's `oauth2_app_client_id` + scopes from workspace metadata (drives auto-exchange)
+- `_compat.py` — MCP SDK major-version shim (see below)
 - `__init__.py` — re-exports `DatabricksAuth`
 
 The proxy bridges an MCP stdio transport to a remote Streamable HTTP MCP server, injecting Databricks OAuth tokens on every request via `DatabricksAuth`.
 
+### SDK compatibility
+
+Both MCP SDK 1.x and 2.x are supported. 2.0 changed three things the proxy
+depends on: the HTTP client library (`httpx` → `httpx2`), the
+`streamable_http_client` yield (dropped the third `get_session_id` element),
+and `SessionMessage.message` (no longer wrapped in a `JSONRPCMessage` root
+model).
+
+The supported dependency ranges are deliberately bounded: `mcp>=1.24,<3`
+matches the API floor and SDK majors exercised by the compatibility shim, while
+`httpx>=0.28,<0.29` and `httpx2>=2.5,<3` prevent untested HTTP-client releases
+from silently changing the retry ordering or redirect-extension behavior. The
+SDK-major CI matrix runs the same retry-invariant tests against the HTTP module
+actually selected by each installed MCP major.
+
+`_compat.py` resolves the HTTP library by reading it back off
+`mcp.client.streamable_http` rather than importing a guessed name — both
+libraries install side by side, so `try: import httpx2` would hand an SDK 1.x
+transport a client built from a library that SDK never imported. Because the
+proxy owns the client it passes to the transport, that mismatch would surface
+as a type error deep in the SDK's request path rather than at import time.
+
+Import `httpx` from `_compat` (tests: from `tests/support.py`) for anything
+that crosses the SDK boundary — the client handed to `streamable_http_client`,
+and every request, response and transport object the SDK's hooks will see.
+
+Code that owns a client end-to-end and never passes it to the SDK may import
+`httpx` directly; `token_exchange.py` does this deliberately, since its RFC
+8693 call to the token endpoint is the proxy's own HTTP request and has
+nothing to do with which library the SDK bound.
+
 ### Error handling
 
 The proxy owns the `httpx.AsyncClient` it hands to the MCP SDK, so that client's

pyproject.toml

@@ -4,16 +4,16 @@ version = "0.5.1"
 description = "MCP stdio-to-Streamable-HTTP proxy with Databricks OAuth"
 requires-python = ">=3.10"
 dependencies = [
-    "mcp>=1.8,<2",
+    "mcp>=1.24,<3",
     "databricks-sdk>=0.30.0",
-    # Bounded deliberately. The auth retry depends on two httpx behaviours that
-    # are implementation details rather than documented contracts: response
-    # event hooks running strictly before the auth flow is handed the response,
-    # and ``extensions`` being copied per redirect hop. A test pins both, but
-    # that only protects CI -- an install that resolves a newer httpx would get
-    # a silently broken retry and never run the suite. Raise the cap only after
-    # ``test_response_hook_precedes_auth_flow`` passes on the newer version.
+    # Bounded deliberately. The auth retry depends on two behaviours shared by
+    # httpx (MCP 1.x) and httpx2 (MCP 2.x) that are implementation details rather
+    # than documented contracts: response hooks running strictly before the auth
+    # flow receives the response, and ``extensions`` being copied per redirect
+    # hop. The SDK-major CI matrix runs the invariant tests against each actual
+    # library. Raise either cap only after those tests pass on the newer major.
     "httpx>=0.28,<0.29",
+    "httpx2>=2.5,<3",
     "anyio",
 ]
 license = "MIT"

src/uc_mcp_proxy/__main__.py

@@ -8,18 +8,17 @@
 import sys
 import time
 from collections.abc import AsyncGenerator, Callable, Generator
-from typing import Any, NoReturn
+from typing import Any, NoReturn, cast
 from urllib.parse import urljoin, urlsplit
 
 import anyio
-import httpx
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from databricks.sdk import WorkspaceClient
 from mcp.client.streamable_http import streamable_http_client
 from mcp.server.stdio import stdio_server
 from mcp.shared.message import SessionMessage
 from mcp.types import JSONRPCRequest
 
+from uc_mcp_proxy._compat import MessageReceiveStream, MessageSendStream, httpx, jsonrpc_payload
 from uc_mcp_proxy.app_discovery import (
     AppDiscoveryError,
     discover_app,
@@ -239,7 +238,7 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
         yield request
 
 
-async def copy_stream(source: MemoryObjectReceiveStream[Any], dest: MemoryObjectSendStream[Any]) -> None:
+async def copy_stream(source: MessageReceiveStream, dest: MessageSendStream) -> None:
     """Copy all messages from source to dest, closing dest when source is exhausted."""
     try:
         async for message in source:
@@ -263,7 +262,7 @@ def inject_meta(
     # for meta params today; _meta is valid on any request per MCP spec.
     if isinstance(message, Exception):
         return message
-    root = message.message.root
+    root = jsonrpc_payload(message)
     if not isinstance(root, JSONRPCRequest) or root.method != "tools/call":
         return message
     if root.params is None:
@@ -281,8 +280,8 @@ def inject_meta(
 
 
 async def inject_meta_stream(
-    source: MemoryObjectReceiveStream[Any],
-    dest: MemoryObjectSendStream[Any],
+    source: MessageReceiveStream,
+    dest: MessageSendStream,
     meta: dict[str, str],
 ) -> None:
     """Like copy_stream, but applies inject_meta to each forwarded message."""
@@ -294,10 +293,10 @@ async def inject_meta_stream(
 
 
 async def bridge(
-    stdio_read: MemoryObjectReceiveStream[Any],
-    stdio_write: MemoryObjectSendStream[Any],
-    http_read: MemoryObjectReceiveStream[Any],
-    http_write: MemoryObjectSendStream[Any],
+    stdio_read: MessageReceiveStream,
+    stdio_write: MessageSendStream,
+    http_read: MessageReceiveStream,
+    http_write: MessageSendStream,
     meta: dict[str, str] | None = None,
 ) -> None:
     """Bidirectional bridge between stdio and HTTP stream pairs.
@@ -649,11 +648,18 @@ async def run(
             ) as httpx_client,
             streamable_http_client(
                 resolved_url,
-                http_client=httpx_client,
+                # Cast because the SDK names its own HTTP library in this
+                # signature, and which library that is varies by SDK major.
+                # ``_compat`` guarantees the client was built from the module
+                # this very SDK imported; mypy only ever sees one of them.
+                http_client=cast(Any, httpx_client),
             ) as (
                 http_read,
                 http_write,
-                _get_session_id,
+                # SDK 1.x yielded a third element, ``get_session_id``, which
+                # this proxy never used; 2.0 dropped it. Starred so the same
+                # unpack accepts either shape.
+                *_,
             ),
         ):
             try:

src/uc_mcp_proxy/_compat.py

@@ -0,0 +1,88 @@
+"""Compatibility across MCP SDK major versions.
+
+SDK 2.0 changed three things this proxy depends on:
+
+1. the HTTP client library moved from ``httpx`` to ``httpx2``;
+2. ``streamable_http_client`` yields ``(read, write)`` where 1.x yielded
+   ``(read, write, get_session_id)``;
+3. ``SessionMessage.message`` is the JSON-RPC model itself, where 1.x wrapped
+   it in the ``JSONRPCMessage`` pydantic root model.
+
+Only (1) and (3) need a shim here -- (2) is absorbed by a starred unpack at
+the one call site in ``__main__``.
+
+The HTTP library is read back off the SDK module rather than imported by a
+guessed name. ``httpx`` and ``httpx2`` install side by side: httpx2 does not
+replace httpx, and databricks-sdk and others still pull httpx in. So
+``try: import httpx2`` would hand an SDK 1.x transport a client built from a
+library that SDK never imported -- and because the proxy owns the client it
+passes to ``streamable_http_client``, that mismatch surfaces as a type error
+deep inside the SDK's request path rather than at import time. Asking the SDK
+which module it bound is the only answer that cannot drift.
+"""
+
+from __future__ import annotations
+
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, Protocol
+
+from mcp.client import streamable_http as _sdk_streamable_http
+
+__all__ = ["MessageReceiveStream", "MessageSendStream", "httpx", "jsonrpc_payload"]
+
+#: Module names to look for on the SDK transport, newest SDK first.
+_HTTPX_MODULE_NAMES = ("httpx2", "httpx")
+
+
+def _resolve_httpx() -> ModuleType:
+    """Return the HTTP client module the installed MCP SDK builds clients from."""
+    for name in _HTTPX_MODULE_NAMES:
+        module = getattr(_sdk_streamable_http, name, None)
+        if isinstance(module, ModuleType):
+            return module
+    raise RuntimeError(
+        "uc-mcp-proxy: cannot tell which HTTP client library this MCP SDK uses. "
+        f"Expected mcp.client.streamable_http to import one of {_HTTPX_MODULE_NAMES}."
+    )
+
+
+if TYPE_CHECKING:
+    # httpx2 mirrors the httpx API for every name the proxy touches, so the
+    # httpx stubs describe both. Only the runtime object has to match the SDK.
+    import httpx
+else:
+    httpx = _resolve_httpx()
+
+
+def jsonrpc_payload(message: Any) -> Any:
+    """Return the JSON-RPC model carried by a ``SessionMessage``.
+
+    SDK 1.x wraps it in the ``JSONRPCMessage`` root model; 2.0 stores the
+    ``JSONRPCRequest``/``JSONRPCNotification``/... directly. Both are reached
+    through the attribute that exists, so neither version is special-cased.
+    """
+    payload = message.message
+    return getattr(payload, "root", payload)
+
+
+class MessageReceiveStream(Protocol):
+    """The read half of a transport, as the bridge actually uses it.
+
+    Structural on purpose. SDK 1.x hands out anyio ``MemoryObjectReceiveStream``
+    objects and 2.0 hands out its own context-carrying wrappers; naming either
+    concrete class here would type-check against one SDK and fail on the other.
+    """
+
+    def __aiter__(self) -> Any: ...
+
+    async def __anext__(self) -> Any: ...
+
+    async def aclose(self) -> None: ...
+
+
+class MessageSendStream(Protocol):
+    """The write half of a transport, as the bridge actually uses it."""
+
+    async def send(self, item: Any, /) -> None: ...
+
+    async def aclose(self) -> None: ...

src/uc_mcp_proxy/errors.py

@@ -29,7 +29,8 @@
 from collections.abc import Sequence
 
 import anyio
-import httpx
+
+from uc_mcp_proxy._compat import httpx
 
 
 class ProxyFatalError(Exception):

tests/integration/test_proxy.py

@@ -7,10 +7,11 @@
 from unittest.mock import patch
 
 import anyio
-import httpx
 import pytest
 from mcp.client.streamable_http import streamable_http_client as _real_streamable_http_client
 
+from tests.support import httpx
+
 pytestmark = [pytest.mark.integration, pytest.mark.anyio]
 
 

tests/support.py

@@ -0,0 +1,33 @@
+"""Version-agnostic helpers for building SDK objects in tests.
+
+SDK 2.0 turned ``JSONRPCMessage`` from a pydantic root model into a plain
+union alias, so ``JSONRPCMessage(...)`` is no longer callable and the
+``SessionMessage.message`` it produced no longer has ``.root``. Tests build
+messages by validating a raw wire payload through the SDK's own schema, which
+accepts both shapes and keeps the JSON that actually goes on the wire as the
+source of truth rather than a hand-built object graph.
+
+``httpx`` is re-exported from the proxy's compat module so tests construct
+requests, responses and transports from the same library the installed SDK
+uses -- otherwise ``isinstance`` checks inside the proxy would compare objects
+from two different HTTP libraries.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from mcp.shared.message import SessionMessage
+from mcp.types import JSONRPCMessage
+from pydantic import TypeAdapter
+
+from uc_mcp_proxy._compat import httpx, jsonrpc_payload
+
+__all__ = ["httpx", "jsonrpc_payload", "session_message"]
+
+_JSONRPC_ADAPTER: TypeAdapter[Any] = TypeAdapter(JSONRPCMessage)
+
+
+def session_message(payload: dict[str, Any]) -> SessionMessage:
+    """Build a ``SessionMessage`` from a raw JSON-RPC ``payload`` dict."""
+    return SessionMessage(_JSONRPC_ADAPTER.validate_python(payload))

tests/unit/test_auth.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
-import httpx
 import pytest
 
+from tests.support import httpx
+
 pytestmark = pytest.mark.unit
 
 

tests/unit/test_cli.py

@@ -7,9 +7,10 @@
 from unittest.mock import MagicMock, patch
 
 import anyio
-import httpx
 import pytest
 
+from tests.support import httpx
+
 pytestmark = pytest.mark.unit