From 4abbe8370fe99b5b32775f4cf7e7529a6aeba237 Mon Sep 17 00:00:00 2001
From: Yug <yugborana000@gmail.com>
Date: Fri, 1 May 2026 10:48:05 +0530
Subject: [PATCH] resolve

---
 .../proxy/_experimental/mcp_server/sampling_handler.py | 10 +++++++++-
 tests/mcp_tests/test_mcp_server.py                     |  6 ++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py
index 8d87ee6d528..31e8c5a7887 100644
--- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py
+++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py
@@ -472,8 +472,10 @@ async def handle_sampling_create_message(
         if openai_tool_choice is not None:
             completion_kwargs["tool_choice"] = openai_tool_choice
         # 5. Add metadata for tracking
+        completion_kwargs["metadata"] = {}
         if params.metadata:
-            completion_kwargs["metadata"] = params.metadata
+            # We nest MCP metadata to avoid collisions with internal LiteLLM auth keys
+            completion_kwargs["metadata"]["mcp_metadata"] = params.metadata
 
         # 6. Inject auth context for cost tracking
         if user_api_key_auth:
@@ -536,10 +538,14 @@ async def handle_sampling_create_message(
         from litellm.exceptions import (
             AuthenticationError,
             BudgetExceededError,
+            ContextWindowExceededError,
             PermissionDeniedError,
             RateLimitError,
+            ServiceUnavailableError,
         )
 
+        # Re-raise known LiteLLM errors so they can be handled by the proxy's
+        # global exception handlers or retry logic if applicable.
         if isinstance(
             e,
             (
@@ -547,6 +553,8 @@ async def handle_sampling_create_message(
                 RateLimitError,
                 AuthenticationError,
                 PermissionDeniedError,
+                ContextWindowExceededError,
+                ServiceUnavailableError,
             ),
         ):
             raise
diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py
index cacd2c171f9..7b2730efabe 100644
--- a/tests/mcp_tests/test_mcp_server.py
+++ b/tests/mcp_tests/test_mcp_server.py
@@ -471,6 +471,11 @@ async def test_sse_mcp_handler_mock():
     mock_server.run = AsyncMock()
     mock_server.create_initialization_options = MagicMock()
 
+    # Mock the gateway request scope as an async context manager
+    mock_scope_manager = MagicMock()
+    mock_scope_manager.__aenter__ = AsyncMock()
+    mock_scope_manager.__aexit__ = AsyncMock()
+
     with (
         patch(
             "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
@@ -490,6 +495,7 @@ async def test_sse_mcp_handler_mock():
         ),
         patch(
             "litellm.proxy._experimental.mcp_server.server._gateway_initialize_instructions_request_scope",
+            return_value=mock_scope_manager,
         ),
         patch(
             "litellm.proxy._experimental.mcp_server.server.set_auth_context",