diff --git a/backend/packages/harness/deerflow/mcp/oauth.py b/backend/packages/harness/deerflow/mcp/oauth.py index dad074e94..19b1d18d3 100644 --- a/backend/packages/harness/deerflow/mcp/oauth.py +++ b/backend/packages/harness/deerflow/mcp/oauth.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import logging +import threading from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any @@ -28,7 +29,17 @@ class OAuthTokenManager: def __init__(self, oauth_by_server: dict[str, McpOAuthConfig]): self._oauth_by_server = oauth_by_server self._tokens: dict[str, _OAuthToken] = {} - self._locks: dict[str, asyncio.Lock] = {name: asyncio.Lock() for name in oauth_by_server} + # A plain threading.Lock, not asyncio.Lock: the embedded/TUI sync tool-call + # path (DeerFlowClient.stream() -> LangGraph ToolNode._func -> a + # ThreadPoolExecutor -> deerflow.tools.sync.make_sync_tool_wrapper's + # per-call asyncio.run()) invokes get_authorization_header from a fresh + # event loop on a fresh OS thread for every concurrent tool call. An + # asyncio.Lock binds to whichever loop first contends on it; a second + # caller's release/wake-up crossing loops without call_soon_threadsafe + # either deadlocks silently or raises "bound to a different event loop". + # threading.Lock has no loop affinity, so it is safe to share across + # however many event loops/threads call into the same server's lock. + self._locks: dict[str, threading.Lock] = {name: threading.Lock() for name in oauth_by_server} @classmethod def from_extensions_config(cls, extensions_config: ExtensionsConfig) -> OAuthTokenManager: @@ -54,7 +65,40 @@ class OAuthTokenManager: return f"{token.token_type} {token.access_token}" lock = self._locks[server_name] - async with lock: + # Acquire the OS-level lock off-thread so a blocking wait never blocks this + # event loop, then release it synchronously (release() never blocks). This + # keeps the de-duplication behavior of the old `async with lock:` (only one + # concurrent caller per server actually fetches a token) while remaining + # safe when callers are on different event loops/threads. + # + # The acquisition itself runs as an explicit Task, shielded from this + # coroutine's own cancellation. A bare `await asyncio.to_thread(lock.acquire)` + # cannot be safely cancelled: once the executor thread has started running + # lock.acquire(), Python has no way to stop it, so a cancellation delivered + # at that await would still let the thread go on to acquire the lock later + # (whenever the current holder releases it) with this coroutine already + # gone and nobody left to call release() -- the lock would stay locked + # forever and every later call for this server would block permanently at + # this same line. Shielding the acquisition task means a cancelled caller + # can instead wait for that (unstoppable) acquisition to actually land and + # release the lock immediately, rather than leaking ownership of it. + acquire_task = asyncio.create_task(asyncio.to_thread(lock.acquire), name=f"oauth-lock-acquire:{server_name}") + try: + await asyncio.shield(acquire_task) + except asyncio.CancelledError: + # Keep waiting -- shielded on every retry -- until the acquisition + # actually finishes, even if this coroutine is cancelled again while + # cleaning up: the underlying thread cannot be interrupted, so this is + # the only way to learn when the lock becomes ours and release it + # right away instead of leaving it locked forever. + while not acquire_task.done(): + try: + await asyncio.shield(acquire_task) + except asyncio.CancelledError: + continue + lock.release() + raise + try: token = self._tokens.get(server_name) if token and not self._is_expiring(token, oauth): return f"{token.token_type} {token.access_token}" @@ -63,6 +107,8 @@ class OAuthTokenManager: self._tokens[server_name] = fresh logger.info(f"Refreshed OAuth access token for MCP server: {server_name}") return f"{fresh.token_type} {fresh.access_token}" + finally: + lock.release() @staticmethod def _is_expiring(token: _OAuthToken, oauth: McpOAuthConfig) -> bool: diff --git a/backend/tests/test_mcp_oauth.py b/backend/tests/test_mcp_oauth.py index 0ad233aaf..da9b5f6db 100644 --- a/backend/tests/test_mcp_oauth.py +++ b/backend/tests/test_mcp_oauth.py @@ -3,8 +3,11 @@ from __future__ import annotations import asyncio +import threading from typing import Any +import pytest + from deerflow.config.extensions_config import ExtensionsConfig from deerflow.mcp.oauth import OAuthTokenManager, build_oauth_tool_interceptor, get_initial_oauth_headers @@ -332,3 +335,229 @@ def test_oauth_refresh_token_rotation_persists_rotated_value(monkeypatch): assert second == "Bearer at-1" assert len(post_calls) == 2 assert post_calls[1]["data"]["refresh_token"] == "rt-rotated-1" + + +def test_get_authorization_header_concurrent_threads_no_deadlock(monkeypatch): + """Concurrent callers on different event loops/threads must not deadlock. + + The embedded/TUI sync tool-call path (``DeerFlowClient.stream()`` -> + LangGraph's ``ToolNode._func`` -> a ``ThreadPoolExecutor`` -> + ``deerflow.tools.sync.make_sync_tool_wrapper``'s per-call ``asyncio.run()``) + invokes ``get_authorization_header`` from a fresh event loop on a fresh OS + thread for every concurrent tool call. A per-server ``asyncio.Lock`` binds + to whichever loop first contends on it; when a caller on a *different* + loop later releases/wakes a waiter, it does so without + ``call_soon_threadsafe``, so the waiting loop's selector is never woken + and that caller hangs forever with no exception (a silent hang). A third + concurrent caller instead hits a synchronous ``RuntimeError: ... is bound + to a different event loop``. Both failure modes are reproducible with the + old ``asyncio.Lock``-per-server implementation. + + This test uses a bounded thread-join timeout so that a regression back to + the old behavior fails this test quickly instead of hanging the whole + suite. + """ + post_calls: list[dict[str, Any]] = [] + post_calls_guard = threading.Lock() + holder_in_critical_section = threading.Event() + + class _SlowMockAsyncClient: + def __init__(self, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def post(self, url: str, data: dict[str, Any]): + with post_calls_guard: + post_calls.append({"url": url, "data": data}) + # Signal that this call is inside the critical section (the lock + # is held) and stay there briefly so the other threads have time + # to reach their own acquire() and genuinely contend, rather than + # racing to also take an uncontended fast path. + holder_in_critical_section.set() + await asyncio.sleep(0.3) + return _MockResponse( + { + "access_token": "concurrent-token", + "token_type": "Bearer", + "expires_in": 3600, + } + ) + + monkeypatch.setattr("httpx.AsyncClient", _SlowMockAsyncClient) + + config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "secure-http": { + "enabled": True, + "type": "http", + "url": "https://api.example.com/mcp", + "oauth": { + "enabled": True, + "token_url": "https://auth.example.com/oauth/token", + "grant_type": "client_credentials", + "client_id": "client-id", + "client_secret": "client-secret", + }, + } + } + } + ) + + manager = OAuthTokenManager.from_extensions_config(config) + results: dict[str, Any] = {} + + def run_in_own_loop(name: str, wait_for_holder: bool) -> None: + if wait_for_holder: + # Only start once another thread is confirmed to be holding the + # lock, guaranteeing this call contends instead of racing for + # the uncontended fast path itself. + assert holder_in_critical_section.wait(timeout=5), "holder thread never entered critical section" + try: + results[name] = asyncio.run(manager.get_authorization_header("secure-http")) + except BaseException as exc: # noqa: BLE001 - captured to assert absence below + results[name] = exc + + threads = [ + threading.Thread(target=run_in_own_loop, args=("holder", False), name="holder", daemon=True), + threading.Thread(target=run_in_own_loop, args=("waiter-1", True), name="waiter-1", daemon=True), + threading.Thread(target=run_in_own_loop, args=("waiter-2", True), name="waiter-2", daemon=True), + ] + + for t in threads: + t.start() + + # Bounded timeout: under the old per-server asyncio.Lock, at least one of + # these threads would never return. Joining with a timeout keeps a + # regression from hanging the test suite forever; it fails fast instead. + for t in threads: + t.join(timeout=5) + + still_alive = [t.name for t in threads if t.is_alive()] + assert not still_alive, f"deadlock: thread(s) still blocked after bounded timeout: {still_alive}" + + for name, result in results.items(): + assert not isinstance(result, BaseException), f"{name} raised instead of completing: {result!r}" + assert result == "Bearer concurrent-token" + + # De-duplication must be preserved: three concurrent callers racing for + # the same (initially uncached) server must still only perform ONE real + # token fetch, not one per caller. + assert len(post_calls) == 1 + + +def test_get_authorization_header_cancelled_while_waiting_does_not_leak_lock(monkeypatch): + """A caller cancelled while waiting on the per-server lock must not leak it. + + ``get_authorization_header`` runs ``lock.acquire()`` on a real OS thread via + ``asyncio.to_thread`` so a blocking wait never blocks the event loop. Once that + thread has actually started running ``lock.acquire()``, Python cannot interrupt + it: cancelling the *caller* only stops the caller from continuing, it does not + stop the thread. If cancellation at that await let the thread go on to acquire + the lock unobserved (nobody left holding a reference that will call + ``release()`` for it), the lock would stay held forever and every subsequent + call for this server would block permanently at the same line -- the very + cross-thread deadlock this file's lock was introduced to fix, reintroduced via + a different path. + + This test holds the per-server lock (simulating another in-flight caller), + starts a second caller that has to wait for it, cancels that waiter while it + is genuinely blocked in its executor thread, releases the original holder, and + then asserts a third caller completes within a bounded timeout and performs + exactly one token fetch. Every potentially-hanging await is wrapped in a + bounded timeout so a regression fails this test quickly instead of hanging the + suite. + """ + post_calls: list[dict[str, Any]] = [] + + def _client_factory(*args, **kwargs): + return _MockAsyncClient( + payload={ + "access_token": "after-cancel-token", + "token_type": "Bearer", + "expires_in": 3600, + }, + post_calls=post_calls, + **kwargs, + ) + + monkeypatch.setattr("httpx.AsyncClient", _client_factory) + + config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "secure-http": { + "enabled": True, + "type": "http", + "url": "https://api.example.com/mcp", + "oauth": { + "enabled": True, + "token_url": "https://auth.example.com/oauth/token", + "grant_type": "client_credentials", + "client_id": "client-id", + "client_secret": "client-secret", + }, + } + } + } + ) + + manager = OAuthTokenManager.from_extensions_config(config) + lock = manager._locks["secure-http"] + + async def scenario() -> None: + # Simulate another in-flight caller already holding the per-server lock + # (uncontended, so this succeeds immediately without blocking). + lock.acquire() + try: + waiter = asyncio.create_task(manager.get_authorization_header("secure-http")) + + # Let the waiter's asyncio.to_thread(lock.acquire) actually get + # scheduled onto an executor thread and start genuinely blocking on + # the real lock before cancelling it -- otherwise the cancellation + # could land before the thread even starts, which would not exercise + # the bug. + await asyncio.sleep(0.2) + + waiter.cancel() + # The original holder finishes its own work and releases *before* we + # wait on the cancelled waiter: a correct fix must keep the lock's + # eventual acquisition shielded from this coroutine's cancellation and + # wait for it to actually land before releasing, so awaiting the + # cancelled waiter can legitimately block until the lock is free + # either way. + lock.release() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(waiter, timeout=5) + + # The crux of the regression: under the bug, the waiter's abandoned + # executor thread went on to acquire the lock with nobody left to + # release it, so this third call would block forever. Bound it so a + # regression fails fast instead of hanging the test itself. + third = await asyncio.wait_for(manager.get_authorization_header("secure-http"), timeout=5) + assert third == "Bearer after-cancel-token" + finally: + # Test-only safety net, independent of the assertions above: under + # the bug, the lock is left permanently locked with a background + # thread (from whichever caller's orphaned acquisition landed last) + # still parked on a *subsequent* acquire() that will now never + # return. asyncio.run()'s own teardown joins every thread the + # default executor ever created before it returns, so leaving that + # thread stuck would hang this test process at interpreter/loop + # shutdown even after the failure above is already reported. Forcing + # the lock open here lets any such thread finish so the process can + # exit; it is a no-op once the fix keeps the lock correctly balanced. + if lock.locked(): + lock.release() + + asyncio.run(scenario()) + + # Exactly one real token fetch: the cancelled waiter must never reach + # _fetch_token, so the third call is the only one that performs it. + assert len(post_calls) == 1