mirror of
https://github.com/langchain-ai/deepagents.git
synced 2026-07-21 17:25:26 -04:00
fix(code): serialize MCP OAuth token refreshes to prevent reuse revocation (#4565)
Fixed frequent LangSmith MCP OAuth timeouts/re-auth prompts caused by concurrent token refreshes invalidating the session; refreshes are now serialized across dcode processes and provider instances. --- dcode stores MCP OAuth tokens per-server under `~/.deepagents/.state/mcp-tokens/`. The MCP SDK's `OAuthClientProvider` serializes refreshes with an in-memory `OAuthContext.lock`, which only covers a single provider instance. Two dcode processes — or two provider instances in one process (e.g. a throwaway discovery session created during tool loading alongside the cached runtime session) — can each read the same refresh token from disk and `POST` the `refresh_token` grant concurrently. The LangSmith OAuth server (in `langchainplus`) rotates refresh tokens and, on replay of an already-rotated token, treats it as reuse and revokes every active refresh token for that identity+client. That is why the LangSmith MCP worked for a few minutes (one access-token lifetime) and then hung until a full re-auth. This adds a cross-process lock (`filelock`) keyed on a sibling `.lock` file next to each token file. Before refreshing, `_ExpiryAwareOAuthClientProvider` acquires the lock, reloads tokens from disk (so a peer's rotation is observed and the refresh is skipped when the token is already valid), and only then performs and persists the SDK refresh grant. The lock wait is bounded (falls back to a best-effort refresh on timeout so a crashed peer can't hang tool calls), and all blocking file IO runs in a worker thread to keep the `blockbuster`-guarded server loop responsive. The lock file holds no token material. `filelock` is added as a direct dependency (`>=3.12,<4.0.0`); it was already resolved transitively (via `huggingface-hub`) at `3.29.4`. It is MIT-licensed and actively maintained. Made by [Open SWE](https://openswe.vercel.app/agents/2c15284a-7a9d-8076-66e9-3c97ab30976c) --------- Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
This commit is contained in:
@@ -11,6 +11,7 @@ server X") rather than the token itself.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import contextlib
|
||||
import contextvars
|
||||
@@ -24,11 +25,12 @@ import secrets
|
||||
import stat
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Any, Literal, TypedDict
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
from filelock import FileLock, Timeout
|
||||
from mcp.client.auth import OAuthClientProvider, TokenStorage
|
||||
from mcp.client.auth.utils import (
|
||||
build_oauth_authorization_server_metadata_discovery_urls,
|
||||
@@ -129,6 +131,9 @@ token itself is expired/revoked — i.e. the expected re-auth cases our hint
|
||||
replaces. Transient failures (`429`, `5xx`, gateway timeouts) must stay
|
||||
visible so a provider outage isn't silently relabeled as "go re-login".
|
||||
"""
|
||||
_EXPECTED_REAUTH_REFRESH_STATUS_CODES = frozenset(
|
||||
int(status) for status in _EXPECTED_REAUTH_REFRESH_STATUSES
|
||||
)
|
||||
|
||||
|
||||
class _ExpectedReauthLogFilter(logging.Filter):
|
||||
@@ -175,6 +180,16 @@ rejected as expired by the server — without the margin, a 401 sends the SDK
|
||||
into the full re-auth (browser) flow instead of the cheaper refresh grant.
|
||||
"""
|
||||
|
||||
_REFRESH_LOCK_TIMEOUT_SECONDS = 60.0
|
||||
"""Longest a provider waits for the cross-process token-refresh lock.
|
||||
|
||||
Bounds the wait so a live-but-stuck peer (e.g. one whose refresh network call
|
||||
hangs) can't block tool calls indefinitely. A peer that outright crashes is not
|
||||
the concern: the OS releases an `fcntl` lock when the holding process exits. On
|
||||
timeout the provider reloads tokens from disk and avoids using any still-stale
|
||||
refresh token (see `_acquire_refresh_lock`).
|
||||
"""
|
||||
|
||||
|
||||
def resolve_headers(
|
||||
headers: dict[str, str],
|
||||
@@ -287,6 +302,18 @@ class FileTokenStorage(TokenStorage):
|
||||
stem = _token_file_stem(self._server_name, self._server_url)
|
||||
return _tokens_dir() / f"{stem}.json"
|
||||
|
||||
@property
|
||||
def refresh_lock_path(self) -> Path:
|
||||
"""Sibling lock file that serializes token refreshes across processes.
|
||||
|
||||
A dedicated `.lock` file (never the token file itself) lets `filelock`
|
||||
coordinate refreshes between dcode processes and provider instances
|
||||
without ever holding an exclusive lock on the credential file. It holds
|
||||
no token material.
|
||||
"""
|
||||
path = self.path
|
||||
return path.with_name(f"{path.name}.lock")
|
||||
|
||||
async def get_tokens(self) -> OAuthToken | None:
|
||||
"""Return the stored `OAuthToken`, or `None` if none is persisted."""
|
||||
data = self._read()
|
||||
@@ -1154,6 +1181,19 @@ class _ExpiryAwareOAuthClientProvider(OAuthClientProvider):
|
||||
# fail loudly rather than silently regress to the 401-on-restart
|
||||
# bug this class exists to prevent.
|
||||
await super()._initialize()
|
||||
await self._apply_stored_expiry()
|
||||
|
||||
async def _apply_stored_expiry(self) -> None:
|
||||
"""Seed `context.token_expiry_time` from the persisted sidecar.
|
||||
|
||||
Upstream `_initialize` loads stored tokens but leaves the expiry unset,
|
||||
so a token whose access portion expired long ago still reports as valid
|
||||
and is sent stale. Restoring the absolute expiry recorded beside the
|
||||
token lets `is_token_valid` return `False` in time for the cheaper
|
||||
refresh grant to fire. Also caches persisted OAuth metadata so the
|
||||
refresh uses the advertised token endpoint. Safe to call repeatedly, so
|
||||
it doubles as the post-reload expiry refresh.
|
||||
"""
|
||||
if self.context.oauth_metadata is None:
|
||||
get_oauth_metadata = getattr(
|
||||
self.context.storage,
|
||||
@@ -1197,6 +1237,94 @@ class _ExpiryAwareOAuthClientProvider(OAuthClientProvider):
|
||||
)
|
||||
self.context.token_expiry_time = expires_at - _REFRESH_SAFETY_MARGIN_SECONDS
|
||||
|
||||
async def _reload_tokens_from_storage(self) -> None:
|
||||
"""Re-read persisted tokens so a peer's refresh is observed.
|
||||
|
||||
Another dcode process (or a separate provider instance in this process)
|
||||
may have rotated the refresh token on disk while this provider held a
|
||||
now-stale copy in memory. Re-reading before deciding to refresh keeps
|
||||
this provider from replaying an already-rotated refresh token, which
|
||||
the LangSmith OAuth server treats as reuse and punishes by revoking the
|
||||
whole identity+client token family.
|
||||
"""
|
||||
self.context.current_tokens = await self.context.storage.get_tokens()
|
||||
client_info = await self.context.storage.get_client_info()
|
||||
if client_info is not None:
|
||||
self.context.client_info = client_info
|
||||
await self._apply_stored_expiry()
|
||||
|
||||
async def _acquire_refresh_lock(self, lock: FileLock) -> bool:
|
||||
"""Wait for the cross-process refresh lock off the event loop.
|
||||
|
||||
`lock.acquire` blocks for up to `_REFRESH_LOCK_TIMEOUT_SECONDS` while a
|
||||
peer finishes its refresh, so it runs in a worker thread to avoid
|
||||
stalling the event loop for that long.
|
||||
|
||||
Args:
|
||||
lock: The `filelock.FileLock` serializing refreshes for this server
|
||||
(backed by the sidecar `.lock` file, not the token file).
|
||||
|
||||
Returns:
|
||||
`True` when the lock was acquired; `False` when the wait timed out
|
||||
or the lock could not be created, signalling the caller to avoid
|
||||
using the possibly in-flight refresh token after reloading.
|
||||
"""
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
lock.acquire,
|
||||
timeout=_REFRESH_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
except Timeout:
|
||||
# A timeout means a peer may still be mid-refresh with this same
|
||||
# token. Do not refresh unlocked: rotating-token servers can treat
|
||||
# the second grant as reuse and revoke the whole token family.
|
||||
logger.warning(
|
||||
"Timed out after %.0fs waiting for the MCP token refresh lock "
|
||||
"for %s; skipping refresh to avoid refresh-token reuse.",
|
||||
_REFRESH_LOCK_TIMEOUT_SECONDS,
|
||||
self.context.server_url,
|
||||
)
|
||||
return False
|
||||
except OSError as exc:
|
||||
# Creating/locking the sidecar can fail (read-only or missing
|
||||
# tokens dir, permission denial on a hardened host). Avoid an
|
||||
# unlocked refresh so we do not replay a rotating refresh token if a
|
||||
# peer did manage to take the lock.
|
||||
logger.warning(
|
||||
"Could not acquire the MCP token refresh lock for %s (%s); "
|
||||
"skipping refresh to avoid refresh-token reuse.",
|
||||
self.context.server_url,
|
||||
type(exc).__name__,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def _refresh_lock_guard(self, lock_path: Path) -> AsyncIterator[bool]:
|
||||
"""Hold the cross-process refresh lock across the serialized refresh.
|
||||
|
||||
Acquires the lock (waiting up to `_REFRESH_LOCK_TIMEOUT_SECONDS`; on
|
||||
timeout it yields `False` so the caller can avoid the refresh grant) and
|
||||
always releases it on exit. Release is gated on `lock.is_locked` rather
|
||||
than the acquire result, so a cancellation that lands *after* the worker
|
||||
thread took the lock still frees it instead of orphaning it until GC.
|
||||
|
||||
Args:
|
||||
lock_path: Sibling `.lock` path from `FileTokenStorage`.
|
||||
|
||||
Yields:
|
||||
Whether the refresh lock was acquired.
|
||||
"""
|
||||
# `thread_local=False` because acquire and release run in different
|
||||
# `asyncio.to_thread` worker threads; the default would refuse the
|
||||
# cross-thread release and leak the OS lock until process exit.
|
||||
lock = FileLock(str(lock_path), thread_local=False)
|
||||
try:
|
||||
yield await self._acquire_refresh_lock(lock)
|
||||
finally:
|
||||
if lock.is_locked:
|
||||
await asyncio.to_thread(lock.release)
|
||||
|
||||
async def _persist_oauth_metadata(self) -> None:
|
||||
"""Persist discovered public OAuth metadata when storage supports it."""
|
||||
if self.context.oauth_metadata is None:
|
||||
@@ -1210,6 +1338,31 @@ class _ExpiryAwareOAuthClientProvider(OAuthClientProvider):
|
||||
await super()._handle_token_response(response)
|
||||
await self._persist_oauth_metadata()
|
||||
|
||||
async def _handle_locked_refresh_response(self, response: httpx.Response) -> bool:
|
||||
"""Handle a serialized refresh without bypassing SDK re-auth fallback.
|
||||
|
||||
Args:
|
||||
response: Refresh endpoint response returned through the auth generator.
|
||||
|
||||
Returns:
|
||||
`True` when refresh succeeded, otherwise `False` so the caller can
|
||||
continue into the delegated SDK flow.
|
||||
"""
|
||||
try:
|
||||
return bool(await self._handle_refresh_response(response))
|
||||
except Exception:
|
||||
if response.status_code not in _EXPECTED_REAUTH_REFRESH_STATUS_CODES:
|
||||
raise
|
||||
logger.debug(
|
||||
"Locked MCP token refresh for %s failed with %s; "
|
||||
"deferring to the SDK re-auth flow.",
|
||||
self.context.server_url,
|
||||
response.status_code,
|
||||
)
|
||||
self.context.clear_tokens()
|
||||
self._initialized = False
|
||||
return False
|
||||
|
||||
async def async_auth_flow(
|
||||
self,
|
||||
request: httpx.Request,
|
||||
@@ -1281,6 +1434,45 @@ class _ExpiryAwareOAuthClientProvider(OAuthClientProvider):
|
||||
type(exc).__name__,
|
||||
)
|
||||
|
||||
if (
|
||||
not self.context.is_token_valid()
|
||||
and self.context.can_refresh_token()
|
||||
and isinstance(self.context.storage, FileTokenStorage)
|
||||
):
|
||||
# Serialize the refresh across processes and provider instances.
|
||||
# Without this, two holders of the same token file can both
|
||||
# replay the same refresh token; the LangSmith OAuth server
|
||||
# rotates refresh tokens and revokes the entire token family on
|
||||
# reuse, which surfaces as requests hanging until a full
|
||||
# re-auth. `self.context.lock` only guards this one provider,
|
||||
# so a file lock is required for the cross-process case.
|
||||
async with self._refresh_lock_guard(
|
||||
self.context.storage.refresh_lock_path
|
||||
) as refresh_lock_acquired:
|
||||
# A peer may have rotated the token while we waited for the
|
||||
# lock; reload so a now-valid token skips the refresh.
|
||||
await self._reload_tokens_from_storage()
|
||||
if (
|
||||
not self.context.is_token_valid()
|
||||
and self.context.can_refresh_token()
|
||||
):
|
||||
if refresh_lock_acquired:
|
||||
# ASYNC119: the refresh lock must stay held across this
|
||||
# yield — the request/response round-trip is the
|
||||
# critical section being serialized. Release is safe
|
||||
# because httpx deterministically drives and
|
||||
# `aclose()`s this generator (see the delegation note
|
||||
# below), so the guard's `finally` runs rather than
|
||||
# deferring cleanup to GC.
|
||||
refresh_response = yield await self._refresh_token() # noqa: ASYNC119
|
||||
await self._handle_locked_refresh_response(refresh_response)
|
||||
else:
|
||||
# The delegated SDK flow has its own refresh branch;
|
||||
# clear only in-memory tokens so this request falls
|
||||
# through to re-auth instead of replaying the refresh
|
||||
# token while another process may still be using it.
|
||||
self.context.clear_tokens()
|
||||
|
||||
# Delegate to the SDK flow by manually pumping the inner generator so
|
||||
# the HTTP responses httpx feeds back via `auth_flow.asend(response)`
|
||||
# are forwarded into it. A plain `async for` would advance the inner
|
||||
|
||||
@@ -77,6 +77,8 @@ dependencies = [
|
||||
|
||||
# MCP
|
||||
"langchain-mcp-adapters>=0.3.0,<1.0.0",
|
||||
# Cross-process lock serializing MCP OAuth token refreshes.
|
||||
"filelock>=3.12,<4.0.0",
|
||||
|
||||
# ACP
|
||||
"deepagents-acp>=0.0.8,<1.0.0",
|
||||
|
||||
@@ -638,6 +638,343 @@ class TestExpiryAwareOAuthClientProvider:
|
||||
await flow.aclose()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("fake_home")
|
||||
class TestRefreshTokenSerialization:
|
||||
"""Cross-process-safe refresh serialization to avoid refresh-token reuse.
|
||||
|
||||
The LangSmith OAuth server rotates refresh tokens and revokes the entire
|
||||
identity+client token family when an already-rotated token is replayed, so
|
||||
the provider must reload the on-disk token under a lock before refreshing.
|
||||
"""
|
||||
|
||||
async def test_refresh_lock_path_is_sibling_of_token_file(self) -> None:
|
||||
"""The lock lives beside the token file and never replaces it."""
|
||||
storage = FileTokenStorage("notion", server_url="https://mcp.notion.com/mcp")
|
||||
assert storage.refresh_lock_path == storage.path.with_name(
|
||||
f"{storage.path.name}.lock"
|
||||
)
|
||||
assert storage.refresh_lock_path.parent == storage.path.parent
|
||||
assert storage.refresh_lock_path != storage.path
|
||||
|
||||
async def test_skips_refresh_when_peer_already_rotated_on_disk(self) -> None:
|
||||
"""A peer's fresh token is reloaded and used instead of refreshing.
|
||||
|
||||
Guards the reuse fix: if this provider still has a stale token in
|
||||
memory but disk already holds a peer's rotated token, it must attach
|
||||
the reloaded token rather than replay its own (now-revoked) refresh
|
||||
token.
|
||||
"""
|
||||
from deepagents_code.mcp_auth import build_oauth_provider
|
||||
|
||||
storage = FileTokenStorage("notion")
|
||||
await storage.set_client_info(_make_client_info())
|
||||
await storage.set_oauth_metadata(_make_oauth_metadata())
|
||||
await storage.set_tokens(_make_tokens(access_token="stale"))
|
||||
# Backdate the sidecar so the loaded token reports as expired.
|
||||
path = storage.path
|
||||
data = json.loads(path.read_text())
|
||||
data["expires_at"] = time.time() - 60
|
||||
path.write_text(json.dumps(data))
|
||||
|
||||
provider = build_oauth_provider(
|
||||
server_name="notion",
|
||||
server_url="https://mcp.notion.com/mcp",
|
||||
storage=storage,
|
||||
interactive=False,
|
||||
)
|
||||
# Initialize with the stale token held in memory.
|
||||
await provider._initialize()
|
||||
assert provider.context.is_token_valid() is False
|
||||
|
||||
# A peer rotates the token on disk: fresh access token, future expiry.
|
||||
await storage.set_tokens(
|
||||
OAuthToken(
|
||||
access_token="peer-rotated",
|
||||
token_type="Bearer",
|
||||
refresh_token="rt-new",
|
||||
expires_in=3600,
|
||||
)
|
||||
)
|
||||
|
||||
flow = provider.async_auth_flow(
|
||||
httpx.Request("POST", "https://mcp.notion.com/mcp")
|
||||
)
|
||||
first_request = await anext(flow)
|
||||
# No refresh round-trip: the reloaded token is attached directly, and
|
||||
# the first yielded request is the actual server call.
|
||||
assert first_request.headers["Authorization"] == "Bearer peer-rotated"
|
||||
assert str(first_request.url) == "https://mcp.notion.com/mcp"
|
||||
await flow.aclose()
|
||||
|
||||
async def test_performs_locked_refresh_and_persists_rotation(self) -> None:
|
||||
"""A still-stale token triggers exactly one refresh, then persists it."""
|
||||
from deepagents_code.mcp_auth import build_oauth_provider
|
||||
|
||||
storage = FileTokenStorage("notion")
|
||||
await storage.set_client_info(_make_client_info())
|
||||
await storage.set_oauth_metadata(
|
||||
_make_oauth_metadata("https://auth.example/token")
|
||||
)
|
||||
await storage.set_tokens(_make_tokens(access_token="stale"))
|
||||
path = storage.path
|
||||
data = json.loads(path.read_text())
|
||||
data["expires_at"] = time.time() - 60
|
||||
path.write_text(json.dumps(data))
|
||||
|
||||
provider = build_oauth_provider(
|
||||
server_name="notion",
|
||||
server_url="https://mcp.notion.com/mcp",
|
||||
storage=storage,
|
||||
interactive=False,
|
||||
)
|
||||
flow = provider.async_auth_flow(
|
||||
httpx.Request("POST", "https://mcp.notion.com/mcp")
|
||||
)
|
||||
|
||||
refresh_request = await anext(flow)
|
||||
assert refresh_request.method == "POST"
|
||||
assert str(refresh_request.url) == "https://auth.example/token"
|
||||
body = refresh_request.content.decode()
|
||||
assert "grant_type=refresh_token" in body
|
||||
assert "refresh_token=rt" in body
|
||||
|
||||
token_response = httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": "at-rotated",
|
||||
"token_type": "Bearer",
|
||||
"refresh_token": "rt-rotated",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
request=refresh_request,
|
||||
)
|
||||
actual_request = await flow.asend(token_response)
|
||||
assert actual_request.headers["Authorization"] == "Bearer at-rotated"
|
||||
assert str(actual_request.url) == "https://mcp.notion.com/mcp"
|
||||
await flow.aclose()
|
||||
|
||||
# The rotated pair is persisted so the next process reads it too.
|
||||
persisted = await storage.get_tokens()
|
||||
assert persisted is not None
|
||||
assert persisted.access_token == "at-rotated"
|
||||
assert persisted.refresh_token == "rt-rotated"
|
||||
|
||||
async def test_locked_refresh_rejection_delegates_to_sdk_reauth_fallback(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A rejected locked refresh does not bypass the SDK re-auth fallback."""
|
||||
from mcp.client.auth import OAuthClientProvider
|
||||
|
||||
from deepagents_code.mcp_auth import build_oauth_provider
|
||||
|
||||
storage = FileTokenStorage("notion")
|
||||
await storage.set_client_info(_make_client_info())
|
||||
await storage.set_oauth_metadata(
|
||||
_make_oauth_metadata("https://auth.example/token")
|
||||
)
|
||||
await storage.set_tokens(_make_tokens(access_token="stale"))
|
||||
path = storage.path
|
||||
data = json.loads(path.read_text())
|
||||
data["expires_at"] = time.time() - 60
|
||||
path.write_text(json.dumps(data))
|
||||
|
||||
delegated: dict[str, bool] = {}
|
||||
|
||||
async def fake_sdk_flow(
|
||||
self: OAuthClientProvider,
|
||||
request: httpx.Request,
|
||||
):
|
||||
delegated["entered"] = True
|
||||
assert self.context.current_tokens is None
|
||||
_ = yield request
|
||||
|
||||
monkeypatch.setattr(OAuthClientProvider, "async_auth_flow", fake_sdk_flow)
|
||||
|
||||
provider = build_oauth_provider(
|
||||
server_name="notion",
|
||||
server_url="https://mcp.notion.com/mcp",
|
||||
storage=storage,
|
||||
interactive=False,
|
||||
)
|
||||
|
||||
async def raise_refresh_failure(response: httpx.Response) -> bool:
|
||||
msg = "refresh token rejected"
|
||||
raise httpx.HTTPStatusError(
|
||||
msg,
|
||||
request=response.request,
|
||||
response=response,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
provider,
|
||||
"_handle_refresh_response",
|
||||
raise_refresh_failure,
|
||||
)
|
||||
|
||||
flow = provider.async_auth_flow(
|
||||
httpx.Request("POST", "https://mcp.notion.com/mcp")
|
||||
)
|
||||
|
||||
refresh_request = await anext(flow)
|
||||
assert str(refresh_request.url) == "https://auth.example/token"
|
||||
|
||||
actual_request = await flow.asend(httpx.Response(401, request=refresh_request))
|
||||
|
||||
assert delegated == {"entered": True}
|
||||
assert str(actual_request.url) == "https://mcp.notion.com/mcp"
|
||||
await flow.aclose()
|
||||
|
||||
async def _build_stale_refreshable_provider(
|
||||
self,
|
||||
) -> tuple[Any, FileTokenStorage]:
|
||||
"""Build a provider whose on-disk token is expired but refreshable.
|
||||
|
||||
Shared setup for the lock-behavior tests: client info and OAuth
|
||||
metadata are persisted and the sidecar expiry is backdated so the
|
||||
refresh branch fires against `https://auth.example/token`.
|
||||
"""
|
||||
from deepagents_code.mcp_auth import build_oauth_provider
|
||||
|
||||
storage = FileTokenStorage("notion")
|
||||
await storage.set_client_info(_make_client_info())
|
||||
await storage.set_oauth_metadata(
|
||||
_make_oauth_metadata("https://auth.example/token")
|
||||
)
|
||||
await storage.set_tokens(_make_tokens(access_token="stale"))
|
||||
path = storage.path
|
||||
data = json.loads(path.read_text())
|
||||
data["expires_at"] = time.time() - 60
|
||||
path.write_text(json.dumps(data))
|
||||
|
||||
provider = build_oauth_provider(
|
||||
server_name="notion",
|
||||
server_url="https://mcp.notion.com/mcp",
|
||||
storage=storage,
|
||||
interactive=False,
|
||||
)
|
||||
return provider, storage
|
||||
|
||||
async def test_lock_timeout_skips_refresh_token_reuse(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""A peer-held lock timeout avoids replaying the refresh token.
|
||||
|
||||
The wait is shrunk so the contended acquire times out promptly; the
|
||||
provider must not fire either its locked refresh or the delegated SDK
|
||||
refresh while a peer may still be using the same rotating refresh token.
|
||||
"""
|
||||
from filelock import FileLock
|
||||
|
||||
from deepagents_code import mcp_auth
|
||||
|
||||
provider, storage = await self._build_stale_refreshable_provider()
|
||||
monkeypatch.setattr(mcp_auth, "_REFRESH_LOCK_TIMEOUT_SECONDS", 0.1)
|
||||
|
||||
# A peer holds the refresh lock for the duration of this flow, forcing
|
||||
# the provider's acquire to time out.
|
||||
peer_lock = FileLock(str(storage.refresh_lock_path), thread_local=False)
|
||||
peer_lock.acquire()
|
||||
try:
|
||||
caplog.set_level(logging.WARNING, logger="deepagents_code.mcp_auth")
|
||||
flow = provider.async_auth_flow(
|
||||
httpx.Request("POST", "https://mcp.notion.com/mcp")
|
||||
)
|
||||
actual_request = await anext(flow)
|
||||
assert str(actual_request.url) == "https://mcp.notion.com/mcp"
|
||||
assert "Authorization" not in actual_request.headers
|
||||
await flow.aclose()
|
||||
finally:
|
||||
peer_lock.release()
|
||||
|
||||
assert any(
|
||||
"skipping refresh to avoid refresh-token reuse" in record.getMessage()
|
||||
for record in caplog.records
|
||||
)
|
||||
|
||||
async def test_refresh_waits_for_peer_holding_lock_then_proceeds(self) -> None:
|
||||
"""The flow blocks on the refresh lock until the peer releases it.
|
||||
|
||||
Proves the serialization the PR exists for: while a peer holds the file
|
||||
lock the provider cannot reach the refresh, and it proceeds only once
|
||||
the lock is free. Guards against a regression that drops the file lock
|
||||
(e.g. reverting to the per-provider `context.lock`).
|
||||
"""
|
||||
from filelock import FileLock
|
||||
|
||||
provider, storage = await self._build_stale_refreshable_provider()
|
||||
|
||||
async def drive_to_refresh() -> httpx.Request:
|
||||
# Keep the whole generator lifecycle on one task: its inner
|
||||
# `anyio` lock is task-affine, so driving `anext` here and
|
||||
# `aclose` on another task would raise on release.
|
||||
flow = provider.async_auth_flow(
|
||||
httpx.Request("POST", "https://mcp.notion.com/mcp")
|
||||
)
|
||||
try:
|
||||
return await anext(flow)
|
||||
finally:
|
||||
await flow.aclose()
|
||||
|
||||
peer_lock = FileLock(str(storage.refresh_lock_path), thread_local=False)
|
||||
peer_lock.acquire()
|
||||
task = asyncio.ensure_future(drive_to_refresh())
|
||||
try:
|
||||
# While the peer holds the lock, the flow can't reach the refresh.
|
||||
await asyncio.sleep(0.2)
|
||||
assert not task.done()
|
||||
|
||||
peer_lock.release()
|
||||
|
||||
# Once the lock frees, the provider acquires it and yields the refresh.
|
||||
refresh_request = await asyncio.wait_for(task, timeout=5)
|
||||
assert str(refresh_request.url) == "https://auth.example/token"
|
||||
finally:
|
||||
if peer_lock.is_locked:
|
||||
peer_lock.release()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
async def test_refresh_lock_released_after_successful_refresh(self) -> None:
|
||||
"""A completed locked refresh frees the lock for the next caller.
|
||||
|
||||
Defends against a self-deadlock regression: if the guard failed to
|
||||
release, the next process to refresh this server would hang.
|
||||
"""
|
||||
from filelock import FileLock, Timeout
|
||||
|
||||
provider, storage = await self._build_stale_refreshable_provider()
|
||||
|
||||
flow = provider.async_auth_flow(
|
||||
httpx.Request("POST", "https://mcp.notion.com/mcp")
|
||||
)
|
||||
refresh_request = await anext(flow)
|
||||
token_response = httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": "at-rotated",
|
||||
"token_type": "Bearer",
|
||||
"refresh_token": "rt-rotated",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
request=refresh_request,
|
||||
)
|
||||
await flow.asend(token_response)
|
||||
await flow.aclose()
|
||||
|
||||
# The guard must have released; a fresh holder takes the lock at once.
|
||||
probe = FileLock(str(storage.refresh_lock_path), thread_local=False)
|
||||
try:
|
||||
await asyncio.to_thread(probe.acquire, timeout=0)
|
||||
except Timeout:
|
||||
pytest.fail("refresh lock was not released after a successful refresh")
|
||||
else:
|
||||
probe.release()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("fake_home")
|
||||
class TestBasicAuthClientIdStripping:
|
||||
"""Tests for dropping the duplicate body `client_id` under HTTP Basic auth."""
|
||||
|
||||
Generated
+2
@@ -1117,6 +1117,7 @@ dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
{ name = "deepagents" },
|
||||
{ name = "deepagents-acp" },
|
||||
{ name = "filelock" },
|
||||
{ name = "httpx" },
|
||||
{ name = "langchain" },
|
||||
{ name = "langchain-anthropic" },
|
||||
@@ -1290,6 +1291,7 @@ requires-dist = [
|
||||
{ name = "deepagents-acp", specifier = ">=0.0.8,<1.0.0" },
|
||||
{ name = "deepagents-code", extras = ["agentcore", "daytona", "modal", "runloop", "vercel"], marker = "extra == 'all-sandboxes'" },
|
||||
{ name = "deepagents-code", extras = ["anthropic", "baseten", "bedrock", "cohere", "deepseek", "fireworks", "google-genai", "groq", "huggingface", "ibm", "litellm", "mistralai", "nvidia", "ollama", "openai", "openrouter", "perplexity", "together", "vertex", "xai"], marker = "extra == 'all-providers'" },
|
||||
{ name = "filelock", specifier = ">=3.12,<4.0.0" },
|
||||
{ name = "httpx", specifier = ">=0.28.1,<1.0.0" },
|
||||
{ name = "langchain", specifier = ">=1.3.11,<2.0.0" },
|
||||
{ name = "langchain-agentcore-codeinterpreter", marker = "extra == 'agentcore'", specifier = ">=0.0.5,<1.0.0" },
|
||||
|
||||
Reference in New Issue
Block a user