kaxil commented on code in PR #74198: URL: https://github.com/apache/airflow/pull/74198#discussion_r4340344290
########## providers/databricks/docs/toolsets/unity_mcp.rst: ########## @@ -0,0 +1,121 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/toolset:DatabricksUnityMCPToolset: + +Unity Gateway MCP Services +========================== + +Use :class:`~airflow.providers.databricks.toolsets.unity_mcp.DatabricksUnityMCPToolset` to give an +agent the tools of a Unity Gateway MCP Service: either one Databricks provides for workspace tools +and SaaS applications, such as ``system.ai.google_calendar``, or an external MCP server registered as +an MCP Service in Unity Catalog. The toolset works with +:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, ``@task.agent``, and the LangChain +bridge of the Common AI provider. + +The Dag names the service by its three-level Unity Catalog name, ``catalog.schema.service``, and the +:ref:`Databricks connection <howto/connection:databricks>` to use. The toolset builds the service URL, +``https://<workspace host>/ai-gateway/mcp-services/<catalog.schema.service>``, from the connection's +host, so neither the gateway URL nor a token appears in Dag code, and the connection's token is only +sent to that workspace, over HTTPS. A connection whose ``schema`` is ``http`` is rejected unless its +host is a loopback address. Requests go through the proxy in the connection's ``proxies`` extra, if +set. Service names may contain only ASCII letters, digits and ``_`` in each part; any other name is +rejected before a request is made. + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_unity_mcp.py + :language: python + :start-after: [START howto_toolset_databricks_unity_mcp] + :end-before: [END howto_toolset_databricks_unity_mcp] + +The service name and connection ID are templated when the toolset is passed to ``AgentOperator`` or +``@task.agent``. + +Caller identity and authentication +---------------------------------- + +MCP Services need a workspace enabled for Unity Catalog, in a region where Model Serving is +supported. + +The gateway runs every tool call as the identity of the connection's credentials. That identity needs +``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its parent catalog and schema +(``EXECUTE`` alone is not enough), and an assignment to the workspace. It needs no privilege on the +Unity Catalog connection behind the service. On the built-in ``system.ai`` services, account users +hold these privileges by default (see `MCP Services +<https://docs.databricks.com/aws/en/agents/mcp-tools/mcp-services>`__). The gateway exposes only the +tools selected for the service, and the service's policies apply. Grant the identity only the services +the agent should use. + +Built-in services that act on a user's own data, such as ``system.ai.google_calendar`` or +``system.ai.gmail``, need that identity to complete a one-time OAuth login first, for example by +opening the service in Catalog Explorer and clicking **Login**. + +The toolset sends the connection's token as a bearer token, so the connection must use one of these +authentication modes of the Databricks connection: + +* a personal access token (the identity is the token's user or service principal); +* service principal OAuth (``service_principal_oauth``); +* Azure AD: a service principal, a managed identity, or ``DefaultAzureCredential``; +* workload identity federation (Kubernetes, AWS IAM, or a supplied token provider). + +For service principal OAuth, the OAuth secret must allow the ``all-apis`` scope, which the connection +requests; a secret restricted to narrower scopes fails (see `OAuth for service principals +<https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m>`__). Username and password +authentication is not supported. A token is fetched from the connection for +every request, so OAuth and Azure AD tokens are refreshed during a long agent run and when the toolset +reconnects. Tokens minted for the toolset are masked in task logs. + +Errors and retries +------------------ + +When the MCP server answers a tool call with an error, such as invalid arguments, the error goes to the +model, which can correct the call and try again. + +When the gateway gives no answer of its own, a tool call that was already sent may or may not have +run, and repeating a tool that changes data could apply the change twice. The toolset never retries +it: the task fails with one of these exceptions from :mod:`airflow.providers.databricks.exceptions`, Review Comment: Through the LangChain bridge the task fails with `ToolCallError` instead: `airflow_toolset_to_langchain_tools` wraps each tool in `AirflowTool.call`, which re-raises any exception as `ToolCallError` carrying only its message. So the exception classes, `retry_after` and `http_status_code` only apply to Pydantic AI runs, which this section could say. ########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,607 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import logging +import re +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, mask_secret +from airflow.providers.databricks.exceptions import ( + DatabricksUnityMCPAccessDeniedError, + DatabricksUnityMCPError, + DatabricksUnityMCPThrottledError, + DatabricksUnityMCPTransportError, +) +from airflow.providers.databricks.hooks.databricks import DatabricksHook +from airflow.providers.databricks.version_compat import AIRFLOW_V_3_2_PLUS + +try: + # httpx2 and fastmcp come with pydantic-ai's mcp extra, which the common.ai extra installs. + import httpx2 + from fastmcp.client.transports import StreamableHttpTransport + from fastmcp.exceptions import McpError + from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset + + from airflow.providers.common.ai.utils.toolset_base import AirflowToolset +except ImportError as e: + raise AirflowOptionalProviderFeatureException( + "DatabricksUnityMCPToolset needs the 'common.ai' extra of the databricks provider: " + "pip install 'apache-airflow-providers-databricks[common.ai]'" + ) from e + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Sequence + + from pydantic_ai._run_context import RunContext + from pydantic_ai.toolsets.abstract import ToolsetTool + + from airflow.providers.common.ai.tools import AirflowTool + +log = logging.getLogger(__name__) + +GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services" + +# Restricting each part to these characters keeps the name from adding a path, query or +# fragment to the URL, so the request can only reach the service path on the connection's host. +_SERVICE_NAME_PART = r"[A-Za-z0-9_]+" +_SERVICE_NAME = re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}") + +# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the caller may not invoke +# the service, including when no service has that name. +_NOT_AUTHORIZED_CODE = -32007 + +# Errors the MCP client makes up when the gateway gave no JSON-RPC answer, such as an HTTP error +# with a plain body or a response that never arrived (see mcp.client.streamable_http). Every other +# MCP error is the server's own answer. Matched on code and message, the only things they carry. +_CLIENT_ERRORS = { + (-32603, "Server returned an error response"), + (-32601, "Not Found"), + (-32600, "Session terminated"), + (-32000, "Connection closed"), +} +# What fastmcp raises, in place of the TimeoutError it loses, when the MCP handshake outlasts +# init_timeout (see fastmcp.client.client). +_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session" +# What the MCP client raises when a request outlasts its read timeout (see +# mcp.shared.jsonrpc_dispatcher). A server may use the same code, so the message is matched too. +_REQUEST_TIMEOUT_CODE = -32001 +_REQUEST_TIMEOUT_MESSAGE = re.compile(r"Request '[^']*' timed out") +_CLIENT_ERROR_PREFIXES = ( + "Unexpected content type:", + "Failed to parse JSON response:", + "Failed to parse SSE message:", + "SSE stream ", + "server answered a request with 202 Accepted", + "Redirect to ", +) + + +def validate_service_name(service_name: str) -> None: + """ + Raise ``ValueError`` unless ``service_name`` is a three-level ``catalog.schema.service`` name. + + Each part may contain only ASCII letters, digits and underscores. + """ + if not isinstance(service_name, str) or not _SERVICE_NAME.fullmatch(service_name): + raise ValueError( + f"Invalid Unity Gateway MCP Service name {service_name!r}: expected " + "'catalog.schema.service', each part made of ASCII letters, digits or '_'." + ) + + +def _is_loopback(host: str) -> bool: + if host == "localhost": + return True + try: + return ipaddress.ip_address(host.strip("[]")).is_loopback + except ValueError: + return False + + +def _parse_retry_after(value: str | None) -> float | None: + if value is None: + return None + try: + seconds = float(value) + except ValueError: + # The HTTP-date form is not worth parsing: the caller only uses this as a hint. + return None + return seconds if seconds >= 0 else None + + +def _walk_exceptions(exc: BaseException) -> Iterator[BaseException]: + """Yield ``exc``, its causes and the exceptions inside any exception groups, each once.""" + seen: set[int] = set() + pending: list[BaseException] = [exc] + while pending: + current = pending.pop(0) + if id(current) in seen: + continue + seen.add(id(current)) + yield current + # Duck-typed: the ExceptionGroup builtin is Python 3.11+, and anyio uses the backport on 3.10. + if isinstance(grouped := getattr(current, "exceptions", None), (list, tuple)): + pending.extend(e for e in grouped if isinstance(e, BaseException)) + pending.extend(e for e in (current.__cause__, current.__context__) if e is not None) + + +def _describe(error: BaseException) -> str: + return str(error) or type(error).__name__ + + +def _find(exc: BaseException, types: type[BaseException] | tuple[type[BaseException], ...]) -> Any: + """Return the first exception of ``types`` in ``exc``, its causes and any exception groups.""" + return next((e for e in _walk_exceptions(exc) if isinstance(e, types)), None) + + +def _is_handshake_timeout(error: BaseException) -> bool: + return any( + isinstance(e, TimeoutError) + or (isinstance(e, RuntimeError) and str(e).endswith(_HANDSHAKE_TIMEOUT_MESSAGE)) + for e in _walk_exceptions(error) + ) + + +def _is_request_timeout(error: McpError | None) -> bool: + return ( + error is not None + and error.error.code == _REQUEST_TIMEOUT_CODE + and _REQUEST_TIMEOUT_MESSAGE.fullmatch(error.error.message) is not None + ) + + +def _is_client_error(error: McpError) -> bool: + code, message = error.error.code, error.error.message + return (code, message) in _CLIENT_ERRORS or message.startswith(_CLIENT_ERROR_PREFIXES) + + +class _TokenUnavailableError(Exception): + """Raised, for this module only, when no token could be fetched for a gateway request.""" + + +@dataclass(eq=False) +class _Operation: + """ + One toolset operation (connecting, listing tools or calling a tool) and the gateway errors seen while it ran. + + The MCP client does not say which HTTP response made a call fail, so the error statuses are + only attributed to the operation when no other operation of the toolset ran at the same time. + Request failures are recorded on every operation running then and kept whatever the overlap: + they only explain a call that saw no answer, and a request that got no response usually ends + the session for all of them. An overlapping call that fails for its own reason may still get + another call's failure, which changes its class and message but not that it fails. + """ + + overlapped: bool = False + errors: list[tuple[int, float | None]] = field(default_factory=list) + failures: list[Exception] = field(default_factory=list) + + @property + def http_error(self) -> tuple[int, float | None] | None: + """Return the status and ``Retry-After`` of the last error response it got, if it is known.""" + return self.errors[-1] if self.errors and not self.overlapped else None + + +class _DatabricksTokenAuth(httpx2.Auth): + """ + Authenticate each gateway request with a token from the Databricks connection. + + Asking the hook on every request, rather than once, lets OAuth tokens refresh during a long + agent run and on reconnection; the hook caches tokens until they are about to expire. + """ + + def __init__(self, hook: DatabricksHook) -> None: + self._hook = hook + self._token_lock = threading.Lock() + self._masked_token: str | None = None + + def fetch_token(self) -> str: + with self._token_lock: + token = self._hook._get_token(raise_error=False) + if not token: + raise ValueError( + f"Connection {self._hook.databricks_conn_id!r} has no token-based authentication " + "configured. This toolset sends a bearer token: use a personal access token, " + "service principal OAuth, Azure AD, or workload identity federation. Username and " + "password authentication is not supported." + ) + return token + + def fetch_and_mask_token(self) -> None: + """Fetch a token and mask it; only call it through ``AirflowToolset.run_blocking``.""" + self._mask(self.fetch_token()) + + def _mask(self, token: str) -> None: + # A personal access token is masked when the connection is fetched; mask minted + # OAuth tokens too, so they never reach task logs. + mask_secret(token) + self._masked_token = token + + async def async_auth_flow( + self, request: httpx2.Request + ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + try: + if AIRFLOW_V_3_2_PLUS: + # Not AirflowToolset.run_blocking: its lock is shared by every toolset, so a long + # SQL query in another toolset would hold up each request here. The hook is this + # toolset's own, so a lock of its own is enough. + token = await asyncio.to_thread(self.fetch_token) + else: + # Reading the connection's extras masks them, which messages the supervisor over + # a socket that has no lock of its own before Airflow 3.2. + token = await AirflowToolset.run_blocking(self.fetch_token) + except Exception as e: + raise _TokenUnavailableError(str(e)) from e + if token != self._masked_token: + # mask_secret messages the supervisor too. + await AirflowToolset.run_blocking(self._mask, token) + request.headers["Authorization"] = f"Bearer {token}" + yield request + + +class _RecordingStream(httpx2.AsyncByteStream): + """A response body that records a transport error raised while it is read.""" + + def __init__(self, stream: httpx2.AsyncByteStream, record: Callable[[Exception], None]) -> None: + self._stream = stream + self._record = record + + async def __aiter__(self) -> AsyncIterator[bytes]: + try: + async for chunk in self._stream: + yield chunk + except httpx2.TransportError as e: + self._record(e) + raise + + async def aclose(self) -> None: + await self._stream.aclose() + + +class _GatewayClient(httpx2.AsyncClient): + """ + Record, on the toolset's running operations, why each JSON-RPC request to the gateway failed. + + Once the MCP session is open, requests are sent from the MCP client's own tasks, and a call + waiting for an answer only learns that the connection closed. These records say why. + """ + + def __init__(self, *args: Any, operations: set[_Operation], **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._operations = operations + + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if not _is_jsonrpc_request(request): + return await super().send(request, **kwargs) + + def record(failure: Exception) -> None: + for operation in self._operations: + operation.failures.append(failure) + + try: + response = await super().send(request, **kwargs) + except (httpx2.TransportError, _TokenUnavailableError) as e: + record(e) + raise + if response.status_code >= 400: + retry_after = _parse_retry_after(response.headers.get("Retry-After")) + for operation in self._operations: + operation.errors.append((response.status_code, retry_after)) + if isinstance(response.stream, httpx2.AsyncByteStream): + response.stream = _RecordingStream(response.stream, record) + return response + + +def _is_jsonrpc_request(request: httpx2.Request) -> bool: + """ + Whether ``request`` sends a JSON-RPC request, which expects an answer. + + Not a notification or a GET for server events: the MCP client tolerates errors on those, so + they say nothing about why an operation failed. + """ + if request.method != "POST": + return False + try: + return "id" in json.loads(request.content) + except (ValueError, TypeError): + return False + + +class DatabricksUnityMCPToolset(AirflowToolset): + """ + Give an agent the tools of a Unity Gateway MCP Service, authenticated as the connection's identity. + + The service is named by its three-level Unity Catalog name, ``catalog.schema.service``, and + reached at ``https://<workspace host>/ai-gateway/mcp-services/<catalog.schema.service>``. The + workspace host and the credentials both come from the Databricks connection, so Dag code holds + neither a gateway URL nor a token, and the token is only ever sent to that workspace, over + HTTPS. The connection's ``proxies`` extra applies to gateway requests. + + The gateway runs every tool call as the identity of the connection's credentials (a user's + personal access token, a service principal, or an Azure AD / federated identity). That identity + needs ``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its catalog and + schema, and an assignment to the workspace. It sees only the tools selected for the service, and + the service's policies apply. + + Tokens are fetched from the connection for each request, so OAuth tokens refresh during a long + agent run and on reconnection. An error the server returns for a tool call reaches the model, + which can correct its arguments and try again. When the gateway gives no answer of its own, a + tool call may or may not have run, so it is never retried; it fails the task with + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPAccessDeniedError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`, or + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`, as do failures + while connecting and listing tools. + + .. code-block:: python + + from airflow.providers.common.ai.operators.agent import AgentOperator + from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + + AgentOperator( + task_id="ask_mcp_service", + prompt="Which tools do you have?", + llm_conn_id="pydanticai_default", + toolsets=[DatabricksUnityMCPToolset("main.default.my_mcp", databricks_conn_id="databricks")], + ) + + :param service_name: Three-level name of the MCP Service, ``catalog.schema.service``. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param databricks_conn_id: Databricks connection whose host and credentials are used. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param tool_prefix: Optional prefix prepended to tool names. + :param init_timeout: Seconds to wait for the MCP handshake when connecting. Unity Gateway forwards + a response only once the MCP server has finished sending it, so a server that keeps its + response stream open needs more than pydantic-ai's default of 5 seconds. + """ + + # Rendered, on a copy, by AgentOperator. + agent_template_fields: Sequence[str] = ("_databricks_conn_id", "_service_name") + + def __init__( + self, + service_name: str, + *, + databricks_conn_id: str = DatabricksHook.default_conn_name, + tool_prefix: str | None = None, + init_timeout: float = 30, + ) -> None: + if not init_timeout > 0: + raise ValueError(f"init_timeout must be a positive number of seconds, got {init_timeout!r}.") + self._databricks_conn_id = databricks_conn_id + self._service_name = service_name + self._tool_prefix = tool_prefix + self._init_timeout = init_timeout + self._server: Any = None + self._operations: set[_Operation] = set() + # A templated name is checked once it has been rendered, when the toolset connects. + if "{{" not in service_name: + validate_service_name(service_name) + + @property + def id(self) -> str: + return f"databricks-unity-mcp-{self._databricks_conn_id}-{self._service_name}" + + def get_service_url(self, hook: DatabricksHook) -> str: + """Return the gateway URL of the MCP Service on the connection's workspace.""" + validate_service_name(self._service_name) + if not hook.host: + raise ValueError(f"Connection {self._databricks_conn_id!r} has no workspace host.") + url = hook._endpoint_url(f"{GATEWAY_MCP_SERVICES_PATH}/{self._service_name}") + parts = urlsplit(url) + if parts.scheme != "https" and not _is_loopback(parts.hostname or ""): + raise ValueError( + f"Connection {self._databricks_conn_id!r} uses the {parts.scheme!r} scheme. Unity Gateway " + "requests carry a bearer token, so they are only sent over HTTPS." + ) + return url + + def _get_server(self) -> Any: + hook = DatabricksHook(self._databricks_conn_id, caller=type(self).__name__) + url = self.get_service_url(hook) + auth = _DatabricksTokenAuth(hook) + # Fail here, with a clear message, rather than inside the MCP client, which reports + # any failure as a generic connection error. + try: + auth.fetch_and_mask_token() + except ValueError: + raise + except Exception as e: + raise _TokenUnavailableError(str(e)) from e + proxy = (hook.proxies or {}).get(urlsplit(url).scheme) + transport = StreamableHttpTransport( + url, + headers=hook.user_agent_header, + auth=auth, + httpx_client_factory=_build_client_factory(self._operations, proxy), + ) + toolset = PydanticAIMCPToolset(transport, init_timeout=self._init_timeout) + return toolset.prefixed(self._tool_prefix) if self._tool_prefix else toolset + + async def _resolve_server(self) -> Any: + if self._server is None: + # Resolving the connection talks to the supervisor, so it takes the blocking-call lock. + self._server = await self.run_blocking(self._get_server) + return self._server + + def _translate_error( + self, error: Exception, operation: _Operation, *, during_tool_call: bool + ) -> Exception | None: + """Return the provider exception for a gateway failure, or ``None`` to re-raise ``error`` as is.""" + service = self._service_name + mcp_error = _find(error, McpError) + # When a request fails once the session is open, the call only sees an error the MCP + # client makes up; what the request ran into is in the operation's failures. + saw_no_answer = mcp_error is not None and _is_client_error(mcp_error) + failure = operation.failures[0] if saw_no_answer and operation.failures else None Review Comment: Nothing pins the `saw_no_answer` check here: changing this to `operation.failures[0] if operation.failures else None` keeps all 73 tests green. Without it, a server's `-32602` answer on a call that overlapped one whose request dropped comes back as `DatabricksUnityMCPTransportError` instead of reaching the model. A direct `_translate_error` test with a recorded `httpx2.ConnectError` and a non-client `McpError` would cover it. ########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,607 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import logging +import re +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, mask_secret +from airflow.providers.databricks.exceptions import ( + DatabricksUnityMCPAccessDeniedError, + DatabricksUnityMCPError, + DatabricksUnityMCPThrottledError, + DatabricksUnityMCPTransportError, +) +from airflow.providers.databricks.hooks.databricks import DatabricksHook +from airflow.providers.databricks.version_compat import AIRFLOW_V_3_2_PLUS + +try: + # httpx2 and fastmcp come with pydantic-ai's mcp extra, which the common.ai extra installs. + import httpx2 + from fastmcp.client.transports import StreamableHttpTransport + from fastmcp.exceptions import McpError + from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset + + from airflow.providers.common.ai.utils.toolset_base import AirflowToolset +except ImportError as e: + raise AirflowOptionalProviderFeatureException( + "DatabricksUnityMCPToolset needs the 'common.ai' extra of the databricks provider: " + "pip install 'apache-airflow-providers-databricks[common.ai]'" + ) from e + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Sequence + + from pydantic_ai._run_context import RunContext + from pydantic_ai.toolsets.abstract import ToolsetTool + + from airflow.providers.common.ai.tools import AirflowTool + +log = logging.getLogger(__name__) + +GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services" + +# Restricting each part to these characters keeps the name from adding a path, query or +# fragment to the URL, so the request can only reach the service path on the connection's host. +_SERVICE_NAME_PART = r"[A-Za-z0-9_]+" +_SERVICE_NAME = re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}") + +# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the caller may not invoke +# the service, including when no service has that name. +_NOT_AUTHORIZED_CODE = -32007 + +# Errors the MCP client makes up when the gateway gave no JSON-RPC answer, such as an HTTP error +# with a plain body or a response that never arrived (see mcp.client.streamable_http). Every other +# MCP error is the server's own answer. Matched on code and message, the only things they carry. +_CLIENT_ERRORS = { + (-32603, "Server returned an error response"), + (-32601, "Not Found"), + (-32600, "Session terminated"), + (-32000, "Connection closed"), +} +# What fastmcp raises, in place of the TimeoutError it loses, when the MCP handshake outlasts +# init_timeout (see fastmcp.client.client). +_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session" +# What the MCP client raises when a request outlasts its read timeout (see +# mcp.shared.jsonrpc_dispatcher). A server may use the same code, so the message is matched too. +_REQUEST_TIMEOUT_CODE = -32001 +_REQUEST_TIMEOUT_MESSAGE = re.compile(r"Request '[^']*' timed out") +_CLIENT_ERROR_PREFIXES = ( + "Unexpected content type:", + "Failed to parse JSON response:", + "Failed to parse SSE message:", + "SSE stream ", + "server answered a request with 202 Accepted", + "Redirect to ", +) + + +def validate_service_name(service_name: str) -> None: + """ + Raise ``ValueError`` unless ``service_name`` is a three-level ``catalog.schema.service`` name. + + Each part may contain only ASCII letters, digits and underscores. + """ + if not isinstance(service_name, str) or not _SERVICE_NAME.fullmatch(service_name): + raise ValueError( + f"Invalid Unity Gateway MCP Service name {service_name!r}: expected " + "'catalog.schema.service', each part made of ASCII letters, digits or '_'." + ) + + +def _is_loopback(host: str) -> bool: + if host == "localhost": + return True + try: + return ipaddress.ip_address(host.strip("[]")).is_loopback + except ValueError: + return False + + +def _parse_retry_after(value: str | None) -> float | None: + if value is None: + return None + try: + seconds = float(value) + except ValueError: + # The HTTP-date form is not worth parsing: the caller only uses this as a hint. + return None + return seconds if seconds >= 0 else None + + +def _walk_exceptions(exc: BaseException) -> Iterator[BaseException]: + """Yield ``exc``, its causes and the exceptions inside any exception groups, each once.""" + seen: set[int] = set() + pending: list[BaseException] = [exc] + while pending: + current = pending.pop(0) + if id(current) in seen: + continue + seen.add(id(current)) + yield current + # Duck-typed: the ExceptionGroup builtin is Python 3.11+, and anyio uses the backport on 3.10. + if isinstance(grouped := getattr(current, "exceptions", None), (list, tuple)): + pending.extend(e for e in grouped if isinstance(e, BaseException)) + pending.extend(e for e in (current.__cause__, current.__context__) if e is not None) + + +def _describe(error: BaseException) -> str: + return str(error) or type(error).__name__ + + +def _find(exc: BaseException, types: type[BaseException] | tuple[type[BaseException], ...]) -> Any: + """Return the first exception of ``types`` in ``exc``, its causes and any exception groups.""" + return next((e for e in _walk_exceptions(exc) if isinstance(e, types)), None) + + +def _is_handshake_timeout(error: BaseException) -> bool: + return any( + isinstance(e, TimeoutError) + or (isinstance(e, RuntimeError) and str(e).endswith(_HANDSHAKE_TIMEOUT_MESSAGE)) + for e in _walk_exceptions(error) + ) + + +def _is_request_timeout(error: McpError | None) -> bool: + return ( + error is not None + and error.error.code == _REQUEST_TIMEOUT_CODE + and _REQUEST_TIMEOUT_MESSAGE.fullmatch(error.error.message) is not None + ) + + +def _is_client_error(error: McpError) -> bool: + code, message = error.error.code, error.error.message + return (code, message) in _CLIENT_ERRORS or message.startswith(_CLIENT_ERROR_PREFIXES) + + +class _TokenUnavailableError(Exception): + """Raised, for this module only, when no token could be fetched for a gateway request.""" + + +@dataclass(eq=False) +class _Operation: + """ + One toolset operation (connecting, listing tools or calling a tool) and the gateway errors seen while it ran. + + The MCP client does not say which HTTP response made a call fail, so the error statuses are + only attributed to the operation when no other operation of the toolset ran at the same time. + Request failures are recorded on every operation running then and kept whatever the overlap: + they only explain a call that saw no answer, and a request that got no response usually ends + the session for all of them. An overlapping call that fails for its own reason may still get + another call's failure, which changes its class and message but not that it fails. + """ + + overlapped: bool = False + errors: list[tuple[int, float | None]] = field(default_factory=list) + failures: list[Exception] = field(default_factory=list) + + @property + def http_error(self) -> tuple[int, float | None] | None: + """Return the status and ``Retry-After`` of the last error response it got, if it is known.""" + return self.errors[-1] if self.errors and not self.overlapped else None + + +class _DatabricksTokenAuth(httpx2.Auth): + """ + Authenticate each gateway request with a token from the Databricks connection. + + Asking the hook on every request, rather than once, lets OAuth tokens refresh during a long + agent run and on reconnection; the hook caches tokens until they are about to expire. + """ + + def __init__(self, hook: DatabricksHook) -> None: + self._hook = hook + self._token_lock = threading.Lock() + self._masked_token: str | None = None + + def fetch_token(self) -> str: + with self._token_lock: + token = self._hook._get_token(raise_error=False) + if not token: + raise ValueError( + f"Connection {self._hook.databricks_conn_id!r} has no token-based authentication " + "configured. This toolset sends a bearer token: use a personal access token, " + "service principal OAuth, Azure AD, or workload identity federation. Username and " + "password authentication is not supported." + ) + return token + + def fetch_and_mask_token(self) -> None: + """Fetch a token and mask it; only call it through ``AirflowToolset.run_blocking``.""" + self._mask(self.fetch_token()) + + def _mask(self, token: str) -> None: + # A personal access token is masked when the connection is fetched; mask minted + # OAuth tokens too, so they never reach task logs. + mask_secret(token) + self._masked_token = token + + async def async_auth_flow( + self, request: httpx2.Request + ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + try: + if AIRFLOW_V_3_2_PLUS: + # Not AirflowToolset.run_blocking: its lock is shared by every toolset, so a long + # SQL query in another toolset would hold up each request here. The hook is this + # toolset's own, so a lock of its own is enough. + token = await asyncio.to_thread(self.fetch_token) + else: + # Reading the connection's extras masks them, which messages the supervisor over + # a socket that has no lock of its own before Airflow 3.2. + token = await AirflowToolset.run_blocking(self.fetch_token) + except Exception as e: + raise _TokenUnavailableError(str(e)) from e + if token != self._masked_token: + # mask_secret messages the supervisor too. + await AirflowToolset.run_blocking(self._mask, token) + request.headers["Authorization"] = f"Bearer {token}" + yield request + + +class _RecordingStream(httpx2.AsyncByteStream): + """A response body that records a transport error raised while it is read.""" + + def __init__(self, stream: httpx2.AsyncByteStream, record: Callable[[Exception], None]) -> None: + self._stream = stream + self._record = record + + async def __aiter__(self) -> AsyncIterator[bytes]: + try: + async for chunk in self._stream: + yield chunk + except httpx2.TransportError as e: + self._record(e) + raise + + async def aclose(self) -> None: + await self._stream.aclose() + + +class _GatewayClient(httpx2.AsyncClient): + """ + Record, on the toolset's running operations, why each JSON-RPC request to the gateway failed. + + Once the MCP session is open, requests are sent from the MCP client's own tasks, and a call + waiting for an answer only learns that the connection closed. These records say why. + """ + + def __init__(self, *args: Any, operations: set[_Operation], **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._operations = operations + + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if not _is_jsonrpc_request(request): + return await super().send(request, **kwargs) + + def record(failure: Exception) -> None: + for operation in self._operations: + operation.failures.append(failure) + + try: + response = await super().send(request, **kwargs) + except (httpx2.TransportError, _TokenUnavailableError) as e: + record(e) + raise + if response.status_code >= 400: + retry_after = _parse_retry_after(response.headers.get("Retry-After")) + for operation in self._operations: + operation.errors.append((response.status_code, retry_after)) + if isinstance(response.stream, httpx2.AsyncByteStream): + response.stream = _RecordingStream(response.stream, record) + return response + + +def _is_jsonrpc_request(request: httpx2.Request) -> bool: + """ + Whether ``request`` sends a JSON-RPC request, which expects an answer. + + Not a notification or a GET for server events: the MCP client tolerates errors on those, so + they say nothing about why an operation failed. + """ + if request.method != "POST": + return False + try: + return "id" in json.loads(request.content) + except (ValueError, TypeError): + return False + + +class DatabricksUnityMCPToolset(AirflowToolset): + """ + Give an agent the tools of a Unity Gateway MCP Service, authenticated as the connection's identity. + + The service is named by its three-level Unity Catalog name, ``catalog.schema.service``, and + reached at ``https://<workspace host>/ai-gateway/mcp-services/<catalog.schema.service>``. The + workspace host and the credentials both come from the Databricks connection, so Dag code holds + neither a gateway URL nor a token, and the token is only ever sent to that workspace, over + HTTPS. The connection's ``proxies`` extra applies to gateway requests. + + The gateway runs every tool call as the identity of the connection's credentials (a user's + personal access token, a service principal, or an Azure AD / federated identity). That identity + needs ``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its catalog and + schema, and an assignment to the workspace. It sees only the tools selected for the service, and + the service's policies apply. + + Tokens are fetched from the connection for each request, so OAuth tokens refresh during a long + agent run and on reconnection. An error the server returns for a tool call reaches the model, + which can correct its arguments and try again. When the gateway gives no answer of its own, a + tool call may or may not have run, so it is never retried; it fails the task with + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPAccessDeniedError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`, or + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`, as do failures + while connecting and listing tools. + + .. code-block:: python + + from airflow.providers.common.ai.operators.agent import AgentOperator + from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + + AgentOperator( + task_id="ask_mcp_service", + prompt="Which tools do you have?", + llm_conn_id="pydanticai_default", + toolsets=[DatabricksUnityMCPToolset("main.default.my_mcp", databricks_conn_id="databricks")], + ) + + :param service_name: Three-level name of the MCP Service, ``catalog.schema.service``. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param databricks_conn_id: Databricks connection whose host and credentials are used. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param tool_prefix: Optional prefix prepended to tool names. + :param init_timeout: Seconds to wait for the MCP handshake when connecting. Unity Gateway forwards + a response only once the MCP server has finished sending it, so a server that keeps its + response stream open needs more than pydantic-ai's default of 5 seconds. + """ + + # Rendered, on a copy, by AgentOperator. + agent_template_fields: Sequence[str] = ("_databricks_conn_id", "_service_name") + + def __init__( + self, + service_name: str, + *, + databricks_conn_id: str = DatabricksHook.default_conn_name, + tool_prefix: str | None = None, + init_timeout: float = 30, + ) -> None: + if not init_timeout > 0: + raise ValueError(f"init_timeout must be a positive number of seconds, got {init_timeout!r}.") + self._databricks_conn_id = databricks_conn_id + self._service_name = service_name + self._tool_prefix = tool_prefix + self._init_timeout = init_timeout + self._server: Any = None + self._operations: set[_Operation] = set() + # A templated name is checked once it has been rendered, when the toolset connects. + if "{{" not in service_name: + validate_service_name(service_name) + + @property + def id(self) -> str: + return f"databricks-unity-mcp-{self._databricks_conn_id}-{self._service_name}" + + def get_service_url(self, hook: DatabricksHook) -> str: + """Return the gateway URL of the MCP Service on the connection's workspace.""" + validate_service_name(self._service_name) + if not hook.host: + raise ValueError(f"Connection {self._databricks_conn_id!r} has no workspace host.") + url = hook._endpoint_url(f"{GATEWAY_MCP_SERVICES_PATH}/{self._service_name}") + parts = urlsplit(url) + if parts.scheme != "https" and not _is_loopback(parts.hostname or ""): + raise ValueError( + f"Connection {self._databricks_conn_id!r} uses the {parts.scheme!r} scheme. Unity Gateway " + "requests carry a bearer token, so they are only sent over HTTPS." + ) + return url + + def _get_server(self) -> Any: + hook = DatabricksHook(self._databricks_conn_id, caller=type(self).__name__) + url = self.get_service_url(hook) + auth = _DatabricksTokenAuth(hook) + # Fail here, with a clear message, rather than inside the MCP client, which reports + # any failure as a generic connection error. + try: + auth.fetch_and_mask_token() + except ValueError: + raise + except Exception as e: + raise _TokenUnavailableError(str(e)) from e + proxy = (hook.proxies or {}).get(urlsplit(url).scheme) + transport = StreamableHttpTransport( + url, + headers=hook.user_agent_header, + auth=auth, + httpx_client_factory=_build_client_factory(self._operations, proxy), + ) + toolset = PydanticAIMCPToolset(transport, init_timeout=self._init_timeout) + return toolset.prefixed(self._tool_prefix) if self._tool_prefix else toolset + + async def _resolve_server(self) -> Any: + if self._server is None: + # Resolving the connection talks to the supervisor, so it takes the blocking-call lock. + self._server = await self.run_blocking(self._get_server) + return self._server + + def _translate_error( + self, error: Exception, operation: _Operation, *, during_tool_call: bool + ) -> Exception | None: + """Return the provider exception for a gateway failure, or ``None`` to re-raise ``error`` as is.""" + service = self._service_name + mcp_error = _find(error, McpError) + # When a request fails once the session is open, the call only sees an error the MCP + # client makes up; what the request ran into is in the operation's failures. + saw_no_answer = mcp_error is not None and _is_client_error(mcp_error) + failure = operation.failures[0] if saw_no_answer and operation.failures else None + token_error = _find(error, _TokenUnavailableError) + if token_error is None and isinstance(failure, _TokenUnavailableError): + token_error = failure + if token_error is not None: + return DatabricksUnityMCPError( + f"Could not get a token from connection {self._databricks_conn_id!r} to call MCP Service " + f"{service!r}: {token_error}" + ) + status, retry_after = operation.http_error or (None, None) + if (mcp_error is not None and mcp_error.error.code == _NOT_AUTHORIZED_CODE) or status in (401, 403): + return DatabricksUnityMCPAccessDeniedError( + f"Unity Gateway denied access to MCP Service {service!r}" + f"{f' (HTTP {status})' if status else ''}. Either no service has that name, or the " + f"identity of connection {self._databricks_conn_id!r} lacks EXECUTE on the service or " + "USE CATALOG and USE SCHEMA on its catalog and schema, or its credentials are not valid.", + http_status_code=status, + ) + if status == 429: + hint = f" Retry after {retry_after:g} seconds." if retry_after is not None else "" + return DatabricksUnityMCPThrottledError( + f"Unity Gateway rate-limited calls to MCP Service {service!r} (HTTP 429).{hint}", + http_status_code=status, + retry_after=retry_after, + ) + ambiguous = ( + " The tool call may or may not have run, so it was not retried." if during_tool_call else "" + ) + # The MCP client reports its own timeouts as MCP errors, so they are told apart from the + # server's answers before those are. + if _is_request_timeout(mcp_error): + return DatabricksUnityMCPTransportError( + f"Timed out waiting for Unity Gateway to answer MCP Service {service!r}.{ambiguous}" + ) + transport_error = _find(error, httpx2.TransportError) + if transport_error is None and isinstance(failure, httpx2.TransportError): + transport_error = failure + if transport_error is not None: + return DatabricksUnityMCPTransportError( + f"Could not reach Unity Gateway for MCP Service {service!r}: " + f"{_describe(transport_error)}.{ambiguous}" + ) + if during_tool_call and _is_handshake_timeout(error): + return DatabricksUnityMCPTransportError( + f"Timed out waiting for Unity Gateway to answer MCP Service {service!r}.{ambiguous}" + ) + if _is_handshake_timeout(error): + return DatabricksUnityMCPTransportError( + f"Unity Gateway did not complete the MCP handshake for MCP Service {service!r} within " + f"{self._init_timeout:g} seconds. Raise init_timeout if the service is slow to answer." + ) + if mcp_error is not None: + if not _is_client_error(mcp_error): + # The server's own answer: for a tool call, it reaches the model as a retry. + if during_tool_call: + return None + return DatabricksUnityMCPError( + f"MCP Service {service!r} returned an error: {mcp_error.error.message} " + f"(code {mcp_error.error.code}).", + http_status_code=status, + ) + what = f"HTTP {status}" if status else f"an error without an answer ({mcp_error.error.message})" + return DatabricksUnityMCPError( + f"Unity Gateway returned {what} for MCP Service {service!r}.{ambiguous}", + http_status_code=status, + ) + return None + + @contextlib.asynccontextmanager + async def _gateway_operation(self, *, during_tool_call: bool = False) -> AsyncIterator[None]: + operation = _Operation(overlapped=bool(self._operations)) + for other in self._operations: + other.overlapped = True + self._operations.add(operation) + try: + yield + except Exception as e: + translated = self._translate_error(e, operation, during_tool_call=during_tool_call) + if translated is None: + raise + raise translated from e + finally: + self._operations.discard(operation) + + async def __aenter__(self) -> DatabricksUnityMCPToolset: + async with self._gateway_operation(): + await (await self._resolve_server()).__aenter__() + return self + + async def __aexit__(self, *args: Any) -> bool | None: + if self._server is None: + return None + if args[0] is None: Review Comment: In an agent run this branch fires even after a failed tool call: pydantic-ai's `CombinedToolset.__aexit__` closes its children with `AsyncExitStack.aclose()`, which passes `(None, None, None)`, and fastmcp's `_disconnect` then re-raises the session's error. With a `FunctionModel` agent against the stub gateway, `drop` and `drop body` fail the task with `httpx2.RemoteProtocolError` and a read timeout with a bare `httpx2.ReadTimeout`, so `DatabricksUnityMCPTransportError` is only left as `__context__`. Could `_gateway_operation` remember that it translated a failure and `__aexit__` log close errors whenever it did, with a test that goes through `CombinedToolset([toolset])` rather than `async with toolset`? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
