From bf1a20ecfe59a951a21b34e3948b3bce4ea529d9 Mon Sep 17 00:00:00 2001 From: Arthur Jenoudet Date: Tue, 30 Jun 2026 11:29:01 -0700 Subject: [PATCH 1/3] Support URL elicitation on MCP tool calls Wire a URL-only elicitation callback into DatabricksMCPClient's ClientSession for both list_tools and call_tool. The MCP SDK ships no default way to surface elicitation to users (its default callback just replies "Elicitation not supported" and advertises no capability), so define one internally: - _handle_url_elicitation: validates the server-supplied URL's scheme against an http(s) allowlist, surfaces a security warning + confirmation prompt, opens the browser on accept, and returns accept/decline/cancel. Cancels when there is no interactive stdin. - _url_only_elicitation_callback: routes URL-mode requests to the handler and rejects form-mode with INVALID_REQUEST so servers can fall back. The callback is module-private; the public DatabricksMCPClient constructor is unchanged. Co-authored-by: Isaac --- databricks_mcp/src/databricks_mcp/mcp.py | 95 +++++++++++- databricks_mcp/tests/unit_tests/test_mcp.py | 157 +++++++++++++++++++- 2 files changed, 248 insertions(+), 4 deletions(-) diff --git a/databricks_mcp/src/databricks_mcp/mcp.py b/databricks_mcp/src/databricks_mcp/mcp.py index 08aeb867f..46960ccf7 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 @@ -11,7 +12,15 @@ 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 +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,82 @@ 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"} + + +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 surface a security warning, validate the scheme, and + require explicit confirmation before launching a browser. Returns ``accept`` once the URL + has been opened, ``decline`` if the user refuses (or the scheme is disallowed), and + ``cancel`` when no decision can be obtained (e.g. no interactive stdin). + """ + url = getattr(params, "url", None) + if not url: + logger.error("URL elicitation request did not include a URL; cancelling.") + return ElicitResult(action="cancel") + + 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") + + # 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: {parsed.netloc}\n" + f" Full URL: {url}\n" + f"\n Server's reason:\n {params.message}\n" + f"\n Elicitation ID: {getattr(params, 'elicitationId', None)}\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(str(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 _url_only_elicitation_callback( + context: RequestContext["ClientSession", Any], + params: ElicitRequestParams, +) -> ElicitResult | ErrorData: + """Default elicitation callback that supports URL-mode elicitation only. + + Form-mode elicitation is intentionally unsupported: form input 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: """ @@ -184,7 +269,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=_url_only_elicitation_callback + ) as session: await session.initialize() return (await session.list_tools()).tools @@ -198,7 +285,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=_url_only_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..31d9701da 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, + _url_only_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,142 @@ 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 _url_only_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 _url_only_elicitation_callback(MagicMock(), params) + + assert isinstance(result, ErrorData) + assert result.code == INVALID_REQUEST + assert "URL elicitation" in result.message + + @pytest.mark.asyncio + async def test_call_tools_async_passes_elicitation_callback(self): + mock_result = CallToolResult(content=[TextContent(type="text", text="ok")]) + 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 + + workspace_client = WorkspaceClient(host="https://test.com", token="test-token") + client = DatabricksMCPClient( + "https://test.com/api/2.0/mcp/functions/catalog/schema", workspace_client + ) + await client._call_tools_async("test_tool", {"arg": "value"}) + + _, kwargs = mock_session_class.call_args + assert kwargs["elicitation_callback"] is _url_only_elicitation_callback + + @pytest.mark.asyncio + async def test_get_tools_async_passes_elicitation_callback(self): + mock_session = AsyncMock() + mock_session.initialize = AsyncMock() + 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 + ) + await client._get_tools_async() + + _, kwargs = mock_session_class.call_args + assert kwargs["elicitation_callback"] is _url_only_elicitation_callback From f121aa8d49af175f865fd6dd8870edc63455bada Mon Sep 17 00:00:00 2001 From: Arthur Jenoudet Date: Tue, 30 Jun 2026 11:41:56 -0700 Subject: [PATCH 2/3] fix: require mcp>=1.23.0 for URL elicitation URL elicitation (ElicitRequestURLParams / ElicitRequestFormParams and the mode-tagged ElicitRequestParams union) was introduced in mcp 1.23.0. The prior floor mcp>=1.13.0 let the lowest-direct CI resolution install a release without those symbols, breaking import of the new tests and the params.mode dispatch. Co-authored-by: Isaac --- databricks_mcp/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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" From 68312cf5c495e6ecf0e2c7156d032970291b1670 Mon Sep 17 00:00:00 2001 From: Arthur Jenoudet Date: Mon, 6 Jul 2026 13:41:03 -0700 Subject: [PATCH 3/3] Make URL elicitation embedding-safe: injectable callback, default off databricks-mcp is frequently embedded in custom agents (often headless in Model Serving, or async). Hardcoding an interactive input()/webbrowser elicitation callback was wrong for that: it silently cancels in headless contexts and blocks the event loop from acall_tool. - Add an optional `elicitation_callback` to DatabricksMCPClient; default None means the ClientSession advertises no elicitation capability (safe default). Agents inject their own callback to surface the URL through their channel. - Provide `interactive_url_elicitation_callback` (now exported) for local/CLI use, and run its blocking prompt/browser off the event loop via run_in_executor so it is safe from async code. - Rename the internal callback accordingly; document the model in the README. Co-authored-by: Isaac --- databricks_mcp/README.md | 22 +++++ databricks_mcp/src/databricks_mcp/__init__.py | 9 +- databricks_mcp/src/databricks_mcp/mcp.py | 89 +++++++++++++------ databricks_mcp/tests/unit_tests/test_mcp.py | 65 +++++++++++--- 4 files changed, 143 insertions(+), 42 deletions(-) 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/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 46960ccf7..8c03653f8 100644 --- a/databricks_mcp/src/databricks_mcp/mcp.py +++ b/databricks_mcp/src/databricks_mcp/mcp.py @@ -10,7 +10,7 @@ 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.shared.context import RequestContext from mcp.types import ( @@ -157,24 +157,14 @@ def sync_wrapper(self, *args, **kwargs): _ALLOWED_URL_SCHEMES = {"http", "https"} -async def _handle_url_elicitation(params: ElicitRequestParams) -> ElicitResult: - """Handle a URL-mode elicitation request by confirming with the user and opening the URL. +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. - The server supplies the URL, so we surface a security warning, validate the scheme, and - require explicit confirmation before launching a browser. Returns ``accept`` once the URL - has been opened, ``decline`` if the user refuses (or the scheme is disallowed), and - ``cancel`` when no decision can be obtained (e.g. no interactive stdin). + 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). """ - url = getattr(params, "url", None) - if not url: - logger.error("URL elicitation request did not include a URL; cancelling.") - return ElicitResult(action="cancel") - - 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") - # 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. @@ -184,10 +174,10 @@ async def _handle_url_elicitation(params: ElicitRequestParams) -> ElicitResult: + "=" * 60 + "\n" - f"\n Domain: {parsed.netloc}\n" + f"\n Domain: {netloc}\n" f" Full URL: {url}\n" - f"\n Server's reason:\n {params.message}\n" - f"\n Elicitation ID: {getattr(params, 'elicitationId', None)}\n" + "-" * 60 + "\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: @@ -198,7 +188,7 @@ async def _handle_url_elicitation(params: ElicitRequestParams) -> ElicitResult: if response in ("y", "yes"): try: - webbrowser.open(str(url)) + 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) @@ -210,14 +200,46 @@ async def _handle_url_elicitation(params: ElicitRequestParams) -> ElicitResult: return ElicitResult(action="cancel") -async def _url_only_elicitation_callback( +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: - """Default elicitation callback that supports URL-mode elicitation only. + """Opt-in elicitation callback that supports URL-mode elicitation interactively. - Form-mode elicitation is intentionally unsupported: form input is rejected so the - server can fall back rather than hang waiting for a response it will never get. + 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) @@ -241,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): @@ -270,7 +303,7 @@ async def _get_tools_async(self) -> List[Tool]: auth=DatabricksOAuthClientProvider(self.client), ) as (read_stream, write_stream, _): async with ClientSession( - read_stream, write_stream, elicitation_callback=_url_only_elicitation_callback + read_stream, write_stream, elicitation_callback=self._elicitation_callback ) as session: await session.initialize() return (await session.list_tools()).tools @@ -286,7 +319,7 @@ async def _call_tools_async( auth=DatabricksOAuthClientProvider(self.client), ) as (read_stream, write_stream, _): async with ClientSession( - read_stream, write_stream, elicitation_callback=_url_only_elicitation_callback + 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 31d9701da..fa99223d7 100644 --- a/databricks_mcp/tests/unit_tests/test_mcp.py +++ b/databricks_mcp/tests/unit_tests/test_mcp.py @@ -25,7 +25,7 @@ _handle_url_elicitation, _is_databricks_apps_url, _is_oauth_auth, - _url_only_elicitation_callback, + interactive_url_elicitation_callback, ) @@ -665,7 +665,7 @@ async def test_callback_routes_url_mode(self): patch("databricks_mcp.mcp.input", return_value="y"), patch("databricks_mcp.mcp.webbrowser.open"), ): - result = await _url_only_elicitation_callback(MagicMock(), params) + result = await interactive_url_elicitation_callback(MagicMock(), params) assert isinstance(result, ElicitResult) assert result.action == "accept" @@ -675,18 +675,42 @@ async def test_callback_rejects_form_mode(self): params = ElicitRequestFormParams( mode="form", message="fill this", requestedSchema=ElicitRequestedSchema(properties={}) ) - result = await _url_only_elicitation_callback(MagicMock(), params) + 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 - async def test_call_tools_async_passes_elicitation_callback(self): - mock_result = CallToolResult(content=[TextContent(type="text", text="ok")]) + @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=mock_result) + 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, @@ -700,15 +724,27 @@ async def test_call_tools_async_passes_elicitation_callback(self): client = DatabricksMCPClient( "https://test.com/api/2.0/mcp/functions/catalog/schema", workspace_client ) - await client._call_tools_async("test_tool", {"arg": "value"}) + 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 _url_only_elicitation_callback + assert kwargs["elicitation_callback"] is None @pytest.mark.asyncio - async def test_get_tools_async_passes_elicitation_callback(self): + @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 ( @@ -721,9 +757,14 @@ async def test_get_tools_async_passes_elicitation_callback(self): workspace_client = WorkspaceClient(host="https://test.com", token="test-token") client = DatabricksMCPClient( - "https://test.com/api/2.0/mcp/functions/catalog/schema", workspace_client + "https://test.com/api/2.0/mcp/functions/catalog/schema", + workspace_client, + elicitation_callback=custom_cb, ) - await client._get_tools_async() + 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 _url_only_elicitation_callback + assert kwargs["elicitation_callback"] is custom_cb