diff --git a/databricks_mcp/src/databricks_mcp/oauth_provider.py b/databricks_mcp/src/databricks_mcp/oauth_provider.py index 084244f6b..086e1f8fe 100644 --- a/databricks_mcp/src/databricks_mcp/oauth_provider.py +++ b/databricks_mcp/src/databricks_mcp/oauth_provider.py @@ -1,11 +1,22 @@ +from collections.abc import AsyncGenerator + +import httpx from databricks.sdk import WorkspaceClient from mcp.client.auth import OAuthClientProvider, TokenStorage from mcp.shared.auth import OAuthToken +from typing_extensions import override TOKEN_EXPIRATION_SECONDS = 60 class DatabricksTokenStorage(TokenStorage): + """Read-only token storage that surfaces the active Databricks bearer token. + + Retained for callers that import `DatabricksTokenStorage` directly. Note: + `DatabricksOAuthClientProvider` no longer uses this class (it adds the + Authorization header directly — see the class docstring there for why). + """ + def __init__(self, workspace_client): self.workspace_client = workspace_client @@ -21,8 +32,48 @@ async def get_tokens(self) -> OAuthToken | None: class DatabricksOAuthClientProvider(OAuthClientProvider): """ - An OAuthClientProvider for Databricks. This class extends mcp.client.auth.OAuthClientProvider - and can be used with the `mcp.client.streamable_http` to authorize the MCP Server with Databricks. + An httpx auth provider for Databricks-fronted MCP servers. The credential + is fully managed by `WorkspaceClient.config.authenticate()` (which already + handles OBO / U2M / PAT / SP refresh), so this class is reduced to a thin + bearer-token stamp on each outgoing request. + + Implementation note: earlier versions of this class inherited the full + `mcp.client.auth.OAuthClientProvider.async_auth_flow` behavior, which + acquires `self.context.lock` and holds it **across** the `yield request` + that suspends until httpx receives the HTTP response: + + async with self.context.lock: + ... + response = yield request # lock still held while awaiting response + ... + + `mcp.client.streamable_http.streamablehttp_client` spawns two concurrent + tasks that share the same auth provider: + + * a long-lived GET (`handle_get_stream`) that subscribes to + server-pushed SSE events for the session; against many MCP servers + (e.g. UC-Connection-backed Atlassian / Jira / Confluence) this GET + stays open indefinitely waiting on the channel. + * per-JSON-RPC POSTs (`_handle_post_request`) for `tools/list`, + `tools/call`, … + + Both go through `auth.async_auth_flow`. Because the GET acquires the + auth lock first and never releases it (its response never returns), every + POST queues on that lock and eventually fails with:: + + mcp.shared.exceptions.McpError: Timed out while waiting for response + to ClientRequest. Waited 20.0 seconds. + + even though the server is healthy (a direct `httpx` call returns the + same response in <1s). + + The Databricks scenario doesn't need any of the upstream OAuth-dance + machinery the lock guards (no client registration, no PKCE, no + callback handlers, no refresh-token flow) — the credential is just + whatever `WorkspaceClient` already produced. So this override skips + the parent's `async_auth_flow` and stamps the Authorization header + directly, without ever taking the lock. That eliminates the deadlock + between GET and POST tasks on the streamable-HTTP transport. Usage: .. code-block:: python @@ -31,7 +82,6 @@ class DatabricksOAuthClientProvider(OAuthClientProvider): from mcp.client.streamable_http import streamablehttp_client from mcp.client.session import ClientSession - # Initialize the Databricks workspace client workspace_client = WorkspaceClient() async with streamablehttp_client( @@ -47,6 +97,8 @@ class DatabricksOAuthClientProvider(OAuthClientProvider): def __init__(self, workspace_client: WorkspaceClient): self.workspace_client = workspace_client + # Retained for backward compatibility — some callers reference this + # attribute directly. Not used by `async_auth_flow` anymore. self.databricks_token_storage = DatabricksTokenStorage(workspace_client) super().__init__( @@ -56,3 +108,21 @@ def __init__(self, workspace_client: WorkspaceClient): redirect_handler=None, callback_handler=None, ) + + @override + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + """Stamp the current Databricks bearer token and yield. No lock. + + Overrides `OAuthClientProvider.async_auth_flow`, which holds an + `async with self.context.lock:` across the request yield and would + otherwise serialize all concurrent requests through this provider + (see class docstring for the deadlock this caused on the + streamable-HTTP transport). + """ + headers = self.workspace_client.config.authenticate() + authorization = headers.get("Authorization") + if authorization: + request.headers["Authorization"] = authorization + yield request diff --git a/databricks_mcp/tests/unit_tests/test_oauth_provider.py b/databricks_mcp/tests/unit_tests/test_oauth_provider.py index 3a394e14f..39c218d37 100644 --- a/databricks_mcp/tests/unit_tests/test_oauth_provider.py +++ b/databricks_mcp/tests/unit_tests/test_oauth_provider.py @@ -1,5 +1,7 @@ +import asyncio from unittest.mock import MagicMock, patch +import httpx import pytest from databricks.sdk import WorkspaceClient @@ -18,6 +20,73 @@ async def test_oauth_provider(): assert oauth_token.token_type.lower() == "bearer" +@pytest.mark.asyncio +async def test_async_auth_flow_stamps_bearer_token(): + """`async_auth_flow` must add the active Databricks bearer token to the + request as `Authorization: Bearer `, and otherwise pass through. + """ + workspace_client = WorkspaceClient(host="https://test-databricks.com", token="test-token") + with patch.object(workspace_client.current_user, "me", return_value=MagicMock()): + provider = DatabricksOAuthClientProvider(workspace_client=workspace_client) + request = httpx.Request("POST", "https://test-databricks.com/api/2.0/mcp/some-path") + + flow = provider.async_auth_flow(request) + stamped = await flow.__anext__() + + assert stamped is request + assert stamped.headers["Authorization"] == "Bearer test-token" + + # The generator should complete after the single yield. + with pytest.raises(StopAsyncIteration): + await flow.__anext__() + + +@pytest.mark.asyncio +async def test_async_auth_flow_does_not_serialize_concurrent_requests(): + """Regression test for the streamable-HTTP deadlock. + + `mcp.client.streamable_http.streamablehttp_client` opens a long-lived GET + (`handle_get_stream`) for server-pushed SSE events AND fires per-RPC POSTs + through the same auth provider. The upstream `OAuthClientProvider`'s + `async_auth_flow` acquires `self.context.lock` and holds it *across* the + `yield request` (i.e. for the entire HTTP request lifetime). When the GET + is long-lived (Atlassian / Jira / Confluence — UC-Connection-backed MCPs + keep the SSE channel open indefinitely), the lock is never released and + every POST queues behind it until `client_session_timeout_seconds` fires: + + mcp.shared.exceptions.McpError: Timed out while waiting for response + to ClientRequest. Waited 20.0 seconds. + + Our override skips the parent's locked flow entirely (no OAuth dance is + needed — the credential is fully managed by WorkspaceClient). This test + pins that behavior by interleaving two `async_auth_flow` calls and + asserting neither blocks on the other. + """ + workspace_client = WorkspaceClient(host="https://test-databricks.com", token="test-token") + with patch.object(workspace_client.current_user, "me", return_value=MagicMock()): + provider = DatabricksOAuthClientProvider(workspace_client=workspace_client) + + # Simulate the streamable-HTTP transport's pattern: one long-lived + # request (the GET) holding its auth_flow generator open while a + # second short-lived request (the POST) needs to authenticate. + long_lived_req = httpx.Request("GET", "https://test-databricks.com/api/2.0/mcp/some-path") + short_lived_req = httpx.Request("POST", "https://test-databricks.com/api/2.0/mcp/some-path") + + long_flow = provider.async_auth_flow(long_lived_req) + # Yield once to "send" the GET; do NOT close the generator — this + # mirrors the live GET sitting on an open SSE channel. + await long_flow.__anext__() + + # The POST's auth_flow must complete promptly even while the GET's + # flow is still open. Wrap in a tight timeout: a regression + # (lock-bound flow) would hang here until the GET closes. + short_flow = provider.async_auth_flow(short_lived_req) + stamped = await asyncio.wait_for(short_flow.__anext__(), timeout=1.0) + + assert stamped is short_lived_req + assert stamped.headers["Authorization"] == "Bearer test-token" + + @pytest.mark.asyncio async def test_authenticate_raises_exception(): workspace_client = WorkspaceClient(host="https://test-databricks.com", token="test-token")