mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
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>
This commit is contained in:
parent
1ad8147947
commit
f39cdd6077
24
CHANGELOG.md
24
CHANGELOG.md
@ -1397,6 +1397,29 @@ As of 2026-10-21, the milestone includes **829 merged pull requests**.
|
||||
|
||||
### Fixed
|
||||
|
||||
- **mcp:** MCP servers using the OAuth `refresh_token` grant keep working
|
||||
across Gateway restarts, tool-cache resets, repeated personal-connection runs
|
||||
and multiple Gateway instances when the provider rotates refresh tokens
|
||||
(Auth0, Okta, Google). Tokens lived only in process memory, so every new
|
||||
process or manager started again from the configured refresh token, which
|
||||
the first refresh had already consumed, and got `invalid_grant`. With
|
||||
`database.backend: sqlite` or `postgres`, each connection's access token and
|
||||
latest refresh token now live in the new `mcp_oauth_tokens` table (migration
|
||||
`0047_mcp_oauth_tokens`), encrypted with `DEER_FLOW_CREDENTIALS_KEY` and keyed
|
||||
per deployment server or per personal-connection owner. A refresh re-reads
|
||||
the row first and is claimed with a short lease (renewed while the token
|
||||
request runs) plus a version compare-and-set, so concurrent instances call
|
||||
the token endpoint once and adopt each other's result; a rotation that could
|
||||
not be stored is kept and stored before the next refresh. Changing the configured seed, token endpoint,
|
||||
client, scope or audience starts a new chain from the configuration.
|
||||
`client_credentials` tokens stay in process. Saving an MCP configuration
|
||||
deletes the stored tokens of connections it no longer has, and a rejected
|
||||
refresh token now fails the tool call with an error asking for a new refresh
|
||||
token (adding, when the instance could not decrypt the stored chain, that
|
||||
every instance needs the same key). A declared multi-instance deployment with
|
||||
an enabled `refresh_token` OAuth server now refuses to start without
|
||||
`DEER_FLOW_CREDENTIALS_KEY`. ([#7207])
|
||||
|
||||
- **community:** Reject booleans in numeric LightRAG and RAGFlow
|
||||
`knowledge_search` settings instead of silently treating `true` as a
|
||||
one-second timeout, one result or one character of output. Invalid settings
|
||||
@ -12944,3 +12967,4 @@ with **180 merged pull requests** since the first 2.0 milestone tag.
|
||||
[#7167]: https://github.com/bytedance/deer-flow/pull/7167
|
||||
[#7170]: https://github.com/bytedance/deer-flow/pull/7170
|
||||
[#7205]: https://github.com/bytedance/deer-flow/pull/7205
|
||||
[#7207]: https://github.com/bytedance/deer-flow/pull/7207
|
||||
|
||||
@ -1112,6 +1112,22 @@
|
||||
|
||||
### 修复
|
||||
|
||||
- **MCP:** 使用 OAuth `refresh_token` 授权的 MCP 服务器,在提供方轮换 refresh token
|
||||
(Auth0、Okta、Google)时,跨 Gateway 重启、工具缓存重置、个人连接的多次运行以及多个
|
||||
Gateway 实例都能继续工作。此前令牌只保存在进程内存中,每个新进程或新 manager 都从配置里
|
||||
的 refresh token 重新开始,而该 token 已被第一次刷新消耗,于是得到 `invalid_grant`。
|
||||
当 `database.backend` 为 `sqlite` 或 `postgres` 时,每个连接的 access token 与最新的
|
||||
refresh token 现保存在新表 `mcp_oauth_tokens`(迁移 `0047_mcp_oauth_tokens`)中,用
|
||||
`DEER_FLOW_CREDENTIALS_KEY` 加密,按部署级服务器或个人连接的所有者分别建键。刷新前先重读
|
||||
该行,并以短租约(token 请求进行期间持续续约)加版本号比较交换(compare-and-set)认领,
|
||||
因此并发实例只调用一次 token 端点并采用彼此的结果;未能写入的轮换结果会被保留,并在下次
|
||||
刷新前先写回。修改配置中的种子、token 端点、client、scope 或 audience 会从配置
|
||||
开始新的令牌链。`client_credentials` 令牌仍留在进程内。保存 MCP 配置时会删除配置中已不存在
|
||||
的连接所存的令牌;refresh token 被拒绝时,工具调用会报出要求重新配置 refresh token 的错误
|
||||
(若本实例无法解密其他实例存下的令牌链,还会提示所有实例需要相同的密钥)。声明多实例部署且启用了
|
||||
`refresh_token` OAuth 服务器时,缺少 `DEER_FLOW_CREDENTIALS_KEY` 现在会拒绝启动。
|
||||
([#7207])
|
||||
|
||||
- **社区工具:** 拒绝 LightRAG 与 RAGFlow 的 `knowledge_search` 数值设置中的
|
||||
布尔值,不再静默把 `true` 当作一秒超时、一条结果或一个字符的输出。无效设置返回
|
||||
现有的配置错误。([#7042])
|
||||
@ -10333,3 +10349,4 @@ DeerFlow 2.0 是围绕"超级智能体"框架的彻底重写,核心包含子
|
||||
[#7167]: https://github.com/bytedance/deer-flow/pull/7167
|
||||
[#7170]: https://github.com/bytedance/deer-flow/pull/7170
|
||||
[#7205]: https://github.com/bytedance/deer-flow/pull/7205
|
||||
[#7207]: https://github.com/bytedance/deer-flow/pull/7207
|
||||
|
||||
@ -613,7 +613,7 @@ Administrators can suspend and restore accounts from **Settings → Users** (adm
|
||||
DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser-facing scheme and origin behind a proxy. The bundled nginx sets `X-Forwarded-Proto`, but preserves an upstream HTTPS value and does not overwrite every forwarded header. Configure the outer trusted proxy to replace or strip client-supplied forwarding headers before traffic reaches DeerFlow.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. Kubernetes replicas run one worker per Pod, which the worker count cannot see: declare them with `deployment.multi_instance: true` (or `DEER_FLOW_MULTI_INSTANCE=1`, which deploy tooling such as a Helm chart can set from its replica count) so the startup gate enforces the same prerequisites instead of staying inert. Declared instances that store credentials (`channel_connections`) must also share one at-rest `DEER_FLOW_CREDENTIALS_KEY` (the Helm chart and `make up` generate it; back it up, see [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#credentials-encryption-key-deer_flow_credentials_key)). The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a trimmed reconnect cursor cannot be verified for complete replay, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. The Memory bridge can replay without recovery when the cursor is its exact last evicted event: every successor is still retained. Older or unverifiable evicted cursors still trigger recovery. An already-observed Redis stream removed during replay also signals a gap, so slow subscribers reload durable state instead of heartbeating forever after cleanup. Redis also tracks terminal stream producers with a short lease: if a worker dies after committing a terminal run status but before publishing END, SSE and `/wait` finish after the lease expires instead of waiting for the retained stream's TTL. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE, `/wait`, and internal stream consumers use `stream_bridge.heartbeat_interval_seconds` (default `15`) for idle liveness checks; changing it requires a Gateway restart. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. Failed-login counters and lockouts for `POST /api/v1/auth/login/local` (`auth.local.max_login_attempts` / `lockout_seconds`) are kept in the shared `login_throttle` table whenever the application database is SQLite or Postgres (`auth.local.throttle_storage: auto`, the default), so every replica enforces one lockout per client IP; `memory` keeps the historical per-process counter, which under N replicas hands an attacker N × `max_login_attempts` guesses and logs a startup warning. IM chat-to-thread bindings live in the shared `channel_thread_bindings` table whenever the application database is SQLite or Postgres (an existing `channels/store.json` is imported once at startup and renamed `store.json.migrated`); IM bot credentials entered in the web UI live in the shared `channel_runtime_configs` table, encrypted with `DEER_FLOW_CREDENTIALS_KEY` (an existing `channels/runtime-config.json` is imported once and then deleted), so a save on any instance reaches the provider's connection within one lease renewal; each provider's platform connection is held by one instance through the shared `channel_leases` table — `channels.leader_lease` (default `auto`, on wherever the multi-process signal is) hands it to another instance when the holder dies, and the startup gate refuses `enabled: false` while an IM channel other than GitHub is enabled (see [IM channel leader lease](backend/docs/CONFIGURATION.md#im-channel-leader-lease-channelsleader_lease)); follow-up buffers still need their own multi-worker coordination. To try this topology on one machine, `make dev-multi` starts two Gateways on shared Postgres, Redis and `DEER_FLOW_HOME`, and `make dev-multi-check` runs the cross-instance checks; see the local two-Gateway harness in [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#local-two-gateway-harness).
|
||||
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. Kubernetes replicas run one worker per Pod, which the worker count cannot see: declare them with `deployment.multi_instance: true` (or `DEER_FLOW_MULTI_INSTANCE=1`, which deploy tooling such as a Helm chart can set from its replica count) so the startup gate enforces the same prerequisites instead of staying inert. Declared instances that store credentials (`channel_connections`, or MCP servers using the OAuth `refresh_token` grant, whose rotating tokens they share through the database) must also share one at-rest `DEER_FLOW_CREDENTIALS_KEY` (the Helm chart and `make up` generate it; back it up, see [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#credentials-encryption-key-deer_flow_credentials_key)). The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a trimmed reconnect cursor cannot be verified for complete replay, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. The Memory bridge can replay without recovery when the cursor is its exact last evicted event: every successor is still retained. Older or unverifiable evicted cursors still trigger recovery. An already-observed Redis stream removed during replay also signals a gap, so slow subscribers reload durable state instead of heartbeating forever after cleanup. Redis also tracks terminal stream producers with a short lease: if a worker dies after committing a terminal run status but before publishing END, SSE and `/wait` finish after the lease expires instead of waiting for the retained stream's TTL. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE, `/wait`, and internal stream consumers use `stream_bridge.heartbeat_interval_seconds` (default `15`) for idle liveness checks; changing it requires a Gateway restart. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. Failed-login counters and lockouts for `POST /api/v1/auth/login/local` (`auth.local.max_login_attempts` / `lockout_seconds`) are kept in the shared `login_throttle` table whenever the application database is SQLite or Postgres (`auth.local.throttle_storage: auto`, the default), so every replica enforces one lockout per client IP; `memory` keeps the historical per-process counter, which under N replicas hands an attacker N × `max_login_attempts` guesses and logs a startup warning. IM chat-to-thread bindings live in the shared `channel_thread_bindings` table whenever the application database is SQLite or Postgres (an existing `channels/store.json` is imported once at startup and renamed `store.json.migrated`); IM bot credentials entered in the web UI live in the shared `channel_runtime_configs` table, encrypted with `DEER_FLOW_CREDENTIALS_KEY` (an existing `channels/runtime-config.json` is imported once and then deleted), so a save on any instance reaches the provider's connection within one lease renewal; each provider's platform connection is held by one instance through the shared `channel_leases` table — `channels.leader_lease` (default `auto`, on wherever the multi-process signal is) hands it to another instance when the holder dies, and the startup gate refuses `enabled: false` while an IM channel other than GitHub is enabled (see [IM channel leader lease](backend/docs/CONFIGURATION.md#im-channel-leader-lease-channelsleader_lease)); follow-up buffers still need their own multi-worker coordination. To try this topology on one machine, `make dev-multi` starts two Gateways on shared Postgres, Redis and `DEER_FLOW_HOME`, and `make dev-multi-check` runs the cross-instance checks; see the local two-Gateway harness in [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#local-two-gateway-harness).
|
||||
>
|
||||
> The producer-lease recovery behavior above is gated by `stream_bridge.producer_lease_recovery_enabled`. Keep it disabled while old Gateway binaries or their runs remain; enable it on every instance only after the writer rollout and drain complete.
|
||||
>
|
||||
|
||||
@ -716,9 +716,10 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
logger.warning("Upload staging file cleanup skipped", exc_info=True)
|
||||
|
||||
# Initialize LangGraph runtime components (StreamBridge, RunManager, checkpointer, store)
|
||||
from app.gateway.mcp_oauth_tokens import mcp_oauth_token_persistence
|
||||
from app.gateway.personal_mcp_access import personal_mcp_authority
|
||||
|
||||
async with personal_mcp_authority(), _runtime_with_mcp_pool_shutdown(app, startup_config):
|
||||
async with personal_mcp_authority(), mcp_oauth_token_persistence(startup_config), _runtime_with_mcp_pool_shutdown(app, startup_config):
|
||||
logger.info("LangGraph runtime initialised")
|
||||
|
||||
# Check admin bootstrap state and migrate orphan threads after admin exists.
|
||||
|
||||
@ -296,9 +296,11 @@ def credentials_key_consumers(config: AppConfig) -> list[str]:
|
||||
"""Return the enabled features that store data under ``DEER_FLOW_CREDENTIALS_KEY``.
|
||||
|
||||
Each entry is the setting that enabled the feature, as refusal text spells
|
||||
it. The Gateway loads the key only when this list is non-empty, and the
|
||||
multi-instance gate below refuses a missing key only for these features, so
|
||||
a deployment that stores no credentials keeps booting without one.
|
||||
it. The Gateway loads the key at startup only when this list is non-empty,
|
||||
and the multi-instance gate below refuses a missing key only for these
|
||||
features, so a deployment that stores no credentials keeps booting without
|
||||
one. Personal MCP OAuth connections are saved at runtime, so they cannot be
|
||||
listed here; their store loads the key on first use.
|
||||
"""
|
||||
consumers: list[str] = []
|
||||
connections = getattr(config, "channel_connections", None)
|
||||
@ -306,6 +308,13 @@ def credentials_key_consumers(config: AppConfig) -> list[str]:
|
||||
# ChannelConnectionRepository encrypts per-connection credentials, and
|
||||
# the runtime config store the bot credentials entered in the web UI.
|
||||
consumers.append("channel_connections.enabled=true")
|
||||
extensions = getattr(config, "extensions", None)
|
||||
if extensions is not None:
|
||||
# deerflow.mcp.oauth_store keeps rotating refresh tokens in the database.
|
||||
for name, server in extensions.get_enabled_mcp_servers().items():
|
||||
oauth = server.oauth
|
||||
if oauth is not None and oauth.enabled and oauth.grant_type == "refresh_token":
|
||||
consumers.append(f"mcpServers.{name}.oauth.enabled=true")
|
||||
return consumers
|
||||
|
||||
|
||||
|
||||
63
backend/app/gateway/mcp_oauth_tokens.py
Normal file
63
backend/app/gateway/mcp_oauth_tokens.py
Normal file
@ -0,0 +1,63 @@
|
||||
"""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)
|
||||
@ -13,6 +13,7 @@ from fastapi import APIRouter, HTTPException, Request, status
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator, model_validator
|
||||
|
||||
from app.gateway.deps import require_admin_user
|
||||
from app.gateway.mcp_oauth_tokens import forget_unconfigured_mcp_oauth_tokens
|
||||
from app.gateway.persistent_writes import run_drained_write
|
||||
from deerflow.config.extensions_config import (
|
||||
ExtensionsConfig,
|
||||
@ -1250,6 +1251,8 @@ def _commit_mcp_config_write(
|
||||
# Re-parse so the committed fence sees the ledger we just wrote.
|
||||
candidate = validate_raw_extensions_config(raw_data)
|
||||
atomic_write_extensions_config(config_path, raw_data)
|
||||
# Disabled servers keep their stored OAuth tokens; removed ones cannot use them again.
|
||||
forget_unconfigured_mcp_oauth_tokens(None, candidate.mcp_servers.keys())
|
||||
try:
|
||||
reload_extensions_config()
|
||||
return prepare_mcp_reconciliation(candidate, config_path=config_path)
|
||||
|
||||
@ -8,10 +8,11 @@ from fastapi import APIRouter, HTTPException, Request
|
||||
from pydantic import ValidationError
|
||||
|
||||
from app.gateway.deps import get_current_user_from_request, is_admin_user
|
||||
from app.gateway.mcp_oauth_tokens import forget_unconfigured_mcp_oauth_tokens
|
||||
from app.gateway.routers import mcp
|
||||
from deerflow.capabilities.runtime import ambiguous_installation_ids
|
||||
from deerflow.config.extensions_config import ExtensionsConfig, atomic_write_extensions_config, extensions_config_file_lock, extensions_config_write_lock
|
||||
from deerflow.mcp.user_config import read_user_mcp_config, user_mcp_config_path
|
||||
from deerflow.mcp.user_config import personal_server_name_variants, read_user_mcp_config, user_mcp_config_path
|
||||
from deerflow.utils.file_io import await_drained
|
||||
|
||||
router = APIRouter(prefix="/api/mcp/personal/config", tags=["mcp"])
|
||||
@ -91,6 +92,9 @@ def _mutate(user_id: str, operation: Literal["create", "update", "delete", "stat
|
||||
raise HTTPException(400, "Duplicate MCP installation IDs")
|
||||
atomic_write_extensions_config(path, candidate)
|
||||
path.chmod(0o600)
|
||||
# A removed or edited connection's runtime name never comes back
|
||||
# (toggles do; see personal_server_name_variants), nor does its token row.
|
||||
forget_unconfigured_mcp_oauth_tokens(user_id, {runtime_name for name, definition in servers.items() for runtime_name in personal_server_name_variants(user_id, name, definition)})
|
||||
return candidate
|
||||
|
||||
|
||||
|
||||
@ -765,7 +765,7 @@ Notes:
|
||||
- Login lockouts are shared through the database: `auth.local.throttle_storage` (startup-only) defaults to `auto`, which keeps the per-IP failed-login counters for `POST /api/v1/auth/login/local` in the `login_throttle` table whenever `database.backend` is `sqlite` or `postgres`, so every instance enforces one `max_login_attempts` limit per client IP and a lockout on one instance holds on all of them. `memory` keeps the historical per-process counter (N instances give an attacker N × `max_login_attempts` guesses; a declared multi-instance deployment logs a warning), and `db` forces the table (it falls back to memory with a warning when the database backend is `memory`, and refuses to start when the configured database's engine is unavailable; `auto` falls back with a warning in that case). `max_login_attempts` and `lockout_seconds` stay live-read.
|
||||
- IM chat-to-thread bindings are shared through the database too: with `database.backend: sqlite` or `postgres` the `ChannelManager` keeps the `channel_name:chat_id[:topic_id]` → thread mapping of unbound IM conversations in the `channel_thread_bindings` table (migration `0038_channel_thread_bindings`), so a conversation created on one instance continues on the same thread when its next message lands on another. On the first start after upgrading, an existing `{base_dir}/channels/store.json` is imported once — only into an empty table, with `INSERT … ON CONFLICT DO NOTHING` so concurrently starting instances cannot duplicate a binding — and renamed `store.json.migrated`; a populated table leaves the file untouched. `memory` keeps the per-process JSON file.
|
||||
- With the default DeerMem backend, the declaration also warns when the derived SQLite retrieval index sits inside the shared `storage_path` (the default `{storage_path}/.retrieval`). Set `memory.backend_config.retrieval_index_path` to an instance-local directory (relative values resolve against `storage_path`): SQLite WAL is unsupported on network filesystems, and the index is rebuilt from the Markdown facts at startup and re-synced per user scope after a peer writes. DeerMem's extraction watermark lives in the shared `storage_path` too (`users/{user}/agents/{agent}/.metadata/watermarks.json`), so a thread whose next turn lands on another instance, or on a restarted one, does not send already-extracted turns to the memory LLM again.
|
||||
- A declared multi-instance deployment that enables a feature storing credentials (today `channel_connections.enabled: true`) also needs one shared `DEER_FLOW_CREDENTIALS_KEY`; see [Credentials encryption key](#credentials-encryption-key-deer_flow_credentials_key).
|
||||
- A declared multi-instance deployment that enables a feature storing credentials (`channel_connections.enabled: true`, or an enabled MCP server whose `oauth.grant_type` is `refresh_token`) also needs one shared `DEER_FLOW_CREDENTIALS_KEY`; see [Credentials encryption key](#credentials-encryption-key-deer_flow_credentials_key).
|
||||
- IM channels: the gate refuses `channels.leader_lease.enabled: false` while a provider other than GitHub is enabled (under `channels.*`, or connectable from the web UI through `channel_connections.<provider>.enabled`), because every instance would otherwise open the platform connection with the same bot credentials. With `auto` (the default) or `true` the instances share each provider's connection through the `channel_leases` table; see [IM channel leader lease](#im-channel-leader-lease-channelsleader_lease).
|
||||
- Restart-required: the gate runs once at startup. Restart all Gateway instances together after changing it.
|
||||
|
||||
@ -830,7 +830,7 @@ Every IM channel except GitHub connects outbound to its platform (Telegram long
|
||||
|
||||
### Credentials encryption key (`DEER_FLOW_CREDENTIALS_KEY`)
|
||||
|
||||
DeerFlow encrypts credentials it stores at rest (today: per-connection IM channel credentials in the `channel_credentials` table, and the IM bot credentials an admin enters in the web UI, in `channel_runtime_configs`) with a deployment key read **only** from the environment — there is no `config.yaml` key, and it is never derived from `AUTH_JWT_SECRET`.
|
||||
DeerFlow encrypts credentials it stores at rest (today: per-connection IM channel credentials in the `channel_credentials` table, the IM bot credentials an admin enters in the web UI, in `channel_runtime_configs`, and the rotating tokens of MCP servers using the OAuth `refresh_token` grant, in `mcp_oauth_tokens` — see [MCP_SERVER.md](MCP_SERVER.md#refresh-token-rotation-across-restarts-and-replicas)) with a deployment key read **only** from the environment — there is no `config.yaml` key, and it is never derived from `AUTH_JWT_SECRET`.
|
||||
|
||||
```bash
|
||||
# Generate a key (urlsafe base64 of 32 random bytes, a Fernet key):
|
||||
@ -840,7 +840,7 @@ export DEER_FLOW_CREDENTIALS_KEY=<key>
|
||||
|
||||
- **Rotation:** the value may be a comma-separated list. The first key encrypts new values; every listed key decrypts. Prepend the new key (`new,old`), restart, and keep the old key listed until the values it encrypted have been rewritten. Stored values carry a `fernet:v2:` prefix; `fernet:v1:` values written by earlier builds stay readable under the same key.
|
||||
- **Single instance:** when the variable is unset or blank, the Gateway generates `{DEER_FLOW_HOME}/.credentials_key` (mode `0600`) on first use and reuses it. The file is created exclusively, so uvicorn workers sharing the runtime home converge on one key. It is only created when a feature that stores credentials is enabled.
|
||||
- **Multi-instance:** instances that do not share a runtime home would each generate their own file and could not decrypt what their peers stored. Startup therefore refuses a declared multi-instance deployment (`deployment.multi_instance: true` / `DEER_FLOW_MULTI_INSTANCE=1`) that enables `channel_connections` without the variable. Multiple workers of one process tree (`GATEWAY_WORKERS > 1` without the declaration) share the file and are not refused. Deployments that store no credentials keep booting without a key.
|
||||
- **Multi-instance:** instances that do not share a runtime home would each generate their own file and could not decrypt what their peers stored. Startup therefore refuses a declared multi-instance deployment (`deployment.multi_instance: true` / `DEER_FLOW_MULTI_INSTANCE=1`) that enables `channel_connections`, or an MCP server with a `refresh_token` OAuth block, without the variable. Personal MCP OAuth connections are saved at runtime, after that check, so set the variable whenever multi-instance users may get one. Multiple workers of one process tree (`GATEWAY_WORKERS > 1` without the declaration) share the file and are not refused. Deployments that store no credentials keep booting without a key.
|
||||
- **Malformed keys** (wrong length, standard rather than urlsafe base64, any bad entry in the list) stop startup with an error that never repeats the key material.
|
||||
- **Backup:** losing the key makes every stored credential unreadable. Unreadable values are treated as missing (logged), never as a crash — Slack, for example, falls back to the deployment bot token, and a bot configured from the web UI reads as not configured until its credentials are saved again. If no key can be loaded at all (an unreadable key file), saving bot credentials from the web UI answers `503` instead of storing them in plaintext.
|
||||
- **Deployment tooling:** the Helm chart generates the key into its app Secret and preserves it across upgrades; `make up` (`scripts/deploy.sh`) honors a shell or `.env` value and otherwise generates and persists `$DEER_FLOW_HOME/.credentials_key` — the same file the Gateway would generate, so whichever side created it first wins.
|
||||
|
||||
@ -412,9 +412,11 @@ Tool-cache resets and rediscovery keep that state. Token rotation does not
|
||||
modify the parsed configuration or trigger the startup-snapshot drift guard;
|
||||
real operator configuration changes still require a restart.
|
||||
|
||||
This state is process-local and is never written back to `extensions_config.json`
|
||||
or an environment variable. A restarted Gateway needs a valid configured
|
||||
refresh token, and separate Gateway workers do not coordinate token rotation.
|
||||
This state is never written back to `extensions_config.json` or an environment
|
||||
variable. With a SQL database, a `refresh_token` connection's tokens are also kept
|
||||
in the shared `mcp_oauth_tokens` table, so a restarted Gateway and its peers
|
||||
continue the rotated chain (see
|
||||
[Refresh-token rotation](#refresh-token-rotation-across-restarts-and-replicas)).
|
||||
Personal MCP connections remain owner-scoped and do not use this deployment
|
||||
state, even if their runtime names match a deployment server. A personal
|
||||
server's runtime name is derived from its complete definition, so editing,
|
||||
@ -481,6 +483,62 @@ Example:
|
||||
}
|
||||
```
|
||||
|
||||
### Refresh-token rotation across restarts and replicas
|
||||
|
||||
Providers such as Auth0, Okta and Google rotate refresh tokens: every refresh
|
||||
returns a new one and invalidates the old one. With `database.backend: sqlite`
|
||||
or `postgres`, the Gateway keeps each `refresh_token` connection's latest access
|
||||
token and refresh token in the `mcp_oauth_tokens` table (migration
|
||||
`0047_mcp_oauth_tokens`), encrypted with `DEER_FLOW_CREDENTIALS_KEY` (see
|
||||
[CONFIGURATION.md](CONFIGURATION.md#credentials-encryption-key-deer_flow_credentials_key)):
|
||||
|
||||
- A restart, a tool-cache reset, or a later run of a personal connection
|
||||
continues from the stored chain instead of the configured `refresh_token`.
|
||||
- Gateway replicas share one chain. Before refreshing, a replica re-reads the
|
||||
row and uses a peer's still-valid access token. Otherwise it claims the
|
||||
refresh with a 30-second lease, so only one replica calls the token endpoint
|
||||
and the others wait for its result. The holder renews the lease while the
|
||||
token request runs (HTTP timeouts bound each read, not the whole response)
|
||||
and abandons a request after 120 seconds. No row lock is held during the
|
||||
token request.
|
||||
- If the new tokens cannot be stored (the database is briefly unavailable),
|
||||
the process keeps using them, retries storing them every few seconds, and
|
||||
stores them before its next refresh, so nobody falls back to the refresh
|
||||
token it already consumed.
|
||||
- A stored access token is reused only by processes that request it the same
|
||||
way (`extra_token_params` such as `resource`, and the response-field
|
||||
settings). After such a change the next request refreshes from the stored
|
||||
refresh token instead, so the chain is kept.
|
||||
- The configured `refresh_token` is only the seed. The row records which
|
||||
`token_url`, `client_id`, `grant_type`, `scope`, `audience` and seed it
|
||||
continues (as a SHA-256 fingerprint). After you change any of them, the next
|
||||
refresh starts again from the configured value. To recover from a revoked
|
||||
chain (`invalid_grant`), configure a new `refresh_token`; changing only
|
||||
`client_secret` keeps the stored chain.
|
||||
- Personal connections get one row per owner, keyed by their runtime name.
|
||||
Editing a personal connection changes that name, so it starts again from its
|
||||
configured refresh token. Saving a configuration deletes the rows of the
|
||||
connections it no longer has: an edited or deleted personal connection, or a
|
||||
deployment server removed through the MCP configuration API (a server removed
|
||||
by hand-editing `extensions_config.json` is dropped on the next API save).
|
||||
Disabling a connection keeps its row, so re-enabling it continues the chain.
|
||||
- When the provider rejects the refresh token (`invalid_grant`), the tool call
|
||||
fails with an error that names the server and asks for a new refresh token.
|
||||
If the instance fell back to the seed because it could not decrypt the chain
|
||||
another instance stored, the error also says that every instance needs the
|
||||
same `DEER_FLOW_CREDENTIALS_KEY`.
|
||||
- `client_credentials` tokens are not stored; each process requests its own.
|
||||
- Tokens stay in process memory, as before, with `database.backend: memory`,
|
||||
in the embedded `DeerFlowClient` (no Gateway), or when the credentials key
|
||||
cannot be loaded. If the database cannot be read or written, the refresh
|
||||
fails instead of consuming a refresh token that nothing records; a still-valid
|
||||
access token keeps working.
|
||||
- A declared multi-instance deployment (`deployment.multi_instance: true`) with
|
||||
an enabled `refresh_token` OAuth server refuses to start without a shared
|
||||
`DEER_FLOW_CREDENTIALS_KEY`. Personal OAuth connections are saved at runtime,
|
||||
so that startup check cannot see them; set the variable on every
|
||||
multi-instance deployment that uses them.
|
||||
|
||||
## Request-Scoped Headers (HTTP/SSE MCP Servers)
|
||||
|
||||
When the credential is chosen by the *caller* rather than by the operator —
|
||||
|
||||
@ -6,7 +6,8 @@ The optional background-task `status` filter flows through `McpTaskSubmitter.lis
|
||||
- **OAuth expiry**: Missing, non-convertible or overflowing `expires_in` values use the one-hour default. Guard both `timedelta` construction and addition to one captured UTC `now`; avoid float-derived range ceilings. `tests/test_mcp_oauth.py` pins both overflow stages and the exact `datetime` boundary with a fixed clock.
|
||||
- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and separate consecutive poll/delivery error counters; `app/mcp_tasks/McpTaskService` performs status, cancellation, and notification work outside the Agent/LLM loop. Notification retries keep their idempotency attempt separate from the delivery-failure count, use capped exponential backoff, and stop after five failures; strict existing-thread admission dead-letters a deleted/mismatched target immediately. A status result is applied only when the worker still owns an unexpired lease, so a stale result cannot be written after expiry even before another worker reclaims the row. Every claim carries a fresh per-claim token, and every poll/cancel/notification release or snapshot apply is fenced by the owner plus that token, so a release that completes after the same service reclaims the row cannot clear the newer lease. Poll, cancel, and notification claims also have separate phase-level single-flight owners that protect only an ambiguous claim outcome: after an uncancelled timeout or caller cancellation, later scans do not start another claim while the outcome is unknown. Once the claim resolves, the phase owner is released immediately, and any returned rows are released through bounded, service-owned background work; per-claim token fencing prevents a late release from mutating a newer claim generation, so a stuck release does not stall the whole phase. While the claim outcome remains unresolved, later scans skip the phase and emit a warning rather than overlapping an ambiguous database claim. Once the claim resolves, the handoff may continue releasing returned rows in the background without blocking later scans. Lease expiry remains the cross-process crash fallback rather than the normal recovery for a live service's late claim. Poll timestamps and retry schedules are based on the remote call's completion time rather than the scan start. If submission succeeds but persistence fails or the caller is cancelled while persistence is in flight, the service best-effort cancels the remote task and preserves the original error or cancellation if that compensation also fails. The exact `uq_mcp_tasks_user_server_remote` conflict is different: an existing durable row already owns the remote handle, so the conflict surfaces without cancelling that tracked task. Unexpected per-task poll failures are isolated from sibling claims and remain recoverable through lease expiry; Gateway shutdown cancels the poller so a hung external status call cannot block process exit. Routine cancellation releases clear only the lease/token and preserve any pre-existing `last_cancel_error` or `notification_error`; cancellation itself is not persisted as a task failure diagnostic. Cancelling an in-flight poll releases only the owner- and per-claim-token-fenced lease and preserves its preclaim schedule and poll-failure state; real poll failures retain exponential backoff and tracking-degradation behavior. `input_required` and terminal states stop polling and become `notification_status=pending` for later Agent/UI delivery. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend leaves the repository/service unavailable. The runtime is startup-configured by `mcp_tasks` and disabled by default until a concrete driver is registered; this foundation does not alter ordinary MCP tool behavior on its own.
|
||||
- **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts.
|
||||
- **Task OAuth lifetime**: `tasks/runtime.py` owns shared OAuth state only for enabled deployment servers with `task_toolsets`, from `set_mcp_task_config_snapshot(config)` until its `None` teardown. Use `_task_enabled_servers` for both snapshot and OAuth selection. Validate the frozen snapshot before reusing that state for discovery, ordinary tools, or durable calls. Cache resets must not discard it. Keep tokens and the cross-loop refresh lock together; rotate a private config copy, never the parsed snapshot or config file. Non-task servers remain reloadable and personal callers use independent managers, even on a name collision. This is not cross-process or cross-restart credential storage.
|
||||
- **Task OAuth lifetime**: `tasks/runtime.py` owns shared OAuth state only for enabled deployment servers with `task_toolsets`, from `set_mcp_task_config_snapshot(config)` until its `None` teardown. Use `_task_enabled_servers` for both snapshot and OAuth selection. Validate the frozen snapshot before reusing that state for discovery, ordinary tools, or durable calls. Cache resets must not discard it. Keep tokens and the cross-loop refresh lock together; rotate a private config copy, never the parsed snapshot or config file. Non-task servers remain reloadable and personal callers use independent managers, even on a name collision.
|
||||
- **OAuth token store**: `oauth_store.py` persists only `refresh_token` states in `mcp_oauth_tokens` (encrypted, key `(scope, owner_user_id or '', server_key)`); personal managers must pass `owner_user_id`. `app/gateway/mcp_oauth_tokens.py` installs the store around the runtime so run drains still persist. Under the process lock: re-read (a mismatched fingerprint means the seed wins), adopt a valid peer token (same `token_request_fingerprint` only), else claim via lease + version CAS, fetch while renewing the lease (`REFRESH_MAX_SECONDS` cap), write under the lease, release on failure. An unstored rotation stays pending (`unsaved_base_version`) and is written back before anything older is used. Never hold a row lock across the HTTP call or touch the store on the loop. DB errors fail the refresh; an unloadable key falls back to process memory. `invalid_grant` raises `McpOAuthReauthorizationRequired`, the owner's tool error. Both MCP config writers drop rows of connections a save removes (`forget_unconfigured_mcp_oauth_tokens`; personal toggles keep both runtime-name variants). Tests: `tests/test_mcp_oauth_token_store.py`, `tests/blocking_io/test_mcp_oauth_token_store.py`.
|
||||
- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers.<server>.task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same canonical `(server_name, versioned user/thread/incarnation scope)` stdio session as ordinary calls inside the persisted deployment/personal ownership domain; non-NULL incarnations use an unambiguous versioned encoding, while an explicitly captured legacy NULL retains the pre-activation `user_id:thread_id` scope. A runtime with a missing/invalid incarnation key fails closed. HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count; a task whose persisted connection no longer resolves (a personal definition edit/disable/delete or a deployment startup-snapshot removal) fails permanently instead, and its pending cancellation fails locally with that error. It derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. A task server whose discovery fails or times out is skipped like any other server (`load_server_tools` returns `None`, not `[]`); only a server that answered without its configured raw tools raises `McpTaskConfigurationError`. Users still cannot submit an answer back to an `input_required` remote task.
|
||||
- **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior.
|
||||
- **Remote task bring-up**: One `session_init_timeout` deadline covers transport context entry and MCP initialization. Report the server and configured bound on expiry; disable it before calling the tool. Enter and exit the adapter's AnyIO context managers in the same task; tests include real SSE endpoint/initialization stalls and connection cleanup.
|
||||
|
||||
@ -5,19 +5,67 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Collection
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
from deerflow.config.extensions_config import ExtensionsConfig, McpOAuthConfig
|
||||
from deerflow.mcp import oauth_store
|
||||
from deerflow.mcp.headers import apply_header_overrides, header_spellings, illegal_header_value_reason
|
||||
from deerflow.mcp.oauth_store import (
|
||||
McpOAuthTokenKey,
|
||||
McpOAuthTokenKeyUnavailable,
|
||||
McpOAuthTokenStore,
|
||||
McpOAuthTokenStoreUnavailable,
|
||||
get_mcp_oauth_token_store,
|
||||
oauth_config_fingerprint,
|
||||
oauth_token_request_fingerprint,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Used when the token endpoint omits ``expires_in`` or returns a value the
|
||||
# manager cannot turn into a lifetime.
|
||||
_DEFAULT_TOKEN_LIFETIME_SECONDS = 3600
|
||||
# How often a still-valid token retries storing a rotation the store missed.
|
||||
_UNSAVED_ROTATION_RETRY_SECONDS = 5.0
|
||||
|
||||
|
||||
class McpOAuthReauthorizationRequired(RuntimeError):
|
||||
"""The token endpoint rejected the refresh token (``invalid_grant``); only a new one fixes the connection.
|
||||
|
||||
The message names the server and the remedy, never a token: on the tool
|
||||
interceptor path it is the model-visible tool error the owner reads.
|
||||
"""
|
||||
|
||||
|
||||
class _RefreshTokenRejected(Exception):
|
||||
"""Internal: ``invalid_grant`` for a refresh_token request, before the server name is known."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__("invalid_grant")
|
||||
# Set when this process fell back to the seed because it could not
|
||||
# decrypt the chain a peer stored: a key mismatch, not a lapsed grant.
|
||||
self.stored_chain_unreadable = False
|
||||
|
||||
|
||||
def _is_invalid_grant(response: Any) -> bool:
|
||||
if getattr(response, "status_code", None) not in (400, 401):
|
||||
return False
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError:
|
||||
return False
|
||||
return isinstance(payload, dict) and payload.get("error") == "invalid_grant"
|
||||
|
||||
|
||||
def _reauthorization_message(server_name: str, *, stored_chain_unreadable: bool) -> str:
|
||||
message = f"The OAuth provider rejected the refresh token for MCP server '{server_name}' (invalid_grant). Re-authorize this connection by configuring a new refresh token for it."
|
||||
if stored_chain_unreadable:
|
||||
message += " This Gateway instance could not decrypt the token chain another instance stored, so every instance needs the same DEER_FLOW_CREDENTIALS_KEY; ask the operator to set it."
|
||||
return message
|
||||
|
||||
|
||||
@dataclass
|
||||
@ -37,12 +85,29 @@ class _OAuthState:
|
||||
token: _OAuthToken | None = None
|
||||
# Sync tool wrappers and the Gateway can refresh from different loops.
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
# Shared-store row and token-chain identity; only the refresh_token grant
|
||||
# has a rotating credential worth sharing across processes.
|
||||
store_key: McpOAuthTokenKey | None = None
|
||||
fingerprint: str | None = None
|
||||
request_fingerprint: str | None = None
|
||||
# Set when this process's latest refresh could not be stored: the row
|
||||
# version it refreshed from. Peers still see the consumed refresh token,
|
||||
# so the live one in ``config`` is written back before anything older is
|
||||
# adopted or refreshed with.
|
||||
unsaved_base_version: int | None = None
|
||||
unsaved_retry_at: float = 0.0
|
||||
|
||||
|
||||
class OAuthTokenManager:
|
||||
"""Acquire/cache/refresh OAuth tokens for MCP servers."""
|
||||
|
||||
def __init__(self, oauth_by_server: dict[str, McpOAuthConfig]):
|
||||
def __init__(
|
||||
self,
|
||||
oauth_by_server: dict[str, McpOAuthConfig],
|
||||
*,
|
||||
owner_user_id: str | None = None,
|
||||
token_store: McpOAuthTokenStore | None = None,
|
||||
):
|
||||
# Refresh-token rotation belongs to runtime state, not the parsed config
|
||||
# used by cache/snapshot validation or custom interceptor builders.
|
||||
# A plain threading.Lock, not asyncio.Lock: the embedded/TUI sync tool-call
|
||||
@ -55,7 +120,21 @@ class OAuthTokenManager:
|
||||
# either deadlocks silently or raises "bound to a different event loop".
|
||||
# threading.Lock has no loop affinity, so it is safe to share across
|
||||
# however many event loops/threads call into the same server's lock.
|
||||
self._states = {name: _OAuthState(config=oauth.model_copy(deep=True)) for name, oauth in oauth_by_server.items()}
|
||||
#
|
||||
# ``owner_user_id`` marks a personal connection's manager; ``token_store``
|
||||
# overrides the host-installed store (``oauth_store``) for refresh_token
|
||||
# connections, which then survive restarts and are shared by replicas.
|
||||
self._token_store = token_store
|
||||
self._states = {name: self._new_state(name, oauth, owner_user_id) for name, oauth in oauth_by_server.items()}
|
||||
|
||||
@staticmethod
|
||||
def _new_state(server_name: str, oauth: McpOAuthConfig, owner_user_id: str | None) -> _OAuthState:
|
||||
state = _OAuthState(config=oauth.model_copy(deep=True))
|
||||
if oauth.grant_type == "refresh_token":
|
||||
state.store_key = McpOAuthTokenKey.for_server(server_name, owner_user_id)
|
||||
state.fingerprint = oauth_config_fingerprint(oauth)
|
||||
state.request_fingerprint = oauth_token_request_fingerprint(oauth)
|
||||
return state
|
||||
|
||||
@classmethod
|
||||
def from_extensions_config(
|
||||
@ -63,17 +142,20 @@ class OAuthTokenManager:
|
||||
extensions_config: ExtensionsConfig,
|
||||
*,
|
||||
shared_manager: OAuthTokenManager | None = None,
|
||||
owner_user_id: str | None = None,
|
||||
token_store: McpOAuthTokenStore | None = None,
|
||||
) -> OAuthTokenManager:
|
||||
"""Build a manager, optionally reusing already-validated connection state.
|
||||
|
||||
The durable runtime supplies ``shared_manager`` only after validating
|
||||
the deployment snapshot. Personal connections must never use it.
|
||||
the deployment snapshot. Personal connections must never use it; they
|
||||
pass their ``owner_user_id`` instead, which keys their stored tokens.
|
||||
"""
|
||||
oauth_by_server: dict[str, McpOAuthConfig] = {}
|
||||
for server_name, server_config in extensions_config.get_enabled_mcp_servers().items():
|
||||
if server_config.oauth and server_config.oauth.enabled:
|
||||
oauth_by_server[server_name] = server_config.oauth
|
||||
manager = cls(oauth_by_server)
|
||||
manager = cls(oauth_by_server, owner_user_id=owner_user_id, token_store=token_store)
|
||||
if shared_manager is not None:
|
||||
for name in manager._states.keys() & shared_manager._states.keys():
|
||||
manager._states[name] = shared_manager._states[name]
|
||||
@ -92,7 +174,7 @@ class OAuthTokenManager:
|
||||
|
||||
oauth = state.config
|
||||
token = state.token
|
||||
if token and not self._is_expiring(token, oauth):
|
||||
if token and not self._is_expiring(token, oauth) and state.unsaved_base_version is None:
|
||||
return self._authorization_value(token, server_name)
|
||||
|
||||
lock = state.lock
|
||||
@ -130,17 +212,195 @@ class OAuthTokenManager:
|
||||
lock.release()
|
||||
raise
|
||||
try:
|
||||
store = self._token_store or get_mcp_oauth_token_store()
|
||||
token = state.token
|
||||
if token and not self._is_expiring(token, oauth):
|
||||
if state.unsaved_base_version is not None and store is not None:
|
||||
await self._retry_unsaved_rotation(state, store)
|
||||
return self._authorization_value(token, server_name)
|
||||
|
||||
fresh = await self._fetch_token(oauth)
|
||||
fresh = None
|
||||
try:
|
||||
if store is not None and state.store_key is not None:
|
||||
try:
|
||||
fresh = await self._refresh_shared(state, store, server_name)
|
||||
except McpOAuthTokenKeyUnavailable:
|
||||
logger.warning("Keeping OAuth tokens for MCP server '%s' in this process only: the credentials key cannot be loaded", server_name)
|
||||
if fresh is None:
|
||||
fresh = await self._fetch_token(oauth)
|
||||
logger.info(f"Refreshed OAuth access token for MCP server: {server_name}")
|
||||
except _RefreshTokenRejected as exc:
|
||||
logger.warning("OAuth refresh token for MCP server '%s' was rejected (invalid_grant); the connection needs a new refresh token", server_name)
|
||||
raise McpOAuthReauthorizationRequired(_reauthorization_message(server_name, stored_chain_unreadable=exc.stored_chain_unreadable)) from exc.__cause__
|
||||
state.token = fresh
|
||||
logger.info(f"Refreshed OAuth access token for MCP server: {server_name}")
|
||||
return self._authorization_value(fresh, server_name)
|
||||
finally:
|
||||
lock.release()
|
||||
|
||||
async def _refresh_shared(self, state: _OAuthState, store: McpOAuthTokenStore, server_name: str) -> _OAuthToken:
|
||||
"""Adopt a peer's token from the shared row, or refresh once for every process.
|
||||
|
||||
Runs under the connection's process lock. The row is re-read first,
|
||||
since a peer (another replica, or this process before a restart) may
|
||||
already hold a newer token. A refresh is claimed with a lease plus a
|
||||
compare-and-set on the row's version and never holds a row lock across
|
||||
the token request; the lease is renewed while the request runs, because
|
||||
HTTP timeouts bound each I/O wait rather than the whole request. While
|
||||
a peer holds the lease, the row is polled until its token lands or the
|
||||
lease expires. The refreshed tokens, including a rotated refresh token,
|
||||
are written back under the claim; a rotation the store missed stays
|
||||
pending and is written back before anything older is used.
|
||||
"""
|
||||
key, fingerprint, request_fingerprint, oauth = state.store_key, state.fingerprint, state.request_fingerprint, state.config
|
||||
assert key is not None and fingerprint is not None and request_fingerprint is not None
|
||||
# A waiter must outlast a live holder, which may renew its lease for
|
||||
# up to REFRESH_MAX_SECONDS; a dead holder's lease expires sooner.
|
||||
deadline = time.monotonic() + oauth_store.REFRESH_MAX_SECONDS + oauth_store.REFRESH_LEASE_SECONDS
|
||||
while True:
|
||||
version, stored, unreadable = await asyncio.to_thread(store.load, key, fingerprint=fingerprint, request_fingerprint=request_fingerprint)
|
||||
if state.unsaved_base_version is not None and version != state.unsaved_base_version:
|
||||
# A peer stored a newer chain after this process's rotation.
|
||||
state.unsaved_base_version = None
|
||||
if state.unsaved_base_version is not None:
|
||||
# The row still holds the refresh token this process consumed;
|
||||
# store the live one first, then continue from the row.
|
||||
if await self._store_unsaved_rotation(state, store, version):
|
||||
continue
|
||||
else:
|
||||
if stored is not None:
|
||||
# The row continues this chain; without one (or under another
|
||||
# fingerprint) the configured seed is the refresh token.
|
||||
oauth.refresh_token = stored.refresh_token
|
||||
if stored.access_token and stored.expires_at is not None:
|
||||
token = _OAuthToken(access_token=stored.access_token, token_type=stored.token_type or oauth.default_token_type, expires_at=stored.expires_at)
|
||||
if not self._is_expiring(token, oauth):
|
||||
return token
|
||||
lease_owner = await asyncio.to_thread(store.claim, key, fingerprint=fingerprint, expected_version=version)
|
||||
if lease_owner is not None:
|
||||
break
|
||||
if time.monotonic() >= deadline:
|
||||
raise McpOAuthTokenStoreUnavailable(f"Timed out waiting for another DeerFlow process to refresh the OAuth token for MCP server '{server_name}'")
|
||||
await asyncio.sleep(oauth_store.REFRESH_POLL_SECONDS)
|
||||
|
||||
# A claim on a missing row inserts it at version 0.
|
||||
base_version = 0 if version is None else version
|
||||
used_refresh_token = oauth.refresh_token
|
||||
renewal = asyncio.create_task(self._renew_lease(store, key, lease_owner, server_name), name=f"oauth-lease-renewal:{server_name}")
|
||||
try:
|
||||
async with asyncio.timeout(oauth_store.REFRESH_MAX_SECONDS):
|
||||
fresh = await self._fetch_token(oauth)
|
||||
except asyncio.CancelledError:
|
||||
# A cancelled caller cannot await the release; hand it to a worker
|
||||
# rather than make every peer wait out the lease.
|
||||
asyncio.get_running_loop().run_in_executor(None, self._release_quietly, store, key, lease_owner)
|
||||
raise
|
||||
except Exception as exc:
|
||||
await asyncio.to_thread(self._release_quietly, store, key, lease_owner)
|
||||
if isinstance(exc, _RefreshTokenRejected):
|
||||
exc.stored_chain_unreadable = unreadable
|
||||
if isinstance(exc, TimeoutError):
|
||||
raise RuntimeError(f"The OAuth token request for MCP server '{server_name}' did not finish within {oauth_store.REFRESH_MAX_SECONDS:g} seconds") from None
|
||||
raise
|
||||
finally:
|
||||
renewal.cancel()
|
||||
logger.info(f"Refreshed OAuth access token for MCP server: {server_name}")
|
||||
try:
|
||||
saved = await asyncio.to_thread(
|
||||
store.save,
|
||||
key,
|
||||
lease_owner=lease_owner,
|
||||
fingerprint=fingerprint,
|
||||
request_fingerprint=request_fingerprint,
|
||||
access_token=fresh.access_token,
|
||||
token_type=fresh.token_type,
|
||||
expires_at=fresh.expires_at,
|
||||
refresh_token=oauth.refresh_token or "",
|
||||
rotated=oauth.refresh_token != used_refresh_token,
|
||||
)
|
||||
except McpOAuthTokenStoreUnavailable:
|
||||
# Logged by the store. Free the claim so the pending rotation (or a
|
||||
# peer) need not wait out the lease once the database answers again.
|
||||
saved = False
|
||||
await asyncio.to_thread(self._release_quietly, store, key, lease_owner)
|
||||
if saved:
|
||||
state.unsaved_base_version = None
|
||||
else:
|
||||
# The token is valid for this process, but peers still read the
|
||||
# refresh token it consumed: keep the new one pending.
|
||||
state.unsaved_base_version = base_version
|
||||
state.unsaved_retry_at = time.monotonic() + _UNSAVED_ROTATION_RETRY_SECONDS
|
||||
logger.warning("The new OAuth token for MCP server '%s' could not be stored; this process keeps it and retries storing it before other processes refresh again", server_name)
|
||||
return fresh
|
||||
|
||||
async def _store_unsaved_rotation(self, state: _OAuthState, store: McpOAuthTokenStore, version: int | None) -> bool:
|
||||
"""Write the pending rotation over the row it was refreshed from; ``False`` while a peer holds the lease."""
|
||||
key, oauth, token = state.store_key, state.config, state.token
|
||||
assert key is not None and state.fingerprint is not None and state.request_fingerprint is not None
|
||||
if token is None:
|
||||
state.unsaved_base_version = None
|
||||
return True
|
||||
lease_owner = await asyncio.to_thread(store.claim, key, fingerprint=state.fingerprint, expected_version=version)
|
||||
if lease_owner is None:
|
||||
return False
|
||||
try:
|
||||
saved = await asyncio.to_thread(
|
||||
store.save,
|
||||
key,
|
||||
lease_owner=lease_owner,
|
||||
fingerprint=state.fingerprint,
|
||||
request_fingerprint=state.request_fingerprint,
|
||||
access_token=token.access_token,
|
||||
token_type=token.token_type,
|
||||
expires_at=token.expires_at,
|
||||
refresh_token=oauth.refresh_token or "",
|
||||
rotated=True,
|
||||
)
|
||||
except McpOAuthTokenStoreUnavailable:
|
||||
await asyncio.to_thread(self._release_quietly, store, key, lease_owner)
|
||||
raise
|
||||
if saved:
|
||||
state.unsaved_base_version = None
|
||||
return saved
|
||||
|
||||
async def _retry_unsaved_rotation(self, state: _OAuthState, store: McpOAuthTokenStore) -> None:
|
||||
"""While the local token is still valid, retry storing a missed rotation now and then.
|
||||
|
||||
Until it is stored, a peer whose token expires first would refresh with
|
||||
the consumed refresh token. Best effort: failures wait for the next try.
|
||||
"""
|
||||
if time.monotonic() < state.unsaved_retry_at:
|
||||
return
|
||||
state.unsaved_retry_at = time.monotonic() + _UNSAVED_ROTATION_RETRY_SECONDS
|
||||
key = state.store_key
|
||||
assert key is not None and state.fingerprint is not None and state.request_fingerprint is not None
|
||||
try:
|
||||
version, _stored, _unreadable = await asyncio.to_thread(store.load, key, fingerprint=state.fingerprint, request_fingerprint=state.request_fingerprint)
|
||||
if version != state.unsaved_base_version:
|
||||
state.unsaved_base_version = None
|
||||
return
|
||||
await self._store_unsaved_rotation(state, store, version)
|
||||
except McpOAuthTokenStoreUnavailable:
|
||||
pass # Logged by the store.
|
||||
|
||||
async def _renew_lease(self, store: McpOAuthTokenStore, key: McpOAuthTokenKey, lease_owner: str, server_name: str) -> None:
|
||||
"""Keep the claim's lease alive while its token request runs (cancelled when it ends)."""
|
||||
while True:
|
||||
await asyncio.sleep(oauth_store.REFRESH_LEASE_SECONDS / 3)
|
||||
try:
|
||||
renewed = await asyncio.to_thread(store.renew, key, lease_owner)
|
||||
except McpOAuthTokenStoreUnavailable:
|
||||
continue # Logged by the store; the next tick retries.
|
||||
if not renewed:
|
||||
logger.warning("The refresh lease for MCP server '%s' expired while its token request was running; another process may refresh too", server_name)
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
def _release_quietly(store: McpOAuthTokenStore, key: McpOAuthTokenKey, lease_owner: str) -> None:
|
||||
try:
|
||||
store.release(key, lease_owner)
|
||||
except McpOAuthTokenStoreUnavailable:
|
||||
pass # Logged by the store; the lease expires on its own.
|
||||
|
||||
@staticmethod
|
||||
def _authorization_value(token: _OAuthToken, server_name: str) -> str:
|
||||
"""Render the Authorization value, refusing one the transport would echo.
|
||||
@ -210,6 +470,11 @@ class OAuthTokenManager:
|
||||
|
||||
async with httpx.AsyncClient(timeout=15.0) as client:
|
||||
response = await client.post(oauth.token_url, data=data)
|
||||
if oauth.grant_type == "refresh_token" and _is_invalid_grant(response):
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise _RefreshTokenRejected() from exc
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
|
||||
@ -219,9 +484,10 @@ class OAuthTokenManager:
|
||||
|
||||
# Persist a rotated refresh_token so subsequent refreshes use the latest
|
||||
# value. This updates our private runtime copy only — neither the parsed
|
||||
# extensions config nor 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.
|
||||
# extensions config nor extensions_config.json; the shared token store,
|
||||
# when installed, receives it from ``_refresh_shared``. 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:
|
||||
|
||||
261
backend/packages/harness/deerflow/mcp/oauth_store.py
Normal file
261
backend/packages/harness/deerflow/mcp/oauth_store.py
Normal file
@ -0,0 +1,261 @@
|
||||
"""Encrypted, cross-process storage for rotating MCP OAuth tokens.
|
||||
|
||||
``OAuthTokenManager`` keeps tokens in process memory. For the ``refresh_token``
|
||||
grant that is not enough: providers that rotate refresh tokens (Auth0, Okta,
|
||||
Google) consume the old one on every refresh, so a second Gateway replica, a
|
||||
restart, a tool-cache reset or a new personal-tool manager that starts again
|
||||
from the configured seed gets ``invalid_grant``. This store keeps each
|
||||
connection's latest tokens in ``mcp_oauth_tokens`` (see
|
||||
``deerflow.persistence.mcp_oauth_tokens``), encrypted with the deployment
|
||||
credentials key (``DEER_FLOW_CREDENTIALS_KEY``), and arbitrates refreshes
|
||||
between processes with a short lease.
|
||||
|
||||
The host installs one store per process (:func:`set_mcp_oauth_token_store`);
|
||||
without one, tokens stay in process memory as before. Every method blocks
|
||||
(SQL, possibly the first key-file read): call it through ``asyncio.to_thread``.
|
||||
``client_credentials`` tokens are never stored; each process can mint its own.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Collection
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from typing import Literal, NamedTuple
|
||||
|
||||
from deerflow.config.credentials_key import CredentialsCipher, CredentialsDecryptError
|
||||
from deerflow.config.extensions_config import McpOAuthConfig
|
||||
from deerflow.persistence.mcp_oauth_tokens import SqlMcpOAuthTokenRepository
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: How long a claimed refresh stays reserved without a renewal. The holder
|
||||
#: renews it every third of this while its token request runs, because HTTP
|
||||
#: timeouts bound each I/O wait, not the whole request.
|
||||
REFRESH_LEASE_SECONDS = 30.0
|
||||
#: The longest a token request may hold the lease; it is abandoned after this.
|
||||
REFRESH_MAX_SECONDS = 120.0
|
||||
#: How often a process waiting on a peer's refresh re-reads the row.
|
||||
REFRESH_POLL_SECONDS = 0.2
|
||||
|
||||
|
||||
class McpOAuthTokenStoreUnavailable(RuntimeError):
|
||||
"""The store cannot be read or written; refreshing anyway could consume a token nobody keeps."""
|
||||
|
||||
|
||||
class McpOAuthTokenKeyUnavailable(McpOAuthTokenStoreUnavailable):
|
||||
"""The credentials key cannot be loaded, so nothing can be stored encrypted."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class McpOAuthTokenKey:
|
||||
"""Row key: personal connections never share a row with a same-named deployment one."""
|
||||
|
||||
scope: Literal["deployment", "personal"]
|
||||
owner_user_id: str
|
||||
server_key: str
|
||||
|
||||
@classmethod
|
||||
def for_server(cls, server_name: str, owner_user_id: str | None) -> McpOAuthTokenKey:
|
||||
if owner_user_id is None:
|
||||
return cls(scope="deployment", owner_user_id="", server_key=server_name)
|
||||
return cls(scope="personal", owner_user_id=owner_user_id, server_key=server_name)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PersistedOAuthToken:
|
||||
"""A usable stored token chain: written under the caller's fingerprint and decrypted."""
|
||||
|
||||
access_token: str | None
|
||||
token_type: str | None
|
||||
expires_at: datetime | None
|
||||
refresh_token: str
|
||||
|
||||
|
||||
class StoredTokenLookup(NamedTuple):
|
||||
"""A row read for one token chain: ``version`` is ``None`` without a row; ``token`` is ``None`` when unusable."""
|
||||
|
||||
version: int | None
|
||||
token: PersistedOAuthToken | None
|
||||
#: The row matched this chain but no configured key could decrypt it.
|
||||
undecryptable: bool = False
|
||||
|
||||
|
||||
def _server_subject(key: McpOAuthTokenKey) -> str:
|
||||
return f"MCP server '{key.server_key}'"
|
||||
|
||||
|
||||
def oauth_config_fingerprint(oauth: McpOAuthConfig) -> str:
|
||||
"""Identify the configured token chain, so a changed seed wins over a stored rotation.
|
||||
|
||||
Covers what makes an issued refresh token valid: endpoint, client, grant,
|
||||
scope, audience and the configured seed. ``client_secret`` is left out
|
||||
because rotating it keeps the issued chain valid. Only the SHA-256 digest
|
||||
is stored, never the seed itself.
|
||||
"""
|
||||
identity = [oauth.token_url, oauth.client_id, oauth.grant_type, oauth.scope, oauth.audience, oauth.refresh_token]
|
||||
return hashlib.sha256(json.dumps(identity, separators=(",", ":")).encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def oauth_token_request_fingerprint(oauth: McpOAuthConfig) -> str:
|
||||
"""Identify how an access token is requested and read, beyond the token chain.
|
||||
|
||||
``extra_token_params`` (for example a ``resource`` or provider-specific
|
||||
audience) select what the issued access token is for, and the response
|
||||
fields select which value is used. A stored access token is reused only
|
||||
by a process whose fingerprint matches; otherwise it refreshes from the
|
||||
stored refresh token, which stays valid across such changes.
|
||||
"""
|
||||
identity = [sorted(oauth.extra_token_params.items()), oauth.token_field, oauth.token_type_field, oauth.expires_in_field, oauth.default_token_type]
|
||||
return hashlib.sha256(json.dumps(identity, separators=(",", ":")).encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
class McpOAuthTokenStore:
|
||||
"""Encrypt and arbitrate one process's view of ``mcp_oauth_tokens``.
|
||||
|
||||
``cipher`` defaults to the process-wide credentials cipher, loaded on first
|
||||
use so deployments without OAuth connections never create a key file.
|
||||
"""
|
||||
|
||||
def __init__(self, repository: SqlMcpOAuthTokenRepository, *, cipher: CredentialsCipher | None = None) -> None:
|
||||
self._repo = repository
|
||||
self._explicit_cipher = cipher
|
||||
|
||||
def _cipher(self) -> CredentialsCipher:
|
||||
if self._explicit_cipher is not None:
|
||||
return self._explicit_cipher
|
||||
from deerflow.config.credentials_key import get_credentials_cipher
|
||||
|
||||
try:
|
||||
return get_credentials_cipher()
|
||||
except Exception:
|
||||
logger.exception("Failed to load the credentials key for MCP OAuth tokens")
|
||||
raise McpOAuthTokenKeyUnavailable("The credentials key for stored MCP OAuth tokens cannot be loaded") from None
|
||||
|
||||
def _call(self, subject: str, action: str, operation: Callable[[], object]):
|
||||
try:
|
||||
return operation()
|
||||
except McpOAuthTokenStoreUnavailable:
|
||||
raise
|
||||
except Exception:
|
||||
# Database errors render the statement and its parameters; keep
|
||||
# them in the log, never in the model-visible tool error.
|
||||
logger.exception("Failed to %s the stored OAuth tokens for %s", action, subject)
|
||||
raise McpOAuthTokenStoreUnavailable(f"OAuth token storage is unavailable for {subject}; try again later") from None
|
||||
|
||||
def load(self, key: McpOAuthTokenKey, *, fingerprint: str, request_fingerprint: str) -> StoredTokenLookup:
|
||||
"""Read the chain's row: no row has no version; a row this chain cannot use has no token.
|
||||
|
||||
The access token is withheld (``None``) when it was requested under
|
||||
another ``request_fingerprint``; the refresh token is still returned.
|
||||
"""
|
||||
cipher = self._cipher()
|
||||
row = self._call(_server_subject(key), "read", lambda: self._repo.get(key.scope, key.owner_user_id, key.server_key))
|
||||
if row is None:
|
||||
return StoredTokenLookup(None, None)
|
||||
if row.config_fingerprint != fingerprint or row.refresh_token_enc is None:
|
||||
# Written for another seed/endpoint, or claimed but never filled.
|
||||
return StoredTokenLookup(row.version, None)
|
||||
try:
|
||||
refresh_token = cipher.decrypt_text(row.refresh_token_enc)
|
||||
reusable = row.access_token_enc is not None and row.token_request_fingerprint == request_fingerprint
|
||||
access_token = cipher.decrypt_text(row.access_token_enc) if reusable else None
|
||||
except CredentialsDecryptError:
|
||||
logger.warning(
|
||||
"Stored OAuth token for MCP server '%s' cannot be decrypted with the configured credentials key(s); using the configured refresh token. Every instance needs the same DEER_FLOW_CREDENTIALS_KEY.",
|
||||
key.server_key,
|
||||
)
|
||||
return StoredTokenLookup(row.version, None, undecryptable=True)
|
||||
expires_at = datetime.fromtimestamp(row.access_expires_at, UTC) if row.access_expires_at is not None else None
|
||||
return StoredTokenLookup(row.version, PersistedOAuthToken(access_token=access_token, token_type=row.token_type, expires_at=expires_at, refresh_token=refresh_token))
|
||||
|
||||
def claim(self, key: McpOAuthTokenKey, *, fingerprint: str, expected_version: int | None) -> str | None:
|
||||
"""Claim the next refresh; returns the lease token, or ``None`` while a peer holds it or wrote since the read."""
|
||||
lease_owner = str(uuid.uuid4())
|
||||
claimed = self._call(
|
||||
_server_subject(key),
|
||||
"claim",
|
||||
lambda: self._repo.claim(
|
||||
key.scope,
|
||||
key.owner_user_id,
|
||||
key.server_key,
|
||||
config_fingerprint=fingerprint,
|
||||
expected_version=expected_version,
|
||||
lease_owner=lease_owner,
|
||||
lease_seconds=REFRESH_LEASE_SECONDS,
|
||||
now=time.time(),
|
||||
),
|
||||
)
|
||||
return lease_owner if claimed else None
|
||||
|
||||
def save(
|
||||
self,
|
||||
key: McpOAuthTokenKey,
|
||||
*,
|
||||
lease_owner: str,
|
||||
fingerprint: str,
|
||||
request_fingerprint: str,
|
||||
access_token: str,
|
||||
token_type: str,
|
||||
expires_at: datetime,
|
||||
refresh_token: str,
|
||||
rotated: bool,
|
||||
) -> bool:
|
||||
"""Store a refreshed token under the claim; ``False`` when the lease was lost meanwhile."""
|
||||
cipher = self._cipher()
|
||||
access_token_enc = cipher.encrypt_text(access_token)
|
||||
refresh_token_enc = cipher.encrypt_text(refresh_token)
|
||||
return self._call(
|
||||
_server_subject(key),
|
||||
"write",
|
||||
lambda: self._repo.write(
|
||||
key.scope,
|
||||
key.owner_user_id,
|
||||
key.server_key,
|
||||
lease_owner=lease_owner,
|
||||
config_fingerprint=fingerprint,
|
||||
access_token_enc=access_token_enc,
|
||||
token_request_fingerprint=request_fingerprint,
|
||||
token_type=token_type,
|
||||
access_expires_at=expires_at.timestamp(),
|
||||
refresh_token_enc=refresh_token_enc,
|
||||
rotated=rotated,
|
||||
now=time.time(),
|
||||
),
|
||||
)
|
||||
|
||||
def renew(self, key: McpOAuthTokenKey, lease_owner: str) -> bool:
|
||||
"""Extend the claim's lease by ``REFRESH_LEASE_SECONDS``; ``False`` once it is no longer held."""
|
||||
return self._call(_server_subject(key), "renew", lambda: self._repo.renew(key.scope, key.owner_user_id, key.server_key, lease_owner=lease_owner, lease_seconds=REFRESH_LEASE_SECONDS))
|
||||
|
||||
def release(self, key: McpOAuthTokenKey, lease_owner: str) -> bool:
|
||||
return self._call(_server_subject(key), "release", lambda: self._repo.release(key.scope, key.owner_user_id, key.server_key, lease_owner=lease_owner))
|
||||
|
||||
def forget_unconfigured(self, owner_user_id: str | None, configured_server_keys: Collection[str]) -> int:
|
||||
"""Delete the rows of one owner's connections (``None``: the deployment's) that a saved configuration no longer has.
|
||||
|
||||
A removed connection's key can never be derived again, so its row would
|
||||
only keep refresh-token ciphertext around. Returns how many were deleted.
|
||||
"""
|
||||
scope_key = McpOAuthTokenKey.for_server("", owner_user_id)
|
||||
subject = "deployment MCP servers" if owner_user_id is None else f"personal MCP connections of user '{owner_user_id}'"
|
||||
return self._call(subject, "prune", lambda: self._repo.delete_except(scope_key.scope, scope_key.owner_user_id, configured_server_keys))
|
||||
|
||||
|
||||
_store: McpOAuthTokenStore | None = None
|
||||
|
||||
|
||||
def set_mcp_oauth_token_store(store: McpOAuthTokenStore | None) -> McpOAuthTokenStore | None:
|
||||
"""Install the host's store for managers built without one, returning the previous store."""
|
||||
global _store
|
||||
previous, _store = _store, store
|
||||
return previous
|
||||
|
||||
|
||||
def get_mcp_oauth_token_store() -> McpOAuthTokenStore | None:
|
||||
return _store
|
||||
@ -167,7 +167,7 @@ class McpTaskToolCaller:
|
||||
if caller is None:
|
||||
# A personal name can equal a deployment name. Never consult
|
||||
# the deployment OAuth state or startup snapshot for this path.
|
||||
caller = McpTaskToolCaller(snapshot.config, oauth_token_manager=OAuthTokenManager.from_extensions_config(snapshot.config))
|
||||
caller = McpTaskToolCaller(snapshot.config, oauth_token_manager=OAuthTokenManager.from_extensions_config(snapshot.config, owner_user_id=user_id))
|
||||
self._personal_callers[user_id] = (snapshot, caller)
|
||||
return caller
|
||||
|
||||
|
||||
@ -1042,6 +1042,7 @@ class _PreparedServerDiscovery:
|
||||
bindings: dict[str, ServerBinding]
|
||||
ownership_domain: MCPPoolDomain
|
||||
connection_scope: str
|
||||
personal_user_id: str | None = None
|
||||
|
||||
|
||||
def _prepare_server_client(
|
||||
@ -1321,6 +1322,7 @@ def _prepare_server_discovery(
|
||||
bindings=server_bindings,
|
||||
ownership_domain=ownership_domain,
|
||||
connection_scope="personal" if personal_user_id is not None else "deployment",
|
||||
personal_user_id=personal_user_id,
|
||||
)
|
||||
|
||||
|
||||
@ -1331,7 +1333,7 @@ async def _discover_and_wrap_prepared(
|
||||
"""Run the OAuth/discovery phase, then wrap successful results in order."""
|
||||
client_cls, load_mcp_tools_fn = adapter_symbols
|
||||
extensions_config = prepared.config
|
||||
oauth_token_manager = OAuthTokenManager.from_extensions_config(extensions_config) if prepared.connection_scope == "personal" else get_mcp_task_oauth_token_manager(extensions_config)
|
||||
oauth_token_manager = OAuthTokenManager.from_extensions_config(extensions_config, owner_user_id=prepared.personal_user_id) if prepared.connection_scope == "personal" else get_mcp_task_oauth_token_manager(extensions_config)
|
||||
|
||||
# Only selected, valid connections reach this point. Selected-but-invalid
|
||||
# servers are intentionally not contacted for OAuth headers (F1 contract).
|
||||
|
||||
@ -44,6 +44,15 @@ def personal_server_name(user_id: str, name: str, configuration: dict[str, Any])
|
||||
return PERSONAL_SERVER_PREFIX + hashlib.sha256(identity.encode()).hexdigest()[:32]
|
||||
|
||||
|
||||
def personal_server_name_variants(user_id: str, name: str, configuration: dict[str, Any]) -> set[str]:
|
||||
"""Runtime names one saved connection takes as stored and with ``enabled`` toggled either way.
|
||||
|
||||
Toggling a connection off and on again restores its earlier runtime name,
|
||||
so state keyed by that name (stored OAuth tokens) must survive the toggle.
|
||||
"""
|
||||
return {personal_server_name(user_id, name, variant) for variant in (configuration, {**configuration, "enabled": True}, {**configuration, "enabled": False})}
|
||||
|
||||
|
||||
def load_user_mcp_config(user_id: str) -> ExtensionsConfig:
|
||||
"""Read literal user values; personal files cannot resolve host $ENV secrets."""
|
||||
raw = read_user_mcp_config(user_id)
|
||||
|
||||
@ -0,0 +1,10 @@
|
||||
"""MCP OAuth tokens shared by every Gateway instance (``mcp_oauth_tokens``)."""
|
||||
|
||||
from deerflow.persistence.mcp_oauth_tokens.model import McpOAuthTokenRow, StoredMcpOAuthToken
|
||||
from deerflow.persistence.mcp_oauth_tokens.sql import SqlMcpOAuthTokenRepository
|
||||
|
||||
__all__ = [
|
||||
"McpOAuthTokenRow",
|
||||
"SqlMcpOAuthTokenRepository",
|
||||
"StoredMcpOAuthToken",
|
||||
]
|
||||
@ -0,0 +1,80 @@
|
||||
"""ORM model for MCP OAuth tokens shared by every Gateway instance (``mcp_oauth_tokens``).
|
||||
|
||||
One row per OAuth connection that uses the ``refresh_token`` grant, keyed by
|
||||
``(scope, owner_user_id, server_key)``:
|
||||
|
||||
- ``scope='deployment'``: ``owner_user_id`` is ``''`` and ``server_key`` the
|
||||
server name in ``extensions_config.json``. The empty string stands in for
|
||||
"no owner" because primary-key columns cannot be NULL on PostgreSQL, and
|
||||
SQLite would treat NULL owners as distinct keys.
|
||||
- ``scope='personal'``: ``owner_user_id`` is the owner and ``server_key`` the
|
||||
``personal_<hash>`` runtime name (``deerflow.mcp.user_config``), so a
|
||||
personal connection never shares a row with a same-named deployment one.
|
||||
|
||||
The row holds only what the token endpoint returned -- the access token, its
|
||||
expiry and the latest (possibly rotated) refresh token -- as
|
||||
``DEER_FLOW_CREDENTIALS_KEY`` ciphertext. ``config_fingerprint`` identifies the
|
||||
configured token chain (``deerflow.mcp.oauth_store.oauth_config_fingerprint``):
|
||||
a row written under another fingerprint is ignored, so a changed seed wins.
|
||||
``token_request_fingerprint`` identifies how the stored access token was
|
||||
requested and read (``extra_token_params`` such as ``resource``, and the
|
||||
response fields): a process configured differently refreshes from the stored
|
||||
refresh token instead of reusing that access token.
|
||||
|
||||
``version`` grows by one on every token write; a refresh is claimed with a
|
||||
compare-and-set on it plus the ``refresh_lease_*`` columns, so peers that read
|
||||
the same expired token refresh it once. A claim before the first write leaves
|
||||
a row at version 0 with no tokens. Times are epoch seconds: the lease against
|
||||
the database clock, the rest from the writing instance.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sqlalchemy import BigInteger, Double, PrimaryKeyConstraint, String, Text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from deerflow.persistence.base import Base
|
||||
|
||||
#: ``'deployment'`` or ``'personal'``; plain text so a new scope needs no migration.
|
||||
MCP_OAUTH_TOKEN_SCOPE_LENGTH = 16
|
||||
MCP_OAUTH_TOKEN_OWNER_LENGTH = 64
|
||||
#: Deployment server names are operator-chosen; personal runtime names are 41 characters.
|
||||
MCP_OAUTH_TOKEN_SERVER_KEY_LENGTH = 255
|
||||
MCP_OAUTH_TOKEN_FINGERPRINT_LENGTH = 64
|
||||
MCP_OAUTH_TOKEN_LEASE_OWNER_LENGTH = 36
|
||||
|
||||
|
||||
class McpOAuthTokenRow(Base):
|
||||
__tablename__ = "mcp_oauth_tokens"
|
||||
|
||||
scope: Mapped[str] = mapped_column(String(MCP_OAUTH_TOKEN_SCOPE_LENGTH), nullable=False)
|
||||
owner_user_id: Mapped[str] = mapped_column(String(MCP_OAUTH_TOKEN_OWNER_LENGTH), nullable=False)
|
||||
server_key: Mapped[str] = mapped_column(String(MCP_OAUTH_TOKEN_SERVER_KEY_LENGTH), nullable=False)
|
||||
config_fingerprint: Mapped[str] = mapped_column(String(MCP_OAUTH_TOKEN_FINGERPRINT_LENGTH), nullable=False)
|
||||
access_token_enc: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
token_request_fingerprint: Mapped[str | None] = mapped_column(String(MCP_OAUTH_TOKEN_FINGERPRINT_LENGTH), nullable=True)
|
||||
token_type: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
access_expires_at: Mapped[float | None] = mapped_column(Double, nullable=True)
|
||||
refresh_token_enc: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
version: Mapped[int] = mapped_column(BigInteger, nullable=False)
|
||||
refresh_lease_owner: Mapped[str | None] = mapped_column(String(MCP_OAUTH_TOKEN_LEASE_OWNER_LENGTH), nullable=True)
|
||||
refresh_lease_until: Mapped[float | None] = mapped_column(Double, nullable=True)
|
||||
rotated_at: Mapped[float | None] = mapped_column(Double, nullable=True)
|
||||
updated_at: Mapped[float] = mapped_column(Double, nullable=False)
|
||||
|
||||
__table_args__ = (PrimaryKeyConstraint("scope", "owner_user_id", "server_key", name="pk_mcp_oauth_tokens"),)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StoredMcpOAuthToken:
|
||||
"""One connection's row as stored: the token columns are opaque ciphertext here."""
|
||||
|
||||
config_fingerprint: str
|
||||
access_token_enc: str | None
|
||||
token_request_fingerprint: str | None
|
||||
token_type: str | None
|
||||
access_expires_at: float | None
|
||||
refresh_token_enc: str | None
|
||||
version: int
|
||||
@ -0,0 +1,189 @@
|
||||
"""Synchronous SQL repository for ``mcp_oauth_tokens`` (SQLite or PostgreSQL).
|
||||
|
||||
MCP OAuth headers are requested from whichever event loop runs the tool call:
|
||||
the Gateway loop, a sync tool wrapper's fresh loop on a worker thread, or the
|
||||
durable task runtime. The async engine is bound to the Gateway loop, so this
|
||||
repository uses a synchronous session factory over the application database
|
||||
(``deerflow.persistence.agents.sql.get_sync_sessionmaker``) and every call is
|
||||
made through ``asyncio.to_thread``.
|
||||
|
||||
A refresh never holds a row lock across the token request (up to 15 s).
|
||||
Instead it is claimed and fenced in single statements:
|
||||
|
||||
- ``claim`` takes the refresh lease with a compare-and-set: ``UPDATE ... WHERE
|
||||
version = :seen AND (refresh_lease_owner IS NULL OR refresh_lease_until <
|
||||
now)``. A peer that wrote a token since the caller's read changed
|
||||
``version``, so the claim fails and the caller re-reads instead of
|
||||
refreshing with a consumed refresh token. A missing row is claimed by
|
||||
inserting it at version 0 with ``ON CONFLICT DO NOTHING``.
|
||||
- ``renew`` extends a live lease while the claim's token request runs: HTTP
|
||||
timeouts bound each I/O wait, not the whole request, so a slow response must
|
||||
not let a peer take over and resend the same refresh token.
|
||||
- ``write`` stores the new tokens, bumps ``version`` and clears the lease, only
|
||||
while ``refresh_lease_owner`` is still the claim's own token.
|
||||
- ``release`` clears the lease under the same fence after a failed refresh.
|
||||
- ``delete_except`` drops one owner's rows for connections a saved
|
||||
configuration no longer has; their keys can never be derived again.
|
||||
|
||||
The lease is judged against the database clock (``db_now``), never an
|
||||
instance's wall clock. Outcomes come from the affected-row count, except the
|
||||
insert, which counts ``RETURNING`` rows (psycopg reports rowcount -1 for ``ON
|
||||
CONFLICT DO NOTHING``). Ciphertext is opaque here; encryption lives in
|
||||
``deerflow.mcp.oauth_store``. Database errors propagate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Collection
|
||||
|
||||
from sqlalchemy import delete, or_, update
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from deerflow.persistence.channel_leases.sql import db_now
|
||||
from deerflow.persistence.mcp_oauth_tokens.model import McpOAuthTokenRow, StoredMcpOAuthToken
|
||||
|
||||
|
||||
def _dialect(session: Session) -> str:
|
||||
name = session.get_bind().dialect.name
|
||||
if name not in ("postgresql", "sqlite"):
|
||||
raise ValueError(f"Unsupported MCP OAuth token database dialect: {name}")
|
||||
return name
|
||||
|
||||
|
||||
def _key(scope: str, owner_user_id: str, server_key: str):
|
||||
return (McpOAuthTokenRow.scope == scope, McpOAuthTokenRow.owner_user_id == owner_user_id, McpOAuthTokenRow.server_key == server_key)
|
||||
|
||||
|
||||
class SqlMcpOAuthTokenRepository:
|
||||
"""Persistence facade for ``mcp_oauth_tokens``; every call blocks, so run it off the event loop."""
|
||||
|
||||
def __init__(self, session_factory: sessionmaker[Session]) -> None:
|
||||
self._sf = session_factory
|
||||
|
||||
def get(self, scope: str, owner_user_id: str, server_key: str) -> StoredMcpOAuthToken | None:
|
||||
with self._sf() as session:
|
||||
row = session.get(McpOAuthTokenRow, (scope, owner_user_id, server_key))
|
||||
if row is None:
|
||||
return None
|
||||
return StoredMcpOAuthToken(
|
||||
config_fingerprint=row.config_fingerprint,
|
||||
access_token_enc=row.access_token_enc,
|
||||
token_request_fingerprint=row.token_request_fingerprint,
|
||||
token_type=row.token_type,
|
||||
access_expires_at=row.access_expires_at,
|
||||
refresh_token_enc=row.refresh_token_enc,
|
||||
version=int(row.version),
|
||||
)
|
||||
|
||||
def claim(
|
||||
self,
|
||||
scope: str,
|
||||
owner_user_id: str,
|
||||
server_key: str,
|
||||
*,
|
||||
config_fingerprint: str,
|
||||
expected_version: int | None,
|
||||
lease_owner: str,
|
||||
lease_seconds: float,
|
||||
now: float,
|
||||
) -> bool:
|
||||
"""Take the refresh lease if the row is still at ``expected_version`` (``None``: no row yet)."""
|
||||
with self._sf() as session:
|
||||
clock = db_now(_dialect(session))
|
||||
if expected_version is None:
|
||||
insert = pg_insert if _dialect(session) == "postgresql" else sqlite_insert
|
||||
stmt = (
|
||||
insert(McpOAuthTokenRow)
|
||||
.values(
|
||||
scope=scope,
|
||||
owner_user_id=owner_user_id,
|
||||
server_key=server_key,
|
||||
config_fingerprint=config_fingerprint,
|
||||
version=0,
|
||||
refresh_lease_owner=lease_owner,
|
||||
refresh_lease_until=clock + lease_seconds,
|
||||
updated_at=now,
|
||||
)
|
||||
.on_conflict_do_nothing(index_elements=[McpOAuthTokenRow.scope, McpOAuthTokenRow.owner_user_id, McpOAuthTokenRow.server_key])
|
||||
.returning(McpOAuthTokenRow.version)
|
||||
)
|
||||
claimed = len(session.execute(stmt).all()) == 1
|
||||
else:
|
||||
stmt = (
|
||||
update(McpOAuthTokenRow)
|
||||
.where(
|
||||
*_key(scope, owner_user_id, server_key),
|
||||
McpOAuthTokenRow.version == expected_version,
|
||||
or_(McpOAuthTokenRow.refresh_lease_owner.is_(None), McpOAuthTokenRow.refresh_lease_until < clock),
|
||||
)
|
||||
.values(refresh_lease_owner=lease_owner, refresh_lease_until=clock + lease_seconds)
|
||||
)
|
||||
claimed = session.execute(stmt).rowcount == 1
|
||||
session.commit()
|
||||
return claimed
|
||||
|
||||
def write(
|
||||
self,
|
||||
scope: str,
|
||||
owner_user_id: str,
|
||||
server_key: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
config_fingerprint: str,
|
||||
access_token_enc: str,
|
||||
token_request_fingerprint: str,
|
||||
token_type: str,
|
||||
access_expires_at: float,
|
||||
refresh_token_enc: str,
|
||||
rotated: bool,
|
||||
now: float,
|
||||
) -> bool:
|
||||
"""Store a refreshed token and release the lease; ``False`` when the lease is no longer this claim's."""
|
||||
values = {
|
||||
"config_fingerprint": config_fingerprint,
|
||||
"access_token_enc": access_token_enc,
|
||||
"token_request_fingerprint": token_request_fingerprint,
|
||||
"token_type": token_type,
|
||||
"access_expires_at": access_expires_at,
|
||||
"refresh_token_enc": refresh_token_enc,
|
||||
"version": McpOAuthTokenRow.version + 1,
|
||||
"refresh_lease_owner": None,
|
||||
"refresh_lease_until": None,
|
||||
"updated_at": now,
|
||||
}
|
||||
if rotated:
|
||||
values["rotated_at"] = now
|
||||
with self._sf() as session:
|
||||
stmt = update(McpOAuthTokenRow).where(*_key(scope, owner_user_id, server_key), McpOAuthTokenRow.refresh_lease_owner == lease_owner).values(**values)
|
||||
written = session.execute(stmt).rowcount == 1
|
||||
session.commit()
|
||||
return written
|
||||
|
||||
def renew(self, scope: str, owner_user_id: str, server_key: str, *, lease_owner: str, lease_seconds: float) -> bool:
|
||||
"""Extend the claim's lease while it is still live; ``False`` once it expired or changed hands."""
|
||||
with self._sf() as session:
|
||||
clock = db_now(_dialect(session))
|
||||
stmt = update(McpOAuthTokenRow).where(*_key(scope, owner_user_id, server_key), McpOAuthTokenRow.refresh_lease_owner == lease_owner, McpOAuthTokenRow.refresh_lease_until > clock).values(refresh_lease_until=clock + lease_seconds)
|
||||
renewed = session.execute(stmt).rowcount == 1
|
||||
session.commit()
|
||||
return renewed
|
||||
|
||||
def release(self, scope: str, owner_user_id: str, server_key: str, *, lease_owner: str) -> bool:
|
||||
"""Give up a claimed lease without writing a token."""
|
||||
with self._sf() as session:
|
||||
stmt = update(McpOAuthTokenRow).where(*_key(scope, owner_user_id, server_key), McpOAuthTokenRow.refresh_lease_owner == lease_owner).values(refresh_lease_owner=None, refresh_lease_until=None)
|
||||
released = session.execute(stmt).rowcount == 1
|
||||
session.commit()
|
||||
return released
|
||||
|
||||
def delete_except(self, scope: str, owner_user_id: str, keep_server_keys: Collection[str]) -> int:
|
||||
"""Delete ``(scope, owner_user_id)``'s rows whose ``server_key`` is not kept; returns how many."""
|
||||
stmt = delete(McpOAuthTokenRow).where(McpOAuthTokenRow.scope == scope, McpOAuthTokenRow.owner_user_id == owner_user_id)
|
||||
if keep_server_keys:
|
||||
stmt = stmt.where(McpOAuthTokenRow.server_key.not_in(sorted(keep_server_keys)))
|
||||
with self._sf() as session:
|
||||
deleted = session.execute(stmt).rowcount
|
||||
session.commit()
|
||||
return int(deleted)
|
||||
@ -34,7 +34,7 @@ The empty-DB path keeps using `create_all` because `Base.metadata` is the only a
|
||||
`0039_user_disabled` → `0040_scheduled_group_notice` →
|
||||
`0041_im_notice_preferences` → `0042_channel_leases` →
|
||||
`0043_runtime_im_notices` → `0044_channel_runtime_configs` →
|
||||
`0045_completed_run_evidence` → `0046_skill_mutations` (current head). The preference
|
||||
`0045_completed_run_evidence` → `0046_skill_mutations` → `0047_mcp_oauth_tokens` (current head). The preference
|
||||
revision adds a separate owner/key table with a cascading users foreign key and
|
||||
does not alter users; the project-documents revision adds a new owner-scoped
|
||||
shelf table, and the MCP lease-token revision adds two nullable token columns to
|
||||
@ -202,4 +202,5 @@ reenabling a provider never reopens notifications canceled before downgrade.
|
||||
Tests: `test_migration_0041_im_notice_preferences.py`.
|
||||
|
||||
- `0045_completed_run_evidence` adds nullable evidence sealing fields and completed-run snapshots after `0044_channel_runtime_configs`.
|
||||
- `0046_skill_mutations` adds durable owner enrollment, asset revisions, proposals, publication operations and scan leases. The single-head pin lives in `tests/test_migration_0046_skill_mutations.py`.
|
||||
- `0046_skill_mutations` adds durable owner enrollment, asset revisions, proposals, publication operations and scan leases.
|
||||
- `0047_mcp_oauth_tokens` creates the table-guarded `mcp_oauth_tokens` (PK `scope`, `owner_user_id` (`''` for deployment), `server_key`; encrypted token columns, chain and request fingerprints, `version`, refresh lease) with no backfill. The single-head pin lives in `tests/test_migration_0047_mcp_oauth_tokens.py`.
|
||||
|
||||
@ -0,0 +1,69 @@
|
||||
"""Shared, encrypted MCP OAuth tokens.
|
||||
|
||||
Revision ID: 0047_mcp_oauth_tokens
|
||||
Revises: 0046_skill_mutations
|
||||
|
||||
Creates ``mcp_oauth_tokens``: one row per MCP OAuth connection that uses the
|
||||
``refresh_token`` grant, keyed by ``(scope, owner_user_id, server_key)``
|
||||
(``owner_user_id`` is ``''`` for deployment servers). It holds the access
|
||||
token and the latest, possibly rotated, refresh token as
|
||||
``DEER_FLOW_CREDENTIALS_KEY`` ciphertext, the fingerprints of the configured
|
||||
token chain and of the access-token request, a ``version`` bumped on every
|
||||
token write, and the refresh lease
|
||||
that lets one Gateway process refresh at a time
|
||||
(``deerflow.mcp.oauth_store``). Tokens used to live only in process memory,
|
||||
so this revision moves no data: every connection starts from its configured
|
||||
seed once, then continues from the stored chain.
|
||||
|
||||
``_helpers.py`` only guards columns, so the table is guarded with a fresh
|
||||
inspector in both directions: an upgrade that finds the table (an interrupted
|
||||
earlier attempt, or an empty-database bootstrap that already ran
|
||||
``create_all``) completes, and a downgrade after the table is gone still
|
||||
moves the version row.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision: str = "0047_mcp_oauth_tokens"
|
||||
down_revision: str | Sequence[str] | None = "0046_skill_mutations"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_TABLE = "mcp_oauth_tokens"
|
||||
|
||||
|
||||
def _has_table() -> bool:
|
||||
return sa.inspect(op.get_bind()).has_table(_TABLE)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if _has_table():
|
||||
return
|
||||
op.create_table(
|
||||
_TABLE,
|
||||
sa.Column("scope", sa.String(length=16), nullable=False),
|
||||
sa.Column("owner_user_id", sa.String(length=64), nullable=False),
|
||||
sa.Column("server_key", sa.String(length=255), nullable=False),
|
||||
sa.Column("config_fingerprint", sa.String(length=64), nullable=False),
|
||||
sa.Column("access_token_enc", sa.Text(), nullable=True),
|
||||
sa.Column("token_request_fingerprint", sa.String(length=64), nullable=True),
|
||||
sa.Column("token_type", sa.String(length=64), nullable=True),
|
||||
sa.Column("access_expires_at", sa.Double(), nullable=True),
|
||||
sa.Column("refresh_token_enc", sa.Text(), nullable=True),
|
||||
sa.Column("version", sa.BigInteger(), nullable=False),
|
||||
sa.Column("refresh_lease_owner", sa.String(length=36), nullable=True),
|
||||
sa.Column("refresh_lease_until", sa.Double(), nullable=True),
|
||||
sa.Column("rotated_at", sa.Double(), nullable=True),
|
||||
sa.Column("updated_at", sa.Double(), nullable=False),
|
||||
sa.PrimaryKeyConstraint("scope", "owner_user_id", "server_key", name="pk_mcp_oauth_tokens"),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _has_table():
|
||||
op.drop_table(_TABLE)
|
||||
@ -27,6 +27,7 @@ from deerflow.persistence.channel_thread_bindings.model import ChannelThreadBind
|
||||
from deerflow.persistence.feedback.model import FeedbackRow
|
||||
from deerflow.persistence.login_throttle.model import LoginThrottleRow
|
||||
from deerflow.persistence.managed_subagents.model import ManagedSubagentRow
|
||||
from deerflow.persistence.mcp_oauth_tokens.model import McpOAuthTokenRow
|
||||
from deerflow.persistence.mcp_tasks.model import McpTaskRow
|
||||
from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow
|
||||
from deerflow.persistence.notification_deliveries.model import NotificationDeliveryRow
|
||||
@ -55,6 +56,7 @@ __all__ = [
|
||||
"ChannelThreadBindingRow",
|
||||
"FeedbackRow",
|
||||
"LoginThrottleRow",
|
||||
"McpOAuthTokenRow",
|
||||
"McpTaskRow",
|
||||
"ManagedSubagentRow",
|
||||
"NotificationDeliveryRow",
|
||||
|
||||
101
backend/tests/blocking_io/test_mcp_oauth_token_store.py
Normal file
101
backend/tests/blocking_io/test_mcp_oauth_token_store.py
Normal file
@ -0,0 +1,101 @@
|
||||
"""Regression anchor: MCP OAuth refreshes through the shared token store must not block the event loop.
|
||||
|
||||
``OAuthTokenManager.get_authorization_header`` runs on whichever loop calls the
|
||||
tool. With the Gateway's ``McpOAuthTokenStore`` installed, a refresh reads the
|
||||
connection's ``mcp_oauth_tokens`` row, claims the refresh lease, encrypts and
|
||||
writes the new tokens -- all synchronous SQL plus Fernet -- and a process
|
||||
waiting on a peer's lease polls the row. Every one of those calls goes through
|
||||
``asyncio.to_thread``; if one regresses onto the loop, the strict Blockbuster
|
||||
gate raises ``BlockingError`` and these tests fail.
|
||||
|
||||
Database setup and inspection are offloaded with ``asyncio.to_thread``; the
|
||||
token endpoint is an in-memory ``httpx.MockTransport``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from urllib.parse import parse_qs
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from cryptography.fernet import Fernet
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from deerflow.config.credentials_key import CredentialsCipher
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.mcp import oauth_store
|
||||
from deerflow.mcp.oauth import OAuthTokenManager
|
||||
from deerflow.mcp.oauth_store import McpOAuthTokenKey, McpOAuthTokenStore, oauth_config_fingerprint
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.mcp_oauth_tokens import McpOAuthTokenRow, SqlMcpOAuthTokenRepository
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
CONFIG = ExtensionsConfig.model_validate(
|
||||
{
|
||||
"mcpServers": {
|
||||
"reports": {
|
||||
"type": "http",
|
||||
"url": "https://mcp.example.invalid/mcp",
|
||||
"oauth": {"token_url": "https://auth.example.invalid/token", "grant_type": "refresh_token", "refresh_token": "refresh-0"},
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def token_endpoint(monkeypatch):
|
||||
issued: list[str] = []
|
||||
|
||||
def handle(request: httpx.Request) -> httpx.Response:
|
||||
issued.append(parse_qs(request.content.decode())["refresh_token"][0])
|
||||
return httpx.Response(200, json={"access_token": f"access-{len(issued)}", "refresh_token": f"refresh-{len(issued)}", "expires_in": 3600})
|
||||
|
||||
original_client = httpx.AsyncClient
|
||||
|
||||
class OfflineClient(original_client):
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs["transport"] = httpx.MockTransport(handle)
|
||||
kwargs["trust_env"] = False
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", OfflineClient)
|
||||
return issued
|
||||
|
||||
|
||||
def _session_factory(path) -> sessionmaker:
|
||||
engine = sa.create_engine(f"sqlite:///{path.as_posix()}")
|
||||
Base.metadata.create_all(engine, tables=[McpOAuthTokenRow.__table__])
|
||||
return sessionmaker(engine, expire_on_commit=False)
|
||||
|
||||
|
||||
def _store(session_factory) -> McpOAuthTokenStore:
|
||||
return McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=CredentialsCipher.from_keys([Fernet.generate_key()]))
|
||||
|
||||
|
||||
async def test_refresh_and_adoption_through_the_store_do_not_block_the_event_loop(tmp_path, token_endpoint) -> None:
|
||||
session_factory = await asyncio.to_thread(_session_factory, tmp_path / "oauth.db")
|
||||
store = _store(session_factory)
|
||||
|
||||
first = OAuthTokenManager.from_extensions_config(CONFIG, token_store=store)
|
||||
assert await first.get_authorization_header("reports") == "Bearer access-1"
|
||||
# A second process adopts the stored token without a refresh.
|
||||
second = OAuthTokenManager.from_extensions_config(CONFIG, token_store=store)
|
||||
assert await second.get_authorization_header("reports") == "Bearer access-1"
|
||||
assert token_endpoint == ["refresh-0"]
|
||||
|
||||
|
||||
async def test_waiting_on_a_peer_lease_does_not_block_the_event_loop(tmp_path, token_endpoint, monkeypatch) -> None:
|
||||
monkeypatch.setattr(oauth_store, "REFRESH_LEASE_SECONDS", 0.3)
|
||||
session_factory = await asyncio.to_thread(_session_factory, tmp_path / "oauth.db")
|
||||
store = _store(session_factory)
|
||||
key = McpOAuthTokenKey.for_server("reports", None)
|
||||
# A peer claimed the refresh and never finished it.
|
||||
assert await asyncio.to_thread(store.claim, key, fingerprint=oauth_config_fingerprint(CONFIG.mcp_servers["reports"].oauth), expected_version=None)
|
||||
|
||||
manager = OAuthTokenManager.from_extensions_config(CONFIG, token_store=store)
|
||||
assert await manager.get_authorization_header("reports") == "Bearer access-1"
|
||||
assert token_endpoint == ["refresh-0"]
|
||||
@ -202,6 +202,34 @@ def test_personal_mcp_authority_spans_runtime_startup_and_shutdown():
|
||||
assert personal_access._admin_checker is previous
|
||||
|
||||
|
||||
def test_mcp_oauth_token_store_spans_runtime_startup_and_shutdown():
|
||||
"""Runs drain inside the runtime's exit; a refresh there must still reach the shared store."""
|
||||
from deerflow.mcp import oauth_store
|
||||
|
||||
events = []
|
||||
store = object()
|
||||
previous = oauth_store.get_mcp_oauth_token_store()
|
||||
|
||||
@asynccontextmanager
|
||||
async def checked_runtime(_app, _config):
|
||||
assert oauth_store.get_mcp_oauth_token_store() is store
|
||||
events.append("start")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
assert oauth_store.get_mcp_oauth_token_store() is store
|
||||
events.append("stop")
|
||||
|
||||
with (
|
||||
patch(f"{__name__}._noop_langgraph_runtime", checked_runtime),
|
||||
patch("app.gateway.mcp_oauth_tokens.build_mcp_oauth_token_store", return_value=store),
|
||||
):
|
||||
asyncio.run(_run_lifespan_with_upload_staging_cleanup())
|
||||
|
||||
assert events == ["start", "stop"]
|
||||
assert oauth_store.get_mcp_oauth_token_store() is previous
|
||||
|
||||
|
||||
async def _run_lifespan_with_mcp_task_config_snapshot() -> None:
|
||||
from app.gateway.app import lifespan
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
|
||||
@ -194,3 +194,29 @@ def test_skills_writer_preserves_the_lifecycle_ledger(
|
||||
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"})]
|
||||
|
||||
675
backend/tests/test_mcp_oauth_token_store.py
Normal file
675
backend/tests/test_mcp_oauth_token_store.py
Normal file
@ -0,0 +1,675 @@
|
||||
"""MCP OAuth tokens shared through ``mcp_oauth_tokens`` (restarts and Gateway replicas).
|
||||
|
||||
A ``refresh_token`` connection's access token and rotated refresh token live in
|
||||
one encrypted row per connection. Every manager below stands for a separate
|
||||
Gateway process: its own ``McpOAuthTokenStore`` and in-process token state over
|
||||
one shared database. The token endpoint rotates on every refresh and rejects a
|
||||
consumed refresh token with ``invalid_grant``, like Auth0/Okta/Google, so any
|
||||
replica or restart that falls back to an already-used seed fails loudly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import threading
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from urllib.parse import parse_qs
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from cryptography.fernet import Fernet
|
||||
from sqlalchemy.engine import make_url
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from deerflow.config.credentials_key import CredentialsCipher
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.mcp import oauth_store
|
||||
from deerflow.mcp.oauth import McpOAuthReauthorizationRequired, OAuthTokenManager
|
||||
from deerflow.mcp.oauth_store import (
|
||||
McpOAuthTokenStore,
|
||||
McpOAuthTokenStoreUnavailable,
|
||||
StoredTokenLookup,
|
||||
oauth_config_fingerprint,
|
||||
oauth_token_request_fingerprint,
|
||||
set_mcp_oauth_token_store,
|
||||
)
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.mcp_oauth_tokens import McpOAuthTokenRow, SqlMcpOAuthTokenRepository
|
||||
from deerflow.persistence.postgres_schema import dsn_with_search_path
|
||||
|
||||
SERVER = "reports"
|
||||
REQUEST = "r" * 64
|
||||
|
||||
|
||||
class _Clock:
|
||||
current = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
|
||||
@classmethod
|
||||
def now(cls, _tz):
|
||||
return cls.current
|
||||
|
||||
|
||||
class _RotatingTokenEndpoint:
|
||||
"""Issue ``access-N``/``refresh-N`` and accept only the latest refresh token."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.refresh_token = "refresh-0"
|
||||
self.generation = 0
|
||||
self.refresh_requests: list[str] = []
|
||||
self.forms: list[dict[str, list[str]]] = []
|
||||
self.fail_with: int | None = None
|
||||
|
||||
def handle(self, request: httpx.Request) -> httpx.Response:
|
||||
assert request.url.path == "/token"
|
||||
form = parse_qs(request.content.decode())
|
||||
submitted = form["refresh_token"][0]
|
||||
self.refresh_requests.append(submitted)
|
||||
self.forms.append(form)
|
||||
if self.fail_with is not None:
|
||||
return httpx.Response(self.fail_with, json={"error": "server_error"})
|
||||
if submitted != self.refresh_token:
|
||||
return httpx.Response(400, json={"error": "invalid_grant"})
|
||||
self.generation += 1
|
||||
self.refresh_token = f"refresh-{self.generation}"
|
||||
return httpx.Response(200, json={"access_token": f"access-{self.generation}", "refresh_token": self.refresh_token, "expires_in": 120})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def endpoint(monkeypatch):
|
||||
server = _RotatingTokenEndpoint()
|
||||
original_client = httpx.AsyncClient
|
||||
|
||||
class OfflineClient(original_client):
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs["transport"] = httpx.MockTransport(lambda request: server.handle(request))
|
||||
kwargs["trust_env"] = False
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", OfflineClient)
|
||||
monkeypatch.setattr("deerflow.mcp.oauth.datetime", _Clock)
|
||||
monkeypatch.setattr(_Clock, "current", datetime(2026, 1, 1, tzinfo=UTC))
|
||||
return server
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_installed_store():
|
||||
previous = set_mcp_oauth_token_store(None)
|
||||
yield
|
||||
set_mcp_oauth_token_store(previous)
|
||||
|
||||
|
||||
@pytest.fixture(params=["sqlite", "postgres"])
|
||||
def session_factory(request, tmp_path) -> sessionmaker[Session]:
|
||||
"""The synchronous session factory the Gateway builds from ``app_sync_sqlalchemy_url``, on each dialect."""
|
||||
schema = None
|
||||
if request.param == "sqlite":
|
||||
engine = sa.create_engine(f"sqlite:///{(tmp_path / 'oauth.db').as_posix()}")
|
||||
else:
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI (real Postgres mcp_oauth_tokens)")
|
||||
schema = f"mcp_oauth_tokens_{uuid.uuid4().hex}"
|
||||
admin = sa.create_engine(make_url(uri).set(drivername="postgresql+psycopg"))
|
||||
with admin.begin() as conn:
|
||||
conn.execute(sa.text(f'CREATE SCHEMA "{schema}"'))
|
||||
admin.dispose()
|
||||
engine = sa.create_engine(make_url(dsn_with_search_path(uri, schema)).set(drivername="postgresql+psycopg"))
|
||||
Base.metadata.create_all(engine, tables=[McpOAuthTokenRow.__table__])
|
||||
yield sessionmaker(engine, expire_on_commit=False)
|
||||
if schema is not None:
|
||||
with engine.begin() as conn:
|
||||
conn.execute(sa.text(f'DROP SCHEMA "{schema}" CASCADE'))
|
||||
engine.dispose()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cipher() -> CredentialsCipher:
|
||||
return CredentialsCipher.from_keys([Fernet.generate_key()])
|
||||
|
||||
|
||||
def _config(*, seed: str = "refresh-0", grant_type: str = "refresh_token", name: str = SERVER, token_url: str = "https://auth.example.invalid/token", extra_token_params: dict[str, str] | None = None) -> ExtensionsConfig:
|
||||
oauth = {"token_url": token_url, "grant_type": grant_type, "refresh_skew_seconds": 0, "extra_token_params": extra_token_params or {}}
|
||||
if grant_type == "refresh_token":
|
||||
oauth["refresh_token"] = seed
|
||||
else:
|
||||
oauth.update(client_id="client", client_secret="secret")
|
||||
return ExtensionsConfig.model_validate({"mcpServers": {name: {"type": "http", "url": "https://mcp.example.invalid/mcp", "oauth": oauth}}})
|
||||
|
||||
|
||||
def _replica(session_factory, cipher, config: ExtensionsConfig | None = None, *, owner_user_id: str | None = None) -> OAuthTokenManager:
|
||||
"""One Gateway process: a fresh store and fresh in-process token state over the shared database."""
|
||||
store = McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=cipher)
|
||||
return OAuthTokenManager.from_extensions_config(config or _config(), owner_user_id=owner_user_id, token_store=store)
|
||||
|
||||
|
||||
def _rows(session_factory) -> list[McpOAuthTokenRow]:
|
||||
with session_factory() as session:
|
||||
return list(session.scalars(sa.select(McpOAuthTokenRow).order_by(McpOAuthTokenRow.scope, McpOAuthTokenRow.owner_user_id, McpOAuthTokenRow.server_key)))
|
||||
|
||||
|
||||
def _expire_access_tokens() -> None:
|
||||
_Clock.current = _Clock.current + timedelta(seconds=300)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_peer_uses_the_rotated_refresh_token_without_calling_the_token_endpoint(session_factory, cipher, endpoint):
|
||||
a = _replica(session_factory, cipher)
|
||||
b = _replica(session_factory, cipher)
|
||||
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-1"
|
||||
# B has never refreshed, but A's still-valid access token is in the row.
|
||||
assert await b.get_authorization_header(SERVER) == "Bearer access-1"
|
||||
assert endpoint.refresh_requests == ["refresh-0"]
|
||||
|
||||
_expire_access_tokens()
|
||||
# B refreshes with A's rotated token, not the consumed seed ...
|
||||
assert await b.get_authorization_header(SERVER) == "Bearer access-2"
|
||||
# ... and A adopts B's result instead of refreshing with its stale copy.
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-2"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restart_does_not_reseed_from_the_consumed_seed(session_factory, cipher, endpoint):
|
||||
assert await _replica(session_factory, cipher).get_authorization_header(SERVER) == "Bearer access-1"
|
||||
_expire_access_tokens()
|
||||
|
||||
restarted = _replica(session_factory, cipher)
|
||||
|
||||
assert await restarted.get_authorization_header(SERVER) == "Bearer access-2"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_changed_seed_wins_over_the_stored_rotation(session_factory, cipher, endpoint):
|
||||
assert await _replica(session_factory, cipher).get_authorization_header(SERVER) == "Bearer access-1"
|
||||
# The operator revokes the chain and configures a new seed.
|
||||
endpoint.refresh_token = "replacement-seed"
|
||||
changed = _config(seed="replacement-seed")
|
||||
|
||||
assert await _replica(session_factory, cipher, changed).get_authorization_header(SERVER) == "Bearer access-2"
|
||||
_expire_access_tokens()
|
||||
# The new chain is what the next process continues from.
|
||||
assert await _replica(session_factory, cipher, changed).get_authorization_header(SERVER) == "Bearer access-3"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "replacement-seed", "refresh-2"]
|
||||
(row,) = _rows(session_factory)
|
||||
assert row.config_fingerprint == oauth_config_fingerprint(changed.mcp_servers[SERVER].oauth)
|
||||
|
||||
|
||||
def test_concurrent_refresh_across_replicas_makes_one_token_request(session_factory, cipher, endpoint, monkeypatch):
|
||||
a = _replica(session_factory, cipher)
|
||||
b = _replica(session_factory, cipher)
|
||||
a_fetching = threading.Event()
|
||||
b_waiting = threading.Event()
|
||||
release = threading.Event()
|
||||
handle = endpoint.handle
|
||||
|
||||
async def held(request):
|
||||
a_fetching.set()
|
||||
assert await asyncio.to_thread(release.wait, 10), "test did not release the token response"
|
||||
return handle(request)
|
||||
|
||||
monkeypatch.setattr(endpoint, "handle", held)
|
||||
b_store = b._token_store
|
||||
original_claim = b_store.claim
|
||||
|
||||
def observed_claim(*args, **kwargs):
|
||||
claimed = original_claim(*args, **kwargs)
|
||||
if claimed is None:
|
||||
b_waiting.set()
|
||||
return claimed
|
||||
|
||||
monkeypatch.setattr(b_store, "claim", observed_claim)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
first = executor.submit(lambda: asyncio.run(a.get_authorization_header(SERVER)))
|
||||
try:
|
||||
assert a_fetching.wait(10), "replica A never reached the token endpoint"
|
||||
second = executor.submit(lambda: asyncio.run(b.get_authorization_header(SERVER)))
|
||||
assert b_waiting.wait(10), "replica B did not find A's refresh lease"
|
||||
finally:
|
||||
release.set()
|
||||
assert first.result(timeout=10) == "Bearer access-1"
|
||||
assert second.result(timeout=10) == "Bearer access-1"
|
||||
assert endpoint.refresh_requests == ["refresh-0"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tokens_are_encrypted_at_rest_and_unreadable_rows_fall_back_to_the_seed(session_factory, cipher, endpoint, caplog):
|
||||
assert await _replica(session_factory, cipher).get_authorization_header(SERVER) == "Bearer access-1"
|
||||
(row,) = _rows(session_factory)
|
||||
assert (row.scope, row.owner_user_id, row.server_key, row.version) == ("deployment", "", SERVER, 1)
|
||||
for column in (row.access_token_enc, row.refresh_token_enc):
|
||||
assert column.startswith("fernet:v2:")
|
||||
assert "access-1" not in column and "refresh-1" not in column
|
||||
assert cipher.decrypt_text(row.access_token_enc) == "access-1"
|
||||
assert cipher.decrypt_text(row.refresh_token_enc) == "refresh-1"
|
||||
|
||||
# A process with another key cannot read the row; it falls back to the
|
||||
# seed (here already consumed) instead of failing on the ciphertext, and
|
||||
# the rejection tells the connection's user why and who can fix it.
|
||||
other = _replica(session_factory, CredentialsCipher.from_keys([Fernet.generate_key()]))
|
||||
with caplog.at_level("WARNING", logger="deerflow.mcp.oauth_store"):
|
||||
with pytest.raises(McpOAuthReauthorizationRequired) as rejected:
|
||||
await other.get_authorization_header(SERVER)
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-0"]
|
||||
assert "cannot be decrypted" in caplog.text
|
||||
assert "refresh-1" not in caplog.text
|
||||
message = str(rejected.value)
|
||||
assert f"'{SERVER}'" in message and "invalid_grant" in message and "DEER_FLOW_CREDENTIALS_KEY" in message
|
||||
assert "refresh-0" not in message and "refresh-1" not in message
|
||||
assert isinstance(rejected.value.__cause__, httpx.HTTPStatusError)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_and_deployment_connections_with_one_name_keep_separate_rows(session_factory, cipher, endpoint):
|
||||
assert await _replica(session_factory, cipher).get_authorization_header(SERVER) == "Bearer access-1"
|
||||
endpoint.refresh_token = "personal-seed"
|
||||
personal = _config(seed="personal-seed")
|
||||
assert await _replica(session_factory, cipher, personal, owner_user_id="user-1").get_authorization_header(SERVER) == "Bearer access-2"
|
||||
_expire_access_tokens()
|
||||
|
||||
# Each owner continues from its own chain after a restart.
|
||||
assert await _replica(session_factory, cipher, personal, owner_user_id="user-1").get_authorization_header(SERVER) == "Bearer access-3"
|
||||
assert [(row.scope, row.owner_user_id, row.server_key) for row in _rows(session_factory)] == [("deployment", "", SERVER), ("personal", "user-1", SERVER)]
|
||||
assert endpoint.refresh_requests == ["refresh-0", "personal-seed", "refresh-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_credentials_tokens_stay_in_process(session_factory, cipher, monkeypatch):
|
||||
requests: list[str] = []
|
||||
|
||||
def handle(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(parse_qs(request.content.decode())["grant_type"][0])
|
||||
return httpx.Response(200, json={"access_token": f"cc-{len(requests)}", "expires_in": 120})
|
||||
|
||||
original_client = httpx.AsyncClient
|
||||
|
||||
class OfflineClient(original_client):
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs["transport"] = httpx.MockTransport(handle)
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", OfflineClient)
|
||||
config = _config(grant_type="client_credentials")
|
||||
|
||||
assert await _replica(session_factory, cipher, config).get_authorization_header(SERVER) == "Bearer cc-1"
|
||||
assert await _replica(session_factory, cipher, config).get_authorization_header(SERVER) == "Bearer cc-2"
|
||||
assert _rows(session_factory) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_refresh_releases_the_lease_for_the_next_caller(session_factory, cipher, endpoint):
|
||||
endpoint.fail_with = 503
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
await _replica(session_factory, cipher).get_authorization_header(SERVER)
|
||||
(row,) = _rows(session_factory)
|
||||
assert row.refresh_lease_owner is None
|
||||
|
||||
endpoint.fail_with = None
|
||||
# The seed was never consumed, so a peer retries it at once, not after the lease runs out.
|
||||
assert await asyncio.wait_for(_replica(session_factory, cipher).get_authorization_header(SERVER), 5) == "Bearer access-1"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-0"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_abandoned_refresh_lease_is_taken_over_after_it_expires(session_factory, cipher, endpoint, monkeypatch):
|
||||
monkeypatch.setattr(oauth_store, "REFRESH_LEASE_SECONDS", 0.5)
|
||||
crashed = McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=cipher)
|
||||
key = oauth_store.McpOAuthTokenKey(scope="deployment", owner_user_id="", server_key=SERVER)
|
||||
fingerprint = oauth_config_fingerprint(_config().mcp_servers[SERVER].oauth)
|
||||
# A process claimed the first refresh and died before writing a token.
|
||||
assert crashed.claim(key, fingerprint=fingerprint, expected_version=None) is not None
|
||||
|
||||
assert await asyncio.wait_for(_replica(session_factory, cipher).get_authorization_header(SERVER), 10) == "Bearer access-1"
|
||||
(row,) = _rows(session_factory)
|
||||
assert row.refresh_lease_owner is None
|
||||
assert row.version == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_errors_fail_the_refresh_without_leaking_sql_parameters(cipher, endpoint, tmp_path, caplog):
|
||||
engine = sa.create_engine(f"sqlite:///{(tmp_path / 'missing-table.db').as_posix()}")
|
||||
manager = _replica(sessionmaker(engine), cipher)
|
||||
try:
|
||||
with caplog.at_level("ERROR"):
|
||||
with pytest.raises(McpOAuthTokenStoreUnavailable) as raised:
|
||||
await manager.get_authorization_header(SERVER)
|
||||
finally:
|
||||
engine.dispose()
|
||||
assert SERVER in str(raised.value)
|
||||
assert "mcp_oauth_tokens" not in str(raised.value)
|
||||
# Never refresh without a place to put the rotated token.
|
||||
assert endpoint.refresh_requests == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_installed_store_is_used_by_managers_built_without_one(session_factory, cipher, endpoint):
|
||||
set_mcp_oauth_token_store(McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=cipher))
|
||||
|
||||
assert await OAuthTokenManager.from_extensions_config(_config()).get_authorization_header(SERVER) == "Bearer access-1"
|
||||
_expire_access_tokens()
|
||||
assert await OAuthTokenManager.from_extensions_config(_config()).get_authorization_header(SERVER) == "Bearer access-2"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_a_store_tokens_stay_in_process(endpoint):
|
||||
assert await OAuthTokenManager.from_extensions_config(_config()).get_authorization_header(SERVER) == "Bearer access-1"
|
||||
_expire_access_tokens()
|
||||
# Today's behaviour: a new manager starts again from the configured seed,
|
||||
# and the rejection names the remedy instead of a bare HTTP 400.
|
||||
with pytest.raises(McpOAuthReauthorizationRequired, match="configuring a new refresh token") as rejected:
|
||||
await OAuthTokenManager.from_extensions_config(_config()).get_authorization_header(SERVER)
|
||||
assert "DEER_FLOW_CREDENTIALS_KEY" not in str(rejected.value)
|
||||
|
||||
|
||||
def test_fingerprint_follows_the_token_chain_identity_not_the_client_secret():
|
||||
base = _config().mcp_servers[SERVER].oauth
|
||||
fingerprint = oauth_config_fingerprint(base)
|
||||
for field, value in {"token_url": "https://other.example.invalid/token", "client_id": "other", "scope": "read", "audience": "api", "refresh_token": "other-seed"}.items():
|
||||
assert oauth_config_fingerprint(base.model_copy(update={field: value})) != fingerprint, field
|
||||
# A rotated client secret keeps the issued refresh-token chain valid.
|
||||
assert oauth_config_fingerprint(base.model_copy(update={"client_secret": "rotated"})) == fingerprint
|
||||
assert "refresh-0" not in fingerprint
|
||||
|
||||
|
||||
def test_a_claim_fails_once_a_peer_wrote_since_the_read(session_factory, cipher):
|
||||
"""The version compare-and-set keeps a reader from refreshing with a refresh token a peer just consumed."""
|
||||
key = oauth_store.McpOAuthTokenKey(scope="deployment", owner_user_id="", server_key=SERVER)
|
||||
fingerprint = oauth_config_fingerprint(_config().mcp_servers[SERVER].oauth)
|
||||
a = McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=cipher)
|
||||
b = McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=cipher)
|
||||
expires_at = datetime(2030, 1, 1, tzinfo=UTC)
|
||||
|
||||
lease = a.claim(key, fingerprint=fingerprint, expected_version=None)
|
||||
assert a.save(key, lease_owner=lease, fingerprint=fingerprint, request_fingerprint=REQUEST, access_token="access-1", token_type="Bearer", expires_at=expires_at, refresh_token="refresh-1", rotated=True)
|
||||
assert b.load(key, fingerprint=fingerprint, request_fingerprint=REQUEST)[0] == 1
|
||||
# Both read version 1; A refreshes first.
|
||||
lease = a.claim(key, fingerprint=fingerprint, expected_version=1)
|
||||
assert lease is not None
|
||||
assert b.claim(key, fingerprint=fingerprint, expected_version=1) is None, "a live lease must block the peer"
|
||||
assert a.save(key, lease_owner=lease, fingerprint=fingerprint, request_fingerprint=REQUEST, access_token="access-2", token_type="Bearer", expires_at=expires_at, refresh_token="refresh-2", rotated=True)
|
||||
|
||||
assert b.claim(key, fingerprint=fingerprint, expected_version=1) is None, "the lease is free, but the row moved on"
|
||||
version, stored, _unreadable = b.load(key, fingerprint=fingerprint, request_fingerprint=REQUEST)
|
||||
assert (version, stored.refresh_token, stored.access_token) == (2, "refresh-2", "access-2")
|
||||
# A write under a lease that is no longer held is refused.
|
||||
assert not b.save(key, lease_owner="not-the-holder", fingerprint=fingerprint, request_fingerprint=REQUEST, access_token="x", token_type="Bearer", expires_at=expires_at, refresh_token="y", rotated=True)
|
||||
assert b.load(key, fingerprint=fingerprint, request_fingerprint=REQUEST)[1].refresh_token == "refresh-2"
|
||||
assert b.load(key, fingerprint="another-chain", request_fingerprint=REQUEST) == StoredTokenLookup(2, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_unloadable_credentials_key_keeps_tokens_in_process(session_factory, endpoint, monkeypatch, caplog):
|
||||
def broken_key():
|
||||
raise oauth_store.McpOAuthTokenKeyUnavailable("no key")
|
||||
|
||||
store = McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory))
|
||||
monkeypatch.setattr(store, "_cipher", broken_key)
|
||||
manager = OAuthTokenManager.from_extensions_config(_config(), token_store=store)
|
||||
|
||||
with caplog.at_level("WARNING", logger="deerflow.mcp.oauth"):
|
||||
assert await manager.get_authorization_header(SERVER) == "Bearer access-1"
|
||||
assert "in this process only" in caplog.text
|
||||
assert _rows(session_factory) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["memory", "sqlite"])
|
||||
def test_gateway_shares_tokens_only_through_a_sql_database(backend, tmp_path):
|
||||
from app.gateway.mcp_oauth_tokens import build_mcp_oauth_token_store
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
|
||||
database = DatabaseConfig(backend=backend, sqlite_dir=str(tmp_path))
|
||||
store = build_mcp_oauth_token_store(SimpleNamespace(database=database))
|
||||
if backend == "memory":
|
||||
assert store is None
|
||||
else:
|
||||
assert isinstance(store, McpOAuthTokenStore)
|
||||
assert build_mcp_oauth_token_store(SimpleNamespace()) is None
|
||||
|
||||
|
||||
def test_forgetting_unconfigured_connections_deletes_only_that_owners_other_rows(session_factory, cipher):
|
||||
store = McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=cipher)
|
||||
for owner, name in ((None, "kept"), (None, "removed"), ("alice", "kept"), ("alice", "removed"), ("bob", "removed")):
|
||||
assert store.claim(oauth_store.McpOAuthTokenKey.for_server(name, owner), fingerprint="f" * 64, expected_version=None)
|
||||
|
||||
assert store.forget_unconfigured("alice", {"kept"}) == 1
|
||||
assert store.forget_unconfigured(None, {"kept", "never-stored"}) == 1
|
||||
assert [(row.scope, row.owner_user_id, row.server_key) for row in _rows(session_factory)] == [
|
||||
("deployment", "", "kept"),
|
||||
("personal", "alice", "kept"),
|
||||
("personal", "bob", "removed"),
|
||||
]
|
||||
# An owner with no connections left loses every row.
|
||||
assert store.forget_unconfigured("bob", set()) == 1
|
||||
assert [row.owner_user_id for row in _rows(session_factory)] == ["", "alice"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_revoked_stored_chain_asks_to_reauthorize_without_blaming_the_key(session_factory, cipher, endpoint):
|
||||
assert await _replica(session_factory, cipher).get_authorization_header(SERVER) == "Bearer access-1"
|
||||
endpoint.refresh_token = "revoked-by-the-provider"
|
||||
_expire_access_tokens()
|
||||
|
||||
with pytest.raises(McpOAuthReauthorizationRequired) as rejected:
|
||||
await _replica(session_factory, cipher).get_authorization_header(SERVER)
|
||||
assert "configuring a new refresh token" in str(rejected.value)
|
||||
assert "DEER_FLOW_CREDENTIALS_KEY" not in str(rejected.value)
|
||||
(row,) = _rows(session_factory)
|
||||
assert row.refresh_lease_owner is None, "a rejected refresh releases its lease"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
|
||||
|
||||
class _SlowStreamingTokenServer:
|
||||
"""A real loopback token endpoint that streams its JSON body in slow chunks.
|
||||
|
||||
HTTPX timeouts bound each read, not the whole response, so a body that
|
||||
keeps trickling in can outlive any fixed lease.
|
||||
"""
|
||||
|
||||
def __init__(self, *, chunks: int, chunk_delay: float) -> None:
|
||||
import http.server
|
||||
import json as json_module
|
||||
import socketserver
|
||||
import time as time_module
|
||||
|
||||
self.requests: list[str] = []
|
||||
self.started = threading.Event()
|
||||
endpoint = self
|
||||
current = {"refresh_token": "refresh-0", "generation": 0}
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
body = self.rfile.read(int(self.headers["Content-Length"])).decode()
|
||||
submitted = parse_qs(body)["refresh_token"][0]
|
||||
endpoint.requests.append(submitted)
|
||||
endpoint.started.set()
|
||||
if submitted != current["refresh_token"]:
|
||||
payload = json_module.dumps({"error": "invalid_grant"}).encode()
|
||||
self.send_response(400)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
return
|
||||
current["generation"] += 1
|
||||
current["refresh_token"] = f"refresh-{current['generation']}"
|
||||
payload = json_module.dumps({"access_token": f"access-{current['generation']}", "refresh_token": current["refresh_token"], "expires_in": 120}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Transfer-Encoding", "chunked")
|
||||
self.end_headers()
|
||||
size = -(-len(payload) // chunks)
|
||||
for offset in range(0, len(payload), size):
|
||||
piece = payload[offset : offset + size]
|
||||
self.wfile.write(f"{len(piece):x}\r\n".encode() + piece + b"\r\n")
|
||||
self.wfile.flush()
|
||||
time_module.sleep(chunk_delay)
|
||||
self.wfile.write(b"0\r\n\r\n")
|
||||
|
||||
self._server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), Handler)
|
||||
self._server.daemon_threads = True
|
||||
self.url = f"http://127.0.0.1:{self._server.server_address[1]}/token"
|
||||
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
||||
|
||||
def __enter__(self):
|
||||
self._thread.start()
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc):
|
||||
self._server.shutdown()
|
||||
self._server.server_close()
|
||||
|
||||
|
||||
def test_a_slow_token_response_keeps_its_lease_until_the_token_lands(session_factory, cipher, monkeypatch):
|
||||
"""The holder renews its lease while the body streams, so a peer waits instead of resending the consumed token."""
|
||||
for variable in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"):
|
||||
monkeypatch.delenv(variable, raising=False)
|
||||
monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost")
|
||||
monkeypatch.setattr(oauth_store, "REFRESH_LEASE_SECONDS", 0.6)
|
||||
monkeypatch.setattr(oauth_store, "REFRESH_POLL_SECONDS", 0.05)
|
||||
|
||||
# About 1.8 s of streaming: three times the lease.
|
||||
with _SlowStreamingTokenServer(chunks=6, chunk_delay=0.3) as server:
|
||||
config = _config(token_url=server.url)
|
||||
a = _replica(session_factory, cipher, config)
|
||||
b = _replica(session_factory, cipher, config)
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
first = executor.submit(lambda: asyncio.run(a.get_authorization_header(SERVER)))
|
||||
assert server.started.wait(10), "replica A never reached the token endpoint"
|
||||
second = executor.submit(lambda: asyncio.run(b.get_authorization_header(SERVER)))
|
||||
assert first.result(timeout=20) == "Bearer access-1"
|
||||
assert second.result(timeout=20) == "Bearer access-1"
|
||||
assert server.requests == ["refresh-0"]
|
||||
(row,) = _rows(session_factory)
|
||||
assert (row.version, row.refresh_lease_owner) == (1, None)
|
||||
assert cipher.decrypt_text(row.refresh_token_enc) == "refresh-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_token_request_that_never_finishes_gives_up_and_releases_the_lease(session_factory, cipher, endpoint, monkeypatch):
|
||||
monkeypatch.setattr(oauth_store, "REFRESH_MAX_SECONDS", 0.3)
|
||||
handle = endpoint.handle
|
||||
|
||||
async def stalled(request):
|
||||
await asyncio.sleep(5)
|
||||
return handle(request)
|
||||
|
||||
monkeypatch.setattr(endpoint, "handle", stalled)
|
||||
with pytest.raises(RuntimeError, match="did not finish within"):
|
||||
await asyncio.wait_for(_replica(session_factory, cipher).get_authorization_header(SERVER), 5)
|
||||
(row,) = _rows(session_factory)
|
||||
assert row.refresh_lease_owner is None
|
||||
|
||||
|
||||
def _fail_next_writes(store: McpOAuthTokenStore, monkeypatch, count: int = 1) -> None:
|
||||
"""Make the next ``count`` token writes fail like a database outage."""
|
||||
original = store._repo.write
|
||||
remaining = {"n": count}
|
||||
|
||||
def write(*args, **kwargs):
|
||||
if remaining["n"] > 0:
|
||||
remaining["n"] -= 1
|
||||
raise sa.exc.OperationalError("UPDATE mcp_oauth_tokens", {}, Exception("database is unavailable"))
|
||||
return original(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(store._repo, "write", write)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_rotation_the_store_missed_is_stored_before_the_next_refresh(session_factory, cipher, endpoint, monkeypatch):
|
||||
a = _replica(session_factory, cipher)
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-1"
|
||||
_expire_access_tokens()
|
||||
_fail_next_writes(a._token_store, monkeypatch)
|
||||
# refresh-1 -> refresh-2, but the row keeps the consumed refresh-1.
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-2"
|
||||
_expire_access_tokens()
|
||||
|
||||
# The live refresh-2 is stored first instead of being replaced by refresh-1.
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-3"
|
||||
_expire_access_tokens()
|
||||
assert await _replica(session_factory, cipher).get_authorization_header(SERVER) == "Bearer access-4"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-1", "refresh-2", "refresh-3"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_missed_rotation_reaches_peers_while_the_access_token_is_still_valid(session_factory, cipher, endpoint, monkeypatch):
|
||||
monkeypatch.setattr("deerflow.mcp.oauth._UNSAVED_ROTATION_RETRY_SECONDS", 0.0)
|
||||
a = _replica(session_factory, cipher)
|
||||
_fail_next_writes(a._token_store, monkeypatch)
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-1"
|
||||
assert _rows(session_factory)[0].refresh_token_enc is None, "the first write failed"
|
||||
|
||||
# The next call, still within the token's lifetime, stores the rotation.
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-1"
|
||||
_expire_access_tokens()
|
||||
assert await _replica(session_factory, cipher).get_authorization_header(SERVER) == "Bearer access-2"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_newer_peer_chain_supersedes_a_missed_rotation(session_factory, cipher, endpoint, monkeypatch):
|
||||
a = _replica(session_factory, cipher)
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-1"
|
||||
_expire_access_tokens()
|
||||
_fail_next_writes(a._token_store, monkeypatch)
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer access-2"
|
||||
|
||||
# Meanwhile a peer stored a newer chain over version 1 (e.g. a provider with a reuse grace period).
|
||||
peer = McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=cipher)
|
||||
key = oauth_store.McpOAuthTokenKey.for_server(SERVER, None)
|
||||
oauth = _config().mcp_servers[SERVER].oauth
|
||||
lease = peer.claim(key, fingerprint=oauth_config_fingerprint(oauth), expected_version=1)
|
||||
assert lease is not None
|
||||
assert peer.save(
|
||||
key,
|
||||
lease_owner=lease,
|
||||
fingerprint=oauth_config_fingerprint(oauth),
|
||||
request_fingerprint=oauth_token_request_fingerprint(oauth),
|
||||
access_token="peer-access",
|
||||
token_type="Bearer",
|
||||
expires_at=_Clock.current + timedelta(seconds=900),
|
||||
refresh_token="peer-refresh",
|
||||
rotated=True,
|
||||
)
|
||||
_expire_access_tokens()
|
||||
|
||||
assert await a.get_authorization_header(SERVER) == "Bearer peer-access"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_changed_token_request_parameters_refresh_from_the_stored_chain_instead_of_reusing_the_access_token(session_factory, cipher, endpoint):
|
||||
assert await _replica(session_factory, cipher, _config(extra_token_params={"resource": "https://old-api"})).get_authorization_header(SERVER) == "Bearer access-1"
|
||||
|
||||
# A hot-reloaded server asks for another resource: the old access token is
|
||||
# for the wrong audience, but its refresh-token chain is still the live one.
|
||||
changed = _config(extra_token_params={"resource": "https://new-api"})
|
||||
assert await _replica(session_factory, cipher, changed).get_authorization_header(SERVER) == "Bearer access-2"
|
||||
assert endpoint.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
assert endpoint.forms[-1]["resource"] == ["https://new-api"]
|
||||
# Processes with the new parameters now share the new access token.
|
||||
assert await _replica(session_factory, cipher, changed).get_authorization_header(SERVER) == "Bearer access-2"
|
||||
assert len(endpoint.refresh_requests) == 2
|
||||
|
||||
|
||||
def test_request_fingerprint_follows_token_request_parameters_not_the_chain():
|
||||
base = _config(extra_token_params={"resource": "https://api"}).mcp_servers[SERVER].oauth
|
||||
fingerprint = oauth_token_request_fingerprint(base)
|
||||
for update in ({"extra_token_params": {"resource": "https://other"}}, {"extra_token_params": {}}, {"token_field": "id_token"}):
|
||||
changed = base.model_copy(update=update)
|
||||
assert oauth_token_request_fingerprint(changed) != fingerprint, update
|
||||
assert oauth_config_fingerprint(changed) == oauth_config_fingerprint(base), "the refresh-token chain is unaffected"
|
||||
assert oauth_token_request_fingerprint(base.model_copy(update={"extra_token_params": {"resource": "https://api"}})) == fingerprint
|
||||
@ -556,7 +556,13 @@ async def test_personal_discovery_uses_personal_pool_domain_and_task_scope(monke
|
||||
AsyncMock(return_value=config),
|
||||
)
|
||||
monkeypatch.setattr("deerflow.mcp.tasks.runtime.is_mcp_task_runtime_available", lambda: True)
|
||||
monkeypatch.setattr("deerflow.mcp.tools.OAuthTokenManager.from_extensions_config", lambda config: personal_manager)
|
||||
oauth_owners: list[str | None] = []
|
||||
|
||||
def personal_manager_for(config, *, owner_user_id=None):
|
||||
oauth_owners.append(owner_user_id)
|
||||
return personal_manager
|
||||
|
||||
monkeypatch.setattr("deerflow.mcp.tools.OAuthTokenManager.from_extensions_config", personal_manager_for)
|
||||
monkeypatch.setattr("deerflow.mcp.tools.get_mcp_task_oauth_token_manager", deployment_manager)
|
||||
monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", AsyncMock(return_value={}))
|
||||
monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", lambda *args, **kwargs: [])
|
||||
@ -574,6 +580,8 @@ async def test_personal_discovery_uses_personal_pool_domain_and_task_scope(monke
|
||||
assert pool.ensure_binding.call_args.kwargs == {"domain": "personal"}
|
||||
assert ownership_domains == ["personal"]
|
||||
assert connection_scopes == ["personal"]
|
||||
# Stored OAuth tokens are keyed by the owner of a personal connection.
|
||||
assert oauth_owners == ["alice"]
|
||||
deployment_manager.assert_not_called()
|
||||
|
||||
|
||||
|
||||
@ -14,14 +14,19 @@ from urllib.parse import parse_qs
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from cryptography.fernet import Fernet
|
||||
from langchain_core.messages import AIMessage
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph
|
||||
from langgraph.prebuilt import ToolNode
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from deerflow.config.credentials_key import CredentialsCipher
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.config.paths import Paths
|
||||
from deerflow.mcp.cache import reset_mcp_tools_cache
|
||||
from deerflow.mcp.oauth import get_initial_oauth_headers
|
||||
from deerflow.mcp.oauth_store import McpOAuthTokenStore, set_mcp_oauth_token_store
|
||||
from deerflow.mcp.task_tool_caller import McpTaskToolCaller
|
||||
from deerflow.mcp.tasks import OrdinaryMcpTaskDriver, TaskReference, TaskStatus
|
||||
from deerflow.mcp.tasks import runtime as task_runtime
|
||||
@ -33,6 +38,8 @@ from deerflow.mcp.tasks.runtime import (
|
||||
)
|
||||
from deerflow.mcp.tools import get_mcp_tools
|
||||
from deerflow.mcp.user_config import load_user_mcp_config, personal_server_name, user_mcp_config_path
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.mcp_oauth_tokens import McpOAuthTokenRow, SqlMcpOAuthTokenRepository
|
||||
|
||||
|
||||
class _Clock:
|
||||
@ -146,6 +153,34 @@ def task_config(tmp_path, monkeypatch):
|
||||
set_mcp_task_config_snapshot(None)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def shared_token_store(tmp_path):
|
||||
"""Install the Gateway's shared token store over a fresh SQLite database."""
|
||||
engine = sa.create_engine(f"sqlite:///{(tmp_path / 'oauth-tokens.db').as_posix()}")
|
||||
Base.metadata.create_all(engine, tables=[McpOAuthTokenRow.__table__])
|
||||
session_factory = sessionmaker(engine, expire_on_commit=False)
|
||||
store = McpOAuthTokenStore(SqlMcpOAuthTokenRepository(session_factory), cipher=CredentialsCipher.from_keys([Fernet.generate_key()]))
|
||||
previous = set_mcp_oauth_token_store(store)
|
||||
try:
|
||||
yield session_factory
|
||||
finally:
|
||||
set_mcp_oauth_token_store(previous)
|
||||
engine.dispose()
|
||||
|
||||
|
||||
@pytest.fixture(params=["process", "shared"])
|
||||
def token_persistence(request):
|
||||
"""Run a contract with tokens in process memory and with the shared store installed."""
|
||||
if request.param == "shared":
|
||||
request.getfixturevalue("shared_token_store")
|
||||
return request.param
|
||||
|
||||
|
||||
def _stored_keys(session_factory) -> list[tuple[str, str, str]]:
|
||||
with session_factory() as session:
|
||||
return [tuple(row) for row in session.execute(sa.select(McpOAuthTokenRow.scope, McpOAuthTokenRow.owner_user_id, McpOAuthTokenRow.server_key).order_by(McpOAuthTokenRow.scope))]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discovery_submit_poll_and_rediscovery_share_rotating_oauth(task_config, rotating_server):
|
||||
path, startup = task_config
|
||||
@ -347,7 +382,7 @@ async def test_cancelled_waiter_does_not_strand_shared_refresh_lock(task_config,
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_task_runtime_reset_drops_rotated_oauth_state(task_config, rotating_server):
|
||||
async def test_task_runtime_reset_drops_rotated_oauth_state(task_config, rotating_server, token_persistence):
|
||||
_path, startup = task_config
|
||||
first = task_runtime.get_mcp_task_oauth_token_manager(ExtensionsConfig.from_file())
|
||||
assert await first.get_authorization_header("reports") == "Bearer access-1"
|
||||
@ -363,7 +398,39 @@ async def test_task_runtime_reset_drops_rotated_oauth_state(task_config, rotatin
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hot_reloadable_servers_do_not_reuse_task_oauth_state(task_config, rotating_server):
|
||||
async def test_task_runtime_restart_with_unchanged_seed_continues_the_stored_chain(task_config, rotating_server, shared_token_store):
|
||||
_path, startup = task_config
|
||||
first = task_runtime.get_mcp_task_oauth_token_manager(ExtensionsConfig.from_file())
|
||||
assert await first.get_authorization_header("reports") == "Bearer access-1"
|
||||
set_mcp_task_config_snapshot(None)
|
||||
_Clock.current = _Clock.current + timedelta(seconds=300)
|
||||
# The counterpart of the reset contract above: with the seed unchanged, a
|
||||
# new Gateway lifetime continues from the stored rotation, not refresh-0.
|
||||
set_mcp_task_config_snapshot(startup)
|
||||
second = task_runtime.get_mcp_task_oauth_token_manager(ExtensionsConfig.from_file())
|
||||
assert await second.get_authorization_header("reports") == "Bearer access-2"
|
||||
assert rotating_server.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
assert _stored_keys(shared_token_store) == [("deployment", "", "reports")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hot_reloaded_server_with_unchanged_seed_continues_the_stored_chain(task_config, rotating_server, shared_token_store):
|
||||
_path, startup = task_config
|
||||
other = startup.mcp_servers["reports"].model_copy(deep=True)
|
||||
other.task_toolsets = []
|
||||
startup.mcp_servers["other"] = other
|
||||
set_mcp_task_config_snapshot(startup)
|
||||
first = task_runtime.get_mcp_task_oauth_token_manager(startup)
|
||||
assert await first.get_authorization_header("other") == "Bearer access-1"
|
||||
_Clock.current = _Clock.current + timedelta(seconds=300)
|
||||
# A tool-cache reset builds a fresh manager for non-task servers.
|
||||
second = task_runtime.get_mcp_task_oauth_token_manager(startup.model_copy(deep=True))
|
||||
assert await second.get_authorization_header("other") == "Bearer access-2"
|
||||
assert rotating_server.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hot_reloadable_servers_do_not_reuse_task_oauth_state(task_config, rotating_server, token_persistence):
|
||||
_path, startup = task_config
|
||||
other = startup.mcp_servers["reports"].model_copy(deep=True)
|
||||
other.task_toolsets = []
|
||||
@ -402,7 +469,7 @@ async def test_snapshot_and_oauth_sharing_follow_the_same_server_selection(task_
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["discovery", "durable_call"])
|
||||
async def test_personal_connection_cannot_reuse_same_named_deployment_oauth(task_config, rotating_server, tmp_path, monkeypatch, operation):
|
||||
async def test_personal_connection_cannot_reuse_same_named_deployment_oauth(task_config, rotating_server, tmp_path, monkeypatch, operation, token_persistence):
|
||||
_path, startup = task_config
|
||||
monkeypatch.setattr("deerflow.mcp.user_config.get_paths", lambda: Paths(base_dir=tmp_path))
|
||||
personal = startup.mcp_servers["reports"].model_dump(mode="json")
|
||||
@ -429,3 +496,29 @@ async def test_personal_connection_cannot_reuse_same_named_deployment_oauth(task
|
||||
assert rotating_server.refresh_requests == ["refresh-0", "personal-seed"]
|
||||
assert rotating_server.calls[-1][1] == "Bearer access-2"
|
||||
assert path.read_bytes() == before
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_connection_continues_its_rotated_chain_on_the_next_run(task_config, rotating_server, tmp_path, monkeypatch, shared_token_store):
|
||||
_path, startup = task_config
|
||||
monkeypatch.setattr("deerflow.mcp.user_config.get_paths", lambda: Paths(base_dir=tmp_path))
|
||||
personal = startup.mcp_servers["reports"].model_dump(mode="json")
|
||||
personal["oauth"]["refresh_token"] = "personal-seed"
|
||||
personal["personal_public_network"] = True
|
||||
path = user_mcp_config_path("user-1")
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps({"mcpServers": {"reports": personal}}), encoding="utf-8")
|
||||
name = personal_server_name("user-1", "reports", personal)
|
||||
rotating_server.refresh_token = "personal-seed"
|
||||
set_mcp_task_submitter(SimpleNamespace())
|
||||
|
||||
# Discovery for one run builds its own manager ...
|
||||
tools = await get_mcp_tools(load_user_mcp_config("user-1"), personal_user_id="user-1")
|
||||
assert await tools[1].ainvoke({})
|
||||
_Clock.current = _Clock.current + timedelta(seconds=300)
|
||||
# ... and a later durable poll another one; both use the owner's row.
|
||||
await McpTaskToolCaller(startup).call_tool(server_name=name, tool_name="status_report", arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", connection_scope="personal")
|
||||
|
||||
assert rotating_server.refresh_requests == ["personal-seed", "refresh-1"]
|
||||
assert rotating_server.calls[-1][1] == "Bearer access-2"
|
||||
assert _stored_keys(shared_token_store) == [("personal", "user-1", name)]
|
||||
|
||||
@ -2,7 +2,7 @@
|
||||
|
||||
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_0046_skill_mutations``.
|
||||
chain-head pin lives in ``test_migration_0047_mcp_oauth_tokens``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
"""Additive Skill host schema and current migration-chain head."""
|
||||
"""Additive Skill host schema; the chain-head pin lives in ``test_migration_0047_mcp_oauth_tokens``."""
|
||||
|
||||
import asyncio
|
||||
|
||||
@ -15,9 +15,10 @@ REVISION = "0046_skill_mutations"
|
||||
TABLES = {"skill_mutation_assets", "skill_mutation_owners", "skill_mutation_proposals", "skill_mutation_operations", "skill_mutation_scan_attempts"}
|
||||
|
||||
|
||||
def test_0046_is_the_chain_head():
|
||||
assert bootstrap._get_head_revision() == REVISION
|
||||
def test_0046_remains_in_the_single_migration_chain():
|
||||
script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR))
|
||||
assert len(script.get_heads()) == 1
|
||||
assert script.get_revision("0047_mcp_oauth_tokens").down_revision == REVISION
|
||||
assert script.get_revision(REVISION).down_revision == "0045_completed_run_evidence"
|
||||
assert script.get_revision("0045_completed_run_evidence").down_revision == "0044_channel_runtime_configs"
|
||||
assert script.get_revision("0044_channel_runtime_configs").down_revision == "0043_runtime_im_notices"
|
||||
|
||||
186
backend/tests/test_migration_0047_mcp_oauth_tokens.py
Normal file
186
backend/tests/test_migration_0047_mcp_oauth_tokens.py
Normal file
@ -0,0 +1,186 @@
|
||||
"""Migration 0047: the shared, encrypted ``mcp_oauth_tokens`` table, and the current migration-chain head."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from alembic import command
|
||||
from alembic.autogenerate import compare_metadata
|
||||
from alembic.migration import MigrationContext
|
||||
from alembic.script import ScriptDirectory
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.persistence import bootstrap
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.mcp_oauth_tokens import McpOAuthTokenRow
|
||||
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
|
||||
|
||||
REVISION = "0047_mcp_oauth_tokens"
|
||||
PREVIOUS = "0046_skill_mutations"
|
||||
TABLE = "mcp_oauth_tokens"
|
||||
COLUMNS = {
|
||||
"scope",
|
||||
"owner_user_id",
|
||||
"server_key",
|
||||
"config_fingerprint",
|
||||
"access_token_enc",
|
||||
"token_request_fingerprint",
|
||||
"token_type",
|
||||
"access_expires_at",
|
||||
"refresh_token_enc",
|
||||
"version",
|
||||
"refresh_lease_owner",
|
||||
"refresh_lease_until",
|
||||
"rotated_at",
|
||||
"updated_at",
|
||||
}
|
||||
NULLABLE = {"access_token_enc", "token_request_fingerprint", "token_type", "access_expires_at", "refresh_token_enc", "refresh_lease_owner", "refresh_lease_until", "rotated_at"}
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
async def test_0047_is_the_chain_head():
|
||||
assert bootstrap._get_head_revision() == REVISION
|
||||
script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR))
|
||||
assert script.get_heads() == [REVISION]
|
||||
assert script.get_revision(REVISION).down_revision == PREVIOUS
|
||||
assert script.get_revision(PREVIOUS).down_revision == "0045_completed_run_evidence"
|
||||
# alembic_version.version_num is VARCHAR(32).
|
||||
assert len(REVISION) <= 32
|
||||
|
||||
|
||||
def _engine(tmp_path, backend):
|
||||
if backend == "sqlite":
|
||||
return create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}"), None
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI (real Postgres migration 0047)")
|
||||
schema = f"mcp_oauth_tokens_{uuid.uuid4().hex}"
|
||||
return create_async_engine(asyncpg_test_url(uri), connect_args=build_asyncpg_connect_args(schema)), schema
|
||||
|
||||
|
||||
async def _shape(engine) -> dict:
|
||||
async with engine.connect() as conn:
|
||||
|
||||
def read(sync):
|
||||
inspector = sa.inspect(sync)
|
||||
tables = set(inspector.get_table_names())
|
||||
if TABLE not in tables:
|
||||
return {"tables": tables}
|
||||
return {
|
||||
"tables": tables,
|
||||
"columns": {col["name"]: col for col in inspector.get_columns(TABLE)},
|
||||
"pk": inspector.get_pk_constraint(TABLE),
|
||||
}
|
||||
|
||||
return await conn.run_sync(read)
|
||||
|
||||
|
||||
async def _orm_diff(engine) -> list:
|
||||
"""Alembic's own drift check, restricted to the table this revision owns."""
|
||||
|
||||
def only_oauth_tokens(obj, name, type_, reflected, compare_to):
|
||||
if type_ == "table":
|
||||
return name == TABLE
|
||||
return getattr(getattr(obj, "table", None), "name", TABLE) == TABLE
|
||||
|
||||
async with engine.connect() as conn:
|
||||
|
||||
def diff(sync):
|
||||
context = MigrationContext.configure(sync, opts={"include_object": only_oauth_tokens, "compare_type": True, "compare_server_default": True})
|
||||
return compare_metadata(context, Base.metadata)
|
||||
|
||||
return await conn.run_sync(diff)
|
||||
|
||||
|
||||
async def _drop_schema(engine, schema):
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
|
||||
|
||||
def _row(**overrides) -> dict:
|
||||
values = {"scope": "deployment", "owner_user_id": "", "server_key": "reports", "config_fingerprint": "f" * 64, "version": 0, "updated_at": 1_700_000_000.0}
|
||||
values.update(overrides)
|
||||
return values
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["sqlite", "postgres"])
|
||||
async def test_0047_creates_the_table_matching_the_orm_and_downgrades_cleanly(tmp_path, backend):
|
||||
engine, schema = _engine(tmp_path, backend)
|
||||
cfg = bootstrap._get_alembic_config(engine, postgres_schema=schema or "")
|
||||
try:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f'CREATE SCHEMA "{schema}"'))
|
||||
await asyncio.to_thread(command.upgrade, cfg, PREVIOUS)
|
||||
assert TABLE not in (await _shape(engine))["tables"]
|
||||
|
||||
await asyncio.to_thread(command.upgrade, cfg, REVISION)
|
||||
shape = await _shape(engine)
|
||||
assert set(shape["columns"]) == COLUMNS
|
||||
assert shape["pk"]["constrained_columns"] == ["scope", "owner_user_id", "server_key"]
|
||||
assert shape["pk"]["name"] in ("pk_mcp_oauth_tokens", None) # SQLite does not report constraint names
|
||||
# Token columns are empty while the first refresh is claimed; the key never is.
|
||||
assert {name for name, column in shape["columns"].items() if column["nullable"]} == NULLABLE
|
||||
# ``make migrate-rev`` would propose nothing: the revision matches the ORM model.
|
||||
assert await _orm_diff(engine) == []
|
||||
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.insert(McpOAuthTokenRow).values(**_row()))
|
||||
# A personal connection may reuse a deployment server's name.
|
||||
await conn.execute(sa.insert(McpOAuthTokenRow).values(**_row(scope="personal", owner_user_id="user-1")))
|
||||
with pytest.raises(sa.exc.IntegrityError):
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.insert(McpOAuthTokenRow).values(**_row(version=3)))
|
||||
|
||||
await asyncio.to_thread(command.downgrade, cfg, PREVIOUS)
|
||||
assert TABLE not in (await _shape(engine))["tables"]
|
||||
await asyncio.to_thread(command.upgrade, cfg, REVISION)
|
||||
assert TABLE in (await _shape(engine))["tables"]
|
||||
finally:
|
||||
await _drop_schema(engine, schema)
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["sqlite", "postgres"])
|
||||
async def test_0047_completes_after_create_all_and_downgrades_after_a_manual_drop(tmp_path, backend):
|
||||
"""The table already exists (an interrupted earlier attempt): upgrade still lands, and vice versa."""
|
||||
engine, schema = _engine(tmp_path, backend)
|
||||
cfg = bootstrap._get_alembic_config(engine, postgres_schema=schema or "")
|
||||
try:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f'CREATE SCHEMA "{schema}"'))
|
||||
await asyncio.to_thread(command.upgrade, cfg, PREVIOUS)
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all, tables=[McpOAuthTokenRow.__table__])
|
||||
|
||||
await asyncio.to_thread(command.upgrade, cfg, REVISION)
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.text("SELECT version_num FROM alembic_version"))).scalar_one() == REVISION
|
||||
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f"DROP TABLE {TABLE}"))
|
||||
await asyncio.to_thread(command.downgrade, cfg, PREVIOUS)
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.text("SELECT version_num FROM alembic_version"))).scalar_one() == PREVIOUS
|
||||
finally:
|
||||
await _drop_schema(engine, schema)
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_bootstrap_provisions_the_table_on_an_empty_database(tmp_path):
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'bootstrap.db'}")
|
||||
try:
|
||||
await bootstrap.bootstrap_schema(engine, backend="sqlite")
|
||||
await bootstrap.bootstrap_schema(engine, backend="sqlite")
|
||||
assert set((await _shape(engine))["columns"]) == COLUMNS
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.text("SELECT version_num FROM alembic_version"))).scalar_one() == REVISION
|
||||
finally:
|
||||
await engine.dispose()
|
||||
@ -31,6 +31,7 @@ from app.gateway.routers.browser import _browser_tools_enabled
|
||||
from deerflow.config.credentials_key import CREDENTIALS_KEY_ENV_VAR, GENERATE_KEY_COMMAND
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.config.deployment_config import MULTI_INSTANCE_ENV_VAR, DeploymentConfig, multi_instance_declaration
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.config.run_ownership_config import RunOwnershipConfig
|
||||
from deerflow.config.stream_bridge_config import StreamBridgeConfig
|
||||
|
||||
@ -876,6 +877,32 @@ def test_credentials_key_consumers_name_the_enabled_features():
|
||||
assert credentials_key_consumers(_with_channel_connections(_cluster_ready())) == ["channel_connections.enabled=true"]
|
||||
|
||||
|
||||
def _with_mcp_oauth(config, *, grant_type: str = "refresh_token", oauth_enabled: bool = True, server_enabled: bool = True):
|
||||
oauth = {"enabled": oauth_enabled, "token_url": "https://auth.example.invalid/token", "grant_type": grant_type, "refresh_token": "seed", "client_id": "client", "client_secret": "secret"}
|
||||
config.extensions = ExtensionsConfig.model_validate({"mcpServers": {"reports": {"enabled": server_enabled, "type": "http", "url": "https://mcp.example.invalid/mcp", "oauth": oauth}}})
|
||||
return config
|
||||
|
||||
|
||||
def test_rotating_mcp_oauth_servers_store_tokens_under_the_credentials_key():
|
||||
"""mcp_oauth_tokens keeps refresh_token connections' tokens encrypted; client_credentials stay in process."""
|
||||
assert credentials_key_consumers(_with_mcp_oauth(_cluster_ready())) == ["mcpServers.reports.oauth.enabled=true"]
|
||||
assert credentials_key_consumers(_with_mcp_oauth(_cluster_ready(), grant_type="client_credentials")) == []
|
||||
assert credentials_key_consumers(_with_mcp_oauth(_cluster_ready(), oauth_enabled=False)) == []
|
||||
assert credentials_key_consumers(_with_mcp_oauth(_cluster_ready(), server_enabled=False)) == []
|
||||
|
||||
|
||||
def test_declared_multi_instance_with_rotating_mcp_oauth_requires_the_credentials_key(monkeypatch):
|
||||
"""A peer that cannot decrypt the stored chain falls back to the consumed seed and gets invalid_grant."""
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
_enforce_credentials_key(_with_mcp_oauth(_cluster_ready(deployment_multi_instance=True)))
|
||||
msg = str(exc_info.value)
|
||||
assert CREDENTIALS_KEY_ENV_VAR in msg
|
||||
assert "mcpServers.reports.oauth.enabled=true" in msg
|
||||
assert "or set mcpServers.reports.oauth.enabled=false" in msg
|
||||
monkeypatch.setenv(CREDENTIALS_KEY_ENV_VAR, _VALID_KEY)
|
||||
_enforce_credentials_key(_with_mcp_oauth(_cluster_ready(deployment_multi_instance=True)))
|
||||
|
||||
|
||||
def test_declared_multi_instance_with_channel_connections_requires_the_credentials_key():
|
||||
"""Each instance would otherwise generate its own key file and could not read its peers' credentials."""
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
|
||||
@ -16,7 +16,7 @@ from langchain_core.tools import ToolException
|
||||
|
||||
from app.gateway.routers import personal_mcp
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.mcp.user_config import load_user_mcp_config, read_user_mcp_config, user_mcp_config_path
|
||||
from deerflow.mcp.user_config import load_user_mcp_config, personal_server_name, read_user_mcp_config, user_mcp_config_path
|
||||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
||||
|
||||
|
||||
@ -620,3 +620,82 @@ async def test_unreachable_personal_task_server_keeps_other_personal_tools(perso
|
||||
|
||||
notes_name = next(name for name, server in personal.mcp_servers.items() if server.url == "https://notes.example.com/mcp")
|
||||
assert [tool.name for tool in tools] == [f"{notes_name}_search"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stored_oauth_rows(tmp_path):
|
||||
"""Install a shared OAuth token store; yields a ``(seed, keys)`` pair over its rows."""
|
||||
import sqlalchemy as sa
|
||||
from cryptography.fernet import Fernet
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from deerflow.config.credentials_key import CredentialsCipher
|
||||
from deerflow.mcp.oauth_store import McpOAuthTokenStore, set_mcp_oauth_token_store
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.mcp_oauth_tokens import McpOAuthTokenRow, SqlMcpOAuthTokenRepository
|
||||
|
||||
engine = sa.create_engine(f"sqlite:///{(tmp_path / 'oauth-tokens.db').as_posix()}")
|
||||
Base.metadata.create_all(engine, tables=[McpOAuthTokenRow.__table__])
|
||||
repo = SqlMcpOAuthTokenRepository(sessionmaker(engine, expire_on_commit=False))
|
||||
previous = set_mcp_oauth_token_store(McpOAuthTokenStore(repo, cipher=CredentialsCipher.from_keys([Fernet.generate_key()])))
|
||||
|
||||
def seed(owner, server_key, scope="personal"):
|
||||
assert repo.claim(scope, owner, server_key, config_fingerprint="f" * 64, expected_version=None, lease_owner="seed", lease_seconds=1.0, now=0.0)
|
||||
|
||||
def keys(owner, scope="personal"):
|
||||
with engine.connect() as conn:
|
||||
return set(conn.scalars(sa.select(McpOAuthTokenRow.server_key).where(McpOAuthTokenRow.scope == scope, McpOAuthTokenRow.owner_user_id == owner)))
|
||||
|
||||
try:
|
||||
yield seed, keys
|
||||
finally:
|
||||
set_mcp_oauth_token_store(previous)
|
||||
engine.dispose()
|
||||
|
||||
|
||||
def test_removed_or_edited_connections_drop_their_stored_oauth_tokens(personal_client, stored_oauth_rows):
|
||||
"""A removed or edited connection's runtime name never returns, so its token row would only hold ciphertext forever."""
|
||||
seed, keys = stored_oauth_rows
|
||||
client = personal_client
|
||||
create(client, "alice", name="github")
|
||||
create(client, "alice", name="linear")
|
||||
create(client, "bob", name="github")
|
||||
alice = set(load_user_mcp_config("alice").mcp_servers)
|
||||
bob = set(load_user_mcp_config("bob").mcp_servers)
|
||||
for name in alice:
|
||||
seed("alice", name)
|
||||
for name in bob:
|
||||
seed("bob", name)
|
||||
seed("alice", "personal_left-by-an-earlier-edit")
|
||||
seed("", "github", scope="deployment")
|
||||
github = personal_server_name("alice", "github", read_user_mcp_config("alice")["mcpServers"]["github"])
|
||||
assert github in alice
|
||||
|
||||
# Turning a connection off and on restores its runtime name: its chain survives.
|
||||
for enabled in (False, True):
|
||||
result = client.patch("/api/mcp/personal/config", headers={"test-user": "alice"}, json={"server_name": "github", "enabled": enabled})
|
||||
assert result.status_code == 200, result.text
|
||||
assert keys("alice") == alice, "a toggle must keep the stored chain; an earlier orphan is swept"
|
||||
|
||||
result = client.put("/api/mcp/personal/config/server", headers={"test-user": "alice"}, json={"server_name": "linear", "server": {"type": "http", "url": "https://example.com/linear-v2", "headers": {"Authorization": "***"}}})
|
||||
assert result.status_code == 200, result.text
|
||||
assert keys("alice") == {github}, "the edited connection's old row is gone"
|
||||
|
||||
assert client.delete("/api/mcp/personal/config/servers/github", headers={"test-user": "alice"}).status_code == 200
|
||||
assert keys("alice") == set()
|
||||
assert keys("bob") == bob, "another owner's rows are untouched"
|
||||
assert keys("", scope="deployment") == {"github"}, "deployment rows are untouched"
|
||||
|
||||
|
||||
def test_personal_config_save_survives_an_unavailable_token_store(personal_client, monkeypatch):
|
||||
from deerflow.mcp.oauth_store import McpOAuthTokenStoreUnavailable, set_mcp_oauth_token_store
|
||||
|
||||
store = MagicMock()
|
||||
store.forget_unconfigured.side_effect = McpOAuthTokenStoreUnavailable("down")
|
||||
previous = set_mcp_oauth_token_store(store)
|
||||
try:
|
||||
assert create(personal_client, "alice").status_code == 200
|
||||
finally:
|
||||
set_mcp_oauth_token_store(previous)
|
||||
assert "github" in read_user_mcp_config("alice")["mcpServers"]
|
||||
store.forget_unconfigured.assert_called_once()
|
||||
|
||||
@ -390,7 +390,8 @@ kubectl -n deer-flow exec deploy/deer-flow-provisioner -- curl -s localhost:8002
|
||||
```
|
||||
- **Credentials key.** The app Secret also holds `DEER_FLOW_CREDENTIALS_KEY`,
|
||||
the Fernet key the gateway encrypts stored credentials with (per-connection
|
||||
IM channel credentials and the bot credentials entered in the web UI). The chart generates it once (32 random bytes,
|
||||
IM channel credentials, the bot credentials entered in the web UI, and
|
||||
rotating MCP OAuth refresh tokens). The chart generates it once (32 random bytes,
|
||||
urlsafe base64) and preserves it across upgrades via `lookup`; every gateway
|
||||
Pod reads the same value. **Back it up**: deleting the Secret or the release
|
||||
generates a new key, and credentials stored under the old one become
|
||||
@ -399,7 +400,8 @@ kubectl -n deer-flow exec deploy/deer-flow-provisioner -- curl -s localhost:8002
|
||||
values have been rewritten. An `existingAppSecret` should carry it; the env
|
||||
entry is optional, so without it a single Pod falls back to an auto-generated
|
||||
`.credentials_key` on the home volume, while a multi-instance gateway with
|
||||
`channel_connections.enabled` refuses to start. If you add the key to an
|
||||
`channel_connections.enabled` or a `refresh_token` MCP OAuth server refuses
|
||||
to start. If you add the key to an
|
||||
`existingAppSecret` later, copy that file's value so stored credentials stay
|
||||
readable:
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user