diff --git a/databricks_mcp/pyproject.toml b/databricks_mcp/pyproject.toml index 6abda1d4f..b41bffa1a 100644 --- a/databricks_mcp/pyproject.toml +++ b/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" diff --git a/databricks_mcp/src/databricks_mcp/mcp.py b/databricks_mcp/src/databricks_mcp/mcp.py index 08aeb867f..b132ee6b9 100644 --- a/databricks_mcp/src/databricks_mcp/mcp.py +++ b/databricks_mcp/src/databricks_mcp/mcp.py @@ -2,6 +2,7 @@ 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 @@ -9,9 +10,26 @@ 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,13 +226,10 @@ 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, @@ -194,13 +237,10 @@ async def _call_tools_async( 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.""" diff --git a/databricks_mcp/src/databricks_mcp/oauth_provider.py b/databricks_mcp/src/databricks_mcp/oauth_provider.py index 084244f6b..892e9af01 100644 --- a/databricks_mcp/src/databricks_mcp/oauth_provider.py +++ b/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), diff --git a/databricks_mcp/tests/integration_tests/conftest.py b/databricks_mcp/tests/integration_tests/conftest.py index 52836c719..2dde3525c 100644 --- a/databricks_mcp/tests/integration_tests/conftest.py +++ b/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 diff --git a/databricks_mcp/tests/integration_tests/test_mcp_core.py b/databricks_mcp/tests/integration_tests/test_mcp_core.py index f83f5ad0c..5d679a07c 100644 --- a/databricks_mcp/tests/integration_tests/test_mcp_core.py +++ b/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, _, diff --git a/databricks_mcp/tests/unit_tests/test_mcp.py b/databricks_mcp/tests/unit_tests/test_mcp.py index 685265c46..56cf9f9e0 100644 --- a/databricks_mcp/tests/unit_tests/test_mcp.py +++ b/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,19 +209,11 @@ 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 @@ -204,7 +221,6 @@ async def test_get_tools_async(self): 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 @@ -212,17 +228,9 @@ 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 @@ -230,12 +238,11 @@ async def test_call_tools_async(self): 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()