From b963282f5ed89bd90bec2874db8ab1fd89e048ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=BB=84=E4=BA=91=E9=BE=99?= <76432572+nankingjing@users.noreply.github.com> Date: Sun, 12 Jul 2026 11:59:36 +0800 Subject: [PATCH] fix(mcp): per-server fail-soft OAuth priming + persist rotated refresh_token (#4084) get_initial_oauth_headers() now wraps each server's token fetch in try/except so one broken OAuth server does not abort the entire multi-server MCP tool load (which returned [] from the outer except, dropping every server's tools). _fetch_token() now captures a rotated refresh_token from the token response and updates the in-memory McpOAuthConfig so subsequent refreshes use the latest value instead of the stale original, which fails with invalid_grant on providers that rotate refresh tokens (Auth0, Okta, Google, etc.). Adds regression tests: - One failing OAuth server still yields the healthy server's headers - Rotated refresh_token is posted on the next refresh attempt Co-authored-by: Claude --- .../packages/harness/deerflow/mcp/oauth.py | 22 ++- backend/tests/test_mcp_oauth.py | 143 ++++++++++++++++++ 2 files changed, 164 insertions(+), 1 deletion(-) diff --git a/backend/packages/harness/deerflow/mcp/oauth.py b/backend/packages/harness/deerflow/mcp/oauth.py index b4cc1c1f1..dad074e94 100644 --- a/backend/packages/harness/deerflow/mcp/oauth.py +++ b/backend/packages/harness/deerflow/mcp/oauth.py @@ -107,6 +107,16 @@ class OAuthTokenManager: if not access_token: raise ValueError(f"OAuth token response missing '{oauth.token_field}'") + # Persist a rotated refresh_token so subsequent refreshes use the latest + # value. This is an in-process update only — it is intentionally NOT + # written back to extensions_config.json. Providers that rotate refresh + # tokens (Auth0, Okta, Google, etc.) return a new refresh_token on each + # refresh; discarding it makes the next refresh fail with invalid_grant. + if oauth.grant_type == "refresh_token": + rotated = payload.get("refresh_token") + if isinstance(rotated, str) and rotated: + oauth.refresh_token = rotated + token_type = str(payload.get(oauth.token_type_field, oauth.default_token_type) or oauth.default_token_type) expires_in_raw = payload.get(oauth.expires_in_field, 3600) @@ -145,6 +155,16 @@ async def get_initial_oauth_headers(extensions_config: ExtensionsConfig) -> dict headers: dict[str, str] = {} for server_name in token_manager.oauth_server_names(): - headers[server_name] = await token_manager.get_authorization_header(server_name) or "" + try: + value = await token_manager.get_authorization_header(server_name) + except Exception: + logger.warning( + "Skipping initial OAuth header for MCP server '%s' after token fetch failed", + server_name, + exc_info=True, + ) + continue + if value: + headers[server_name] = value return {name: value for name, value in headers.items() if value} diff --git a/backend/tests/test_mcp_oauth.py b/backend/tests/test_mcp_oauth.py index 27facd465..0ad233aaf 100644 --- a/backend/tests/test_mcp_oauth.py +++ b/backend/tests/test_mcp_oauth.py @@ -189,3 +189,146 @@ def test_get_initial_oauth_headers(monkeypatch): 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"