deer-flow/backend/tests/test_mcp_lifecycle_writers.py
Nan Gao f39cdd6077
fix(mcp): keep rotating OAuth refresh tokens across restarts and Gateway instances (#7207)
* 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>
2026-10-21 18:17:29 +08:00

223 lines
8.1 KiB
Python

"""Gateway writer coverage for the shared ``mcpLifecycle`` generation."""
from __future__ import annotations
import json
import pytest
# Reuse the process-local isolation fixture so these writer tests do not leak
# cache globals or the session-pool singleton.
from test_mcp_cache_reconciliation import reconciler # noqa: F401
from app.gateway.routers import mcp as mcp_router
from app.gateway.routers.mcp import (
McpConfigUpdateRequest,
McpServerConfigResponse,
McpServerConfigUpdateRequest,
McpServerStateUpdateRequest,
)
from deerflow.config.extensions_config import ExtensionsConfig
from deerflow.mcp.lifecycle import _stdio_fingerprints, legacy_token, parse_mcp_lifecycle
def _stdio(command: str) -> dict:
return {"enabled": True, "type": "stdio", "command": command}
def _write(cfg, servers: dict) -> None:
cfg.write_text(json.dumps({"mcpServers": servers, "skills": {}}), encoding="utf-8")
def _ledger(cfg):
return parse_mcp_lifecycle(json.loads(cfg.read_text(encoding="utf-8")))
def _legacy(cfg, name: str) -> str:
fingerprint = _stdio_fingerprints(ExtensionsConfig.from_file(str(cfg)))[name]
return legacy_token(name, fingerprint)
def test_state_disable_and_reenable_advance_the_generation(
reconciler, # noqa: F811 - fixture imported from the sibling suite
monkeypatch,
tmp_path,
):
"""Raw in-place mutation must not hide the enabled -> disabled transition."""
cfg = tmp_path / "extensions_config.json"
_write(cfg, {"A": _stdio("npx")})
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg))
baseline = _legacy(cfg, "A")
mcp_router._apply_mcp_server_state_update(McpServerStateUpdateRequest(server_name="A", enabled=False))
disabled = _ledger(cfg).servers["A"]
assert disabled != baseline
mcp_router._apply_mcp_server_state_update(McpServerStateUpdateRequest(server_name="A", enabled=True))
reenabled = _ledger(cfg).servers["A"]
assert reenabled not in {baseline, disabled}
def test_delete_then_identical_readd_advance_the_generation(
reconciler, # noqa: F811 - fixture imported from the sibling suite
monkeypatch,
tmp_path,
):
"""Delete tombstones the token and an identical re-add must rotate it again."""
cfg = tmp_path / "extensions_config.json"
_write(cfg, {"A": _stdio("npx")})
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg))
mcp_router._apply_mcp_server_delete("A")
tombstone = _ledger(cfg).servers["A"]
mcp_router._apply_mcp_servers_create(
McpConfigUpdateRequest(
mcp_servers={"A": McpServerConfigResponse(enabled=True, type="stdio", command="npx")},
)
)
rebound = _ledger(cfg).servers["A"]
assert rebound != tombstone
def test_commit_helper_requires_previous_raw():
"""A controlled write must not be able to skip the lifecycle commit."""
import inspect
parameters = inspect.signature(mcp_router._commit_mcp_config_write).parameters
assert parameters["previous_raw"].default is inspect.Parameter.empty
def test_bulk_update_rotates_only_the_changed_server_token(
reconciler, # noqa: F811 - fixture imported from the sibling suite
monkeypatch,
tmp_path,
):
"""A connection change rotates only that server's token; $VAR stays raw."""
cfg = tmp_path / "extensions_config.json"
cfg.write_text(
json.dumps(
{
"mcpServers": {
"A": {"enabled": True, "type": "stdio", "command": "npx", "args": ["pkg-v1"]},
"B": {"enabled": True, "type": "stdio", "command": "npx", "args": ["pkg-b"], "env": {"TOKEN": "$PE_VAR"}},
},
"skills": {},
}
),
encoding="utf-8",
)
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg))
monkeypatch.setenv("PE_VAR", "secret-value")
baseline_a = _legacy(cfg, "A")
mcp_router._apply_mcp_config_update(
McpConfigUpdateRequest(
mcp_servers={
"A": McpServerConfigResponse(enabled=True, type="stdio", command="npx", args=["pkg-v2"]),
"B": McpServerConfigResponse(enabled=True, type="stdio", command="npx", args=["pkg-b"], env={"TOKEN": "$PE_VAR"}),
}
)
)
ledger = _ledger(cfg)
assert ledger is not None
assert ledger.servers["A"] != baseline_a # the connection changed
assert "B" not in ledger.servers # unchanged legacy server stays derived
written = cfg.read_text(encoding="utf-8")
assert "$PE_VAR" in written
assert "secret-value" not in written
@pytest.mark.asyncio
async def test_targeted_metadata_update_keeps_token_and_session(
reconciler, # noqa: F811 - fixture imported from the sibling suite
monkeypatch,
tmp_path,
):
"""A metadata-only targeted update must not rotate the token or the session."""
import test_mcp_cache_reconciliation as harness
from deerflow.mcp import session_pool as session_pool_module
from deerflow.mcp.session_pool import MCPSessionPool
cfg = tmp_path / "extensions_config.json"
pool = MCPSessionPool()
session_pool_module._pool = pool
log = harness._session_log()
await harness._initialize(monkeypatch, cfg, {"A": _stdio("npx")}, log)
binding_before = harness._binding(pool, "A")
entry_before = harness._entry(pool, "A")
mcp_router._apply_mcp_server_config_update(
McpServerConfigUpdateRequest(
server_name="A",
server=McpServerConfigResponse(enabled=True, type="stdio", command="npx", description="renamed"),
)
)
ledger = _ledger(cfg)
assert ledger is not None and "A" not in ledger.servers # metadata-only: no token
assert harness._binding(pool, "A") is binding_before
assert harness._entry(pool, "A") is entry_before
assert json.loads(cfg.read_text(encoding="utf-8"))["mcpServers"]["A"]["description"] == "renamed"
def test_skills_writer_preserves_the_lifecycle_ledger(
reconciler, # noqa: F811 - fixture imported from the sibling suite
monkeypatch,
tmp_path,
):
"""A skills-only write must round-trip mcpLifecycle exactly (no recompute)."""
from unittest.mock import MagicMock
from app.gateway.routers import skills as skills_router
cfg = tmp_path / "extensions_config.json"
ledger = {"version": 1, "servers": {"A": "a" * 32, "B": "b" * 32}}
cfg.write_text(
json.dumps(
{
"mcpServers": {"A": {"enabled": True, "type": "stdio", "command": "npx"}},
"skills": {"demo": {"enabled": True}},
"mcpLifecycle": ledger,
}
),
encoding="utf-8",
)
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg))
monkeypatch.setattr("app.gateway.routers.skills.reload_extensions_config", lambda *a, **k: None)
skills_router._write_extensions_skill_state(MagicMock(_app_config=None), "demo", False, rebuild_public_projection=False)
written = json.loads(cfg.read_text(encoding="utf-8"))
assert written["mcpLifecycle"] == ledger # preserved byte-for-value, never recomputed
assert written["skills"]["demo"]["enabled"] is False
def test_removed_servers_drop_their_stored_oauth_tokens_and_disabled_ones_keep_them(
reconciler, # noqa: F811 - fixture imported from the sibling suite
monkeypatch,
tmp_path,
):
"""A removed server's mcp_oauth_tokens row can never be read again; a disabled one may be re-enabled."""
from unittest.mock import MagicMock
from deerflow.mcp.oauth_store import set_mcp_oauth_token_store
cfg = tmp_path / "extensions_config.json"
_write(cfg, {"A": _stdio("npx"), "B": _stdio("uvx")})
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg))
store = MagicMock()
store.forget_unconfigured.return_value = 0
previous = set_mcp_oauth_token_store(store)
try:
mcp_router._apply_mcp_server_state_update(McpServerStateUpdateRequest(server_name="A", enabled=False))
mcp_router._apply_mcp_server_delete("B")
finally:
set_mcp_oauth_token_store(previous)
calls = [(call.args[0], set(call.args[1])) for call in store.forget_unconfigured.call_args_list]
assert calls == [(None, {"A", "B"}), (None, {"A"})]