fix(mcp): cache toolset name lookups in MCPServerManager to avoid per-request DB calls

This commit is contained in:
Ishaan Jaffer 2026-03-21 20:39:55 -07:00
parent 03e4a5c9ee
commit 399f3f81ba
2 changed files with 49 additions and 17 deletions

View file

@ -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]

View file

@ -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(