mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
f39cdd6077
* fix(mcp): keep rotating OAuth refresh tokens across restarts and Gateway instances
MCP servers on the OAuth refresh_token grant lost their rotated refresh
token whenever a new process or manager started (restart, tool-cache
reset for non-task servers, every personal-connection run, a second
Gateway instance) and re-seeded with the consumed configured token,
getting invalid_grant for good.
- New table mcp_oauth_tokens (migration 0047_mcp_oauth_tokens), one row
per refresh_token connection keyed (scope, owner_user_id, server_key):
('deployment', '', name) or ('personal', owner, personal runtime
name). Access and refresh tokens are DEER_FLOW_CREDENTIALS_KEY
ciphertext; a config fingerprint (token_url, client_id, grant, scope,
audience, seed) makes a changed seed win.
- OAuthTokenManager re-reads the row under its process lock, adopts a
peer's valid token, and otherwise claims the refresh with a lease plus
a version compare-and-set (database clock), fetches outside any
transaction, writes under the lease and releases it on failure.
- The Gateway installs the store (sync engine via to_thread) around the
LangGraph runtime for sqlite/postgres; personal managers pass their
owner. client_credentials stays in process; an unloadable key falls
back to process memory; database errors fail the refresh.
- credentials_key_consumers lists refresh_token MCP OAuth servers, so a
declared multi-instance deployment needs the shared key.
Docs: MCP_SERVER.md, CONFIGURATION.md, Helm README, README, mcp and
migrations AGENTS.md, CHANGELOG (en/zh).
* docs(changelog): link the MCP OAuth token store entry to #7207
* fix(mcp): prune removed connections' OAuth token rows and ask owners to re-authorize on invalid_grant
Review follow-up (willem-bd):
- Both MCP config writers drop mcp_oauth_tokens rows of connections a
save removes (deployment: _commit_mcp_config_write; personal:
_mutate, keeping each connection's enabled/disabled runtime-name
variants so a toggle keeps its chain). Best effort; no startup sweep,
since a missing config file would look like every server removed.
- invalid_grant on a refresh_token request raises
McpOAuthReauthorizationRequired: the tool error names the server and
asks for a new refresh token, plus the shared-key hint when this
instance could not decrypt the stored chain.
- Annotate the gateway helpers' config as AppConfig.
* fix(mcp): renew the OAuth refresh lease, keep unstored rotations, and key access tokens by request parameters
Review follow-up (willem-bd, second round):
- HTTPX timeouts bound each I/O wait, not a whole response: the refresh
holder now renews its lease every third of REFRESH_LEASE_SECONDS while
the token request runs and abandons it after REFRESH_MAX_SECONDS (120 s);
waiters outlast a live holder. Regression uses a real loopback endpoint
that streams its body for three lease lengths.
- A refresh whose tokens could not be stored keeps them pending
(unsaved_base_version): they are written back before an older stored
token is adopted or refreshed with, retried every few seconds while the
access token is still valid, and dropped once a peer stored a newer
chain. A failed save releases the lease.
- New token_request_fingerprint (extra_token_params and response fields):
a stored access token is reused only under the same request parameters;
otherwise the stored refresh-token chain is refreshed.
---------
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
64 lines
2.7 KiB
Python
64 lines
2.7 KiB
Python
"""Share MCP OAuth tokens through the application database for one Gateway lifetime."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections.abc import Collection
|
|
from contextlib import asynccontextmanager
|
|
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.mcp.oauth_store import McpOAuthTokenStore, McpOAuthTokenStoreUnavailable, get_mcp_oauth_token_store, set_mcp_oauth_token_store
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def build_mcp_oauth_token_store(app_config: AppConfig) -> McpOAuthTokenStore | None:
|
|
"""The ``mcp_oauth_tokens`` store for a sqlite/postgres database, else ``None`` (tokens stay in process).
|
|
|
|
The store opens the shared synchronous engine
|
|
(``deerflow.persistence.agents.sql.get_sync_sessionmaker``) because token
|
|
refreshes run on whichever loop calls the tool. Its table comes from the
|
|
persistence bootstrap, which runs before any tool call; the credentials
|
|
key loads on the first refresh_token use.
|
|
"""
|
|
database = getattr(app_config, "database", None)
|
|
if getattr(database, "backend", None) not in ("sqlite", "postgres"):
|
|
return None
|
|
from deerflow.persistence.agents.sql import get_sync_sessionmaker
|
|
from deerflow.persistence.mcp_oauth_tokens import SqlMcpOAuthTokenRepository
|
|
|
|
return McpOAuthTokenStore(SqlMcpOAuthTokenRepository(get_sync_sessionmaker(database.app_sync_sqlalchemy_url)))
|
|
|
|
|
|
@asynccontextmanager
|
|
async def mcp_oauth_token_persistence(app_config: AppConfig):
|
|
"""Install the store for every OAuth manager in this process, around the whole runtime.
|
|
|
|
Entered before and left after the LangGraph runtime, whose exit drains
|
|
active runs: a refresh during that drain must still reach the database, or
|
|
the rotated refresh token is lost to the next process.
|
|
"""
|
|
store = await asyncio.to_thread(build_mcp_oauth_token_store, app_config)
|
|
previous = set_mcp_oauth_token_store(store)
|
|
try:
|
|
yield
|
|
finally:
|
|
set_mcp_oauth_token_store(previous)
|
|
|
|
|
|
def forget_unconfigured_mcp_oauth_tokens(owner_user_id: str | None, configured_server_keys: Collection[str]) -> None:
|
|
"""After a saved MCP configuration, drop stored tokens of connections it no longer has.
|
|
|
|
``owner_user_id`` is ``None`` for the deployment's servers. Best effort and
|
|
blocking: a database failure is logged by the store and leaves the rows for
|
|
the next save; it never fails the configuration write that already landed.
|
|
"""
|
|
store = get_mcp_oauth_token_store()
|
|
if store is None:
|
|
return
|
|
try:
|
|
removed = store.forget_unconfigured(owner_user_id, configured_server_keys)
|
|
except McpOAuthTokenStoreUnavailable:
|
|
return
|
|
if removed:
|
|
logger.info("Deleted stored OAuth tokens of %d removed MCP connection(s)", removed)
|