mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 05:06:59 +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>
56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
"""Migration tests for 0026_mcp_task_lease_tokens.
|
|
|
|
Adds the nullable per-claim token columns ``McpTaskRepository`` uses to fence
|
|
poll, cancel, and notification mutations to the exact claim generation. The
|
|
chain-head pin lives in ``test_migration_0047_mcp_oauth_tokens``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
import sqlalchemy as sa
|
|
from alembic import command
|
|
from alembic.script import ScriptDirectory
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
|
|
from deerflow.persistence import bootstrap
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
REVISION = "0026_mcp_task_lease_tokens"
|
|
PREVIOUS = "0025_repair_run_change_seq"
|
|
TOKEN_COLUMNS = {"lease_token", "notification_lease_token"}
|
|
|
|
|
|
async def test_0026_chains_into_the_single_head():
|
|
script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR))
|
|
assert len(script.get_heads()) == 1
|
|
# Later migrations may advance the head without removing this revision.
|
|
assert REVISION in {revision.revision for revision in script.walk_revisions()}
|
|
assert script.get_revision(REVISION).down_revision == PREVIOUS
|
|
|
|
|
|
async def test_0026_adds_nullable_claim_tokens_and_downgrades(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'lease-tokens.db'}")
|
|
cfg = bootstrap._get_alembic_config(engine)
|
|
|
|
async def token_columns() -> dict[str, bool]:
|
|
async with engine.connect() as conn:
|
|
columns = await conn.run_sync(lambda sync: sa.inspect(sync).get_columns("mcp_tasks"))
|
|
return {column["name"]: column["nullable"] for column in columns if column["name"] in TOKEN_COLUMNS}
|
|
|
|
try:
|
|
await asyncio.to_thread(bootstrap._upgrade, cfg, PREVIOUS)
|
|
assert await token_columns() == {}
|
|
|
|
await asyncio.to_thread(bootstrap._upgrade, cfg, "head")
|
|
# Nullable so rows written before this revision stay valid.
|
|
assert await token_columns() == dict.fromkeys(TOKEN_COLUMNS, True)
|
|
|
|
await asyncio.to_thread(command.downgrade, cfg, PREVIOUS)
|
|
assert await token_columns() == {}
|
|
finally:
|
|
await engine.dispose()
|