mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(mcp): cache toolset name lookups in MCPServerManager to avoid per-request DB calls
This commit is contained in:
parent
03e4a5c9ee
commit
399f3f81ba
2 changed files with 49 additions and 17 deletions
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue