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

[sunishsheth2009]

Summary

DatabricksOAuthClientProvider inherits from mcp.client.auth.OAuthClientProvider, whose async_auth_flow holds self.context.lock across the yield request that suspends until the HTTP response arrives. Combined with mcp.client.streamable_http.streamablehttp_client's pattern of running a long-lived GET (handle_get_stream) alongside per-JSON-RPC POSTs against the same auth provider, this deadlocks every POST whenever the GET's SSE channel stays open — which is the common case for UC-Connection-backed MCP servers (Atlassian, Jira, Confluence).

The user-visible symptom is:

mcp.shared.exceptions.McpError: Timed out while waiting for response to ClientRequest. Waited 20.0 seconds.

…on the first tools/list call after __aenter__, even though the same request issued via a bare httpx call (no auth provider, no concurrent GET) returns in <1 second.

Root cause

# mcp/client/auth/oauth2.py — OAuthClientProvider.async_auth_flow
async def async_auth_flow(self, request):
    async with self.context.lock:                # ← acquired here
        ...
        response = yield request                 # ← held across the yield
        ...

# mcp/client/streamable_http.py — streamablehttp_client
tg.start_soon(transport.handle_get_stream, client, read_stream_writer)   # long-lived GET
tg.start_soon(..., read_stream_writer)                                   # per-RPC POSTs

Both handle_get_stream and _handle_post_request route through the same DatabricksOAuthClientProvider instance. The GET enters async_auth_flow, acquires the lock, yields the request, and httpx then awaits the response — which never arrives for an open SSE channel. The lock stays held; every POST blocks on async with self.context.lock: until the session-level anyio.fail_after(client_session_timeout_seconds) fires.

Fix

Override async_auth_flow in DatabricksOAuthClientProvider to stamp the Authorization header from WorkspaceClient.config.authenticate() and yield, without acquiring the parent's lock. The full OAuth dance the lock guards (client registration, PKCE, callback handlers, refresh-token flow) is unused in the Databricks scenario — the credential is fully managed by WorkspaceClient.

Constructor signature and the databricks_token_storage attribute are preserved; existing callers and the existing token-storage tests are unchanged.

Test

Added test_async_auth_flow_does_not_serialize_concurrent_requests: opens an async_auth_flow generator for a "GET" without closing it (mirrors the open SSE channel), then opens a second async_auth_flow generator for a "POST" and asserts it advances within a 1-second timeout. With the old behavior the POST hangs on the lock; with the override it completes immediately.

Also added test_async_auth_flow_stamps_bearer_token to pin the basic behavior (Authorization header is added, generator completes after one yield).

Repro context

Discovered while running an OpenAI Agents SDK agent against an external UC-Connection MCP (Atlassian) through /api/2.0/mcp/external/<connection>. Every first-turn tools/list from agents.mcp.MCPServerStreamableHttp (which inherits the same MCP client transport) failed with McpError: Timed out … 20.0 seconds. A side-by-side direct httpx POST to the same URL (with the same OBO bearer token, no OAuthClientProvider configured) succeeded in <1s and returned all tools, confirming the latency was introduced by the auth provider rather than the proxy or upstream server.

Test plan

• databricks_mcp/tests/unit_tests/test_oauth_provider.py — test_async_auth_flow_stamps_bearer_token
• databricks_mcp/tests/unit_tests/test_oauth_provider.py — test_async_auth_flow_does_not_serialize_concurrent_requests
• Existing test_oauth_provider and test_authenticate_raises_exception (token-storage path) still pass — that code path is unchanged.

This pull request and its description were written by Isaac.