deer-flow/backend/tests/test_mcp_oauth.py
Daoyuan Li 44990ff194
fix(mcp): use threading.Lock for OAuth token refresh to avoid cross-thread deadlock (#4240)
* fix(mcp): use threading.Lock for OAuth token refresh to avoid cross-thread deadlock

OAuthTokenManager created one asyncio.Lock per server for the process
lifetime. The embedded/TUI sync tool-call path (DeerFlowClient.stream()
-> LangGraph's ToolNode._func -> a ThreadPoolExecutor ->
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. asyncio.Lock binds to whichever loop
first contends on it; when a caller on a different loop later releases
or 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 third concurrent caller instead raises a
synchronous RuntimeError ("bound to a different event loop"). Either
way, two concurrent OAuth-protected tool calls (including the very
first cold-start token fetch) can freeze the entire agent turn.
Gateway's async path (ToolNode._afunc) is unaffected.

Replace the asyncio.Lock with a plain threading.Lock, acquired via
asyncio.to_thread so the blocking wait never blocks the event loop,
and released synchronously in a finally block. This keeps the
single-fetch de-duplication the lock provided while making it safe
across however many event loops/threads call into the same server's
lock.

Adds a regression test that runs three threads, each with its own
event loop, calling get_authorization_header concurrently for the same
server, and asserts (with a bounded join timeout so a regression fails
fast instead of hanging the suite) that none hang or raise, and that
only one real token fetch happens.

* fix(mcp): make OAuth lock acquisition cancellation-safe

get_authorization_header acquired the per-server threading.Lock via a
bare `await asyncio.to_thread(lock.acquire)`, with the try/finally that
guarantees release only starting after that await returned. Once the
executor thread had actually started running lock.acquire(), cancelling
the awaiting caller only stopped the caller -- Python cannot interrupt a
running OS thread. CancelledError was still delivered to the caller
immediately, but the thread kept blocking until the current holder
released, then silently acquired the lock with nobody left to call
release() for it. The lock stayed locked forever and every later OAuth
token refresh for that server blocked permanently at the same line --
the exact cross-thread deadlock this lock was introduced to prevent,
reintroduced via a different path under cancellation (e.g. a caller
wrapped in asyncio.wait_for/asyncio.timeout, or task-group cancellation).

Run the acquisition as an explicit asyncio.create_task, awaited via
asyncio.shield() so cancelling the caller no longer cancels the
underlying acquisition task. If the caller is cancelled, keep
(re-)waiting on the still-shielded acquisition task -- tolerating
further cancellation during this cleanup by simply retrying -- until it
actually finishes, release the lock immediately, and only then
re-raise. This guarantees the lock is released regardless of when or
how many times the caller is cancelled: before the acquisition is even
scheduled, while queued, or after it has already been silently granted.

Adds a regression test that holds the per-server lock, starts a second
caller that has to wait for it, cancels that caller while it is
genuinely blocked in its executor thread, releases the original holder,
and asserts a third caller completes within a bounded asyncio.wait_for
and still performs exactly one token fetch. Every potentially-hanging
await is bounded so a regression fails the test quickly instead of
hanging the suite.
2026-07-22 19:58:43 +08:00

564 lines
21 KiB
Python

"""Tests for MCP OAuth support."""
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
class _MockResponse:
def __init__(self, payload: dict[str, Any]):
self._payload = payload
def raise_for_status(self) -> None:
return None
def json(self) -> dict[str, Any]:
return self._payload
class _MockAsyncClient:
def __init__(self, payload: dict[str, Any], post_calls: list[dict[str, Any]], **kwargs):
self._payload = payload
self._post_calls = post_calls
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]):
self._post_calls.append({"url": url, "data": data})
return _MockResponse(self._payload)
def test_oauth_token_manager_fetches_and_caches_token(monkeypatch):
post_calls: list[dict[str, Any]] = []
def _client_factory(*args, **kwargs):
return _MockAsyncClient(
payload={
"access_token": "token-123",
"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)
first = asyncio.run(manager.get_authorization_header("secure-http"))
second = asyncio.run(manager.get_authorization_header("secure-http"))
assert first == "Bearer token-123"
assert second == "Bearer token-123"
assert len(post_calls) == 1
assert post_calls[0]["url"] == "https://auth.example.com/oauth/token"
assert post_calls[0]["data"]["grant_type"] == "client_credentials"
def test_build_oauth_interceptor_injects_authorization_header(monkeypatch):
post_calls: list[dict[str, Any]] = []
def _client_factory(*args, **kwargs):
return _MockAsyncClient(
payload={
"access_token": "token-abc",
"token_type": "Bearer",
"expires_in": 3600,
},
post_calls=post_calls,
**kwargs,
)
monkeypatch.setattr("httpx.AsyncClient", _client_factory)
config = ExtensionsConfig.model_validate(
{
"mcpServers": {
"secure-sse": {
"enabled": True,
"type": "sse",
"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",
},
}
}
}
)
interceptor = build_oauth_tool_interceptor(config)
assert interceptor is not None
class _Request:
def __init__(self):
self.server_name = "secure-sse"
self.headers = {"X-Test": "1"}
def override(self, **kwargs):
updated = _Request()
updated.server_name = self.server_name
updated.headers = kwargs.get("headers")
return updated
captured: dict[str, Any] = {}
async def _handler(request):
captured["headers"] = request.headers
return "ok"
result = asyncio.run(interceptor(_Request(), _handler))
assert result == "ok"
assert captured["headers"]["Authorization"] == "Bearer token-abc"
assert captured["headers"]["X-Test"] == "1"
def test_get_initial_oauth_headers(monkeypatch):
post_calls: list[dict[str, Any]] = []
def _client_factory(*args, **kwargs):
return _MockAsyncClient(
payload={
"access_token": "token-initial",
"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",
},
},
"no-oauth": {
"enabled": True,
"type": "http",
"url": "https://example.com/mcp",
},
}
}
)
headers = asyncio.run(get_initial_oauth_headers(config))
assert headers == {"secure-http": "Bearer token-initial"}
assert len(post_calls) == 1
def test_get_initial_oauth_headers_one_failing_server_does_not_drop_others(monkeypatch):
"""A single OAuth server whose token endpoint fails must not drop headers
(and therefore tools) from healthy servers."""
class _FailingClient:
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]):
raise RuntimeError("token endpoint unreachable")
class _OkClient:
def __init__(self, post_calls: list[dict[str, Any]], **kwargs):
self._post_calls = post_calls
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]):
self._post_calls.append({"url": url, "data": data})
return _MockResponse(
payload={
"access_token": "token-ok",
"token_type": "Bearer",
"expires_in": 3600,
}
)
ok_post_calls: list[dict[str, Any]] = []
def _client_factory(**kwargs):
# The first call is for the failing server, second for the healthy one,
# because OAuthTokenManager iterates _oauth_by_server in dict order
# ('broken-http' < 'secure-http').
if not hasattr(_client_factory, "_count"):
_client_factory._count = 0 # type: ignore[attr-defined]
_client_factory._count += 1 # type: ignore[attr-defined]
if _client_factory._count == 1: # type: ignore[attr-defined]
return _FailingClient()
return _OkClient(post_calls=ok_post_calls)
monkeypatch.setattr("httpx.AsyncClient", _client_factory)
config = ExtensionsConfig.model_validate(
{
"mcpServers": {
"broken-http": {
"enabled": True,
"type": "http",
"url": "https://broken.example.com/mcp",
"oauth": {
"enabled": True,
"token_url": "https://auth.broken.example.com/oauth/token",
"grant_type": "client_credentials",
"client_id": "client-id",
"client_secret": "client-secret",
},
},
"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-2",
"client_secret": "client-secret-2",
},
},
}
}
)
headers = asyncio.run(get_initial_oauth_headers(config))
# The healthy server's header must still be present.
assert headers == {"secure-http": "Bearer token-ok"}
assert len(ok_post_calls) == 1
def test_oauth_refresh_token_rotation_persists_rotated_value(monkeypatch):
"""When a provider rotates the refresh_token, _fetch_token must capture
the new value so the next refresh uses it instead of the stale original."""
post_calls: list[dict[str, Any]] = []
def _client_factory(*args, **kwargs):
return _MockAsyncClient(
payload={
"access_token": "at-1",
"token_type": "Bearer",
"expires_in": 3600,
"refresh_token": "rt-rotated-1",
},
post_calls=post_calls,
**kwargs,
)
monkeypatch.setattr("httpx.AsyncClient", _client_factory)
config = ExtensionsConfig.model_validate(
{
"mcpServers": {
"rotating-srv": {
"enabled": True,
"type": "http",
"url": "https://api.example.com/mcp",
"oauth": {
"enabled": True,
"token_url": "https://auth.example.com/oauth/token",
"grant_type": "refresh_token",
"refresh_token": "rt-original-seed",
},
}
}
}
)
manager = OAuthTokenManager.from_extensions_config(config)
# Force the _is_expiring check to always return True so we hit _fetch_token.
monkeypatch.setattr(OAuthTokenManager, "_is_expiring", lambda self, token, oauth: True)
first = asyncio.run(manager.get_authorization_header("rotating-srv"))
assert first == "Bearer at-1"
assert len(post_calls) == 1
# First call posted the original seed token.
assert post_calls[0]["data"]["refresh_token"] == "rt-original-seed"
# On the second call, the rotated refresh_token from the first response
# must be used.
second = asyncio.run(manager.get_authorization_header("rotating-srv"))
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