mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(mcp): use _mcp_active_toolset_id ContextVar to detect toolset scope, avoiding DB-default false-positive
This commit is contained in:
parent
9c1ae5d149
commit
2a07c3e90d
3 changed files with 25 additions and 15 deletions
16
litellm/proxy/_experimental/mcp_server/mcp_context.py
Normal file
16
litellm/proxy/_experimental/mcp_server/mcp_context.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
"""
|
||||
Shared ContextVars for the MCP server layer.
|
||||
|
||||
Lives in its own module to avoid circular imports between
|
||||
mcp_server_manager.py and server.py.
|
||||
"""
|
||||
|
||||
from contextvars import ContextVar
|
||||
from typing import Optional
|
||||
|
||||
# Set server-side in proxy_server.py route handlers when a request arrives via
|
||||
# /toolset/{name}/mcp or the toolset fallback in dynamic_mcp_route.
|
||||
# Never populated from client-supplied headers.
|
||||
_mcp_active_toolset_id: ContextVar[Optional[str]] = ContextVar(
|
||||
"_mcp_active_toolset_id", default=None
|
||||
)
|
||||
|
|
@ -793,14 +793,15 @@ class MCPServerManager:
|
|||
)
|
||||
combined_servers = set(allowed_mcp_servers)
|
||||
# Only skip allow_all_keys servers when the request is inside a toolset
|
||||
# scope. _apply_toolset_scope marks this by setting mcp_toolsets=[].
|
||||
# For regular keys (mcp_toolsets=None) or keys with configured toolsets
|
||||
# (mcp_toolsets=[...non-empty...]), allow_all_keys servers must still be
|
||||
# included to preserve backward-compatible behavior.
|
||||
op = user_api_key_auth.object_permission if user_api_key_auth else None
|
||||
in_toolset_scope = (
|
||||
op is not None and op.mcp_servers is not None and op.mcp_toolsets == []
|
||||
# scope. toolset_mcp_route / dynamic_mcp_route set _mcp_active_toolset_id
|
||||
# before calling the handler — that ContextVar is the reliable signal.
|
||||
# Using op.mcp_toolsets==[] would false-positive on DB-default rows where
|
||||
# Postgres initialises the column to ARRAY[]::TEXT[].
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import ( # noqa: PLC0415
|
||||
_mcp_active_toolset_id,
|
||||
)
|
||||
|
||||
in_toolset_scope = _mcp_active_toolset_id.get() is not None
|
||||
if not in_toolset_scope:
|
||||
combined_servers.update(allow_all_server_ids)
|
||||
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ import contextlib
|
|||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
from contextvars import ContextVar
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
Any,
|
||||
|
|
@ -38,6 +37,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
|||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
get_request_base_url,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_active_toolset_id
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
LITELLM_MCP_SERVER_DESCRIPTION,
|
||||
|
|
@ -57,13 +57,6 @@ from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
|
|||
from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall
|
||||
from litellm.utils import Rules, client, function_setup
|
||||
|
||||
# ContextVar holding the active toolset ID when a request arrives via
|
||||
# /toolset/{name}/mcp or /{name}/mcp (toolset fallback). Set server-side
|
||||
# in proxy_server.py route handlers; never read from client-supplied headers.
|
||||
_mcp_active_toolset_id: ContextVar[Optional[str]] = ContextVar(
|
||||
"_mcp_active_toolset_id", default=None
|
||||
)
|
||||
|
||||
# Short-lived in-memory cache for BYOK credentials.
|
||||
# Keyed by (user_id, server_id); value is (credential_or_None, monotonic_timestamp).
|
||||
# Storing the credential value (not just a bool) means _get_byok_credential and
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue