diff --git a/databricks_mcp/src/databricks_mcp/external_client.py b/databricks_mcp/src/databricks_mcp/external_client.py new file mode 100644 index 000000000..a36f48fe3 --- /dev/null +++ b/databricks_mcp/src/databricks_mcp/external_client.py @@ -0,0 +1,140 @@ +"""POST-only JSON-RPC MCP client for external/managed UC connections. + +External MCP connections (``/api/2.0/mcp/external/``) are fronted by a +POST-only JSON-REST bridge: it rejects any non-POST method with ``405 Method Not +Allowed`` and replies with ``application/json`` (never ``text/event-stream``). It is a +request/response transcoder with no GET endpoint, no SSE stream, and no server-initiated +messages. + +The standard Streamable HTTP client (``mcp.client.streamable_http.streamablehttp_client``, +used by :class:`~databricks_mcp.mcp.DatabricksMCPClient` for managed servers) opens a +background GET SSE stream for server -> client messages. Against this POST-only bridge that +GET is rejected with 405 and the stream task dies; a later request whose response handling +expects coordination with that (now-dead) stream then blocks until timeout. See +modelcontextprotocol/python-sdk#1941 ("hangs indefinitely when connecting to POST-only MCP +servers"). + +This client speaks plain request/response JSON-RPC over POST instead, which is exactly what +the bridge supports: ``initialize`` -> ``notifications/initialized`` -> ``tools/list`` / +``tools/call``. Responses are parsed as JSON, or, defensively, as a single ``data:`` frame if +the server ever replies with a one-shot ``text/event-stream`` body. +""" + +import json +from typing import Any + +import httpx +from databricks.sdk import WorkspaceClient +from mcp.types import CallToolResult, Tool + +# The bridge negotiates the 2025-03-26 Streamable HTTP spec. Advertise it explicitly rather +# than the SDK's LATEST_PROTOCOL_VERSION, which can be newer than the bridge supports. +_PROTOCOL_VERSION = "2025-03-26" +_DEFAULT_TIMEOUT_SECONDS = 30.0 + + +class PostOnlyMCPClient: + """Minimal request/response JSON-RPC MCP client over HTTP POST. + + Returns the same ``mcp.types`` objects as the Streamable HTTP path so it is a drop-in + for the external-connection branch of :class:`DatabricksMCPClient`. + """ + + def __init__( + self, + server_url: str, + workspace_client: WorkspaceClient, + timeout: float = _DEFAULT_TIMEOUT_SECONDS, + ) -> None: + self.server_url = server_url + self.workspace_client = workspace_client + self._session_id: str | None = None + self._next_id = 1 + self._client = httpx.AsyncClient(timeout=timeout, follow_redirects=False) + + async def __aenter__(self) -> "PostOnlyMCPClient": + await self._initialize() + return self + + async def __aexit__(self, *exc_info: object) -> None: + await self._client.aclose() + + def _headers(self) -> dict[str, str]: + # Resolve a fresh bearer token from the workspace client on each handshake, matching + # how DatabricksMCPClient authenticates elsewhere. + authorization = self.workspace_client.config.authenticate()["Authorization"] + headers = { + "Authorization": authorization, + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": _PROTOCOL_VERSION, + } + if self._session_id: + headers["Mcp-Session-Id"] = self._session_id + return headers + + @staticmethod + def _parse_body(response: httpx.Response) -> dict[str, Any]: + content_type = response.headers.get("Content-Type", "") + if "text/event-stream" in content_type: + for line in response.text.splitlines(): + if line.startswith("data:"): + payload = line[len("data:"):].strip() + if payload: + return json.loads(payload) + raise RuntimeError( + f"No JSON-RPC frame in event-stream response: {response.text[:300]!r}" + ) + return response.json() + + async def _send( + self, + method: str, + params: dict[str, Any] | None = None, + *, + is_notification: bool = False, + ) -> dict[str, Any] | None: + body: dict[str, Any] = {"jsonrpc": "2.0", "method": method} + if not is_notification: + body["id"] = self._next_id + self._next_id += 1 + if params is not None: + body["params"] = params + + response = await self._client.post(self.server_url, headers=self._headers(), json=body) + + if method == "initialize": + self._session_id = response.headers.get("Mcp-Session-Id") + + response.raise_for_status() + if is_notification: + return None + + frame = self._parse_body(response) + if "error" in frame: + raise RuntimeError(f"MCP {method} error: {frame['error']}") + return frame.get("result", {}) + + async def _initialize(self) -> None: + await self._send( + "initialize", + { + "protocolVersion": _PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "databricks-mcp", "version": "1"}, + }, + ) + # The MCP spec requires this notification before the server answers tools/* requests. + await self._send("notifications/initialized", is_notification=True) + + async def list_tools(self) -> list[Tool]: + result = await self._send("tools/list", {}) + return [Tool.model_validate(tool) for tool in (result or {}).get("tools", [])] + + async def call_tool( + self, tool_name: str, arguments: dict[str, Any] | None = None + ) -> CallToolResult: + result = await self._send( + "tools/call", {"name": tool_name, "arguments": arguments or {}} + ) + return CallToolResult.model_validate(result) diff --git a/databricks_mcp/src/databricks_mcp/mcp.py b/databricks_mcp/src/databricks_mcp/mcp.py index 08aeb867f..c617c5120 100644 --- a/databricks_mcp/src/databricks_mcp/mcp.py +++ b/databricks_mcp/src/databricks_mcp/mcp.py @@ -20,6 +20,7 @@ DatabricksVectorSearchIndex, ) +from databricks_mcp.external_client import PostOnlyMCPClient from databricks_mcp.oauth_provider import DatabricksOAuthClientProvider logger = logging.getLogger(__name__) @@ -180,6 +181,12 @@ 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.""" + # External/managed connections are fronted by a POST-only JSON-REST bridge that the + # Streamable HTTP client hangs against (its background GET SSE stream is rejected). + # Use the plain request/response POST client for those. See external_client.py. + if self._get_databricks_managed_mcp_url_type() == EXTERNAL_MCP: + async with PostOnlyMCPClient(self.server_url, self.client) as client: + return await client.list_tools() async with streamablehttp_client( url=self.server_url, auth=DatabricksOAuthClientProvider(self.client), @@ -194,6 +201,9 @@ async def _call_tools_async( arguments: dict[str, Any] | None = None, ) -> CallToolResult: """Call the tool with the given name and input.""" + if self._get_databricks_managed_mcp_url_type() == EXTERNAL_MCP: + async with PostOnlyMCPClient(self.server_url, self.client) as client: + return await client.call_tool(tool_name, arguments) async with streamablehttp_client( url=self.server_url, auth=DatabricksOAuthClientProvider(self.client), diff --git a/databricks_mcp/tests/unit_tests/test_mcp.py b/databricks_mcp/tests/unit_tests/test_mcp.py index 685265c46..dc81276b1 100644 --- a/databricks_mcp/tests/unit_tests/test_mcp.py +++ b/databricks_mcp/tests/unit_tests/test_mcp.py @@ -233,6 +233,52 @@ async def test_call_tools_async(self): mock_session.initialize.assert_called_once() mock_session.call_tool.assert_called_once_with("test_tool", {"arg": "value"}) + @pytest.mark.asyncio + async def test_get_tools_async_external_uses_post_only_client(self): + """External connections use the POST-only client, not the Streamable HTTP client.""" + mock_tools = [Tool(name="ext_tool", description="External tool", inputSchema={})] + mock_post_client = AsyncMock() + mock_post_client.list_tools = AsyncMock(return_value=mock_tools) + + with ( + patch("databricks_mcp.mcp.PostOnlyMCPClient") as mock_post_class, + patch("databricks_mcp.mcp.streamablehttp_client") as mock_streamable, + ): + mock_post_class.return_value.__aenter__.return_value = mock_post_client + + workspace_client = WorkspaceClient(host="https://test.com", token="test-token") + client = DatabricksMCPClient( + "https://test.com/api/2.0/mcp/external/my-connection", workspace_client + ) + tools = await client._get_tools_async() + + assert tools == mock_tools + mock_post_client.list_tools.assert_called_once() + mock_streamable.assert_not_called() + + @pytest.mark.asyncio + async def test_call_tools_async_external_uses_post_only_client(self): + """External connections call tools via the POST-only client.""" + mock_result = CallToolResult(content=[TextContent(type="text", text="ext result")]) + mock_post_client = AsyncMock() + mock_post_client.call_tool = AsyncMock(return_value=mock_result) + + with ( + patch("databricks_mcp.mcp.PostOnlyMCPClient") as mock_post_class, + patch("databricks_mcp.mcp.streamablehttp_client") as mock_streamable, + ): + mock_post_class.return_value.__aenter__.return_value = mock_post_client + + workspace_client = WorkspaceClient(host="https://test.com", token="test-token") + client = DatabricksMCPClient( + "https://test.com/api/2.0/mcp/external/my-connection", workspace_client + ) + result = await client._call_tools_async("ext_tool", {"arg": "value"}) + + assert result == mock_result + mock_post_client.call_tool.assert_called_once_with("ext_tool", {"arg": "value"}) + mock_streamable.assert_not_called() + def test_list_tools(self): """Test synchronous tool listing.""" mock_tools = [Tool(name="test_tool", description="Test tool", inputSchema={})]