fix(mcp): use _mcp_active_toolset_id ContextVar to detect toolset scope, avoiding DB-default false-positive

This commit is contained in:
Ishaan Jaffer 2026-03-23 13:34:44 -07:00
parent 9c1ae5d149
commit 2a07c3e90d
3 changed files with 25 additions and 15 deletions

View 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
)

View file

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

View file

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