fix(mcp): single-flight SDK-driven OAuth refresh across concurrent tool calls (#15040)

backend/onyx/server/features/mcp/oauth.py

@@ -8,7 +8,7 @@ Route orchestration lives in `oauth_flow.py`.
 
 import json
 import time
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from typing import Any, TypedDict
 from urllib.parse import parse_qs, urlencode, urljoin, urlparse
 from uuid import uuid4
@@ -90,6 +90,10 @@ REQUESTED_SCOPE: str | None = None
 OAUTH_HTTP_TIMEOUT_SECONDS = 30.0
 
 
+def _refresh_lock_name(connection_config_id: int) -> str:
+    return f"mcp_token_refresh:{get_current_tenant_id()}:{connection_config_id}"
+
+
 class MCPRefreshLogContext(TypedDict):
     mcp_server_id: int | None
     mcp_server_name: str
@@ -211,7 +215,7 @@ def refresh_mcp_oauth_token_if_expired(
     lock still can't be acquired *and* the stored token is expired does it return
     None; the caller then falls back to its existing header.
     """
-    lock_name = f"mcp_token_refresh:{get_current_tenant_id()}:{connection_config_id}"
+    lock_name = _refresh_lock_name(connection_config_id)
     try:
         with cache_shared_lock(
             lock_name,
@@ -692,6 +696,63 @@ class OnyxOAuthClientProvider(OAuthClientProvider):
         # invalid_grant discard against wiping a concurrently stored grant.
         self.redeemed_refresh_token: str | None = None
 
+    async def async_auth_flow(
+        self, request: httpx.Request
+    ) -> AsyncGenerator[httpx.Request, httpx.Response]:
+        """Single-flight the refresh, then delegate to the SDK flow.
+
+        The SDK serializes refreshes with a per-instance context lock, but every
+        tool call builds its own provider — so concurrent calls each refresh on
+        their own and providers that rotate refresh tokens (Snowflake, GitHub)
+        reject the second redemption with invalid_grant, which can discard the
+        grant entirely. Refresh here under the cross-process cache lock first;
+        the SDK flow below then sees a valid token and skips its own refresh.
+        """
+        storage = self.context.storage
+        connection_config_id = (
+            storage.connection_config_id
+            if isinstance(storage, OnyxTokenStorage)
+            else None
+        )
+        if connection_config_id is not None:
+            await self._initialize()
+            if not self.context.is_token_valid() and self.context.can_refresh_token():
+                try:
+                    with cache_shared_lock(
+                        _refresh_lock_name(connection_config_id),
+                        max_time_lock_held_s=_REFRESH_LOCK_LEASE_S,
+                        wait_for_lock_s=_REFRESH_LOCK_WAIT_S,
+                        logger=logger,
+                    ):
+                        # The lock holder may have already refreshed; re-read
+                        # the stored grant before redeeming the refresh token.
+                        await self._initialize()
+                        if (
+                            not self.context.is_token_valid()
+                            and self.context.can_refresh_token()
+                        ):
+                            refresh_request = await self._refresh_token()
+                            refresh_response = yield refresh_request
+                            await self._handle_refresh_response(refresh_response)
+                except CacheLockAcquisitionError:
+                    # Contention outlasted the wait; the SDK flow below retries
+                    # its own refresh, no worse than the unlocked status quo.
+                    logger.info(
+                        "mcp_token_refresh.lock_contended config_id=%s",
+                        connection_config_id,
+                    )
+
+        sdk_flow = super().async_auth_flow(request)
+        try:
+            pending = await sdk_flow.__anext__()
+            while True:
+                try:
+                    pending = await sdk_flow.asend((yield pending))
+                except StopAsyncIteration:
+                    return
+        finally:
+            await sdk_flow.aclose()
+
     def build_resumable_authorization_request(
         self,
     ) -> MCPPendingOAuthAuthorization:

backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py

@@ -1,12 +1,16 @@
-"""Regression coverage for SSE-transport MCP OAuth refresh: SSE can't use the
-SDK's httpx.Auth refresh (open stream), so refresh_mcp_oauth_token_if_expired
-drives OAuthClientProvider's refresh step directly instead.
+"""Regression coverage for MCP OAuth refresh: SSE can't use the SDK's
+httpx.Auth refresh (open stream), so refresh_mcp_oauth_token_if_expired drives
+OAuthClientProvider's refresh step directly instead; STREAMABLE_HTTP calls
+refresh inside OnyxOAuthClientProvider.async_auth_flow, single-flighted on the
+same cache lock.
 
 Exercises the real OAuthClientProvider/OAuthContext/OnyxTokenStorage code —
 client-auth-method branching, endpoint resolution, persistence — mocking only
 the DB layer and the outbound network call.
 """
 
+import asyncio
+import threading
 import time
 from collections.abc import Iterator
 from contextlib import contextmanager
@@ -22,13 +26,16 @@ from onyx.cache.interface import CacheLockAcquisitionError
 from onyx.db.enums import MCPOAuthProviderMode, MCPTransport
 from onyx.db.models import MCPServer as DbMCPServer
 from onyx.server.features.mcp.models import MCPOAuthKeys
-from onyx.server.features.mcp.oauth import refresh_mcp_oauth_token_if_expired
+from onyx.server.features.mcp.oauth import (
+    make_oauth_provider,
+    refresh_mcp_oauth_token_if_expired,
+)
 
 _TOKEN_ENDPOINT = "https://gitlab.example.com/oauth/token"
 _REDIRECT_URI = "https://onyx.example.com/mcp/oauth/callback"
 
 
-def _server_stub() -> DbMCPServer:
+def _server_stub(transport: MCPTransport = MCPTransport.SSE) -> DbMCPServer:
     # AUTO_DISCOVERY: token endpoint comes from persisted METADATA, matching
     # real DCR-registered servers (KNOWN_PROVIDER never negotiates
     # client_secret_basic — see _build_oauth_admin_config_data).
@@ -38,7 +45,7 @@ def _server_stub() -> DbMCPServer:
             id=1,
             name="gitlab",
             server_url="https://mcp.gitlab.example.com",
-            transport=MCPTransport.SSE,
+            transport=transport,
             oauth_provider_mode=MCPOAuthProviderMode.AUTO_DISCOVERY,
             oauth_authorization_endpoint=None,
             oauth_token_endpoint=None,
@@ -610,3 +617,145 @@ def test_lock_contention_returns_persisted_header_or_none(
 
     assert header == expected_header
     assert "sent_request" not in captured
+
+
+def _expired_grant_config() -> dict[str, Any]:
+    return {
+        "headers": {"Authorization": "Bearer OLD"},
+        MCPOAuthKeys.TOKENS.value: {
+            "access_token": "OLD",
+            "token_type": "Bearer",
+            "refresh_token": "REFRESH_1",
+        },
+        MCPOAuthKeys.TOKEN_EXPIRES_AT.value: time.time() - 60,
+        MCPOAuthKeys.CLIENT_INFO.value: {
+            "client_id": "cid",
+            "redirect_uris": [_REDIRECT_URI],
+        },
+        MCPOAuthKeys.METADATA.value: {
+            "issuer": "https://gitlab.example.com",
+            "authorization_endpoint": "https://gitlab.example.com/oauth/authorize",
+            "token_endpoint": _TOKEN_ENDPOINT,
+        },
+    }
+
+
+def test_sdk_auth_flow_single_flights_refresh_across_concurrent_calls(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """STREAMABLE_HTTP path: each tool call builds its own provider, so the
+    SDK's per-context lock can't serialize their refreshes. Two concurrent
+    async_auth_flows must share one refresh under cache_shared_lock — the
+    loser re-reads the winner's grant instead of re-redeeming a single-use
+    (rotating) refresh token, which providers like Snowflake reject with
+    invalid_grant. Each flow runs on its own thread/event loop, matching
+    production (run_async_sync_no_cancel per tool call)."""
+    config_data = _expired_grant_config()
+    thread_lock = threading.Lock()
+
+    @contextmanager
+    def _serializing_lock(*_args: Any, **_kwargs: Any) -> Iterator[None]:
+        acquired = thread_lock.acquire(timeout=10)
+        if not acquired:
+            raise CacheLockAcquisitionError("held by a concurrent refresher")
+        try:
+            yield
+        finally:
+            thread_lock.release()
+
+    _install_mocks(monkeypatch, config_data, response=None)
+    monkeypatch.setattr(mcp_oauth, "cache_shared_lock", _serializing_lock)
+
+    refresh_posts: list[httpx.Request] = []
+    authorized_requests: list[str | None] = []
+
+    def handle_request(request: httpx.Request) -> httpx.Response:
+        body = parse_qs(request.content.decode()) if request.content else {}
+        if body.get("grant_type") == ["refresh_token"]:
+            refresh_posts.append(request)
+            # A refresh token is single-use: a second redemption must fail.
+            if body["refresh_token"] != ["REFRESH_1"] or len(refresh_posts) > 1:
+                return httpx.Response(
+                    400, json={"error": "invalid_grant"}, request=request
+                )
+            time.sleep(0.05)  # hold the window so the loser arrives mid-refresh
+            return httpx.Response(
+                200,
+                json={
+                    "access_token": "NEW",
+                    "token_type": "Bearer",
+                    "expires_in": 7200,
+                    "refresh_token": "REFRESH_2",
+                },
+                request=request,
+            )
+        authorized_requests.append(request.headers.get("Authorization"))
+        return httpx.Response(200, json={"ok": True}, request=request)
+
+    errors: list[BaseException | None] = [None, None]
+
+    def run_flow(index: int) -> None:
+        provider = make_oauth_provider(
+            _server_stub(MCPTransport.STREAMABLE_HTTP), 42, None
+        )
+
+        async def run() -> None:
+            async with httpx.AsyncClient(
+                auth=provider,
+                transport=httpx.MockTransport(handle_request),
+            ) as client:
+                await client.post("https://mcp.gitlab.example.com/mcp")
+
+        try:
+            asyncio.run(run())
+        except BaseException as error:  # noqa: BLE001
+            errors[index] = error
+
+    threads = [threading.Thread(target=run_flow, args=(index,)) for index in range(2)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert errors == [None, None]
+    assert len(refresh_posts) == 1
+    assert authorized_requests and all(
+        header == "Bearer NEW" for header in authorized_requests
+    )
+    assert config_data[MCPOAuthKeys.TOKENS.value]["refresh_token"] == "REFRESH_2"
+    assert config_data["headers"]["Authorization"] == "Bearer NEW"
+
+
+def test_sdk_auth_flow_rechecks_grant_inside_the_lock(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """After a prior refresh lands, a later call's flow must not refresh again:
+    inside the lock it re-reads the stored grant, sees a valid token, and only
+    sends the original request."""
+    config_data = _expired_grant_config()
+    _install_mocks(monkeypatch, config_data, response=None)
+
+    refresh_posts: list[httpx.Request] = []
+
+    def handle_request(request: httpx.Request) -> httpx.Response:
+        body = parse_qs(request.content.decode()) if request.content else {}
+        if body.get("grant_type") == ["refresh_token"]:
+            refresh_posts.append(request)
+            return _token_response()
+        return httpx.Response(200, json={"ok": True}, request=request)
+
+    provider = make_oauth_provider(_server_stub(MCPTransport.STREAMABLE_HTTP), 42, None)
+
+    async def run() -> None:
+        async with httpx.AsyncClient(
+            auth=provider,
+            transport=httpx.MockTransport(handle_request),
+        ) as client:
+            await client.post("https://mcp.gitlab.example.com/mcp")
+
+    asyncio.run(run())
+    assert len(refresh_posts) == 1
+    assert config_data["headers"]["Authorization"] == "Bearer NEW"
+
+    asyncio.run(run())
+    assert len(refresh_posts) == 1