From 399f3f81ba10bdea6dd5477da07d75a0b3d3d85b Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Sat, 21 Mar 2026 20:39:55 -0700
Subject: [PATCH] fix(mcp): cache toolset name lookups in MCPServerManager to
 avoid per-request DB calls

---
 .../mcp_server/mcp_server_manager.py          | 34 +++++++++++++++++++
 litellm/proxy/proxy_server.py                 | 32 ++++++++---------
 2 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index ebe7d49ffb8..273be3d1879 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -184,6 +184,8 @@ class MCPServerManager:
         """
 
         self._toolset_perm_cache: Dict[str, Tuple[Dict[str, List[str]], float]] = {}
+        # name → (toolset | None, cached_at)  — 60 s TTL
+        self._toolset_name_cache: Dict[str, Tuple[Optional[Any], float]] = {}
 
     def get_registry(self) -> Dict[str, MCPServer]:
         """
@@ -852,12 +854,44 @@ class MCPServerManager:
         """
         if toolset_id is None:
             self._toolset_perm_cache.clear()
+            self._toolset_name_cache.clear()
             return
         keys_to_remove = [
             k for k in self._toolset_perm_cache if toolset_id in k.split(",")
         ]
         for k in keys_to_remove:
             del self._toolset_perm_cache[k]
+        # Also evict any name cache entry that refers to this toolset_id
+        name_keys_to_remove = [
+            name
+            for name, (ts, _) in self._toolset_name_cache.items()
+            if ts is not None and getattr(ts, "toolset_id", None) == toolset_id
+        ]
+        for k in name_keys_to_remove:
+            del self._toolset_name_cache[k]
+
+    async def get_toolset_by_name_cached(
+        self,
+        prisma_client: Any,
+        toolset_name: str,
+    ) -> Optional[Any]:
+        """Return a toolset by name, using a 60 s in-memory TTL cache.
+
+        Avoids a DB hit on every MCP request routed through a named toolset.
+        """
+        cached_entry = self._toolset_name_cache.get(toolset_name)
+        if cached_entry is not None:
+            toolset, cached_at = cached_entry
+            if time.time() - cached_at < 60:
+                return toolset
+
+        from litellm.proxy._experimental.mcp_server.toolset_db import (
+            get_mcp_toolset_by_name,
+        )
+
+        toolset = await get_mcp_toolset_by_name(prisma_client, toolset_name)
+        self._toolset_name_cache[toolset_name] = (toolset, time.time())
+        return toolset
 
     def filter_server_ids_by_ip(
         self, server_ids: List[str], client_ip: Optional[str]
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index aa1fcf5fd8f..f10a62f328b 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -13585,17 +13585,20 @@ async def toolset_mcp_route(toolset_name: str, request: Request):
     Any valid API key can discover and call the toolset's tools here.
     """
     try:
-        from litellm.proxy._experimental.mcp_server.server import (
-            handle_streamable_http_mcp,
+        from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+            global_mcp_server_manager,
         )
-        from litellm.proxy._experimental.mcp_server.toolset_db import (
-            get_mcp_toolset_by_name,
+        from litellm.proxy._experimental.mcp_server.server import (
+            _mcp_active_toolset_id,
+            handle_streamable_http_mcp,
         )
 
         if prisma_client is None:
             raise HTTPException(status_code=503, detail="Database not available")
 
-        toolset = await get_mcp_toolset_by_name(prisma_client, toolset_name)
+        toolset = await global_mcp_server_manager.get_toolset_by_name_cached(
+            prisma_client, toolset_name
+        )
         if toolset is None:
             raise HTTPException(
                 status_code=404,
@@ -13605,10 +13608,6 @@ async def toolset_mcp_route(toolset_name: str, request: Request):
         scope = dict(request.scope)
         scope["path"] = "/mcp"
 
-        from litellm.proxy._experimental.mcp_server.server import (
-            _mcp_active_toolset_id,
-        )
-
         token = _mcp_active_toolset_id.set(toolset.toolset_id)
         try:
             return await _stream_mcp_asgi_response(
@@ -13649,22 +13648,21 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request):
             # Check if this is a toolset name — toolsets are accessible at /{name}/mcp
             # the same way individual servers are, no separate /toolset/ prefix needed.
             if prisma_client is not None:
+                from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+                    global_mcp_server_manager,
+                )
                 from litellm.proxy._experimental.mcp_server.server import (
+                    _mcp_active_toolset_id,
                     handle_streamable_http_mcp,
                 )
-                from litellm.proxy._experimental.mcp_server.toolset_db import (
-                    get_mcp_toolset_by_name,
-                )
 
-                toolset = await get_mcp_toolset_by_name(prisma_client, mcp_server_name)
+                toolset = await global_mcp_server_manager.get_toolset_by_name_cached(
+                    prisma_client, mcp_server_name
+                )
                 if toolset is not None:
                     scope = dict(request.scope)
                     scope["path"] = "/mcp"
 
-                    from litellm.proxy._experimental.mcp_server.server import (
-                        _mcp_active_toolset_id,
-                    )
-
                     token = _mcp_active_toolset_id.set(toolset.toolset_id)
                     try:
                         return await _stream_mcp_asgi_response(