From 4376c2b35b83ad8e053a07760a8b474ce7b6a305 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Sat, 13 Jun 2026 22:43:55 +0000 Subject: [PATCH] [databricks-mcp] Stop holding the OAuth lock across HTTP requests in DatabricksOAuthClientProvider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `DatabricksOAuthClientProvider` extends `mcp.client.auth.OAuthClientProvider`, whose `async_auth_flow` does: async with self.context.lock: ... response = yield request # lock held while httpx awaits the response ... i.e. the lock is held for the entire HTTP request lifetime, not just the moments around the yield. `mcp.client.streamable_http.streamablehttp_client` spawns two concurrent tasks against the same auth provider: * a long-lived GET (`handle_get_stream`) that opens an SSE channel for server-pushed events, * per-JSON-RPC POSTs (`_handle_post_request`) for `tools/list`, `tools/call`, ... For most Databricks-fronted MCP servers reached through the AI Gateway UC-Connection proxy (`/api/2.0/mcp/external/`) — Atlassian, Jira, Confluence, ... — the GET keeps the SSE channel open indefinitely waiting on server events. The auth lock is acquired first by that GET task and is never released, so every POST queues behind it and eventually surfaces as: mcp.shared.exceptions.McpError: Timed out while waiting for response to ClientRequest. Waited 20.0 seconds. even though the upstream MCP server is healthy (a bare `httpx` call — no auth provider, no concurrent GET — returns the same response in <1s). The Databricks scenario doesn't need any of the OAuth-dance machinery the lock guards: there is no client registration, no PKCE, no callback handlers, no refresh-token flow — the credential is whatever `WorkspaceClient.config.authenticate()` already produced. This change overrides `async_auth_flow` in `DatabricksOAuthClientProvider` to stamp the Authorization header and yield, without acquiring the parent's lock. The override keeps the same constructor signature and the `databricks_token_storage` attribute, so existing callers and the existing token-storage tests continue to work unchanged. Added regression test `test_async_auth_flow_does_not_serialize_concurrent_requests` which interleaves two `async_auth_flow` generators (mirroring the GET + POST pattern of the streamable-HTTP transport) and asserts neither blocks on the other. Co-authored-by: Isaac --- .../src/databricks_mcp/oauth_provider.py | 76 ++++++++++++++++++- .../tests/unit_tests/test_oauth_provider.py | 69 +++++++++++++++++ 2 files changed, 142 insertions(+), 3 deletions(-) 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")