From 2a07c3e90de1221e9f42c3ea276593b398c45414 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Mon, 23 Mar 2026 13:34:44 -0700
Subject: [PATCH] fix(mcp): use _mcp_active_toolset_id ContextVar to detect
 toolset scope, avoiding DB-default false-positive

---
 .../_experimental/mcp_server/mcp_context.py      | 16 ++++++++++++++++
 .../mcp_server/mcp_server_manager.py             | 15 ++++++++-------
 litellm/proxy/_experimental/mcp_server/server.py |  9 +--------
 3 files changed, 25 insertions(+), 15 deletions(-)
 create mode 100644 litellm/proxy/_experimental/mcp_server/mcp_context.py

diff --git a/litellm/proxy/_experimental/mcp_server/mcp_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py
new file mode 100644
index 00000000000..12830db1d6a
--- /dev/null
+++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py
@@ -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
+)
diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index 33a8457d626..102bd615738 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -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)
 
diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index eb97e1247ff..166be712d5a 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -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