[databricks-mcp] Fix streamable-HTTP deadlock by not holding OAuth lock across requests

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

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 <token>`, 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")