fix: bound _refresh_locks with LRU eviction to prevent memory leak (#3968)

Author: jlowin

src/fastmcp/server/auth/oauth_proxy/proxy.py

@@ -22,6 +22,7 @@
 import secrets
 import time
 from base64 import urlsafe_b64encode
+from collections import OrderedDict
 from typing import Any, Literal
 from urllib.parse import urlencode, urlparse, urlunparse
 
@@ -92,6 +93,8 @@
 
 logger = get_logger(__name__)
 
+_REFRESH_LOCK_CACHE_SIZE = 10_000
+
 
 def _normalize_resource_url(url: str) -> str:
     """Normalize a resource URL by removing query parameters and trailing slashes.
@@ -576,7 +579,7 @@ def __init__(
         # refresh the same token within a single process. Does not protect
         # against cross-process races in distributed deployments — those are
         # handled by re-reading from storage after refresh failure.
-        self._refresh_locks: dict[str, anyio.Lock] = {}
+        self._refresh_locks: OrderedDict[str, anyio.Lock] = OrderedDict()
 
         logger.debug(
             "Initialized OAuth proxy provider with upstream server %s",
@@ -647,6 +650,18 @@ def _create_upstream_oauth_client(self) -> AsyncOAuth2Client:
             timeout=HTTP_TIMEOUT_SECONDS,
         )
 
+    def _get_refresh_lock(self, token_id: str) -> anyio.Lock:
+        """Get or create a per-token refresh lock, evicting LRU entries when at capacity."""
+        lock = self._refresh_locks.get(token_id)
+        if lock is None:
+            lock = anyio.Lock()
+            self._refresh_locks[token_id] = lock
+            if len(self._refresh_locks) > _REFRESH_LOCK_CACHE_SIZE:
+                self._refresh_locks.popitem(last=False)
+        else:
+            self._refresh_locks.move_to_end(token_id)
+        return lock
+
     # -------------------------------------------------------------------------
     # PKCE Helper Methods
     # -------------------------------------------------------------------------
@@ -1656,9 +1671,7 @@ async def load_access_token(self, token: str) -> AccessToken | None:  # type: ig
 
                     # Advisory lock prevents concurrent requests from racing
                     # to refresh the same upstream token.
-                    if token_id not in self._refresh_locks:
-                        self._refresh_locks[token_id] = anyio.Lock()
-                    lock = self._refresh_locks[token_id]
+                    lock = self._get_refresh_lock(token_id)
 
                     async with lock:
                         # Re-read from storage — another task may have

tests/server/auth/oauth_proxy/test_tokens.py

@@ -907,6 +907,38 @@ async def mock_get(key: str) -> UpstreamTokenSet | None:
         assert result is not None
         assert result.token == "refreshed-upstream-access"
 
+    def test_refresh_lock_cache_bounded(self, proxy, monkeypatch):
+        """_get_refresh_lock never grows beyond _REFRESH_LOCK_CACHE_SIZE."""
+        monkeypatch.setattr(
+            "fastmcp.server.auth.oauth_proxy.proxy._REFRESH_LOCK_CACHE_SIZE", 3
+        )
+        for i in range(10):
+            proxy._get_refresh_lock(f"token-{i}")
+        assert len(proxy._refresh_locks) == 3
+
+    def test_refresh_lock_lru_evicts_least_recently_used(self, proxy, monkeypatch):
+        """Touching an entry promotes it; eviction removes the oldest untouched entry."""
+        monkeypatch.setattr(
+            "fastmcp.server.auth.oauth_proxy.proxy._REFRESH_LOCK_CACHE_SIZE", 3
+        )
+        proxy._get_refresh_lock("a")
+        proxy._get_refresh_lock("b")
+        proxy._get_refresh_lock("c")
+        # Touch "a" to move it to MRU position
+        proxy._get_refresh_lock("a")
+        # Adding "d" should evict "b" (oldest untouched)
+        proxy._get_refresh_lock("d")
+        assert "b" not in proxy._refresh_locks
+        assert "a" in proxy._refresh_locks
+        assert "c" in proxy._refresh_locks
+        assert "d" in proxy._refresh_locks
+
+    def test_refresh_lock_same_token_returns_same_lock(self, proxy):
+        """Requesting the same token ID twice returns the same lock object."""
+        lock1 = proxy._get_refresh_lock("tok")
+        lock2 = proxy._get_refresh_lock("tok")
+        assert lock1 is lock2
+
     async def test_upstream_claims_propagated(self, proxy):
         jwt = await self._setup_session_with_claims(
             proxy, upstream_claims={"sub": "user-123"}