fix(databricks-mcp): support mcp 2.x and drop the temporary <2.0.0 pin

databricks_mcp/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "databricks-mcp"
-version = "0.9.1"
+version = "0.9.2"
 description = "MCP helpers for Databricks"
 authors = [
     { name="Databricks", email="agent-feedback@databricks.com" },
@@ -9,7 +9,7 @@ readme = "README.md"
 license = { text="Apache-2.0" }
 requires-python = ">=3.10"
 dependencies = [
-    "mcp>=1.13.0,<2.0.0",
+    "mcp>=1.13.0",
     "databricks-sdk>=0.49.0",
     "databricks-ai-bridge>=0.4.2",
     "mlflow>=3.1"

databricks_mcp/src/databricks_mcp/mcp.py

@@ -2,16 +2,34 @@
 import json
 import logging
 import re
+from contextlib import asynccontextmanager
 from functools import wraps
 from typing import Any, Callable, List, Optional
 from urllib.parse import urlparse
 
 import requests
 from databricks.sdk import WorkspaceClient
 from databricks_ai_bridge.utils.annotations import experimental
-from mcp.client.session import ClientSession
-from mcp.client.streamable_http import streamablehttp_client
 from mcp.types import CallToolResult, Tool
+
+# mcp 2.0.0 replaced the transport + ``ClientSession`` layering with a single
+# ``Client``, and moved auth onto an ``httpx2.AsyncClient`` handed to
+# ``streamable_http_client`` (the old ``streamablehttp_client`` alias and its
+# ``auth=`` keyword were removed). Detect which major is installed and adapt so
+# this package stays compatible with both mcp 1.x and mcp 2.x.
+try:
+    from mcp import Client  # noqa: F401  (present only in mcp >= 2.0.0)
+
+    _MCP_V2 = True
+except ImportError:  # mcp < 2.0.0
+    _MCP_V2 = False
+
+if _MCP_V2:
+    import httpx2
+    from mcp.client.streamable_http import streamable_http_client
+else:
+    from mcp.client.session import ClientSession
+    from mcp.client.streamable_http import streamablehttp_client  # ty:ignore[unresolved-import]
 from mlflow.models.resources import (
     DatabricksFunction,
     DatabricksGenieSpace,
@@ -25,6 +43,34 @@
 logger = logging.getLogger(__name__)
 
 
+@asynccontextmanager
+async def _open_mcp_session(server_url: str, auth: DatabricksOAuthClientProvider):
+    """Open an authenticated MCP session against ``server_url``.
+
+    Yields an object exposing ``list_tools()`` and ``call_tool()`` and works
+    with both mcp 1.x (transport + ``ClientSession``) and mcp 2.x (``Client``).
+    """
+    if _MCP_V2:
+        # mcp 2.x: auth is attached to an httpx2 client passed to the
+        # transport, and ``Client`` performs the initialize handshake itself.
+        async with httpx2.AsyncClient(auth=auth, follow_redirects=True) as http_client:
+            async with Client(
+                streamable_http_client(server_url, http_client=http_client)
+            ) as session:
+                yield session
+    else:
+        # mcp 1.x: auth is a keyword on the transport and the session must be
+        # explicitly initialized.
+        async with streamablehttp_client(url=server_url, auth=auth) as (
+            read_stream,
+            write_stream,
+            _,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+                yield session
+
+
 def _is_databricks_apps_url(url: str) -> bool:
     """Check if the URL is hosted on Databricks Apps."""
     parsed = urlparse(url)
@@ -180,27 +226,21 @@ def _get_databricks_managed_mcp_url_type(self) -> str | None:
 
     async def _get_tools_async(self) -> List[Tool]:
         """Fetch tools from the MCP endpoint asynchronously."""
-        async with streamablehttp_client(
-            url=self.server_url,
-            auth=DatabricksOAuthClientProvider(self.client),
-        ) as (read_stream, write_stream, _):
-            async with ClientSession(read_stream, write_stream) as session:
-                await session.initialize()
-                return (await session.list_tools()).tools
+        async with _open_mcp_session(
+            self.server_url, DatabricksOAuthClientProvider(self.client)
+        ) as session:
+            return (await session.list_tools()).tools
 
     async def _call_tools_async(
         self,
         tool_name: str,
         arguments: dict[str, Any] | None = None,
     ) -> CallToolResult:
         """Call the tool with the given name and input."""
-        async with streamablehttp_client(
-            url=self.server_url,
-            auth=DatabricksOAuthClientProvider(self.client),
-        ) as (read_stream, write_stream, _):
-            async with ClientSession(read_stream, write_stream) as session:
-                await session.initialize()
-                return await session.call_tool(tool_name, arguments)
+        async with _open_mcp_session(
+            self.server_url, DatabricksOAuthClientProvider(self.client)
+        ) as session:
+            return await session.call_tool(tool_name, arguments)
 
     def _extract_genie_id(self) -> str:
         """Extract the Genie space ID from the URL."""

databricks_mcp/src/databricks_mcp/oauth_provider.py

@@ -22,18 +22,40 @@ async def get_tokens(self) -> OAuthToken | None:
 class DatabricksOAuthClientProvider(OAuthClientProvider):
     """
     An OAuthClientProvider for Databricks. This class extends mcp.client.auth.OAuthClientProvider
-    and can be used with the `mcp.client.streamable_http` to authorize the MCP Server with Databricks.
+    and can be used with `mcp.client.streamable_http` to authorize the MCP Server with Databricks.
+
+    Usage (mcp >= 2.0.0): the provider is an ``httpx2.Auth`` and must be attached to an
+    ``httpx2.AsyncClient`` that is handed to ``streamable_http_client``:
 
-    Usage:
         .. code-block:: python
 
+            import httpx2
+            from mcp import Client
+            from mcp.client.streamable_http import streamable_http_client
+
             from databricks_mcp.oauth_provider import DatabricksOAuthClientProvider
-            from mcp.client.streamable_http import streamablehttp_client
-            from mcp.client.session import ClientSession
 
             # Initialize the Databricks workspace client
             workspace_client = WorkspaceClient()
 
+            async with httpx2.AsyncClient(
+                auth=DatabricksOAuthClientProvider(workspace_client),
+                follow_redirects=True,
+            ) as http_client:
+                async with Client(
+                    streamable_http_client("https://mcp-server-url", http_client=http_client)
+                ) as session:
+                    tools = await session.list_tools()
+
+    Usage (mcp < 2.0.0): the provider is passed directly as the transport ``auth`` keyword:
+
+        .. code-block:: python
+
+            from mcp.client.streamable_http import streamablehttp_client
+            from mcp.client.session import ClientSession
+
+            from databricks_mcp.oauth_provider import DatabricksOAuthClientProvider
+
             async with streamablehttp_client(
                 url="https://mcp-server-url",
                 auth=DatabricksOAuthClientProvider(workspace_client),

databricks_mcp/tests/integration_tests/conftest.py

@@ -21,7 +21,11 @@
 import os
 
 import pytest
-from mcp.shared.exceptions import McpError
+
+try:
+    from mcp import MCPError as McpError  # mcp >= 2.0.0 (renamed from McpError)
+except ImportError:  # mcp < 2.0.0
+    from mcp.shared.exceptions import McpError  # ty:ignore[unresolved-import]
 
 from databricks_mcp import DatabricksMCPClient
 

databricks_mcp/tests/integration_tests/test_mcp_core.py

@@ -12,25 +12,48 @@
 
 import pytest
 from conftest import _skip_if_not_found
-from mcp.shared.exceptions import McpError
 from mcp.types import CallToolResult
 
+try:
+    from mcp import MCPError as McpError  # mcp >= 2.0.0 (renamed from McpError)
+except ImportError:  # mcp < 2.0.0
+    from mcp.shared.exceptions import McpError  # ty:ignore[unresolved-import]
+
 
 @asynccontextmanager
 async def raw_mcp_session(url, workspace_client):
-    """Create a raw MCP ClientSession using streamable_http_client with Databricks OAuth."""
-    import httpx
-    from mcp import ClientSession
-    from mcp.client.streamable_http import streamable_http_client
+    """Create a raw MCP session with Databricks OAuth, using the raw MCP SDK.
 
+    Works with both mcp 1.x (transport + ``ClientSession``) and mcp 2.x
+    (``Client`` + ``httpx2``).
+    """
     from databricks_mcp import DatabricksOAuthClientProvider
 
-    async with httpx.AsyncClient(
-        auth=DatabricksOAuthClientProvider(workspace_client),
-        follow_redirects=True,
-        timeout=httpx.Timeout(120.0, read=120.0),
-    ) as http_client:
-        async with streamable_http_client(url, http_client=http_client) as (
+    auth = DatabricksOAuthClientProvider(workspace_client)
+
+    try:
+        from mcp import Client  # present only in mcp >= 2.0.0
+
+        mcp_v2 = True
+    except ImportError:
+        mcp_v2 = False
+
+    if mcp_v2:
+        import httpx2
+        from mcp.client.streamable_http import streamable_http_client
+
+        async with httpx2.AsyncClient(
+            auth=auth,
+            follow_redirects=True,
+            timeout=httpx2.Timeout(120.0, read=120.0),
+        ) as http_client:
+            async with Client(streamable_http_client(url, http_client=http_client)) as session:
+                yield session
+    else:
+        from mcp import ClientSession
+        from mcp.client.streamable_http import streamablehttp_client  # ty:ignore[unresolved-import]
+
+        async with streamablehttp_client(url=url, auth=auth) as (
             read_stream,
             write_stream,
             _,

databricks_mcp/tests/unit_tests/test_mcp.py

@@ -1,4 +1,5 @@
 import re
+from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -17,6 +18,30 @@
 )
 
 
+def _patch_mcp_session(mock_session):
+    """Patch the version-agnostic ``_open_mcp_session`` helper to yield ``mock_session``.
+
+    Works regardless of whether mcp 1.x or 2.x is installed, since both code
+    paths funnel through ``_open_mcp_session``.
+    """
+
+    @asynccontextmanager
+    async def _fake_session(*args, **kwargs):
+        yield mock_session
+
+    return patch("databricks_mcp.mcp._open_mcp_session", _fake_session)
+
+
+def _make_tool(name, description=""):
+    """Construct an mcp ``Tool`` compatibly across mcp 1.x and 2.x.
+
+    mcp 2.x renamed the ``inputSchema`` field to ``input_schema`` (keeping
+    ``inputSchema`` as a validation alias), so build via ``model_validate`` with
+    the wire name, which is accepted on both versions and avoids a type error.
+    """
+    return Tool.model_validate({"name": name, "description": description, "inputSchema": {}})
+
+
 class TestDatabricksMCPClient:
     """Test cases for DatabricksMCPClient class."""
 
@@ -184,58 +209,40 @@ def test_normalize_tool_name(self, input_name, expected_name):
     @pytest.mark.asyncio
     async def test_get_tools_async(self):
         """Test asynchronous tool fetching."""
-        mock_tools = [Tool(name="test_tool", description="Test tool", inputSchema={})]
+        mock_tools = [_make_tool("test_tool", "Test tool")]
         mock_session = AsyncMock()
-        mock_session.initialize = AsyncMock()
         mock_session.list_tools = AsyncMock(return_value=MagicMock(tools=mock_tools))
 
-        with (
-            patch("databricks_mcp.mcp.streamablehttp_client") as mock_client,
-            patch("databricks_mcp.mcp.ClientSession") as mock_session_class,
-            patch("databricks_mcp.mcp.DatabricksOAuthClientProvider"),
-        ):
-            mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), None)
-            mock_session_class.return_value.__aenter__.return_value = mock_session
-
+        with _patch_mcp_session(mock_session):
             workspace_client = WorkspaceClient(host="https://test.com", token="test-token")
             client = DatabricksMCPClient(
                 "https://test.com/api/2.0/mcp/functions/catalog/schema", workspace_client
             )
             tools = await client._get_tools_async()
 
             assert tools == mock_tools
-            mock_session.initialize.assert_called_once()
             mock_session.list_tools.assert_called_once()
 
     @pytest.mark.asyncio
     async def test_call_tools_async(self):
         """Test asynchronous tool calling."""
         mock_result = CallToolResult(content=[TextContent(type="text", text="test result")])
         mock_session = AsyncMock()
-        mock_session.initialize = AsyncMock()
         mock_session.call_tool = AsyncMock(return_value=mock_result)
 
-        with (
-            patch("databricks_mcp.mcp.streamablehttp_client") as mock_client,
-            patch("databricks_mcp.mcp.ClientSession") as mock_session_class,
-            patch("databricks_mcp.mcp.DatabricksOAuthClientProvider"),
-        ):
-            mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), None)
-            mock_session_class.return_value.__aenter__.return_value = mock_session
-
+        with _patch_mcp_session(mock_session):
             workspace_client = WorkspaceClient(host="https://test.com", token="test-token")
             client = DatabricksMCPClient(
                 "https://test.com/api/2.0/mcp/functions/catalog/schema", workspace_client
             )
             result = await client._call_tools_async("test_tool", {"arg": "value"})
 
             assert result == mock_result
-            mock_session.initialize.assert_called_once()
             mock_session.call_tool.assert_called_once_with("test_tool", {"arg": "value"})
 
     def test_list_tools(self):
         """Test synchronous tool listing."""
-        mock_tools = [Tool(name="test_tool", description="Test tool", inputSchema={})]
+        mock_tools = [_make_tool("test_tool", "Test tool")]
 
         with patch.object(DatabricksMCPClient, "_get_tools_async", return_value=mock_tools):
             workspace_client = WorkspaceClient(host="https://test.com", token="test-token")
@@ -278,9 +285,7 @@ def test_get_databricks_resources_with_tools(
         self, mcp_type, tool_names, expected_resource_names
     ):
         """Test getting Databricks resources for MCP types that require tool listing."""
-        mock_tools = [
-            Tool(name=name, description=f"Tool {name}", inputSchema={}) for name in tool_names
-        ]
+        mock_tools = [_make_tool(name, f"Tool {name}") for name in tool_names]
 
         with (
             patch.object(DatabricksMCPClient, "list_tools", return_value=mock_tools),
@@ -330,7 +335,7 @@ def test_get_databricks_resources_invalid_url(self):
 
     def test_get_databricks_resources_unknown_mcp_type(self):
         """Test getting Databricks resources for unknown MCP type."""
-        mock_tools = [Tool(name="test_tool", description="Test tool", inputSchema={})]
+        mock_tools = [_make_tool("test_tool", "Test tool")]
 
         with (
             patch.object(DatabricksMCPClient, "list_tools", return_value=mock_tools),
@@ -454,7 +459,7 @@ def test_error_decorator_paths(self, status_code, expected_exc, expected_msg, me
             patch.object(client, "_get_databricks_managed_mcp_url_type", return_value=None),
             patch("databricks_mcp.mcp.DatabricksOAuthClientProvider") as mock_auth_provider,
             patch("requests.request") as mock_request,
-            patch("databricks_mcp.mcp.streamablehttp_client") as mock_client,
+            patch("databricks_mcp.mcp._open_mcp_session") as mock_session,
             patch.object(
                 client.client.config,
                 "authenticate",
@@ -466,7 +471,7 @@ def test_error_decorator_paths(self, status_code, expected_exc, expected_msg, me
             mock_request.return_value = mock_response
 
             # Trigger decorator by failing the MCP call
-            mock_client.side_effect = original_error
+            mock_session.side_effect = original_error
 
             method = getattr(client, method_name)
             if expected_exc is Exception:
@@ -508,11 +513,11 @@ def test_error_decorator_managed_server_reraises_original(self):
             patch.object(
                 client, "_get_databricks_managed_mcp_url_type", return_value=UC_FUNCTIONS_MCP
             ),
-            patch("databricks_mcp.mcp.streamablehttp_client") as mock_client,
+            patch("databricks_mcp.mcp._open_mcp_session") as mock_session,
             patch("databricks_mcp.mcp.DatabricksOAuthClientProvider"),
             patch("requests.request") as mock_request,
         ):
-            mock_client.side_effect = original_error
+            mock_session.side_effect = original_error
 
             with pytest.raises(Exception, match="Databricks server error"):
                 client.list_tools()