diff --git a/databricks_mcp/README.md b/databricks_mcp/README.md index 5b1c90732..961a15b9d 100644 --- a/databricks_mcp/README.md +++ b/databricks_mcp/README.md @@ -18,6 +18,28 @@ pip install git+https://git@github.com/databricks/databricks-ai-bridge.git#subdi - **OAuth Provider**: Enables authentication across Databricks Notebooks, Model Serving, and local environments using the Databricks CLI. +## URL elicitation + +An MCP server may ask the user to open a URL mid tool call (for example, to complete an auth flow). `DatabricksMCPClient` supports this, but **elicitation is off by default** — safe for agents embedded in headless or async runtimes (Model Serving, notebooks), where a stdin prompt or browser launch is meaningless and would block the event loop. + +To opt in, pass an `elicitation_callback`. Agents should supply their own callback that surfaces the URL through their own channel (UI, response stream, approval card): + +```python +from databricks_mcp import DatabricksMCPClient + +client = DatabricksMCPClient(server_url, elicitation_callback=my_callback) +``` + +For local or CLI use, a ready-made interactive handler is provided. It prints a security warning, asks for confirmation on stdin, and opens the approved URL in a browser (running the blocking prompt off the event loop so it is safe from async code). Only URL-mode elicitation is supported; form-mode requests are rejected: + +```python +from databricks_mcp import DatabricksMCPClient, interactive_url_elicitation_callback + +client = DatabricksMCPClient( + server_url, elicitation_callback=interactive_url_elicitation_callback +) +``` + --- ## Contribution Guide diff --git a/databricks_mcp/pyproject.toml b/databricks_mcp/pyproject.toml index 868e21d9c..7474e2962 100644 --- a/databricks_mcp/pyproject.toml +++ b/databricks_mcp/pyproject.toml @@ -9,7 +9,7 @@ readme = "README.md" license = { text="Apache-2.0" } requires-python = ">=3.10" dependencies = [ - "mcp>=1.13.0", + "mcp>=1.23.0", "databricks-sdk>=0.49.0", "databricks-ai-bridge>=0.4.2", "mlflow>=3.1" diff --git a/databricks_mcp/src/databricks_mcp/__init__.py b/databricks_mcp/src/databricks_mcp/__init__.py index faa4d4a1e..72bdbfe83 100644 --- a/databricks_mcp/src/databricks_mcp/__init__.py +++ b/databricks_mcp/src/databricks_mcp/__init__.py @@ -1,5 +1,10 @@ from databricks_mcp.connector import register_mcp_server_via_dcr -from databricks_mcp.mcp import DatabricksMCPClient +from databricks_mcp.mcp import DatabricksMCPClient, interactive_url_elicitation_callback from databricks_mcp.oauth_provider import DatabricksOAuthClientProvider -__all__ = ["DatabricksOAuthClientProvider", "DatabricksMCPClient", "register_mcp_server_via_dcr"] +__all__ = [ + "DatabricksOAuthClientProvider", + "DatabricksMCPClient", + "register_mcp_server_via_dcr", + "interactive_url_elicitation_callback", +] diff --git a/databricks_mcp/src/databricks_mcp/mcp.py b/databricks_mcp/src/databricks_mcp/mcp.py index 08aeb867f..8c03653f8 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 +import webbrowser from functools import wraps from typing import Any, Callable, List, Optional from urllib.parse import urlparse @@ -9,9 +10,17 @@ import requests from databricks.sdk import WorkspaceClient from databricks_ai_bridge.utils.annotations import experimental -from mcp.client.session import ClientSession +from mcp.client.session import ClientSession, ElicitationFnT from mcp.client.streamable_http import streamablehttp_client -from mcp.types import CallToolResult, Tool +from mcp.shared.context import RequestContext +from mcp.types import ( + INVALID_REQUEST, + CallToolResult, + ElicitRequestParams, + ElicitResult, + ErrorData, + Tool, +) from mlflow.models.resources import ( DatabricksFunction, DatabricksGenieSpace, @@ -142,6 +151,104 @@ def sync_wrapper(self, *args, **kwargs): return sync_wrapper +# Schemes we are willing to hand to the user's browser. URL elicitation lets the +# server choose the URL, so anything outside http(s) (e.g. file://, javascript:) is +# refused rather than opened. +_ALLOWED_URL_SCHEMES = {"http", "https"} + + +def _blocking_url_prompt(url: str, netloc: str, message: str, elicitation_id: Any) -> ElicitResult: + """Synchronously prompt the user to open a URL and open it on confirmation. + + Uses blocking ``input()`` / ``webbrowser.open()``, so callers must run this off the event + loop (see ``_handle_url_elicitation``). Returns ``accept`` once the URL is opened, + ``decline`` if the user refuses, and ``cancel`` when no decision can be obtained (e.g. no + interactive stdin). + """ + # The server controls the URL, so the confirmation prompt doubles as a security warning. + # It is written to stdout via input()'s prompt (rather than print) so the human sees it + # exactly when they are being asked to act. + prompt = ( + "\n" + "=" * 60 + "\n" + "SECURITY WARNING: the MCP server is requesting you open an external URL\n" + + "=" + * 60 + + "\n" + f"\n Domain: {netloc}\n" + f" Full URL: {url}\n" + f"\n Server's reason:\n {message}\n" + f"\n Elicitation ID: {elicitation_id}\n" + "-" * 60 + "\n" + "\nOpen this URL in your browser? (y/n): " + ) + try: + response = input(prompt).strip().lower() + except EOFError: + # No interactive stdin (e.g. running inside a service); cancel rather than guess. + return ElicitResult(action="cancel") + + if response in ("y", "yes"): + try: + webbrowser.open(url) + except Exception as e: + logger.warning("Failed to open browser (%s); open the URL manually: %s", e, url) + logger.info("Opened browser for URL elicitation; awaiting completion: %s", url) + return ElicitResult(action="accept") + if response in ("n", "no"): + return ElicitResult(action="decline") + + logger.info("Unrecognized response %r to URL elicitation; cancelling.", response) + return ElicitResult(action="cancel") + + +async def _handle_url_elicitation(params: ElicitRequestParams) -> ElicitResult: + """Handle a URL-mode elicitation request by confirming with the user and opening the URL. + + The server supplies the URL, so we validate the scheme and then run the interactive + confirmation off the event loop (``run_in_executor``) so blocking ``input()`` never + freezes the caller's loop when used from ``acall_tool``/``alist_tools``. Returns + ``decline`` if the scheme is disallowed or the URL is missing. + """ + url = getattr(params, "url", None) + if not url: + logger.error("URL elicitation request did not include a URL; declining.") + return ElicitResult(action="decline") + + parsed = urlparse(str(url)) + if parsed.scheme.lower() not in _ALLOWED_URL_SCHEMES: + logger.warning("Refusing URL elicitation with disallowed scheme %r: %s", parsed.scheme, url) + return ElicitResult(action="decline") + + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + _blocking_url_prompt, + str(url), + parsed.netloc, + params.message, + getattr(params, "elicitationId", None), + ) + + +async def interactive_url_elicitation_callback( + context: RequestContext["ClientSession", Any], + params: ElicitRequestParams, +) -> ElicitResult | ErrorData: + """Opt-in elicitation callback that supports URL-mode elicitation interactively. + + Pass this as ``elicitation_callback`` to :class:`DatabricksMCPClient` for local/CLI use: it + prompts the user on stdin and opens the approved URL in a browser. It is not the default, + because agents embedding this client should surface the URL through their own channel by + supplying their own callback. Form-mode elicitation is rejected so the server can fall back + rather than hang waiting for a response it will never get. + """ + if params.mode == "url": + return await _handle_url_elicitation(params) + return ErrorData( + code=INVALID_REQUEST, + message=f"Only URL elicitation is supported; got mode={params.mode!r}", + ) + + @experimental class DatabricksMCPClient: """ @@ -156,9 +263,20 @@ class DatabricksMCPClient: client (databricks.sdk.WorkspaceClient): The Databricks workspace client used for authentication and requests. """ - def __init__(self, server_url: str, workspace_client: Optional[WorkspaceClient] = None): + def __init__( + self, + server_url: str, + workspace_client: Optional[WorkspaceClient] = None, + elicitation_callback: Optional[ElicitationFnT] = None, + ): self.client = workspace_client or WorkspaceClient() self.server_url = server_url + # Elicitation is off by default: with no callback the ClientSession advertises no + # elicitation capability, which is the safe behavior when this client is embedded in a + # headless/async agent (Model Serving, notebooks). Callers opt in by injecting their own + # callback that surfaces the request through the agent's own channel, or by passing the + # exported interactive_url_elicitation_callback for local/CLI use. + self._elicitation_callback = elicitation_callback # Early detection: error if using non-OAuth auth with Databricks Apps if _is_databricks_apps_url(server_url) and not _is_oauth_auth(self.client): @@ -184,7 +302,9 @@ async def _get_tools_async(self) -> List[Tool]: url=self.server_url, auth=DatabricksOAuthClientProvider(self.client), ) as (read_stream, write_stream, _): - async with ClientSession(read_stream, write_stream) as session: + async with ClientSession( + read_stream, write_stream, elicitation_callback=self._elicitation_callback + ) as session: await session.initialize() return (await session.list_tools()).tools @@ -198,7 +318,9 @@ async def _call_tools_async( url=self.server_url, auth=DatabricksOAuthClientProvider(self.client), ) as (read_stream, write_stream, _): - async with ClientSession(read_stream, write_stream) as session: + async with ClientSession( + read_stream, write_stream, elicitation_callback=self._elicitation_callback + ) as session: await session.initialize() return await session.call_tool(tool_name, arguments) diff --git a/databricks_mcp/tests/unit_tests/test_mcp.py b/databricks_mcp/tests/unit_tests/test_mcp.py index 685265c46..fa99223d7 100644 --- a/databricks_mcp/tests/unit_tests/test_mcp.py +++ b/databricks_mcp/tests/unit_tests/test_mcp.py @@ -3,7 +3,17 @@ import pytest from databricks.sdk import WorkspaceClient -from mcp.types import CallToolResult, TextContent, Tool +from mcp.types import ( + INVALID_REQUEST, + CallToolResult, + ElicitRequestedSchema, + ElicitRequestFormParams, + ElicitRequestURLParams, + ElicitResult, + ErrorData, + TextContent, + Tool, +) from databricks_mcp.mcp import ( EXTERNAL_MCP, @@ -12,11 +22,17 @@ UC_FUNCTIONS_MCP, VECTOR_SEARCH_MCP, DatabricksMCPClient, + _handle_url_elicitation, _is_databricks_apps_url, _is_oauth_auth, + interactive_url_elicitation_callback, ) +def _url_params(url="https://example.com/auth", message="please authenticate", eid="elicit-1"): + return ElicitRequestURLParams(mode="url", message=message, url=url, elicitationId=eid) + + class TestDatabricksMCPClient: """Test cases for DatabricksMCPClient class.""" @@ -572,3 +588,183 @@ def test_allows_non_oauth_with_non_databricks_apps(self): mock_client.config.oauth_token.side_effect = ValueError("not available") client = DatabricksMCPClient("https://test.com/api/2.0/mcp/functions/a/b", mock_client) assert client.server_url == "https://test.com/api/2.0/mcp/functions/a/b" + + +class TestUrlElicitation: + """Test cases for URL-mode elicitation support.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("answer", ["y", "yes", "Y", " Yes "]) + async def test_handle_url_elicitation_accepts_and_opens_browser(self, answer): + params = _url_params(url="https://example.com/auth") + with ( + patch("databricks_mcp.mcp.input", return_value=answer), + patch("databricks_mcp.mcp.webbrowser.open") as mock_open, + ): + result = await _handle_url_elicitation(params) + + assert result.action == "accept" + mock_open.assert_called_once_with("https://example.com/auth") + + @pytest.mark.asyncio + @pytest.mark.parametrize("answer", ["n", "no", "N"]) + async def test_handle_url_elicitation_declines(self, answer): + params = _url_params() + with ( + patch("databricks_mcp.mcp.input", return_value=answer), + patch("databricks_mcp.mcp.webbrowser.open") as mock_open, + ): + result = await _handle_url_elicitation(params) + + assert result.action == "decline" + mock_open.assert_not_called() + + @pytest.mark.asyncio + async def test_handle_url_elicitation_cancels_on_invalid_response(self): + params = _url_params() + with ( + patch("databricks_mcp.mcp.input", return_value="maybe"), + patch("databricks_mcp.mcp.webbrowser.open") as mock_open, + ): + result = await _handle_url_elicitation(params) + + assert result.action == "cancel" + mock_open.assert_not_called() + + @pytest.mark.asyncio + async def test_handle_url_elicitation_cancels_without_interactive_stdin(self): + params = _url_params() + with ( + patch("databricks_mcp.mcp.input", side_effect=EOFError), + patch("databricks_mcp.mcp.webbrowser.open") as mock_open, + ): + result = await _handle_url_elicitation(params) + + assert result.action == "cancel" + mock_open.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("url", ["file:///etc/passwd", "ftp://host/x", "javascript:alert(1)"]) + async def test_handle_url_elicitation_rejects_disallowed_scheme(self, url): + params = _url_params(url=url) + # A disallowed scheme is refused before any prompt or browser launch. + with ( + patch("databricks_mcp.mcp.input") as mock_input, + patch("databricks_mcp.mcp.webbrowser.open") as mock_open, + ): + result = await _handle_url_elicitation(params) + + assert result.action == "decline" + mock_input.assert_not_called() + mock_open.assert_not_called() + + @pytest.mark.asyncio + async def test_callback_routes_url_mode(self): + params = _url_params() + with ( + patch("databricks_mcp.mcp.input", return_value="y"), + patch("databricks_mcp.mcp.webbrowser.open"), + ): + result = await interactive_url_elicitation_callback(MagicMock(), params) + + assert isinstance(result, ElicitResult) + assert result.action == "accept" + + @pytest.mark.asyncio + async def test_callback_rejects_form_mode(self): + params = ElicitRequestFormParams( + mode="form", message="fill this", requestedSchema=ElicitRequestedSchema(properties={}) + ) + result = await interactive_url_elicitation_callback(MagicMock(), params) + + assert isinstance(result, ErrorData) + assert result.code == INVALID_REQUEST + assert "URL elicitation" in result.message + + def test_elicitation_off_by_default(self): + """No callback passed -> elicitation is off (safe for headless/async embedding).""" + workspace_client = WorkspaceClient(host="https://test.com", token="test-token") + client = DatabricksMCPClient( + "https://test.com/api/2.0/mcp/functions/catalog/schema", workspace_client + ) + assert client._elicitation_callback is None + + def test_injected_elicitation_callback_is_stored(self): + async def custom_cb(context, params): + return ElicitResult(action="cancel") + + workspace_client = WorkspaceClient(host="https://test.com", token="test-token") + client = DatabricksMCPClient( + "https://test.com/api/2.0/mcp/functions/catalog/schema", + workspace_client, + elicitation_callback=custom_cb, + ) + assert client._elicitation_callback is custom_cb + + @pytest.mark.asyncio + @pytest.mark.parametrize("method", ["_call_tools_async", "_get_tools_async"]) + async def test_default_passes_no_elicitation_callback(self, method): + """With no callback, both session paths pass elicitation_callback=None (off).""" + mock_session = AsyncMock() + mock_session.initialize = AsyncMock() + mock_session.call_tool = AsyncMock( + return_value=CallToolResult(content=[TextContent(type="text", text="ok")]) + ) + mock_session.list_tools = AsyncMock(return_value=MagicMock(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 + + workspace_client = WorkspaceClient(host="https://test.com", token="test-token") + client = DatabricksMCPClient( + "https://test.com/api/2.0/mcp/functions/catalog/schema", workspace_client + ) + if method == "_call_tools_async": + await client._call_tools_async("test_tool", {"arg": "value"}) + else: + await client._get_tools_async() + + _, kwargs = mock_session_class.call_args + assert kwargs["elicitation_callback"] is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("method", ["_call_tools_async", "_get_tools_async"]) + async def test_injected_callback_passed_through(self, method): + """An injected callback is forwarded to the ClientSession on both paths.""" + + async def custom_cb(context, params): + return ElicitResult(action="cancel") + + mock_session = AsyncMock() + mock_session.initialize = AsyncMock() + mock_session.call_tool = AsyncMock( + return_value=CallToolResult(content=[TextContent(type="text", text="ok")]) + ) + mock_session.list_tools = AsyncMock(return_value=MagicMock(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 + + workspace_client = WorkspaceClient(host="https://test.com", token="test-token") + client = DatabricksMCPClient( + "https://test.com/api/2.0/mcp/functions/catalog/schema", + workspace_client, + elicitation_callback=custom_cb, + ) + if method == "_call_tools_async": + await client._call_tools_async("test_tool", {"arg": "value"}) + else: + await client._get_tools_async() + + _, kwargs = mock_session_class.call_args + assert kwargs["elicitation_callback"] is custom_cb