From ba95b204c7e348d4d74a646f718a0e11b5b55879 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaan@berri.ai>
Date: Fri, 1 May 2026 11:24:42 -0700
Subject: [PATCH] fix(mcp): use WeakValueDictionary for token exchange locks to
 prevent memory leak

Per-user lock entries are now eligible for GC once no coroutine holds
a reference, so rotating JWTs across many users don't accumulate entries
indefinitely. Also moved IDP error body to debug log only to avoid
leaking partial token values or internal metadata into server error replies.
---
 .../mcp_server/auth/token_exchange.py         | 22 +++++++++++++++----
 test_obo_config.yaml                          | 20 -----------------
 2 files changed, 18 insertions(+), 24 deletions(-)
 delete mode 100644 test_obo_config.yaml

diff --git a/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py b/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py
index 6750c4359b1..68ff367e594 100644
--- a/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py
+++ b/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py
@@ -10,6 +10,7 @@ See: https://datatracker.ietf.org/doc/html/rfc8693
 
 import asyncio
 import hashlib
+import weakref
 from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
 import httpx
@@ -46,10 +47,18 @@ class TokenExchangeHandler:
             max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
             default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
         )
-        self._locks: Dict[str, asyncio.Lock] = {}
+        # WeakValueDictionary so locks are GC'd once no coroutine holds a reference,
+        # preventing unbounded growth with many rotating user tokens.
+        self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
+            weakref.WeakValueDictionary()
+        )
 
     def _get_lock(self, cache_key: str) -> asyncio.Lock:
-        return self._locks.setdefault(cache_key, asyncio.Lock())
+        lock = self._locks.get(cache_key)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._locks[cache_key] = lock
+        return lock
 
     @staticmethod
     def _cache_key(subject_token: str, server_id: str) -> str:
@@ -131,10 +140,15 @@ class TokenExchangeHandler:
             response = await client.post(endpoint, data=data)
             response.raise_for_status()
         except httpx.HTTPStatusError as exc:
+            verbose_logger.debug(
+                "Token exchange IDP error for MCP server %s (status %d): %s",
+                server.server_id,
+                exc.response.status_code,
+                exc.response.text,
+            )
             raise ValueError(
                 f"Token exchange for MCP server '{server.server_id}' "
-                f"failed with status {exc.response.status_code}: "
-                f"{exc.response.text}"
+                f"failed with status {exc.response.status_code}"
             ) from exc
 
         body = response.json()
diff --git a/test_obo_config.yaml b/test_obo_config.yaml
deleted file mode 100644
index 6008fbb702d..00000000000
--- a/test_obo_config.yaml
+++ /dev/null
@@ -1,20 +0,0 @@
-model_list:
-  - model_name: gpt-4o
-    litellm_params:
-      model: gpt-4o
-
-mcp_servers:
-  test_obo:
-    url: http://localhost:9998/mcp
-    transport: http
-    auth_type: oauth2_token_exchange
-    token_exchange_endpoint: http://localhost:9999/oauth2/token
-    audience: api://test-mcp-server
-    scopes:
-      - mcp.tools.read
-      - mcp.tools.execute
-    client_id: litellm-client
-    client_secret: litellm-secret
-
-general_settings:
-  master_key: sk-1234