Litellm oss staging 040626 (#29671)

* fix(azure): apply api_version fallback chain to image edit URL

`AzureImageEditConfig.get_complete_url` only read `api_version` from
`litellm_params`. When callers configured it via `litellm.api_version`
or `AZURE_API_VERSION`, the constructed URL had no `?api-version=` and
Azure responded `404 Resource not found`.

Apply the same fallback chain the Azure chat path already uses in
`common_utils.py`:

    litellm_params > litellm.api_version > AZURE_API_VERSION env >
    litellm.AZURE_DEFAULT_API_VERSION

Adds 5 unit tests pinning each layer of the chain plus a regression
guard for `api_base` that already carries `?api-version=`.

* feat(mcp): core sampling and elicitation flow with security hardening

- Add sampling_handler.py: full MCP sampling/createMessage flow with
  model selection (hint-based + priority-based), auth enforcement,
  budget checks, route restriction gates, and tag policy pre-auth
- Add elicitation_handler.py: MCP elicitation/create relay with
  downstream client capability detection
- Wire sampling/elicitation callbacks in mcp_server_manager.py
  gated behind allow_sampling/allow_elicitation config flags
- Add allow_sampling/allow_elicitation fields to MCPServer type
- Fix session lock deadlock: skip lock for JSON-RPC response POSTs
  (elicitation/sampling replies) with truncated-body heuristic
- Extend client.py with sampling_callback and elicitation_callback
- Security: RouteChecks gate, tag-budget bypass fix, x-forwarded-for
  spoofing fix, Latin-1 header encoding guard
- Add 4 new test modules (model access, priority selection, request
  builder, tool conversion) + update existing MCP tests

* fix(security): run pre-call guardrails before MCP sampling acompletion

Without this, an upstream MCP server with allow_sampling enabled could
send prompts that bypass every guardrail (content filtering, PII
redaction, prompt-injection detection) configured on /chat/completions.

- Call proxy_logging_obj.pre_call_hook(call_type='acompletion') before
  llm_router.acompletion so guardrails fire for sampling sub-calls
- Add HTTPException to the re-raise list so guardrail rejections
  propagate correctly instead of being swallowed as generic errors

* feat(bedrock_mantle): add Responses API support (/openai/v1/responses) (#29490)

* feat(bedrock_mantle): add Responses API transformation config

* test(bedrock_mantle): cover trailing-slash api_base normalization

* feat(bedrock_mantle): export BedrockMantleResponsesAPIConfig

* feat(bedrock_mantle): register gpt-5.x Responses config (gpt-oss unchanged)

* feat(bedrock_mantle): add gpt-5.5/gpt-5.4 Responses price-map entries

* refactor(bedrock_mantle): exclude gpt-oss instead of allow-listing gpt-5 for Responses routing

Frontier OpenAI models on Bedrock Mantle are Responses-only on /openai/v1/responses;
gpt-oss is the legacy family that also speaks chat-completions. Gate by excluding
gpt-oss (which keeps its chat-completions emulation) and defaulting everything else
to the native Responses config, so future frontier models (gpt-6, etc.) route
correctly without a code change. Verified against the live us-east-2 Mantle endpoint:
gpt-oss 400s on /openai/v1/responses while gpt-5.5 400s on both standard paths.

* test(bedrock_mantle): cover supports_native_websocket opt-out

Closes the one uncovered line flagged by codecov on the Responses config.
The assertion documents that Mantle Responses has no realtime/websocket
transport, so realtime routing must not attempt a socket it cannot serve.

* fix(bedrock_mantle): route file_search through emulation instead of forwarding to Mantle

BedrockMantleResponsesAPIConfig inherited supports_native_file_search()
-> True from OpenAIResponsesAPIConfig but never overrode it. Mantle has no
OpenAI vector stores, so a forwarded file_search tool is rejected with a
400 (verified upstream: Tool type 'file_search' is not supported). Opting
out, like the existing supports_native_websocket override, routes the tool
through LiteLLM's file_search emulation instead.

* fix(bedrock_mantle): only route openai.gpt frontier models to Responses

The previous gate excluded gpt-oss and routed every other model to the
native Responses config. But on Mantle only the OpenAI gpt frontier models
(gpt-5.x) are served on /openai/v1/responses; gpt-oss and the non-OpenAI
families (nvidia, mistral, google, zai, ...) are chat-completions only and
400 on that path. Allow-list the openai.gpt- family (excluding gpt-oss)
instead, so chat-only models fall through to the chat-completions emulation.
Verified against the live us-east-2 endpoint: nvidia.nemotron-nano-9b-v2
returns 400 on /openai/v1/responses and 200 on /v1/chat/completions.

* feat(custom_llm): allow streaming/astreaming to yield ModelResponseStream (#27580)

* fix(custom_llm): allow streaming/astreaming to yield ModelResponseStream directly

* fix(streaming): enhance ModelResponseStream handling for custom LLM providers

* fix(streaming): strip finish_reason from content chunks and ensure tool_calls are preserved

* fix(streaming): add type ignore for finish_reason assignment in CustomStreamWrapper

* fix(proxy): strip stack trace from HTTP 503 responses (CWE-209) (#28330)

* fix(proxy/cwe-209): strip Python traceback from HTTP 503 error responses

The /cache/ping endpoint included a full Python traceback in its 503 error
response body (inside the ProxyException message), leaking internal file
paths, line numbers, and call stacks to any caller. Two MCP route handlers
in proxy_server.py similarly interpolated str(e) into "Internal server
error" detail strings.

Fix: log the traceback server-side via verbose_proxy_logger.exception()
and omit it from the ProxyException payload / HTTPException detail returned
to clients. Tests updated to assert no "traceback" keyword or frame paths
appear in the 503 body, with a new dedicated regression test.

CWE-209: Generation of Error Message Containing Sensitive Information.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(proxy/cwe-209): apply Greptile P2 fixes and add MCP exception-path tests

Greptile 4/5 review identified two remaining gaps and Codecov reported
0% coverage on the two MCP handler exception branches:

1. caching_routes.py — str(e) in "Service Unhealthy ({str(e)})" could
   still leak Redis hostnames/IPs; replaced with static "Service Unhealthy".
   HTTPException is now re-raised before the generic handler so the
   "cache not initialized" 503 still reaches callers with its detail.
   Removed the redundant str(e) arg from verbose_proxy_logger.exception()
   (exception() already appends the traceback automatically).

2. tests — two new unit tests cover the exception paths in
   dynamic_mcp_route and toolset_mcp_route that were previously at 0%:
   - test_dynamic_mcp_route_unexpected_exception_returns_500_without_traceback
   - test_toolset_mcp_route_unexpected_exception_returns_500_without_traceback

All 25 tests pass (9 caching + 16 MCP).

CWE-209: Generation of Error Message Containing Sensitive Information.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* test(caching_routes): restore precise assertion in test_cache_ping_no_cache_initialized

The assertion was weakened to `"Cache not initialized" in str(data)`, which
matches the raw string of the entire response dict and would pass even if the
error moved to an unexpected field or changed structure.

Restore a targeted check on the parsed response: assert the exact string in
the correct field `data["detail"]`, matching FastAPI's HTTPException
serialisation format {"detail": "<message>"}.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* test(caching_routes): restore precise assertion and add CWE-209 no-cache path test

The assertion in test_cache_ping_no_cache_initialized was weakened to
`"Cache not initialized" in str(data)`, which matched against the raw string
representation of the entire response dict. This would pass silently even if
the error message moved to an unexpected field or the structure changed.

Restore a targeted assertion on the parsed field:
  assert data["detail"] == "Cache not initialized. litellm.cache is None"
matching FastAPI's HTTPException serialisation format exactly.

Add test_cache_ping_no_cache_does_not_expose_internals to show the code path
is still working correctly after the CWE-209 fix: verifies that the HTTPException
is re-raised as-is (no traceback, no source paths), and asserts the complete
response structure is exactly {"detail": "Cache not initialized. litellm.cache is None"}.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(caching_routes): restore ProxyException envelope for null-cache 503

The except HTTPException: raise guard (added in the CWE-209 fix) caused
the null-cache HTTPException to escape as FastAPI's {"detail": "..."} shape
instead of the {"error": {...}} ProxyException envelope that callers expect.

Move the null-cache guard before the try block and raise ProxyException
directly so the response structure is consistent with all other /cache/ping
503s, and the except HTTPException: raise guard is only reachable by
unexpected downstream HTTPExceptions.

Update the two no-cache tests to assert the correct ProxyException envelope.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

* Update utils.py (#26609)

* feat(pricing): add Snowflake Cortex REST API model pricing (#26612)

* feat(pricing): add Snowflake Cortex REST API model pricing

## Summary

Adds pricing and context window information for 20+ Snowflake Cortex REST API models to `model_prices_and_context_window.json`.

## What's included

- **7 Claude models** (sonnet-4-5, sonnet-4-6, 4-sonnet, 4-opus, haiku-4-5, 3-7-sonnet, 3-5-sonnet) — with prompt caching rates
- **4 OpenAI models** (gpt-4.1, gpt-5, gpt-5-mini, gpt-5-nano) — with prompt caching rates  
- **5 Llama models** (3.1-8b, 3.1-70b, 3.1-405b, 3.3-70b, 4-maverick)
- **1 DeepSeek model** (deepseek-r1)
- **1 Mistral model** (mistral-large2)
- **1 Snowflake model** (snowflake-llama-3.3-70b)
- **2 Embedding models** (arctic-embed-l-v2.0, arctic-embed-m-v2.0)

Each entry includes `input_cost_per_token`, `output_cost_per_token`, `cache_read_input_token_cost` (where applicable), `max_input_tokens`, `max_output_tokens`, and capability flags (`supports_function_calling`, `supports_vision`, `supports_prompt_caching`, `supports_reasoning`).

## Pricing source

All prices are in USD per token, sourced from the official [Snowflake Service Consumption Table](https://www.snowflake.com/legal-files/CreditConsumptionTable.pdf) — Tables 6(b) (REST API with Prompt Caching) and 6(c) (REST API).

## Context

The existing `snowflake/` provider has zero model entries in the pricing JSON, which means LiteLLM cannot track costs for Snowflake Cortex calls. This PR fills that gap.

## Related

- Existing provider: `litellm/llms/snowflake/`
- Cortex REST API docs: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api

* Update model_prices_and_context_window.json

Fix the JSON parsing error

* Update model_prices_and_context_window.json

Removed the duplicate entry

* fix(utils): copy extra_body before adding unknown params to prevent model config mutation (#29620)

Fixes #29615. In add_provider_specific_params_to_optional_params, the line:

    extra_body = passed_params.pop("extra_body", None) or {}

returns the original dict reference when extra_body is non-empty (truthy).
Subsequent writes like extra_body[k] = passed_params[k] then mutate the
shared model config object held by the router, poisoning /model/info and
all subsequent requests for that deployment.

The or {} short-circuit creates a new dict only when extra_body is falsy
(None or {}), which is why the bug does not reproduce with extra_body: {}.

Fix: wrap in dict() so we always work on a fresh shallow copy.

* fix(vertex_ai): Bake tool_choice into Gemini CachedContent body to prevent silent drop (#29097)

* fix(vertex_ai): bake tool_choice into Gemini CachedContent body to prevent silent drop

* address greptile feedback on tool_choice cache test

* adds test that uses ToolConfig(functionCallingConfig=FunctionCallingConfig(mode=ANY)) instead of a dict literal, mirroring what map_tool_choice_values actually produce

* fix(gemini/veo): move image from parameters into instances[0] (#29501)

* fix(gemini/veo): move image from parameters into instances[0]

Veo's predictLongRunning schema puts image (and prompt) on the
instances element; parameters is for aspectRatio/durationSeconds/etc.
The Gemini path was leaving image in params_copy, so it ended up
nested under parameters and the API silently ignored it.

The Vertex path already builds the instance dict explicitly, so this
just aligns the Gemini path with it.

Fixes #29498

* address greptile: unconditional pop + BytesIO test

- Pop `image` from params_copy unconditionally so it never reaches
  GeminiVideoGenerationParameters even when None, removing implicit
  reliance on Pydantic's extra-field-ignore.
- Add test_transform_video_create_request_image_filelike_goes_to_instance
  covering the BytesIO path (_convert_image_to_gemini_format) — round-trips
  the base64 to confirm encoding.
- Add test_transform_video_create_request_image_none_is_dropped covering
  the new None branch.

* fix(huggingface): handle special token text in embedding usage (#29660)

* fix(guardrails): recompile ToolPermissionGuardrail rules on update_in_memory_litellm_params (#29655)

* fix(guardrails): recompile ToolPermissionGuardrail rules on update_in_memory_litellm_params

ToolPermissionGuardrail builds self.rules and the compiled target/pattern
maps only in __init__. The base update_in_memory_litellm_params re-sets raw
attributes via setattr but never rebuilds those maps, so a guardrail updated
in place (PUT /guardrails, or the immediate in-memory sync) keeps enforcing
the construction-time rules until it is reinitialized (PATCH path, periodic
DB poll, or restart).

Extract the compile step into _load_rules and override
update_in_memory_litellm_params to rebuild from it (dict- and model-safe),
re-normalizing default_action / on_disallowed_action. Mirrors the existing
PresidioGuardrail override of the same method. Adds regression tests.

Fixes #29592.

* fix(guardrails): handle dict params in ToolPermissionGuardrail in-memory update

Delegate to super() only for LitellmParams input (the base setattr loop is
model-only); apply the raw-dict case inline. Fixes the mypy arg-type error
and makes the recompile work when the proxy passes the raw DB dict.

* fix(guardrails): preserve tool-permission rules on a partial in-memory update

A partial update (e.g. a LitellmParams whose rules field is None) ran through
the generic setattr, which set self.rules to None, and the recompile was
skipped, leaving the guardrail with no rules. Snapshot the previous rules and
restore them when the update carries no rules; an explicit empty list still
clears them. Adds a regression test for the rules-absent case.

Addresses the Greptile review note on #29655.

* fix(bedrock): stop base_model label from stripping tools/tool_choice (#29621)

* fix(bedrock): stop base_model label from stripping tools/tool_choice

A Router/proxy Bedrock deployment whose model_info.base_model is a friendly
label (e.g. claude-haiku-4-5) silently lost tools/tool_choice: the outgoing
Converse request was built without toolConfig, so the model behaved as if no
tools were provided. Worked in v1.84.0, regressed in v1.85.0, and with
drop_params=true it failed silently.

Two changes compound into the bug. completion() passed model_info.base_model
as the model argument to get_optional_params, so the real Bedrock model id
never reached supported-param resolution; and get_supported_openai_params
resolved the provider config's params from base_model or model, letting the
label fully replace the real model. For Bedrock the label resolves to no tool
support, so tools/tool_choice were dropped before transformation.

completion() now keeps model as the real deployment model and threads the
resolved base_model (kwarg or model_info) through separately, and
get_supported_openai_params treats base_model as additive: it returns the
union of the params supported by model and by base_model. A hint can only add
capabilities, never strip ones the real model already exposes, which also
preserves the original base_model behavior from #27717 and Azure's base_model
driven model-type detection.

Fixes #29618

* test(main): make base_model param test robust to new parametrize cases

Restore an explicit per-case expected_model_param literal instead of
hardcoding the gemini id, so a future case with a different model can't
produce a misleading assertion failure.

* fix(fireworks_ai): pass response_format json_schema through unchanged (#29606)

FireworksAIConfig.map_openai_params was rewriting the OpenAI strict
`{type: json_schema, json_schema: {name, strict, schema}}` shape into
`{type: json_object, schema: ...}` before sending to Fireworks, dropping
`strict` and `name` and changing the `type`. Per Fireworks' docs json_object
means "force any valid JSON output (no specific schema)", so the schema
constraint was effectively dropped and grammar-guided decoding never ran;
model output silently violated the schema.

The rewrite landed in #7085 (Dec 2024) when Fireworks did not yet accept
native json_schema. Fireworks accepts the OpenAI strict shape natively now,
so the rewrite has become a regression.

Removes the rewrite. Passes response_format through unchanged. Updates the
existing test_map_response_format to assert pass-through. Adds focused
regression tests in tests/test_litellm/ covering preservation of type,
strict, name, and schema body, plus that json_object alone still works.

* fix(types): import Required from typing_extensions in gemini types

* style: reformat sampling_handler.py for py312 black compat

* refactor(mcp-sampling): extract helpers to fix PLR0915 too-many-statements in handle_sampling_create_message

* fix(proxy-server): add explicit ProxyLogging type annotation to proxy_logging_obj to fix mypy inference

* fix(mcp-sampling): suppress mypy assignment error on ImportError fallback for proxy_logging_obj

* fix(test): use .value when comparing LlmProviders enum against string in test_default_api_base

* fix(test): iterate LlmProviders enum in test_default_api_base to avoid str pollution from custom provider registration

litellm.provider_list is a mutable global initialized to list(LlmProviders) but custom_llm_setup() appends plain provider strings to it. When a test_custom_llm.py test runs first in the same xdist worker, provider_list contains a str and calling .value on it raises AttributeError. Iterate the immutable LlmProviders enum instead, which is deterministic and what the check intends.

* fix(mcp): depth-aware JSON-RPC response detection and neutral speed-priority fallback

Replace the flat substring check in the truncated-body routing path with a
top-level-key scan so a JSON-RPC response whose result payload nests a
"method" field is still detected as a response and skips the session lock,
removing a deadlock against the in-flight tool call awaiting it.

Drop the inverse max_output_tokens speed proxy when no model exposes
output_tokens_per_second; context-window size does not track latency, so a
neutral score avoids biasing speedPriority toward the smallest-context model.

* fix(guardrails): make ToolPermission rule reload atomic on invalid regex

_load_rules appended each rule to self.rules before compiling its regex, so an
invalid pattern raised mid-loop after the bad rule was already live but without
a _compiled_rule_targets entry. _matches_regex reads a missing compiled target
as a None pattern and returns True, turning the bad rule into a match-all that
silently applies its decision to every tool. Via update_in_memory_litellm_params
(PUT /guardrails) this corrupted the live guardrail.

Build the parsed rules and compiled maps into locals and swap them in only after
every regex compiles, and restore the previous ruleset if a live update is
rejected, so an invalid regex now fails the update without leaving the guardrail
enforcing a broken policy.

* test(mcp): cover sampling conversion, model resolution, and elicitation relay paths

The MCP sampling and elicitation handlers shipped with partial test
coverage, leaving the response-to-MCP conversion, the model resolution
fallback chain, completion-kwargs assembly, guardrail routing, and the
entire elicitation relay untested. That pulled the PR's diff (patch)
coverage below the codecov threshold even though overall project
coverage rose.

Add focused unit tests for _convert_openai_response_to_mcp_result,
_convert_mcp_tools_to_openai, _convert_mcp_tool_choice_to_openai, image
and audio content conversion, the hint-matching and fallback branches of
_resolve_model_from_preferences, _build_completion_kwargs, the router and
guardrail-rejection paths of _run_guardrails_and_call_llm, the
handle_sampling_create_message success and error-propagation flows, the
marker-hoisting fallback for tool content on unexpected roles, and the
elicitation form/url/generic relay together with its decline paths

---------

Co-authored-by: shin-berri <shin-laptop@berri.ai>
Co-authored-by: yuneng-jiang <yuneng@berri.ai>
Co-authored-by: lengkejun <lengkejun@xd.com>
Co-authored-by: Yug <yugborana000@gmail.com>
Co-authored-by: Kent <72616338+kingdoooo@users.noreply.github.com>
Co-authored-by: tanmay958 <53569547+tanmay958@users.noreply.github.com>
Co-authored-by: DrishnaTrivedi <142084770+DrishnaTrivedi@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Navnit Shukla <Navnit.shukla25@gmail.com>
Co-authored-by: PRABHU KIRAN VANDRANKI <72809214+VANDRANKI@users.noreply.github.com>
Co-authored-by: Adrian Lopez <109683617+adriangomez24@users.noreply.github.com>
Co-authored-by: hcl <chenglunhu@gmail.com>
Co-authored-by: JooHo Lee <96564470+BWAAEEEK@users.noreply.github.com>
Co-authored-by: Dinesh Girbide <85330597+Dinesh-Girbide@users.noreply.github.com>
Co-authored-by: cloudwiz <22098246+andrey-dubnik@users.noreply.github.com>
Co-authored-by: Ahmad Khan <ahmadkhan2508@gmail.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/__init__.py

@@ -1740,6 +1740,9 @@ if TYPE_CHECKING:
     from .llms.openrouter.responses.transformation import (
         OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig,
     )
+    from .llms.bedrock_mantle.responses.transformation import (
+        BedrockMantleResponsesAPIConfig as BedrockMantleResponsesAPIConfig,
+    )
     from .llms.gemini.interactions.transformation import (
         GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
     )

litellm/_lazy_imports_registry.py

@@ -237,6 +237,7 @@ LLM_CONFIG_NAMES = (
     "PerplexityResponsesConfig",
     "DatabricksResponsesAPIConfig",
     "OpenRouterResponsesAPIConfig",
+    "BedrockMantleResponsesAPIConfig",
     "GoogleAIStudioInteractionsConfig",
     "OpenAIOSeriesConfig",
     "AnthropicSkillsConfig",
@@ -958,6 +959,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
         ".llms.openrouter.responses.transformation",
         "OpenRouterResponsesAPIConfig",
     ),
+    "BedrockMantleResponsesAPIConfig": (
+        ".llms.bedrock_mantle.responses.transformation",
+        "BedrockMantleResponsesAPIConfig",
+    ),
     "GoogleAIStudioInteractionsConfig": (
         ".llms.gemini.interactions.transformation",
         "GoogleAIStudioInteractionsConfig",

litellm/experimental_mcp_client/client.py

@@ -4,6 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
 
 import asyncio
 import base64
+import os
 from typing import (
     Any,
     Awaitable,
@@ -16,7 +17,6 @@ from typing import (
     TypeVar,
     Union,
 )
-
 import httpx
 from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
 from mcp.client.sse import sse_client
@@ -42,9 +42,8 @@ from mcp.types import (
 )
 from mcp.types import Tool as MCPTool
 from pydantic import AnyUrl
-
 from litellm._logging import verbose_logger
-from litellm.constants import MCP_CLIENT_TIMEOUT
+from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
 from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
 from litellm.types.llms.custom_http import VerifyTypes
 from litellm.types.mcp import (
@@ -67,7 +66,6 @@ TSessionResult = TypeVar("TSessionResult")
 class MCPSigV4Auth(httpx.Auth):
     """
     httpx Auth class that signs each request with AWS SigV4.
-
     This is used for MCP servers that require AWS SigV4 authentication,
     such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow()
     for every outgoing request, enabling per-request signature computation.
@@ -92,10 +90,8 @@ class MCPSigV4Auth(httpx.Auth):
                 "Missing botocore to use AWS SigV4 authentication. "
                 "Run 'pip install boto3'."
             )
-
         self.service_name = aws_service_name or "bedrock-agentcore"
         self.region_name = aws_region_name or "us-east-1"
-
         # Note: os.environ/ prefixed values are already resolved by
         # ProxyConfig._check_for_os_environ_vars() at config load time.
         # Values arrive here as plain strings.
@@ -143,20 +139,17 @@ class MCPSigV4Auth(httpx.Auth):
         session_name = (
             aws_session_name or f"litellm-mcp-{int(__import__('time').time())}"
         )
-
         sts_kwargs: dict = {"region_name": aws_region_name}
         if aws_access_key_id and aws_secret_access_key:
             sts_kwargs["aws_access_key_id"] = aws_access_key_id
             sts_kwargs["aws_secret_access_key"] = aws_secret_access_key
             if aws_session_token:
                 sts_kwargs["aws_session_token"] = aws_session_token
-
         sts_client = boto3.client("sts", **sts_kwargs)
         sts_response = sts_client.assume_role(
             RoleArn=aws_role_name,
             RoleSessionName=session_name,
         )
-
         sts_creds = sts_response["Credentials"]
         return Credentials(
             access_key=sts_creds["AccessKeyId"],
@@ -178,17 +171,14 @@ class MCPSigV4Auth(httpx.Auth):
             data=request.content,
             headers=dict(request.headers),
         )
-
         # Sign the request — SigV4Auth.add_auth() adds Authorization,
         # X-Amz-Date, and X-Amz-Security-Token (if session token present).
         # Host header is derived automatically from the URL.
         sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name)
         sigv4.add_auth(aws_request)
-
         # Copy SigV4 headers back to the httpx request
         for header_name, header_value in aws_request.headers.items():
             request.headers[header_name] = header_value
-
         yield request
 
 
@@ -198,6 +188,8 @@ class MCPClient:
       SSE and HTTP transports
       Authentication via Bearer token, Basic Auth, or API Key
       Tool calling with error handling and result parsing
+      Sampling callbacks for upstream server LLM requests
+      Elicitation callbacks for upstream server user-input requests
     """
 
     def __init__(
@@ -211,6 +203,9 @@ class MCPClient:
         extra_headers: Optional[Dict[str, str]] = None,
         ssl_verify: Optional[VerifyTypes] = None,
         aws_auth: Optional[httpx.Auth] = None,
+        sampling_callback: Optional[Callable] = None,
+        elicitation_callback: Optional[Callable] = None,
+        logging_callback: Optional[Callable] = None,
     ):
         self.server_url: str = server_url
         self.transport_type: MCPTransport = transport_type
@@ -222,6 +217,9 @@ class MCPClient:
         self.ssl_verify: Optional[VerifyTypes] = ssl_verify
         self._aws_auth: Optional[httpx.Auth] = aws_auth
         self._last_initialize_instructions: Optional[str] = None
+        self._sampling_callback: Optional[Callable] = sampling_callback
+        self._elicitation_callback: Optional[Callable] = elicitation_callback
+        self._logging_callback: Optional[Callable] = logging_callback
         # handle the basic auth value if provided
         if auth_value:
             self.update_auth_value(auth_value)
@@ -231,23 +229,20 @@ class MCPClient:
     ) -> Tuple[Any, Optional[httpx.AsyncClient]]:
         """
         Create the appropriate transport context based on transport type.
-
         Returns:
             Tuple of (transport_context, http_client).
             http_client is only set for HTTP transport and needs cleanup.
         """
         http_client: Optional[httpx.AsyncClient] = None
-
         if self.transport_type == MCPTransport.stdio:
             if not self.stdio_config:
                 raise ValueError("stdio_config is required for stdio transport")
             server_params = StdioServerParameters(
                 command=self.stdio_config.get("command", ""),
                 args=self.stdio_config.get("args", []),
-                env=self.stdio_config.get("env", {}),
+                env=self._get_safe_stdio_env(self.stdio_config.get("env")),
             )
             return stdio_client(server_params), None
-
         if self.transport_type == MCPTransport.sse:
             headers = self._get_auth_headers()
             httpx_client_factory = self._create_httpx_client_factory()
@@ -260,14 +255,12 @@ class MCPClient:
                 ),
                 None,
             )
-
         # HTTP transport (default)
         if streamable_http_client is None:
             raise ImportError(
                 "streamable_http_client is not available. "
                 "Please install mcp with HTTP support."
             )
-
         headers = self._get_auth_headers()
         httpx_client_factory = self._create_httpx_client_factory()
         verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
@@ -281,6 +274,54 @@ class MCPClient:
         )
         return transport_ctx, http_client
 
+    def _get_safe_stdio_env(
+        self, provided_env: Optional[Dict[str, str]]
+    ) -> Optional[Dict[str, str]]:
+        """
+        Return a safe environment for the stdio subprocess.
+
+        If provided_env is set, we use it as-is.
+        If provided_env is None, we return a minimal allowlist from the parent environment
+        to avoid leaking sensitive LiteLLM keys (OPENAI_API_KEY, etc.) to sub-processes.
+        """
+        if provided_env is not None:
+            return provided_env
+
+        # Minimal allowlist of safe/standard environment variables
+        safe_keys = {
+            "PATH",
+            "HOME",
+            "USER",
+            "LOGNAME",
+            "TMPDIR",
+            "TMP",
+            "TEMP",
+            "SHELL",
+            "LANG",
+            "LC_ALL",
+            # Node/Package manager caches
+            "NPM_CONFIG_CACHE",
+            "PNPM_HOME",
+            "XDG_CACHE_HOME",
+            "XDG_CONFIG_HOME",
+            "XDG_DATA_HOME",
+            # System info
+            "SYSTEMROOT",
+            "COMSPEC",
+            "PATHEXT",
+            "WINDIR",
+        }
+
+        safe_env = {}
+        for key in safe_keys:
+            if key in os.environ:
+                safe_env[key] = os.environ[key]
+
+        if "NPM_CONFIG_CACHE" not in safe_env:
+            safe_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR
+
+        return safe_env
+
     async def _execute_session_operation(
         self,
         transport_ctx: Any,
@@ -288,13 +329,23 @@ class MCPClient:
     ) -> TSessionResult:
         """
         Execute an operation within a transport and session context.
-
         Handles entering/exiting contexts and running the operation.
+        Passes sampling/elicitation/logging callbacks to the ClientSession
+        so that upstream MCP servers can request LLM inference (sampling),
+        user input (elicitation), or send log messages.
         """
         transport = await transport_ctx.__aenter__()
         try:
             read_stream, write_stream = transport[0], transport[1]
-            session_ctx = ClientSession(read_stream, write_stream)
+            # Build session kwargs with optional callbacks
+            session_kwargs: Dict[str, Any] = {}
+            if self._sampling_callback is not None:
+                session_kwargs["sampling_callback"] = self._sampling_callback
+            if self._elicitation_callback is not None:
+                session_kwargs["elicitation_callback"] = self._elicitation_callback
+            if self._logging_callback is not None:
+                session_kwargs["logging_callback"] = self._logging_callback
+            session_ctx = ClientSession(read_stream, write_stream, **session_kwargs)
             session = await session_ctx.__aenter__()
             try:
                 init_result = await session.initialize()
@@ -351,7 +402,6 @@ class MCPClient:
     def _get_auth_headers(self) -> dict:
         """Generate authentication headers based on auth type."""
         headers = {}
-
         if self._mcp_auth_value:
             if isinstance(self._mcp_auth_value, str):
                 if self.auth_type == MCPAuth.bearer_token:
@@ -373,17 +423,14 @@ class MCPClient:
         # Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
         # signing (including the body hash), so it uses httpx.Auth flow instead
         # of static headers. See MCPSigV4Auth and _create_httpx_client_factory().
-
         # update the headers with the extra headers
         if self.extra_headers:
             headers.update(self.extra_headers)
-
         return headers
 
     def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]:
         """
         Create a custom httpx client factory that uses LiteLLM's SSL configuration.
-
         This factory follows the same CA bundle path logic as http_handler.py:
         1. Check ssl_verify parameter (can be SSLContext, bool, or path to CA bundle)
         2. Check SSL_VERIFY environment variable
@@ -400,17 +447,14 @@ class MCPClient:
             """Create an httpx.AsyncClient with LiteLLM's SSL configuration."""
             # Get unified SSL configuration using the same logic as http_handler.py
             ssl_config = get_ssl_configuration(self.ssl_verify)
-
             verbose_logger.debug(
                 f"MCP client using SSL configuration: {type(ssl_config).__name__}"
             )
-
             # Use SigV4 auth if configured and no explicit auth provided.
             # The MCP SDK's sse_client and streamable_http_client call this
             # factory without passing auth=, so self._aws_auth is used.
             # For non-SigV4 clients, self._aws_auth is None — no behavior change.
             effective_auth = auth if auth is not None else self._aws_auth
-
             return httpx.AsyncClient(
                 headers=headers,
                 timeout=timeout,
@@ -458,7 +502,6 @@ class MCPClient:
                 f"Server: {self.server_url or 'stdio'}, "
                 f"Transport: {self.transport_type}"
             )
-
             # Check if it's a stream/connection error
             if "BrokenResourceError" in error_type or "Broken" in error_type:
                 verbose_logger.error(
@@ -491,7 +534,6 @@ class MCPClient:
                 f"MCP Tool '{call_tool_request_params.name}' progress: "
                 f"{progress}/{total} ({percentage:.0f}%) - {message or ''}"
             )
-
             # Forward to Host if callback provided
             if host_progress_callback:
                 try:
@@ -521,7 +563,6 @@ class MCPClient:
 
             error_trace = traceback.format_exc()
             verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}")
-
             # Log detailed error information
             error_type = type(e).__name__
             verbose_logger.error(
@@ -532,14 +573,12 @@ class MCPClient:
                 f"Server: {self.server_url or 'stdio'}, "
                 f"Transport: {self.transport_type}"
             )
-
             # Check if it's a stream/connection error
             if "BrokenResourceError" in error_type or "Broken" in error_type:
                 verbose_logger.error(
                     "MCP client detected broken connection/stream - "
                     "the MCP server may have crashed, disconnected, or timed out."
                 )
-
             # Return a default error result instead of raising
             return MCPCallToolResult(
                 content=[
@@ -577,14 +616,12 @@ class MCPClient:
                 f"Server: {self.server_url or 'stdio'}, "
                 f"Transport: {self.transport_type}"
             )
-
             # Check if it's a stream/connection error
             if "BrokenResourceError" in error_type or "Broken" in error_type:
                 verbose_logger.error(
                     "MCP client detected broken connection/stream during list_tools - "
                     "the MCP server may have crashed, disconnected, or timed out"
                 )
-
             # Return empty list instead of raising to allow graceful degradation
             return []
 
@@ -617,7 +654,6 @@ class MCPClient:
 
             error_trace = traceback.format_exc()
             verbose_logger.debug(f"MCP client get_prompt traceback:\n{error_trace}")
-
             # Log detailed error information
             error_type = type(e).__name__
             verbose_logger.error(
@@ -628,14 +664,12 @@ class MCPClient:
                 f"Server: {self.server_url or 'stdio'}, "
                 f"Transport: {self.transport_type}"
             )
-
             # Check if it's a stream/connection error
             if "BrokenResourceError" in error_type or "Broken" in error_type:
                 verbose_logger.error(
                     "MCP client detected broken connection/stream during get_prompt - "
                     "the MCP server may have crashed, disconnected, or timed out."
                 )
-
             raise
 
     async def list_resources(self) -> list[Resource]:
@@ -667,14 +701,12 @@ class MCPClient:
                 f"Server: {self.server_url or 'stdio'}, "
                 f"Transport: {self.transport_type}"
             )
-
             # Check if it's a stream/connection error
             if "BrokenResourceError" in error_type or "Broken" in error_type:
                 verbose_logger.error(
                     "MCP client detected broken connection/stream during list_resources - "
                     "the MCP server may have crashed, disconnected, or timed out"
                 )
-
             # Return empty list instead of raising to allow graceful degradation
             return []
 
@@ -709,14 +741,12 @@ class MCPClient:
                 f"Server: {self.server_url or 'stdio'}, "
                 f"Transport: {self.transport_type}"
             )
-
             # Check if it's a stream/connection error
             if "BrokenResourceError" in error_type or "Broken" in error_type:
                 verbose_logger.error(
                     "MCP client detected broken connection/stream during list_resource_templates - "
                     "the MCP server may have crashed, disconnected, or timed out"
                 )
-
             # Return empty list instead of raising to allow graceful degradation
             return []
 
@@ -742,7 +772,6 @@ class MCPClient:
 
             error_trace = traceback.format_exc()
             verbose_logger.debug(f"MCP client read_resource traceback:\n{error_trace}")
-
             # Log detailed error information
             error_type = type(e).__name__
             verbose_logger.error(
@@ -753,12 +782,10 @@ class MCPClient:
                 f"Server: {self.server_url or 'stdio'}, "
                 f"Transport: {self.transport_type}"
             )
-
             # Check if it's a stream/connection error
             if "BrokenResourceError" in error_type or "Broken" in error_type:
                 verbose_logger.error(
                     "MCP client detected broken connection/stream during read_resource - "
                     "the MCP server may have crashed, disconnected, or timed out."
                 )
-
             raise

litellm/litellm_core_utils/get_supported_openai_params.py

@@ -22,9 +22,11 @@ def get_supported_openai_params(  # noqa: PLR0915
     ```
 
     Args:
-        base_model: For Azure, the true underlying model (e.g. ``"azure/gpt-5.2"``)
-            when the deployment name differs. Used for model-type detection so that
-            non-standard deployment names route to the correct config.
+        base_model: An optional capability hint for deployments whose ``model``
+            label isn't recognized on its own (e.g. an Azure deployment name, or a
+            friendly Bedrock alias). It is additive: the result is the union of the
+            params supported by ``model`` and by ``base_model``, so a hint can only
+            add capabilities, never strip ones the real model already supports.
 
     Returns:
     - List if custom_llm_provider is mapped
@@ -52,7 +54,15 @@ def get_supported_openai_params(  # noqa: PLR0915
         provider_config = None
 
     if provider_config and request_type == "chat_completion":
-        return provider_config.get_supported_openai_params(model=base_model or model)
+        supported_params = provider_config.get_supported_openai_params(model=model)
+        if base_model and base_model != model:
+            base_model_params = provider_config.get_supported_openai_params(
+                model=base_model
+            )
+            supported_params = list(
+                dict.fromkeys([*supported_params, *base_model_params])
+            )
+        return supported_params
 
     if custom_llm_provider == "bedrock":
         return litellm.AmazonConverseConfig().get_supported_openai_params(model=model)

litellm/litellm_core_utils/streaming_handler.py

@@ -1149,6 +1149,32 @@ class CustomStreamWrapper:
             completion_obj: Dict[str, Any] = {"content": ""}
             from litellm.types.utils import GenericStreamingChunk as GChunk
 
+            if (
+                isinstance(chunk, ModelResponseStream)
+                and self.custom_llm_provider is not None
+                and self.custom_llm_provider in litellm._custom_providers
+            ):
+                _has_content = bool(
+                    chunk.choices
+                    and chunk.choices[0].delta is not None
+                    and (
+                        chunk.choices[0].delta.content
+                        or chunk.choices[0].delta.tool_calls
+                    )
+                )
+                if self.received_finish_reason is not None:
+                    if not _has_content:
+                        raise StopIteration
+                if chunk.choices and chunk.choices[0].finish_reason:
+                    self.received_finish_reason = chunk.choices[0].finish_reason
+                    if not _has_content:
+                        return None
+                    # Strip finish_reason from the content chunk so it appears
+                    # only on the trailing empty-delta chunk (OpenAI spec).
+                    # finish_reason_handler() will emit the proper terminal chunk.
+                    chunk.choices[0].finish_reason = None  # type: ignore[assignment]
+                return chunk
+
             if (
                 isinstance(chunk, dict)
                 and generic_chunk_has_all_required_fields(

litellm/llms/azure/image_edit/transformation.py

@@ -97,8 +97,15 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
             )
         original_url = httpx.URL(api_base)
 
-        # Extract api_version or use default
-        api_version = cast(Optional[str], litellm_params.get("api_version"))
+        # Resolve api_version: litellm_params > litellm.api_version > AZURE_API_VERSION env > default.
+        # Mirrors the fallback chain used by the Azure chat path in common_utils.py,
+        # so callers that set a global / env api_version don't get an unversioned URL.
+        api_version = (
+            cast(Optional[str], litellm_params.get("api_version"))
+            or litellm.api_version
+            or get_secret_str("AZURE_API_VERSION")
+            or litellm.AZURE_DEFAULT_API_VERSION
+        )
 
         # Create a new dictionary with existing params
         query_params = dict(original_url.params)

litellm/llms/bedrock_mantle/responses/transformation.py

@@ -0,0 +1,81 @@
+"""
+Amazon Bedrock Mantle - Responses API backend.
+
+gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses`
+path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI
+Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides
+only the endpoint URL and Bearer authentication.
+
+Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the
+standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4.
+"""
+
+from typing import Optional
+
+from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+BEDROCK_MANTLE_DEFAULT_REGION = "us-east-1"
+
+# Checked longest/most-specific first so a full endpoint URL collapses to host
+# in one pass and the appended path never doubles.
+_BASE_SUFFIXES_TO_STRIP = (
+    "/openai/v1/responses",
+    "/v1/responses",
+    "/responses",
+    "/openai/v1",
+    "/v1",
+)
+
+
+class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
+    @property
+    def custom_llm_provider(self) -> LlmProviders:
+        return LlmProviders.BEDROCK_MANTLE
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        litellm_params: dict,
+    ) -> str:
+        region = (
+            get_secret_str("BEDROCK_MANTLE_REGION")
+            or get_secret_str("AWS_REGION")
+            or BEDROCK_MANTLE_DEFAULT_REGION
+        )
+        base = (
+            api_base
+            or get_secret_str("BEDROCK_MANTLE_API_BASE")
+            or f"https://bedrock-mantle.{region}.api.aws"
+        )
+        base = base.rstrip("/")
+        for suffix in _BASE_SUFFIXES_TO_STRIP:
+            if base.endswith(suffix):
+                base = base[: -len(suffix)]
+                break
+        return f"{base}/openai/v1/responses"
+
+    def validate_environment(
+        self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]
+    ) -> dict:
+        litellm_params = litellm_params or GenericLiteLLMParams()
+        api_key = (
+            litellm_params.api_key
+            or get_secret_str("BEDROCK_MANTLE_API_KEY")
+            or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
+        )
+        if not api_key:
+            raise ValueError(
+                "Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY "
+                "(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key."
+            )
+        headers["Authorization"] = f"Bearer {api_key}"
+        return headers
+
+    def supports_native_file_search(self) -> bool:
+        return False
+
+    def supports_native_websocket(self) -> bool:
+        return False

litellm/llms/fireworks_ai/chat/transformation.py

@@ -170,11 +170,6 @@ class FireworksAIConfig(OpenAIGPTConfig):
                         is_response_format_supported=False,
                         enforce_tool_choice=False,  # tools and response_format are both set, don't enforce tool_choice
                     )
-                elif "json_schema" in value:
-                    optional_params["response_format"] = {
-                        "type": "json_object",
-                        "schema": value["json_schema"]["schema"],
-                    }
                 else:
                     optional_params["response_format"] = value
             elif param == "max_completion_tokens":

litellm/llms/gemini/videos/transformation.py

@@ -265,7 +265,11 @@ class GeminiVideoConfig(BaseVideoConfig):
         {
             "instances": [
                 {
-                    "prompt": "A cat playing with a ball of yarn"
+                    "prompt": "A cat playing with a ball of yarn",
+                    "image": {
+                        "bytesBase64Encoded": "...",
+                        "mimeType": "image/jpeg"
+                    }
                 }
             ],
             "parameters": {
@@ -275,13 +279,18 @@ class GeminiVideoConfig(BaseVideoConfig):
             }
         }
         """
-        instance = GeminiVideoGenerationInstance(prompt=prompt)
+        instance: GeminiVideoGenerationInstance = {"prompt": prompt}
 
         params_copy = video_create_optional_request_params.copy()
 
-        if "image" in params_copy and params_copy["image"] is not None:
-            image_data = _convert_image_to_gemini_format(params_copy["image"])
-            params_copy["image"] = image_data
+        if "image" in params_copy:
+            image = params_copy.pop("image")
+            if image is not None:
+                if isinstance(image, dict):
+                    image_data = image
+                else:
+                    image_data = _convert_image_to_gemini_format(image)
+                instance["image"] = image_data
 
         parameters = GeminiVideoGenerationParameters(**params_copy)
 

litellm/llms/huggingface/embedding/handler.py

@@ -239,7 +239,7 @@ class HuggingFaceEmbedding(BaseLLM):
         model_response.model = model
         input_tokens = 0
         for text in input:
-            input_tokens += len(encoding.encode(text))
+            input_tokens += len(encoding.encode(text, disallowed_special=()))
 
         setattr(
             model_response,

litellm/llms/snowflake/utils.py

@@ -25,6 +25,7 @@ class SnowflakeBaseConfig:
             "temperature",
             "max_tokens",
             "top_p",
+            "stream",
             "response_format",
             "tools",
             "tool_choice",

litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py

@@ -337,6 +337,7 @@ class ContextCachingEndpoints(VertexBase):
             return messages, optional_params, None
 
         tools = optional_params.pop("tools", None)
+        tool_choice = optional_params.pop("tool_choice", None)
 
         ## AUTHORIZATION ##
         token, url = self._get_token_and_url_context_caching(
@@ -371,7 +372,7 @@ class ContextCachingEndpoints(VertexBase):
 
         ## CHECK IF CACHED ALREADY
         generated_cache_key = local_cache_obj.get_cache_key(
-            messages=cached_messages, tools=tools, model=model
+            messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
         )
         google_cache_name = self.check_cache(
             cache_key=generated_cache_key,
@@ -402,6 +403,8 @@ class ContextCachingEndpoints(VertexBase):
         )
 
         cached_content_request_body["tools"] = tools
+        if tool_choice is not None:
+            cached_content_request_body["toolConfig"] = tool_choice
 
         ## LOGGING
         logging_obj.pre_call(
@@ -487,6 +490,7 @@ class ContextCachingEndpoints(VertexBase):
             return messages, optional_params, None
 
         tools = optional_params.pop("tools", None)
+        tool_choice = optional_params.pop("tool_choice", None)
 
         ## AUTHORIZATION ##
         token, url = self._get_token_and_url_context_caching(
@@ -518,7 +522,7 @@ class ContextCachingEndpoints(VertexBase):
 
         ## CHECK IF CACHED ALREADY
         generated_cache_key = local_cache_obj.get_cache_key(
-            messages=cached_messages, tools=tools, model=model
+            messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
         )
         google_cache_name = await self.async_check_cache(
             cache_key=generated_cache_key,
@@ -550,6 +554,8 @@ class ContextCachingEndpoints(VertexBase):
         )
 
         cached_content_request_body["tools"] = tools
+        if tool_choice is not None:
+            cached_content_request_body["toolConfig"] = tool_choice
 
         ## LOGGING
         logging_obj.pre_call(

litellm/main.py

@@ -1322,7 +1322,9 @@ def completion(  # type: ignore # noqa: PLR0915
     preset_cache_key = kwargs.get("preset_cache_key", None)
     hf_model_name = kwargs.get("hf_model_name", None)
     supports_system_message = kwargs.get("supports_system_message", None)
-    base_model = kwargs.get("base_model", None)
+    base_model = kwargs.get("base_model", None) or (
+        model_info.get("base_model") if isinstance(model_info, dict) else None
+    )
     ### DISABLE FLAGS ###
     disable_add_transform_inline_image_block = kwargs.get(
         "disable_add_transform_inline_image_block", None
@@ -1534,11 +1536,7 @@ def completion(  # type: ignore # noqa: PLR0915
             "logit_bias": logit_bias,
             "user": user,
             # params to identify the model
-            "model": (
-                model_info.get("base_model")
-                if isinstance(model_info, dict) and model_info.get("base_model")
-                else model
-            ),
+            "model": model,
             "custom_llm_provider": custom_llm_provider,
             "response_format": response_format,
             "seed": seed,

litellm/model_prices_and_context_window_backup.json

@@ -41223,6 +41223,44 @@
         "supports_response_schema": true,
         "supports_tool_choice": true
     },
+    "bedrock_mantle/openai.gpt-5.5": {
+        "input_cost_per_token": 5.5e-06,
+        "cache_read_input_token_cost": 5.5e-07,
+        "output_cost_per_token": 3.3e-05,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 272000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "responses",
+        "supported_endpoints": ["/v1/responses"],
+        "supported_modalities": ["text", "image"],
+        "supported_output_modalities": ["text"],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/openai.gpt-5.4": {
+        "input_cost_per_token": 2.75e-06,
+        "cache_read_input_token_cost": 2.75e-07,
+        "output_cost_per_token": 1.65e-05,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 272000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "responses",
+        "supported_endpoints": ["/v1/responses"],
+        "supported_modalities": ["text", "image"],
+        "supported_output_modalities": ["text"],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "volcengine/doubao-seed-2-0-pro-260215": {
         "litellm_provider": "volcengine",
         "max_input_tokens": 256000,

litellm/proxy/_experimental/mcp_server/elicitation_handler.py

@@ -0,0 +1,163 @@
+"""
+MCP Elicitation Handler
+Handles `elicitation/create` requests from upstream MCP servers by either:
+1. Relaying them to the connected downstream MCP client (if it supports elicitation)
+2. Returning a decline/error response (if no downstream client or unsupported)
+Supports both Form mode (structured data collection) and URL mode (external URL
+navigation for sensitive interactions like OAuth).
+MCP Spec Reference:
+    https://modelcontextprotocol.io/specification/2025-11-25/client/elicitation
+"""
+
+from typing import Any, Optional, Union
+from litellm._logging import verbose_logger
+
+# Guard imports that require the mcp package
+try:
+    from mcp.types import (
+        ElicitRequestFormParams,
+        ElicitRequestParams,
+        ElicitRequestURLParams,
+        ElicitResult,
+        ErrorData,
+    )
+
+    MCP_ELICITATION_AVAILABLE = True
+except ImportError:
+    MCP_ELICITATION_AVAILABLE = False
+
+
+async def handle_elicitation_request(
+    context: Any,
+    params: "ElicitRequestParams",
+    downstream_session: Optional[Any] = None,
+    downstream_capabilities: Optional[Any] = None,
+) -> Union["ElicitResult", "ErrorData"]:
+    """
+    Handle an MCP elicitation/create request from an upstream MCP server.
+    In Gateway mode (Mode A), we relay the elicitation request to the
+    connected downstream client if they declared elicitation capabilities.
+    In Tool Bridge mode (Mode B), there's no persistent downstream MCP
+    client, so we return a decline response.
+    Args:
+        context: MCP RequestContext from the upstream server connection.
+        params: The ElicitRequestParams (either form or URL mode).
+        downstream_session: The ServerSession to the downstream client,
+            if available (for relaying).
+        downstream_capabilities: The downstream client's declared
+            capabilities, used to check elicitation support.
+    Returns:
+        ElicitResult with the user's response, or ErrorData on failure.
+    """
+    if not MCP_ELICITATION_AVAILABLE:
+        return ErrorData(
+            code=-1,
+            message="MCP elicitation is not available (mcp package not installed)",
+        )
+    try:
+        mode = getattr(params, "mode", "form")
+        verbose_logger.info(
+            "MCP elicitation: received request mode=%s, message=%s",
+            mode,
+            getattr(params, "message", ""),
+        )
+        # Check if we have a downstream session to relay to
+        if downstream_session is not None:
+            return await _relay_elicitation_to_downstream(
+                params=params,
+                downstream_session=downstream_session,
+                downstream_capabilities=downstream_capabilities,
+            )
+        # No downstream session — we're in Tool Bridge mode
+        # or the client doesn't support elicitation
+        verbose_logger.info(
+            "MCP elicitation: no downstream session available, declining"
+        )
+        return ElicitResult(
+            action="decline",
+        )
+    except Exception as e:
+        verbose_logger.exception("MCP elicitation handler failed: %s", e)
+        return ErrorData(
+            code=-1,
+            message=f"Elicitation failed: {str(e)}",
+        )
+
+
+async def _relay_elicitation_to_downstream(
+    params: "ElicitRequestParams",
+    downstream_session: Any,
+    downstream_capabilities: Optional[Any] = None,
+) -> Union["ElicitResult", "ErrorData"]:
+    """
+    Relay an elicitation request to the downstream MCP client.
+    Uses the ServerSession's elicit_form() or elicit_url() methods to
+    send the elicitation request back to the connected client.
+    Args:
+        params: The elicitation request parameters.
+        downstream_session: The ServerSession connected to the downstream client.
+        downstream_capabilities: Client capabilities to check support.
+    Returns:
+        ElicitResult from the downstream client.
+    """
+    mode = getattr(params, "mode", "form")
+    # Check if the downstream client supports the requested mode
+    if downstream_capabilities is not None:
+        elicit_caps = getattr(downstream_capabilities, "elicitation", None)
+        if elicit_caps is None:
+            verbose_logger.info(
+                "MCP elicitation: downstream client does not support elicitation"
+            )
+            return ElicitResult(action="decline")
+        if mode == "url":
+            url_cap = getattr(elicit_caps, "url", None)
+            if url_cap is None:
+                verbose_logger.info(
+                    "MCP elicitation: downstream client does not support URL mode"
+                )
+                return ElicitResult(action="decline")
+        if mode == "form":
+            form_cap = getattr(elicit_caps, "form", None)
+            if form_cap is None:
+                verbose_logger.info(
+                    "MCP elicitation: downstream client does not support form mode"
+                )
+                return ElicitResult(action="decline")
+    try:
+        if mode == "url" and isinstance(params, ElicitRequestURLParams):
+            # URL mode: relay URL to client for external navigation
+            verbose_logger.info(
+                "MCP elicitation: relaying URL mode to downstream, url=%s",
+                getattr(params, "url", ""),
+            )
+            result = await downstream_session.elicit_url(
+                message=params.message,
+                url=params.url,
+                elicitation_id=getattr(params, "elicitationId", None),
+            )
+        elif isinstance(params, ElicitRequestFormParams):
+            # Form mode: relay structured form to client
+            verbose_logger.info("MCP elicitation: relaying form mode to downstream")
+            result = await downstream_session.elicit_form(
+                message=params.message,
+                requestedSchema=getattr(params, "requestedSchema", None),
+            )
+        else:
+            # Fallback for generic ElicitRequestParams — pass an empty schema
+            # since elicit() requires requestedSchema as a positional arg.
+            verbose_logger.info(
+                "MCP elicitation: relaying generic elicitation to downstream"
+            )
+            result = await downstream_session.elicit(
+                message=getattr(params, "message", ""),
+                requestedSchema=getattr(params, "requestedSchema", {}),
+            )
+        verbose_logger.info(
+            "MCP elicitation: downstream responded with action=%s",
+            getattr(result, "action", "unknown"),
+        )
+        return result
+    except Exception as e:
+        verbose_logger.warning("MCP elicitation: failed to relay to downstream: %s", e)
+        # If relay fails, decline gracefully
+        return ElicitResult(action="decline")

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -49,6 +49,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
     MCPRequestHandler,
 )
 from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
+from litellm.proxy._experimental.mcp_server.elicitation_handler import (
+    MCP_ELICITATION_AVAILABLE,
+)
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    MCP_SAMPLING_AVAILABLE,
+)
 from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mcp_auth
 from litellm.proxy._experimental.mcp_server.utils import (
     MCP_TOOL_PREFIX_SEPARATOR,
@@ -289,6 +295,82 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]:
         return data
 
 
+def _create_sampling_callback(user_api_key_auth: Optional[Any] = None):
+    """
+    Create a sampling callback for MCP ClientSession.
+    Returns a callable that handles sampling/createMessage requests from
+    upstream MCP servers by routing them through litellm.acompletion().
+    """
+    if not MCP_SAMPLING_AVAILABLE:
+        return None
+
+    async def _sampling_callback(context, params):
+        from litellm.proxy._experimental.mcp_server.sampling_handler import (
+            handle_sampling_create_message,
+        )
+        import litellm
+        from litellm.proxy._experimental.mcp_server.server import (
+            get_active_auth_context,
+        )
+
+        auth_context = get_active_auth_context()
+        resolved_auth = user_api_key_auth or (
+            auth_context.user_api_key_auth if auth_context else None
+        )
+        # Forward original HTTP headers and client IP so that
+        # header-dependent guardrails, tag-based routing, trace
+        # correlation, and forward_llm_provider_auth_headers work
+        # correctly for sampling sub-calls.
+        _raw_headers = getattr(auth_context, "raw_headers", None)
+        _client_ip = getattr(auth_context, "client_ip", None)
+
+        return await handle_sampling_create_message(
+            context=context,
+            params=params,
+            default_model=getattr(litellm, "default_mcp_sampling_model", None),
+            user_api_key_auth=resolved_auth,
+            raw_headers=_raw_headers,
+            client_ip=_client_ip,
+        )
+
+    return _sampling_callback
+
+
+def _create_elicitation_callback():
+    """
+    Create an elicitation callback for MCP ClientSession.
+    Returns a callable that handles elicitation/create requests from
+    upstream MCP servers. In gateway mode, this relays to the downstream
+    client; in tool bridge mode, it returns a decline response.
+    """
+    if not MCP_ELICITATION_AVAILABLE:
+        return None
+
+    async def _elicitation_callback(context, params):
+        from litellm.proxy._experimental.mcp_server.elicitation_handler import (
+            handle_elicitation_request,
+        )
+        from litellm.proxy._experimental.mcp_server.server import get_active_mcp_session
+
+        # In Gateway mode, we relay the elicitation request to the downstream client
+        # that triggered the current operation.
+        downstream_session = get_active_mcp_session()
+        downstream_capabilities = (
+            getattr(downstream_session, "capabilities", None)
+            if downstream_session
+            else None
+        )
+
+        return await handle_elicitation_request(
+            context=context,
+            params=params,
+            downstream_session=downstream_session,
+            downstream_capabilities=downstream_capabilities,
+        )
+
+    return _elicitation_callback
+
+
 class MCPServerManager:
     _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$")
 
@@ -600,6 +682,8 @@ class MCPServerManager:
                     "subject_token_type",
                     "urn:ietf:params:oauth:token-type:access_token",
                 ),
+                allow_sampling=bool(server_config.get("allow_sampling", False)),
+                allow_elicitation=bool(server_config.get("allow_elicitation", False)),
             )
             self._assign_unique_short_prefix(new_server)
             _warn_internal_delegate_pkce_if_applicable(new_server, source="config")
@@ -699,8 +783,7 @@ class MCPServerManager:
             )
 
             verbose_logger.debug(
-                f"Using headers for OpenAPI tools (excluding sensitive values): "
-                f"{list(headers.keys())}"
+                f"Using headers for OpenAPI tools (excluding sensitive values): {list(headers.keys())}"
             )
 
             # Extract and register tools from OpenAPI paths
@@ -1494,6 +1577,7 @@ class MCPServerManager:
         extra_headers: Optional[Dict[str, str]] = None,
         stdio_env: Optional[Dict[str, str]] = None,
         subject_token: Optional[str] = None,
+        user_api_key_auth: Optional[UserAPIKeyAuth] = None,
     ) -> MCPClient:
         """
         Create an MCPClient instance for the given server.
@@ -1510,6 +1594,7 @@ class MCPServerManager:
             extra_headers: Additional headers to forward.
             stdio_env: Environment variables for stdio transport.
             subject_token: Optional user JWT for token exchange (OBO) flow.
+            user_api_key_auth: Optional auth context for sampling callbacks.
 
         Returns:
             Configured MCP client instance.
@@ -1520,23 +1605,44 @@ class MCPServerManager:
 
         transport = server.transport or MCPTransport.sse
 
+        # Create sampling and elicitation callbacks for this client
+        sampling_cb = (
+            _create_sampling_callback(user_api_key_auth=user_api_key_auth)
+            if server.allow_sampling
+            else None
+        )
+        elicitation_cb = (
+            _create_elicitation_callback() if server.allow_elicitation else None
+        )
+
         # Handle stdio transport
         if transport == MCPTransport.stdio:
             resolved_env = (
-                stdio_env if stdio_env is not None else dict(server.env or {})
+                stdio_env
+                if stdio_env is not None
+                else (dict(server.env) if server.env is not None else None)
             )
 
             # Ensure npm-based STDIO MCP servers have a writable cache dir.
             # In containers the default (~/.npm or /app/.npm) may not exist
             # or be read-only, causing npx to fail with ENOENT.
-            if "NPM_CONFIG_CACHE" not in resolved_env:
+            if resolved_env is not None and "NPM_CONFIG_CACHE" not in resolved_env:
                 resolved_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR
             # Defense-in-depth: block commands not in the allowlist.
             # The Pydantic validator blocks new servers; this catches legacy
             # config/DB records predating the allowlist.
             if server.command:
                 base_command = os.path.basename(server.command)
-                if base_command not in MCP_STDIO_ALLOWED_COMMANDS:
+                # Strip .exe/.cmd/.bat/.com suffix for Windows compatibility
+                base_command_no_ext = base_command.lower()
+                for ext in [".exe", ".cmd", ".bat", ".com"]:
+                    if base_command.lower().endswith(ext):
+                        base_command_no_ext = base_command[: -len(ext)].lower()
+                        break
+                if (
+                    base_command.lower() not in MCP_STDIO_ALLOWED_COMMANDS
+                    and base_command_no_ext not in MCP_STDIO_ALLOWED_COMMANDS
+                ):
                     raise HTTPException(
                         status_code=403,
                         detail=f"MCP stdio command '{server.command}' is not in the allowlist ({sorted(MCP_STDIO_ALLOWED_COMMANDS)}). "
@@ -1559,6 +1665,8 @@ class MCPServerManager:
                 timeout=MCP_CLIENT_TIMEOUT,
                 stdio_config=stdio_config,
                 extra_headers=extra_headers,
+                sampling_callback=sampling_cb,
+                elicitation_callback=elicitation_cb,
             )
         else:
             # For HTTP/SSE transports
@@ -1585,6 +1693,8 @@ class MCPServerManager:
                 timeout=MCP_CLIENT_TIMEOUT,
                 extra_headers=extra_headers,
                 aws_auth=aws_auth,
+                sampling_callback=sampling_cb,
+                elicitation_callback=elicitation_cb,
             )
 
     async def _get_tools_from_server(
@@ -1668,6 +1778,7 @@ class MCPServerManager:
                 mcp_auth_header=mcp_auth_header,
                 extra_headers=extra_headers,
                 stdio_env=stdio_env,
+                user_api_key_auth=user_api_key_auth,
             )
 
             ## HANDLE OPENAPI TOOLS
@@ -3030,6 +3141,7 @@ class MCPServerManager:
             extra_headers=extra_headers,
             stdio_env=stdio_env,
             subject_token=subject_token,
+            user_api_key_auth=user_api_key_auth,
         )
 
         call_tool_params = MCPCallToolRequestParams(
@@ -3260,7 +3372,6 @@ class MCPServerManager:
                 )
             )
         else:
-            # For regular MCP servers, use the MCP client
             return await self._call_regular_mcp_tool(
                 mcp_server=mcp_server,
                 original_tool_name=name,

litellm/proxy/_experimental/mcp_server/server.py

@@ -6,6 +6,7 @@ LiteLLM MCP Server Routes
 
 import asyncio
 import contextlib
+import contextvars
 import hashlib
 import json
 import time
@@ -125,6 +126,18 @@ try:
         GetPromptResult,
         ResourceTemplate,
         TextResourceContents,
+        Tool,
+    )
+    from mcp.server.session import ServerSession as _McpServerSession
+    import weakref
+
+    # Robust auth lookup keyed by session_object.
+    _session_obj_auth_storage: (
+        "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]"
+    ) = weakref.WeakKeyDictionary()
+
+    active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = (
+        contextvars.ContextVar("active_mcp_session", default=None)
     )
 except ImportError as e:
     verbose_logger.debug(f"MCP module not found: {e}")
@@ -160,6 +173,60 @@ def _mcp_session_id_from_headers(
     return None
 
 
+def _jsonrpc_text_has_top_level_method(text: str) -> bool:
+    """Whether a (possibly truncated) JSON-RPC envelope has a ``method`` key at
+    the root object's top level.
+
+    Used to tell a request/notification (carries ``method``) apart from a
+    response (carries ``result``/``error`` and no top-level ``method``). A
+    response payload can itself nest a ``method`` field, so only keys at the
+    root object's depth are inspected rather than searching the whole string.
+    Returns ``True`` only when a top-level ``method`` key is positively found;
+    truncation that hides it yields ``False``.
+    """
+    depth = 0
+    in_string = False
+    escaped = False
+    in_object: List[bool] = []
+    reading_key = False
+    expect_key = False
+    key_chars: List[str] = []
+    for ch in text:
+        if in_string:
+            if escaped:
+                escaped = False
+            elif ch == "\\":
+                escaped = True
+            elif ch == '"':
+                in_string = False
+                if reading_key and depth == 1 and "".join(key_chars) == "method":
+                    return True
+            elif reading_key:
+                key_chars.append(ch)
+            continue
+        if ch == '"':
+            in_string = True
+            reading_key = expect_key and depth >= 1 and in_object[-1]
+            key_chars = []
+            expect_key = False
+        elif ch == "{" or ch == "[":
+            depth += 1
+            in_object.append(ch == "{")
+            expect_key = ch == "{"
+        elif ch == "}" or ch == "]":
+            if in_object:
+                in_object.pop()
+            depth -= 1
+            if depth <= 0:
+                break
+            expect_key = False
+        elif ch == ",":
+            expect_key = bool(in_object) and in_object[-1]
+        elif ch == ":":
+            expect_key = False
+    return False
+
+
 if MCP_AVAILABLE:
     from mcp.server import Server
     from mcp.server.lowlevel.server import NotificationOptions
@@ -483,10 +550,18 @@ if MCP_AVAILABLE:
     ########################################################
 
     @server.list_tools()
-    async def list_tools() -> List[MCPTool]:
+    async def handle_list_tools() -> List[Tool]:
         """
-        List all available tools
+        List all available tools.
+        Also captures the active session for propagation to callbacks.
         """
+        from mcp.server.lowlevel.server import request_ctx
+
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
+
         try:
             # Get user authentication from context variable
             (
@@ -497,7 +572,7 @@ if MCP_AVAILABLE:
                 oauth2_headers,
                 raw_headers,
                 _client_ip,
-            ) = get_auth_context()
+            ) = await get_or_extract_auth_context()
             verbose_logger.debug(
                 f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
             )
@@ -528,152 +603,178 @@ if MCP_AVAILABLE:
             # Return empty list instead of failing completely
             # This prevents the HTTP stream from failing and allows the client to get a response
             return []
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.call_tool()
-    async def mcp_server_tool_call(
-        name: str, arguments: Optional[Dict[str, Any]]
+    async def mcp_server_tool_call(  # noqa: PLR0915
+        name: str, arguments: Dict[str, Any] | None
     ) -> CallToolResult:
         """
         Call a specific tool with the provided arguments
-
         Args:
             name (str): Name of the tool to call
             arguments (Dict[str, Any] | None): Arguments to pass to the tool
-
         Returns:
             List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]: Tool execution results
-
         Raises:
             HTTPException: If tool not found or arguments missing
         """
         from fastapi import Request
-
         from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
         from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
         from litellm.proxy.proxy_server import proxy_config
+        from mcp.types import CallToolResult
+        from mcp.server.lowlevel.server import request_ctx
 
-        # Validate arguments
-        (
-            user_api_key_auth,
-            mcp_auth_header,
-            mcp_servers,
-            mcp_server_auth_headers,
-            oauth2_headers,
-            raw_headers,
-            _client_ip,
-        ) = get_auth_context()
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
 
-        verbose_logger.debug(
-            f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
-        )
-        host_progress_callback = None
         try:
-            host_ctx = server.request_context
-            if host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta:
-                host_token = getattr(host_ctx.meta, "progressToken", None)
-                if host_token and hasattr(host_ctx, "session") and host_ctx.session:
-                    host_session = host_ctx.session
-
-                    async def forward_progress(progress: float, total: Optional[float]):
-                        """Forward progress notifications from external MCP to Host"""
-                        try:
-                            await host_session.send_progress_notification(
-                                progress_token=host_token,
-                                progress=progress,
-                                total=total,
-                            )
-                            verbose_logger.debug(
-                                f"Forwarded progress {progress}/{total} to Host"
-                            )
-                        except Exception as e:
-                            verbose_logger.error(
-                                f"Failed to forward progress to Host: {e}"
-                            )
-
-                    host_progress_callback = forward_progress
-                    verbose_logger.debug(
-                        f"Host progressToken captured: {host_token[:8]}..."
-                    )
-        except Exception as e:
-            verbose_logger.warning(f"Could not capture host progress context: {e}")
-        try:
-            # Create a body date for logging
-            body_data = {"name": name, "arguments": arguments}
-            # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A)
-            chain_id = get_chain_id_from_headers(raw_headers)
-            if chain_id:
-                body_data["litellm_trace_id"] = chain_id
-                body_data["litellm_session_id"] = chain_id
-
-            request = Request(
-                scope={
-                    "type": "http",
-                    "method": "POST",
-                    "path": "/mcp/tools/call",
-                    "headers": [(b"content-type", b"application/json")],
-                }
+            # Validate arguments
+            (
+                user_api_key_auth,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+                _client_ip,
+            ) = await get_or_extract_auth_context()
+            verbose_logger.debug(
+                f"MCP mcp_server_tool_call - user_api_key_auth={user_api_key_auth}, user_role={getattr(user_api_key_auth, 'user_role', 'N/A')}"
             )
-            if user_api_key_auth is not None:
-                data = await add_litellm_data_to_request(
-                    data=body_data,
-                    request=request,
-                    user_api_key_dict=user_api_key_auth,
-                    proxy_config=proxy_config,
+
+            verbose_logger.debug(
+                f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
+            )
+            host_progress_callback = None
+            try:
+                host_ctx = server.request_context
+                if host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta:
+                    host_token = getattr(host_ctx.meta, "progressToken", None)
+                    if host_token and hasattr(host_ctx, "session") and host_ctx.session:
+                        host_session = host_ctx.session
+
+                        async def forward_progress(
+                            progress: float, total: Optional[float]
+                        ):
+                            """Forward progress notifications from external MCP to Host"""
+                            try:
+                                await host_session.send_progress_notification(
+                                    progress_token=host_token,
+                                    progress=progress,
+                                    total=total,
+                                )
+                                verbose_logger.debug(
+                                    f"Forwarded progress {progress}/{total} to Host"
+                                )
+                            except Exception as e:
+                                verbose_logger.error(
+                                    f"Failed to forward progress to Host: {e}"
+                                )
+
+                        host_progress_callback = forward_progress
+                        verbose_logger.debug(
+                            f"Host progressToken captured: {host_token[:8]}..."
+                        )
+            except Exception as e:
+                verbose_logger.warning(f"Could not capture host progress context: {e}")
+            try:
+                # Create a body date for logging
+                body_data = {"name": name, "arguments": arguments}
+                # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A)
+                chain_id = get_chain_id_from_headers(raw_headers)
+                if chain_id:
+                    body_data["litellm_trace_id"] = chain_id
+                    body_data["litellm_session_id"] = chain_id
+
+                request = Request(
+                    scope={
+                        "type": "http",
+                        "method": "POST",
+                        "path": "/mcp/tools/call",
+                        "headers": [(b"content-type", b"application/json")],
+                    }
                 )
-            else:
-                data = body_data
-
-            response = await call_mcp_tool(
-                user_api_key_auth=user_api_key_auth,
-                mcp_auth_header=mcp_auth_header,
-                mcp_servers=mcp_servers,
-                mcp_server_auth_headers=mcp_server_auth_headers,
-                oauth2_headers=oauth2_headers,
-                raw_headers=raw_headers,
-                host_progress_callback=host_progress_callback,
-                **data,  # for logging
-            )
-        except BlockedPiiEntityError as e:
-            verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}")
-            return CallToolResult(
-                content=[
-                    TextContent(
-                        text=f"Error: Blocked PII entity detected - {str(e)}",
-                        type="text",
+                if user_api_key_auth is not None:
+                    data = await add_litellm_data_to_request(
+                        data=body_data,
+                        request=request,
+                        user_api_key_dict=user_api_key_auth,
+                        proxy_config=proxy_config,
                     )
-                ],
-                isError=True,
-            )
-        except GuardrailRaisedException as e:
-            verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}")
-            return CallToolResult(
-                content=[
-                    TextContent(
-                        text=f"Error: Guardrail violation - {str(e)}", type="text"
-                    )
-                ],
-                isError=True,
-            )
-        except HTTPException as e:
-            verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}")
-            return CallToolResult(
-                content=[TextContent(text=f"Error: {str(e.detail)}", type="text")],
-                isError=True,
-            )
-        except Exception as e:
-            verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}")
-            return CallToolResult(
-                content=[TextContent(text=f"Error: {str(e)}", type="text")],
-                isError=True,
-            )
+                else:
+                    data = body_data
 
-        return response
+                response = await call_mcp_tool(
+                    user_api_key_auth=user_api_key_auth,
+                    mcp_auth_header=mcp_auth_header,
+                    mcp_servers=mcp_servers,
+                    mcp_server_auth_headers=mcp_server_auth_headers,
+                    oauth2_headers=oauth2_headers,
+                    raw_headers=raw_headers,
+                    host_progress_callback=host_progress_callback,
+                    **data,  # for logging
+                )
+            except BlockedPiiEntityError as e:
+                verbose_logger.error(
+                    f"BlockedPiiEntityError in MCP tool call: {str(e)}"
+                )
+                return CallToolResult(
+                    content=[
+                        TextContent(
+                            text=f"Error: Blocked PII entity detected - {str(e)}",
+                            type="text",
+                        )
+                    ],
+                    isError=True,
+                )
+            except GuardrailRaisedException as e:
+                verbose_logger.error(
+                    f"GuardrailRaisedException in MCP tool call: {str(e)}"
+                )
+                return CallToolResult(
+                    content=[
+                        TextContent(
+                            text=f"Error: Guardrail violation - {str(e)}", type="text"
+                        )
+                    ],
+                    isError=True,
+                )
+            except HTTPException as e:
+                verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}")
+                return CallToolResult(
+                    content=[TextContent(text=f"Error: {str(e.detail)}", type="text")],
+                    isError=True,
+                )
+            except Exception as e:
+                verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}")
+                return CallToolResult(
+                    content=[TextContent(text=f"Error: {str(e)}", type="text")],
+                    isError=True,
+                )
+
+            return response
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.list_prompts()
     async def list_prompts() -> List[Prompt]:
         """
         List all available prompts
         """
+        from mcp.server.lowlevel.server import request_ctx
+
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
+
         try:
             # Get user authentication from context variable
             (
@@ -684,7 +785,7 @@ if MCP_AVAILABLE:
                 oauth2_headers,
                 raw_headers,
                 _client_ip,
-            ) = get_auth_context()
+            ) = await get_or_extract_auth_context()
             verbose_logger.debug(
                 f"MCP list_prompts - User API Key Auth from context: {user_api_key_auth}"
             )
@@ -713,6 +814,9 @@ if MCP_AVAILABLE:
             # Return empty list instead of failing completely
             # This prevents the HTTP stream from failing and allows the client to get a response
             return []
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.get_prompt()
     async def get_prompt(
@@ -730,33 +834,13 @@ if MCP_AVAILABLE:
         """
 
         # Validate arguments
-        (
-            user_api_key_auth,
-            mcp_auth_header,
-            mcp_servers,
-            mcp_server_auth_headers,
-            oauth2_headers,
-            raw_headers,
-            _client_ip,
-        ) = get_auth_context()
+        from mcp.server.lowlevel.server import request_ctx
 
-        verbose_logger.debug(
-            f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
-        )
-        return await mcp_get_prompt(
-            name=name,
-            arguments=arguments,
-            user_api_key_auth=user_api_key_auth,
-            mcp_auth_header=mcp_auth_header,
-            mcp_servers=mcp_servers,
-            mcp_server_auth_headers=mcp_server_auth_headers,
-            oauth2_headers=oauth2_headers,
-            raw_headers=raw_headers,
-        )
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
 
-    @server.list_resources()
-    async def list_resources() -> List[Resource]:
-        """List all available resources."""
         try:
             (
                 user_api_key_auth,
@@ -766,7 +850,45 @@ if MCP_AVAILABLE:
                 oauth2_headers,
                 raw_headers,
                 _client_ip,
-            ) = get_auth_context()
+            ) = await get_or_extract_auth_context()
+
+            verbose_logger.debug(
+                f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
+            )
+            return await mcp_get_prompt(
+                name=name,
+                arguments=arguments,
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=mcp_auth_header,
+                mcp_servers=mcp_servers,
+                mcp_server_auth_headers=mcp_server_auth_headers,
+                oauth2_headers=oauth2_headers,
+                raw_headers=raw_headers,
+            )
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
+
+    @server.list_resources()
+    async def list_resources() -> List[Resource]:
+        """List all available resources."""
+        from mcp.server.lowlevel.server import request_ctx
+
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
+
+        try:
+            (
+                user_api_key_auth,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+                _client_ip,
+            ) = await get_or_extract_auth_context()
             verbose_logger.debug(
                 f"MCP list_resources - User API Key Auth from context: {user_api_key_auth}"
             )
@@ -792,10 +914,20 @@ if MCP_AVAILABLE:
         except Exception as e:
             verbose_logger.exception(f"Error in list_resources endpoint: {str(e)}")
             return []
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.list_resource_templates()
     async def list_resource_templates() -> List[ResourceTemplate]:
         """List all available resource templates."""
+        from mcp.server.lowlevel.server import request_ctx
+
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
+
         try:
             (
                 user_api_key_auth,
@@ -805,7 +937,7 @@ if MCP_AVAILABLE:
                 oauth2_headers,
                 raw_headers,
                 _client_ip,
-            ) = get_auth_context()
+            ) = await get_or_extract_auth_context()
             verbose_logger.debug(
                 f"MCP list_resource_templates - User API Key Auth from context: {user_api_key_auth}"
             )
@@ -825,8 +957,7 @@ if MCP_AVAILABLE:
                 raw_headers=raw_headers,
             )
             verbose_logger.info(
-                "MCP list_resource_templates - Successfully returned "
-                f"{len(resource_templates)} resource templates"
+                f"MCP list_resource_templates - Successfully returned {len(resource_templates)} resource templates"
             )
             return resource_templates
         except Exception as e:
@@ -834,30 +965,44 @@ if MCP_AVAILABLE:
                 f"Error in list_resource_templates endpoint: {str(e)}"
             )
             return []
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.read_resource()
     async def read_resource(url: AnyUrl) -> list[ReadResourceContents]:
-        (
-            user_api_key_auth,
-            mcp_auth_header,
-            mcp_servers,
-            mcp_server_auth_headers,
-            oauth2_headers,
-            raw_headers,
-            _client_ip,
-        ) = get_auth_context()
+        from mcp.server.lowlevel.server import request_ctx
 
-        read_resource_result = await mcp_read_resource(
-            url=url,
-            user_api_key_auth=user_api_key_auth,
-            mcp_auth_header=mcp_auth_header,
-            mcp_servers=mcp_servers,
-            mcp_server_auth_headers=mcp_server_auth_headers,
-            oauth2_headers=oauth2_headers,
-            raw_headers=raw_headers,
-        )
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
 
-        return _normalize_resource_contents(read_resource_result.contents)
+        try:
+            (
+                user_api_key_auth,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+                _client_ip,
+            ) = await get_or_extract_auth_context()
+
+            read_resource_result = await mcp_read_resource(
+                url=url,
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=mcp_auth_header,
+                mcp_servers=mcp_servers,
+                mcp_server_auth_headers=mcp_server_auth_headers,
+                oauth2_headers=oauth2_headers,
+                raw_headers=raw_headers,
+            )
+
+            return _normalize_resource_contents(read_resource_result.contents)
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     ########################################################
     ############ End of MCP Server Routes ##################
@@ -1180,8 +1325,7 @@ if MCP_AVAILABLE:
                 cached_token = await mcp_per_user_token_cache.get(user_id, server_id)
                 if cached_token is not None:
                     verbose_logger.debug(
-                        "_get_user_oauth_extra_headers_from_db: Redis hit for "
-                        "user=%s server=%s",
+                        "_get_user_oauth_extra_headers_from_db: Redis hit for user=%s server=%s",
                         user_id,
                         server_id,
                     )
@@ -1207,8 +1351,7 @@ if MCP_AVAILABLE:
 
             if is_oauth_credential_expired(cred):
                 verbose_logger.debug(
-                    "_get_user_oauth_extra_headers_from_db: token expired for "
-                    "user=%s server=%s — attempting refresh",
+                    "_get_user_oauth_extra_headers_from_db: token expired for user=%s server=%s — attempting refresh",
                     user_id,
                     server_id,
                 )
@@ -1230,8 +1373,7 @@ if MCP_AVAILABLE:
                         )
                     except Exception as refresh_exc:
                         verbose_logger.warning(
-                            "_get_user_oauth_extra_headers_from_db: refresh failed "
-                            "for user=%s server=%s: %s",
+                            "_get_user_oauth_extra_headers_from_db: refresh failed for user=%s server=%s: %s",
                             user_id,
                             server_id,
                             refresh_exc,
@@ -1275,8 +1417,7 @@ if MCP_AVAILABLE:
             return {"Authorization": f"Bearer {access_token}"}
         except Exception as e:
             verbose_logger.warning(
-                "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for "
-                "user=%s server=%s: %s",
+                "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for user=%s server=%s: %s",
                 user_id,
                 server_id,
                 e,
@@ -2485,7 +2626,7 @@ if MCP_AVAILABLE:
                 arguments=arguments or {},
                 server_name=server_name or mcp_server.name,
                 user_api_key_auth=user_api_key_auth,
-                proxy_logging_obj=proxy_logging_obj,
+                proxy_logging_obj=proxy_logging_obj,  # type: ignore[arg-type]
                 server=mcp_server,
                 raw_headers=raw_headers,
             )
@@ -2744,8 +2885,7 @@ if MCP_AVAILABLE:
             raise HTTPException(
                 status_code=400,
                 detail=(
-                    "Multiple MCP servers configured; read_resource currently "
-                    "supports exactly one allowed server."
+                    "Multiple MCP servers configured; read_resource currently supports exactly one allowed server."
                 ),
             )
 
@@ -3124,8 +3264,7 @@ if MCP_AVAILABLE:
                 return False
         except Exception:
             verbose_logger.debug(
-                "Unable to inspect active MCP sessions for '%s'. "
-                "Deferring to session manager.",
+                "Unable to inspect active MCP sessions for '%s'. Deferring to session manager.",
                 _session_id,
             )
             return False
@@ -3136,8 +3275,7 @@ if MCP_AVAILABLE:
         if method == "DELETE":
             _remove_stateful_session_tracking(_session_id)
             verbose_logger.info(
-                "DELETE request for non-existent MCP session '%s'. "
-                "Returning success (idempotent DELETE).",
+                "DELETE request for non-existent MCP session '%s'. Returning success (idempotent DELETE).",
                 _session_id,
             )
             success_response = JSONResponse(
@@ -3615,6 +3753,7 @@ if MCP_AVAILABLE:
                     return
                 session_id = _get_session_id_from_scope(scope)
 
+            body = b""
             if scope.get("method") == "POST":
                 consumed_messages, body = await _read_request_body_for_routing(receive)
                 is_initialize = _is_initialize_request(body)
@@ -3639,8 +3778,7 @@ if MCP_AVAILABLE:
                 )
                 if not await _enforce_stateful_session_cap_for_owner(request_owner):
                     verbose_logger.warning(
-                        "Rejecting MCP initialize: caller already holds the maximum "
-                        "number of active stateful sessions."
+                        "Rejecting MCP initialize: caller already holds the maximum number of active stateful sessions."
                     )
                     too_many_response = JSONResponse(
                         status_code=429,
@@ -3672,9 +3810,56 @@ if MCP_AVAILABLE:
             # POST/DELETE are the methods that actually mutate the shared
             # auth context, so serializing those is sufficient for the
             # clobbering race between concurrent JSON-RPC calls.
-            session_lock: Optional[asyncio.Lock] = None
+            #
+            # Also skip the lock for JSON-RPC *responses* (POSTs that carry
+            # a ``result`` or ``error`` but no ``method``). These are replies
+            # to server-initiated requests such as ``elicitation/create`` or
+            # ``sampling/createMessage``. The in-flight tool-call POST that
+            # triggered the server request already holds the session lock, so
+            # trying to acquire it again for the response POST would deadlock.
+            is_jsonrpc_response = False
             request_method = (scope.get("method") or "").upper()
-            if use_stateful and session_id and request_method in ("POST", "DELETE"):
+            if body and request_method == "POST":
+                try:
+                    _peeked = json.loads(body)
+                    if (
+                        isinstance(_peeked, dict)
+                        and _peeked.get("jsonrpc") == "2.0"
+                        and "id" in _peeked
+                        and "method" not in _peeked
+                        and ("result" in _peeked or "error" in _peeked)
+                    ):
+                        is_jsonrpc_response = True
+                        verbose_logger.debug(
+                            "MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock",
+                            _peeked.get("id"),
+                        )
+                except (json.JSONDecodeError, TypeError):
+                    # Peek cap truncated the body, so it can't be fully parsed.
+                    # Scan the top-level keys (depth-aware) instead of a flat
+                    # substring search: a response's result payload may nest a
+                    # "method" field, and misreading that would acquire the lock
+                    # and deadlock the in-flight tool call awaiting this
+                    # response. A false skip is harmless; a false acquire is not.
+                    _body_str = body.decode("utf-8", errors="replace")
+                    if (
+                        '"jsonrpc"' in _body_str
+                        and ('"result"' in _body_str or '"error"' in _body_str)
+                        and not _jsonrpc_text_has_top_level_method(_body_str)
+                    ):
+                        is_jsonrpc_response = True
+                        verbose_logger.debug(
+                            "MCP: detected truncated JSON-RPC response POST via "
+                            "top-level key scan, skipping session lock to avoid deadlock"
+                        )
+
+            session_lock: Optional[asyncio.Lock] = None
+            if (
+                use_stateful
+                and session_id
+                and request_method in ("POST", "DELETE")
+                and not is_jsonrpc_response
+            ):
                 session_lock = _stateful_session_locks.setdefault(
                     session_id, asyncio.Lock()
                 )
@@ -4099,6 +4284,119 @@ if MCP_AVAILABLE:
             )
         return None, None, None, None, None, None, None
 
+    def _get_current_session():
+        try:
+            from mcp.server.lowlevel.server import request_ctx
+
+            return request_ctx.get().session
+        except (LookupError, ImportError):
+            return None
+
+    def _cache_auth_context_lazily():
+        session = _get_current_session()
+        if session is None:
+            return
+        try:
+            if session in _session_obj_auth_storage:
+                return
+        except TypeError:
+            verbose_logger.debug(
+                "_cache_auth_context_lazily: session object is unhashable (type=%s), cannot cache auth context",
+                type(session).__name__,
+            )
+            return
+
+        auth = auth_context_var.get()
+        if auth and isinstance(auth, MCPAuthenticatedUser):
+            try:
+                _session_obj_auth_storage[session] = auth
+            except TypeError:
+                verbose_logger.debug(
+                    "_cache_auth_context_lazily: could not store auth via "
+                    "session identity — session object is unhashable"
+                )
+
+    def _recover_auth_from_session() -> Optional[MCPAuthenticatedUser]:
+        session = _get_current_session()
+        if session is None:
+            return None
+
+        stored: Optional[MCPAuthenticatedUser] = None
+        try:
+            stored = _session_obj_auth_storage.get(session)
+        except TypeError:
+            verbose_logger.debug(
+                "_recover_auth_from_session: session object is unhashable "
+                "(type=%s), skipping _session_obj_auth_storage lookup",
+                type(session).__name__,
+            )
+
+        return stored
+
+    async def get_or_extract_auth_context() -> Tuple[
+        Optional[UserAPIKeyAuth],
+        Optional[str],
+        Optional[List[str]],
+        Optional[Dict[str, Dict[str, str]]],
+        Optional[Dict[str, str]],
+        Optional[Dict[str, str]],
+        Optional[str],
+    ]:
+        """
+        Get auth context from ContextVar first, then fall back to session
+        storage (which survives cross-task boundaries in the MCP SDK).
+        """
+        (
+            user_api_key_auth,
+            mcp_auth_header,
+            mcp_servers,
+            mcp_server_auth_headers,
+            oauth2_headers,
+            raw_headers,
+            _client_ip,
+        ) = get_auth_context()
+
+        if user_api_key_auth is not None:
+            _cache_auth_context_lazily()
+        else:
+            stored = _recover_auth_from_session()
+
+            if stored:
+                user_api_key_auth = stored.user_api_key_auth
+                mcp_auth_header = stored.mcp_auth_header
+                mcp_servers = stored.mcp_servers
+                mcp_server_auth_headers = stored.mcp_server_auth_headers
+                oauth2_headers = stored.oauth2_headers
+                raw_headers = stored.raw_headers
+                _client_ip = stored.client_ip
+        return (
+            user_api_key_auth,
+            mcp_auth_header,
+            mcp_servers,
+            mcp_server_auth_headers,
+            oauth2_headers,
+            raw_headers,
+            _client_ip,
+        )
+
+    def get_active_mcp_session() -> Optional[_McpServerSession]:
+        """Return the active MCP session captured during handler execution."""
+        session = active_mcp_session_var.get()
+        if session is not None:
+            return session
+        return _get_current_session()
+
+    def get_active_auth_context() -> Optional[MCPAuthenticatedUser]:
+        """Return auth context from ContextVar or session storage."""
+        auth = auth_context_var.get()
+        if auth and isinstance(auth, MCPAuthenticatedUser):
+            return auth
+
+        stored = _recover_auth_from_session()
+        if stored is not None:
+            return stored
+        return None
+
     ########################################################
     ############ End of Auth Context Functions #############
     ########################################################

litellm/proxy/caching_routes.py

@@ -60,11 +60,20 @@ async def cache_ping():
     """
     litellm_cache_params: Dict[str, Any] = {}
     cleaned_cache_params: Dict[str, Any] = {}
+    if litellm.cache is None:
+        raise ProxyException(
+            message=safe_dumps(
+                {
+                    "message": "Cache not initialized. litellm.cache is None",
+                    "litellm_cache_params": "{}",
+                    "health_check_cache_params": "{}",
+                }
+            ),
+            type=ProxyErrorTypes.cache_ping_error,
+            param="cache_ping",
+            code=503,
+        )
     try:
-        if litellm.cache is None:
-            raise HTTPException(
-                status_code=503, detail="Cache not initialized. litellm.cache is None"
-            )
         litellm_cache_params = masker.mask_dict(vars(litellm.cache))
         # remove field that might reference itself
         litellm_cache_params.pop("cache", None)
@@ -97,14 +106,14 @@ async def cache_ping():
                 cache_type=str(litellm.cache.type),
                 litellm_cache_params=safe_dumps(litellm_cache_params),
             )
-    except Exception as e:
-        import traceback
-
+    except HTTPException:
+        raise
+    except Exception:
+        verbose_proxy_logger.exception("Cache health check failed")
         error_message = {
-            "message": f"Service Unhealthy ({str(e)})",
+            "message": "Service Unhealthy",
             "litellm_cache_params": safe_dumps(litellm_cache_params),
             "health_check_cache_params": safe_dumps(cleaned_cache_params),
-            "traceback": traceback.format_exc(),
         }
         raise ProxyException(
             message=safe_dumps(error_message),

litellm/proxy/guardrails/guardrail_hooks/tool_permission.py

@@ -16,7 +16,7 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.callback_utils import (
     add_guardrail_to_applied_guardrails_header,
 )
-from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
 from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
     PermissionError,
     ToolPermissionRule,
@@ -60,53 +60,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
         super().__init__(**kwargs)
 
-        self.rules: List[ToolPermissionRule] = []
-        self._compiled_rule_patterns: Dict[str, Dict[str, re.Pattern]] = {}
-        self._compiled_rule_targets: Dict[str, Dict[str, Optional[re.Pattern]]] = {}
-        if rules:
-            for rule_item in rules:
-                if isinstance(rule_item, ToolPermissionRule):
-                    rule = rule_item
-                else:
-                    rule = ToolPermissionRule(**rule_item)
-                self.rules.append(rule)
-
-                compiled_target_patterns: Dict[str, Optional[re.Pattern]] = {
-                    "tool_name": None,
-                    "tool_type": None,
-                }
-                if rule.tool_name is not None:
-                    try:
-                        compiled_target_patterns["tool_name"] = re.compile(
-                            rule.tool_name
-                        )
-                    except re.error as exc:
-                        raise ValueError(
-                            f"Invalid regex for tool_name in rule '{rule.id}': {exc}"
-                        ) from exc
-                if rule.tool_type is not None:
-                    try:
-                        compiled_target_patterns["tool_type"] = re.compile(
-                            rule.tool_type
-                        )
-                    except re.error as exc:
-                        raise ValueError(
-                            f"Invalid regex for tool_type in rule '{rule.id}': {exc}"
-                        ) from exc
-                self._compiled_rule_targets[rule.id] = compiled_target_patterns
-
-                if rule.allowed_param_patterns:
-                    compiled_patterns: Dict[str, re.Pattern] = {}
-                    for path, pattern in rule.allowed_param_patterns.items():
-                        try:
-                            compiled_patterns[path] = re.compile(pattern)
-                        except re.error as exc:
-                            raise ValueError(
-                                f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}"
-                            ) from exc
-
-                    if compiled_patterns:
-                        self._compiled_rule_patterns[rule.id] = compiled_patterns
+        self._load_rules(rules)
 
         # Normalize to lowercase for case-insensitive handling
         self.default_action = (
@@ -126,6 +80,115 @@ class ToolPermissionGuardrail(CustomGuardrail):
             self.default_action,
         )
 
+    def _load_rules(self, rules: Optional[List[Any]]) -> None:
+        """Parse ``rules`` and (re)build the compiled target/pattern lookups.
+
+        ``self.rules`` plus ``_compiled_rule_targets`` / ``_compiled_rule_patterns``
+        are the state every matching path reads. Centralizing the build here lets
+        both ``__init__`` and ``update_in_memory_litellm_params`` recompile from a
+        single source of truth, so an in-place update (PUT /guardrails, immediate
+        sync) reflects rule changes instead of keeping the construction-time maps.
+        """
+        parsed_rules: List[ToolPermissionRule] = []
+        compiled_targets: Dict[str, Dict[str, Optional[re.Pattern]]] = {}
+        compiled_patterns: Dict[str, Dict[str, re.Pattern]] = {}
+
+        for rule_item in rules or []:
+            rule = (
+                rule_item
+                if isinstance(rule_item, ToolPermissionRule)
+                else ToolPermissionRule(**rule_item)
+            )
+
+            target_patterns: Dict[str, Optional[re.Pattern]] = {
+                "tool_name": None,
+                "tool_type": None,
+            }
+            if rule.tool_name is not None:
+                try:
+                    target_patterns["tool_name"] = re.compile(rule.tool_name)
+                except re.error as exc:
+                    raise ValueError(
+                        f"Invalid regex for tool_name in rule '{rule.id}': {exc}"
+                    ) from exc
+            if rule.tool_type is not None:
+                try:
+                    target_patterns["tool_type"] = re.compile(rule.tool_type)
+                except re.error as exc:
+                    raise ValueError(
+                        f"Invalid regex for tool_type in rule '{rule.id}': {exc}"
+                    ) from exc
+
+            rule_patterns: Dict[str, re.Pattern] = {}
+            for path, pattern in (rule.allowed_param_patterns or {}).items():
+                try:
+                    rule_patterns[path] = re.compile(pattern)
+                except re.error as exc:
+                    raise ValueError(
+                        f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}"
+                    ) from exc
+
+            parsed_rules.append(rule)
+            compiled_targets[rule.id] = target_patterns
+            if rule_patterns:
+                compiled_patterns[rule.id] = rule_patterns
+
+        # Swap in the fully-built maps only after every rule compiles, so an
+        # invalid regex raises without leaving a partially-built ruleset (a
+        # missing compiled target is read as a match-all wildcard).
+        self.rules = parsed_rules
+        self._compiled_rule_targets = compiled_targets
+        self._compiled_rule_patterns = compiled_patterns
+
+    def update_in_memory_litellm_params(
+        self, litellm_params: Union[LitellmParams, dict]
+    ) -> None:
+        """Apply updated params in place, rebuilding the compiled rule state.
+
+        The base implementation only ``setattr``s raw fields, which would leave
+        ``_compiled_rule_targets`` / ``_compiled_rule_patterns`` (built in
+        ``__init__``) stale, so a guardrail updated without reinitialization would
+        keep enforcing the old ruleset. Recompile here so PUT /guardrails and the
+        immediate in-memory sync take effect, mirroring the PresidioGuardrail
+        override of this method.
+        """
+        # ``litellm_params`` may arrive as the raw DB dict (the proxy ``cast()``s
+        # it to ``LitellmParams`` without converting), so handle both shapes. The
+        # base ``setattr`` loop is model-only, so apply the dict case here.
+        previous_rules = self.rules
+        if isinstance(litellm_params, dict):
+            params = litellm_params
+            for key, value in params.items():
+                setattr(self, key, value)
+        else:
+            super().update_in_memory_litellm_params(litellm_params)
+            params = vars(litellm_params)
+
+        # The generic update above sets ``self.rules`` from the incoming value
+        # (None on a partial update that omits rules), but never rebuilds the
+        # compiled maps. Rebuild them when rules are provided; otherwise restore
+        # the previous ruleset so a partial update doesn't silently wipe it. An
+        # explicit empty list still clears the rules.
+        rules = params.get("rules")
+        if rules is not None:
+            try:
+                self._load_rules(rules)
+            except Exception:
+                # The generic update above may have overwritten self.rules with
+                # the raw payload; restore the prior consistent ruleset so a
+                # rejected update can't leave the live guardrail enforcing a
+                # broken policy.
+                self.rules = previous_rules
+                raise
+        else:
+            self.rules = previous_rules
+        default_action = params.get("default_action")
+        if isinstance(default_action, str):
+            self.default_action = default_action.lower()
+        on_disallowed_action = params.get("on_disallowed_action")
+        if isinstance(on_disallowed_action, str):
+            self.on_disallowed_action = on_disallowed_action.lower()
+
     @staticmethod
     def get_config_model():
         from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (

litellm/proxy/proxy_server.py

@@ -1910,7 +1910,7 @@ prompt_injection_detection_obj: Optional[_OPTIONAL_PromptInjectionDetection] = N
 store_model_in_db: bool = False
 open_telemetry_logger: Optional[OpenTelemetry] = None
 ### INITIALIZE GLOBAL LOGGING OBJECT ###
-proxy_logging_obj = ProxyLogging(
+proxy_logging_obj: ProxyLogging = ProxyLogging(
     user_api_key_cache=user_api_key_cache, premium_user=premium_user
 )
 ### REDIS QUEUE ###
@@ -15844,10 +15844,10 @@ async def toolset_mcp_route(toolset_name: str, request: Request):
     except HTTPException as e:
         raise e
     except Exception as e:
-        verbose_proxy_logger.error(
-            f"Error handling toolset MCP route for {toolset_name}: {str(e)}"
+        verbose_proxy_logger.exception(
+            "Error handling toolset MCP route for %s: %s", toolset_name, str(e)
         )
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")
 
 
 async def _mcp_forward_as_path(path_segment: str, request: Request):
@@ -16028,7 +16028,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request):
     except HTTPException as e:
         raise e
     except Exception as e:
-        verbose_proxy_logger.error(
-            f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}"
+        verbose_proxy_logger.exception(
+            "Error handling dynamic MCP route for %s: %s", mcp_server_name, str(e)
         )
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")

litellm/types/llms/gemini.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional
 
-from typing_extensions import TypedDict
+from typing_extensions import Required, TypedDict
 
 from .vertex_ai import (
     GenerationConfig,
@@ -233,10 +233,11 @@ class GeminiImageGenerationResponse(TypedDict):
 
 
 # Video Generation Types
-class GeminiVideoGenerationInstance(TypedDict):
+class GeminiVideoGenerationInstance(TypedDict, total=False):
     """Instance data for Gemini video generation request"""
 
-    prompt: str
+    prompt: Required[str]
+    image: Dict[str, Any]
 
 
 class GeminiVideoGenerationParameters(BaseModel):
@@ -264,11 +265,6 @@ class GeminiVideoGenerationParameters(BaseModel):
     negativePrompt: Optional[str] = None
     """Text describing what not to include in the video."""
 
-    image: Optional[Any] = None
-    """
-    An initial image to animate (Image object).
-    """
-
     lastFrame: Optional[Any] = None
     """
     The final image for interpolation video to transition.

litellm/types/mcp_server/mcp_server_manager.py

@@ -115,6 +115,8 @@ class MCPServer(BaseModel):
     # different ``server_id`` values are bumped deterministically.  Left
     # ``None`` in default-prefix mode.
     short_prefix: Optional[str] = None
+    allow_sampling: bool = False
+    allow_elicitation: bool = False
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @property

litellm/utils.py

@@ -4871,7 +4871,7 @@ def add_provider_specific_params_to_optional_params(
             )
             is False
         ):
-            extra_body = passed_params.pop("extra_body", None) or {}
+            extra_body = dict(passed_params.pop("extra_body", None) or {})
             for k in passed_params.keys():
                 if k not in openai_params and passed_params[k] is not None:
                     extra_body[k] = passed_params[k]
@@ -8909,6 +8909,16 @@ class ProviderConfigManager:
             return litellm.OpenRouterResponsesAPIConfig()
         elif litellm.LlmProviders.HOSTED_VLLM == provider:
             return litellm.HostedVLLMResponsesAPIConfig()
+        elif litellm.LlmProviders.BEDROCK_MANTLE == provider:
+            # Only OpenAI gpt frontier models (gpt-5.x, and future gpt-6 etc.) are
+            # served on the /openai/v1/responses path. gpt-oss and every non-OpenAI
+            # model on Mantle (nvidia, mistral, google, zai, ...) are chat-completions
+            # only and 400 on that path, so they fall through to None to keep the
+            # chat-completions emulation (see litellm/responses/main.py "config is None").
+            model_lower = model.lower() if model else ""
+            if "openai.gpt-" in model_lower and "gpt-oss" not in model_lower:
+                return litellm.BedrockMantleResponsesAPIConfig()
+            return None
         return None
 
     @staticmethod

model_prices_and_context_window.json

@@ -30424,21 +30424,32 @@
         "supports_reasoning": true,
         "source": "https://cloud.sambanova.ai/plans/pricing"
     },
-    "snowflake/claude-3-5-sonnet": {
+      "snowflake/claude-3-5-sonnet": {
         "litellm_provider": "snowflake",
-        "max_input_tokens": 18000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
+        "max_input_tokens": 200000,
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
         "mode": "chat",
-        "supports_computer_use": true
+        "input_cost_per_token": 0.000003,
+        "output_cost_per_token": 0.000015,
+        "cache_read_input_token_cost": 0.0000003,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_vision": true,
+        "supports_prompt_caching": true,
+        "supports_system_messages": true,
+        "supports_response_schema": true
     },
-    "snowflake/deepseek-r1": {
+      "snowflake/deepseek-r1": {
         "litellm_provider": "snowflake",
-        "max_input_tokens": 32768,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
+        "max_input_tokens": 128000,
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
         "mode": "chat",
-        "supports_reasoning": true
+        "input_cost_per_token": 0.00000135,
+        "output_cost_per_token": 0.0000054,
+        "supports_reasoning": true,
+        "supports_system_messages": true
     },
     "snowflake/gemma-7b": {
         "litellm_provider": "snowflake",
@@ -30492,23 +30503,34 @@
     "snowflake/llama3.1-405b": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
+        "mode": "chat",
+        "input_cost_per_token": 0.0000012,
+        "output_cost_per_token": 0.0000012,
+        "supports_function_calling": true,
+        "supports_system_messages": true
     },
     "snowflake/llama3.1-70b": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
+        "mode": "chat",
+        "input_cost_per_token": 0.00000072,
+        "output_cost_per_token": 0.00000072,
+        "supports_function_calling": true,
+        "supports_system_messages": true
     },
     "snowflake/llama3.1-8b": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
+        "mode": "chat",
+        "input_cost_per_token": 0.00000024,
+        "output_cost_per_token": 0.00000024,
+        "supports_system_messages": true
     },
     "snowflake/llama3.2-1b": {
         "litellm_provider": "snowflake",
@@ -30524,13 +30546,17 @@
         "max_tokens": 8192,
         "mode": "chat"
     },
-    "snowflake/llama3.3-70b": {
-        "litellm_provider": "snowflake",
+      "snowflake/llama3.3-70b": {
+        "max_tokens": 16384,
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
-    },
+        "max_output_tokens": 16384,
+        "input_cost_per_token": 0.00000072,
+        "output_cost_per_token": 0.00000072,
+        "litellm_provider": "snowflake",
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_system_messages": true
+      },
     "snowflake/mistral-7b": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 32000,
@@ -30545,12 +30571,17 @@
         "max_tokens": 8192,
         "mode": "chat"
     },
-    "snowflake/mistral-large2": {
+     "snowflake/mistral-large2": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
+        "mode": "chat",
+        "input_cost_per_token": 0.000002,
+        "output_cost_per_token": 0.000006,
+        "supports_function_calling": true,
+        "supports_system_messages": true,
+        "supports_response_schema": true
     },
     "snowflake/mixtral-8x7b": {
         "litellm_provider": "snowflake",
@@ -30587,13 +30618,17 @@
         "max_tokens": 8192,
         "mode": "chat"
     },
-    "snowflake/snowflake-llama-3.3-70b": {
+      "snowflake/snowflake-llama-3.3-70b": {
+        "max_tokens": 16384,
+        "max_input_tokens": 128000,
+        "max_output_tokens": 16384,
+        "input_cost_per_token": 0.00000072,
+        "output_cost_per_token": 0.00000072,
         "litellm_provider": "snowflake",
-        "max_input_tokens": 8000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
-    },
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_system_messages": true
+      },
     "stability/sd3": {
         "litellm_provider": "stability",
         "mode": "image_generation",
@@ -41223,6 +41258,44 @@
         "supports_response_schema": true,
         "supports_tool_choice": true
     },
+    "bedrock_mantle/openai.gpt-5.5": {
+        "input_cost_per_token": 5.5e-06,
+        "cache_read_input_token_cost": 5.5e-07,
+        "output_cost_per_token": 3.3e-05,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 272000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "responses",
+        "supported_endpoints": ["/v1/responses"],
+        "supported_modalities": ["text", "image"],
+        "supported_output_modalities": ["text"],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/openai.gpt-5.4": {
+        "input_cost_per_token": 2.75e-06,
+        "cache_read_input_token_cost": 2.75e-07,
+        "output_cost_per_token": 1.65e-05,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 272000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "responses",
+        "supported_endpoints": ["/v1/responses"],
+        "supported_modalities": ["text", "image"],
+        "supported_output_modalities": ["text"],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "volcengine/doubao-seed-2-0-pro-260215": {
         "litellm_provider": "volcengine",
         "max_input_tokens": 256000,
@@ -41501,5 +41574,180 @@
         "supports_vision": true,
         "supports_native_structured_output": true,
         "supports_pdf_input": true
-    }
-}
+    },
+  "snowflake/claude-sonnet-4-5": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000003,
+    "output_cost_per_token": 0.000015,
+    "cache_read_input_token_cost": 0.0000003,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-sonnet-4-6": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000003,
+    "output_cost_per_token": 0.000015,
+    "cache_read_input_token_cost": 0.0000003,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-4-sonnet": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000003,
+    "output_cost_per_token": 0.000015,
+    "cache_read_input_token_cost": 0.0000003,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-4-opus": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000005,
+    "output_cost_per_token": 0.000025,
+    "cache_read_input_token_cost": 0.0000005,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_reasoning": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-haiku-4-5": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000001,
+    "output_cost_per_token": 0.000005,
+    "cache_read_input_token_cost": 0.0000001,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-3-7-sonnet": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000003,
+    "output_cost_per_token": 0.000015,
+    "cache_read_input_token_cost": 0.0000003,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_reasoning": true,
+    "supports_response_schema": true
+  },
+  "snowflake/openai-gpt-4.1": {
+    "max_tokens": 16384,
+    "max_input_tokens": 300000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000002,
+    "output_cost_per_token": 0.000008,
+    "cache_read_input_token_cost": 0.0000005,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/openai-gpt-5": {
+    "max_tokens": 16384,
+    "max_input_tokens": 300000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.00000125,
+    "output_cost_per_token": 0.00001,
+    "cache_read_input_token_cost": 0.000000125,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_reasoning": true,
+    "supports_response_schema": true
+  },
+  "snowflake/openai-gpt-5-mini": {
+    "max_tokens": 16384,
+    "max_input_tokens": 1000000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.0000003,
+    "output_cost_per_token": 0.0000012,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/openai-gpt-5-nano": {
+    "max_tokens": 16384,
+    "max_input_tokens": 5000000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.00000015,
+    "output_cost_per_token": 0.0000006,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/llama4-maverick": {
+    "max_tokens": 16384,
+    "max_input_tokens": 128000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.00000024,
+    "output_cost_per_token": 0.00000097,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_system_messages": true
+  },
+  "snowflake/snowflake-arctic-embed-l-v2.0": {
+    "max_tokens": 8192,
+    "max_input_tokens": 8192,
+    "input_cost_per_token": 0.00000007,
+    "output_cost_per_token": 0.0,
+    "litellm_provider": "snowflake",
+    "mode": "embedding"
+  },
+  "snowflake/snowflake-arctic-embed-m-v2.0": {
+    "max_tokens": 8192,
+    "max_input_tokens": 8192,
+    "input_cost_per_token": 0.00000007,
+    "output_cost_per_token": 0.0,
+    "litellm_provider": "snowflake",
+    "mode": "embedding"
+  }
+ }
+

tests/llm_translation/test_fireworks_ai_translation.py

@@ -43,12 +43,14 @@ def test_map_openai_params_tool_choice():
 
 def test_map_response_format():
     """
-    Test that the response format is translated correctly.
+    json_schema response_format is passed through to Fireworks unchanged.
 
-    h/t to https://github.com/DaveDeCaprio (@DaveDeCaprio) for the test case
+    Fireworks accepts the OpenAI strict json_schema shape natively. The earlier
+    downgrade to {type: json_object, schema: ...} silently dropped `strict` and
+    `name`, producing a request that Fireworks treats as "any valid JSON" per
+    its docs, disabling grammar-guided decoding.
 
-    Relevant Issue: https://github.com/BerriAI/litellm/issues/6797
-    Fireworks AI Ref: https://docs.fireworks.ai/structured-responses/structured-response-formatting#step-1-import-libraries
+    Ref: https://docs.fireworks.ai/structured-responses/structured-response-formatting
     """
     response_format = {
         "type": "json_schema",
@@ -65,16 +67,7 @@ def test_map_response_format():
     result = fireworks.map_openai_params(
         {"response_format": response_format}, {}, "some_model", drop_params=False
     )
-    assert result == {
-        "response_format": {
-            "type": "json_object",
-            "schema": {
-                "properties": {"result": {"type": "boolean"}},
-                "required": ["result"],
-                "type": "object",
-            },
-        }
-    }
+    assert result == {"response_format": response_format}
 
 
 class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest):

tests/local_testing/test_custom_llm.py

@@ -44,7 +44,14 @@ from litellm import (
     image_generation,
 )
 from litellm.utils import ModelResponseIterator
-from litellm.types.utils import ImageResponse, ImageObject, EmbeddingResponse
+from litellm.types.utils import (
+    ImageResponse,
+    ImageObject,
+    EmbeddingResponse,
+    ModelResponseStream,
+    StreamingChoices,
+    Delta,
+)
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 
 
@@ -644,3 +651,82 @@ async def test_simple_aembedding():
         "embedding": [0.1, 0.2, 0.3],
         "index": 1,
     }
+
+
+# ── Tests for ModelResponseStream passthrough in custom providers (issue #27389) ──
+
+
+class ModelResponseStreamLLM(MyCustomLLM):
+    """Subclass that overrides streaming/astreaming to yield ModelResponseStream directly."""
+
+    def __init__(self, finish_reason: str = "stop"):
+        self._finish_reason = finish_reason
+
+    def streaming(self, *args, **kwargs) -> Iterator[ModelResponseStream]:  # type: ignore
+        yield ModelResponseStream(
+            id="test-stream-id",
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(content="Hello world"),
+                    finish_reason=self._finish_reason,
+                )
+            ],
+        )
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[ModelResponseStream]:  # type: ignore
+        yield ModelResponseStream(
+            id="test-stream-id",
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(content="Hello world"),
+                    finish_reason=self._finish_reason,
+                )
+            ],
+        )
+
+
+@pytest.mark.parametrize(
+    "finish_reason", ["stop", "tool_calls", "length", "content_filter"]
+)
+def test_custom_llm_streaming_model_response_stream(finish_reason):
+    my_custom_llm = ModelResponseStreamLLM(finish_reason=finish_reason)
+    litellm.custom_provider_map = [
+        {"provider": "custom_llm", "custom_handler": my_custom_llm}
+    ]
+    resp = completion(
+        model="custom_llm/my-fake-model",
+        messages=[{"role": "user", "content": "Hello world!"}],
+        stream=True,
+    )
+
+    for chunk in resp:
+        print(chunk)
+        if chunk.choices[0].finish_reason is None:
+            assert isinstance(chunk.choices[0].delta.content, str)
+        else:
+            assert chunk.choices[0].finish_reason == finish_reason
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "finish_reason", ["stop", "tool_calls", "length", "content_filter"]
+)
+async def test_custom_llm_astreaming_model_response_stream(finish_reason):
+    my_custom_llm = ModelResponseStreamLLM(finish_reason=finish_reason)
+    litellm.custom_provider_map = [
+        {"provider": "custom_llm", "custom_handler": my_custom_llm}
+    ]
+    resp = await litellm.acompletion(
+        model="custom_llm/my-fake-model",
+        messages=[{"role": "user", "content": "Hello world!"}],
+        stream=True,
+    )
+
+    async for chunk in resp:
+        print(chunk)
+        if chunk.choices[0].finish_reason is None:
+            assert isinstance(chunk.choices[0].delta.content, str)
+        else:
+            assert chunk.choices[0].finish_reason == finish_reason

tests/local_testing/test_get_llm_provider.py

@@ -131,6 +131,7 @@ def test_default_api_base():
     from litellm.litellm_core_utils.get_llm_provider_logic import (
         _get_openai_compatible_provider_info,
     )
+    from litellm.types.utils import LlmProviders
 
     # Patch environment variable to remove API base if it's set
     with patch.dict(os.environ, {}, clear=True):
@@ -150,13 +151,13 @@ def test_default_api_base():
             if api_base is None:
                 continue
 
-            for other_provider in litellm.provider_list:
-                if other_provider != provider and provider != "{}_chat".format(
+            for other_provider in LlmProviders:
+                if other_provider.value != provider and provider != "{}_chat".format(
                     other_provider.value
                 ):
-                    if provider == "codestral" and other_provider == "mistral":
+                    if provider == "codestral" and other_provider.value == "mistral":
                         continue
-                    elif provider == "github" and other_provider == "azure":
+                    elif provider == "github" and other_provider.value == "azure":
                         continue
                     assert other_provider.value not in api_base.replace("/openai", "")
 

tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py

@@ -0,0 +1,134 @@
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../.."))
+
+from litellm.litellm_core_utils.get_supported_openai_params import (
+    get_supported_openai_params,
+)
+
+BEDROCK_REAL_MODEL = "eu.anthropic.claude-haiku-4-5-20251001-v1:0"
+BEDROCK_LABEL = "claude-haiku-4-5"
+
+
+def test_base_model_label_does_not_strip_bedrock_tools():
+    """Regression for #29618.
+
+    A Bedrock deployment whose ``model_info.base_model`` is a friendly label
+    (``claude-haiku-4-5``) must still advertise ``tools``/``tool_choice``. The label
+    on its own resolves to no tool support, so before the fix it stripped the
+    capability the real model id exposes, silently dropping function calling under
+    ``drop_params``."""
+    params = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL,
+        custom_llm_provider="bedrock",
+        base_model=BEDROCK_LABEL,
+    )
+
+    assert params is not None
+    assert "tools" in params
+    assert "tool_choice" in params
+
+
+def test_base_model_label_alone_lacks_bedrock_tools():
+    """The label by itself does not advertise tools; this is what made the union
+    necessary. Guards against the discrepancy disappearing (and the regression test
+    above silently passing for the wrong reason)."""
+    params = get_supported_openai_params(
+        model=BEDROCK_LABEL, custom_llm_provider="bedrock"
+    )
+
+    assert params is not None
+    assert "tools" not in params
+
+
+def test_base_model_is_additive_not_replacement():
+    """``base_model`` may only add capabilities, never remove ones the real model has.
+
+    Bedrock: real id supports ``tools`` but not the label's reasoning hint; the union
+    must contain the real model's ``tools`` regardless of the label being a subset."""
+    real_only = set(
+        get_supported_openai_params(
+            model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
+        )
+    )
+    label_only = set(
+        get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock")
+    )
+    combined = set(
+        get_supported_openai_params(
+            model=BEDROCK_REAL_MODEL,
+            custom_llm_provider="bedrock",
+            base_model=BEDROCK_LABEL,
+        )
+    )
+
+    assert combined == real_only | label_only
+    assert real_only - label_only  # the label really is a strict subset here
+    assert real_only <= combined
+
+
+def test_base_model_adds_capabilities_the_real_model_lacks():
+    """Regression for #27717 (the behavior the union must preserve).
+
+    ``gemini-3.1-pro`` isn't in the cost map so it advertises no reasoning support,
+    but the registered ``gemini-3.1-pro-preview`` base_model does. The hint must add
+    ``reasoning_effort``/``thinking`` without the call erroring."""
+    real_only = set(
+        get_supported_openai_params(
+            model="gemini-3.1-pro", custom_llm_provider="gemini"
+        )
+    )
+    assert "reasoning_effort" not in real_only
+
+    combined = set(
+        get_supported_openai_params(
+            model="gemini-3.1-pro",
+            custom_llm_provider="gemini",
+            base_model="gemini-3.1-pro-preview",
+        )
+    )
+    assert "reasoning_effort" in combined
+    assert "thinking" in combined
+
+
+def test_no_base_model_is_unchanged():
+    """Omitting ``base_model`` must resolve purely from ``model``."""
+    with_none = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None
+    )
+    plain = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
+    )
+
+    assert with_none == plain
+
+
+def test_base_model_equal_to_model_is_unchanged():
+    """A ``base_model`` identical to ``model`` must not double-resolve or reorder."""
+    plain = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
+    )
+    same = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL,
+        custom_llm_provider="bedrock",
+        base_model=BEDROCK_REAL_MODEL,
+    )
+
+    assert same == plain
+
+
+def test_azure_base_model_detection_preserved():
+    """Azure relies on ``base_model`` for model-type detection when the deployment
+    name is opaque; the union must keep advertising the gpt-5 capabilities."""
+    params = get_supported_openai_params(
+        model="my-opaque-deployment",
+        custom_llm_provider="azure",
+        base_model="azure/gpt-5.2",
+    )
+
+    assert params is not None
+    assert "reasoning_effort" in params
+    assert "tools" in params

tests/test_litellm/litellm_core_utils/test_streaming_handler.py

@@ -2118,3 +2118,172 @@ def test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum():
         f"Expected 'tool_calls' but got {final.choices[0].finish_reason!r}. "
         "STOP enum was not normalised through map_finish_reason()."
     )
+
+
+@pytest.mark.parametrize(
+    "finish_reason", ["stop", "tool_calls", "length", "content_filter"]
+)
+def test_chunk_creator_passes_through_model_response_stream(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+    finish_reason: str,
+):
+    """
+    chunk_creator must pass ModelResponseStream chunks from custom providers
+    straight through and preserve finish_reason exactly — not force-cast to GChunk.
+    Regression test for issue #27389.
+    """
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    litellm._custom_providers.append("my-custom-provider")
+
+    chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(content="Hello", role="assistant"),
+                finish_reason=finish_reason,
+            )
+        ],
+    )
+
+    result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+    assert result is not None
+    assert initialized_custom_stream_wrapper.received_finish_reason == finish_reason
+
+
+def test_chunk_creator_drops_empty_finish_chunk(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+):
+    """
+    A ModelResponseStream chunk with finish_reason but no content should return
+    None so finish_reason_handler() synthesises the final chunk — mirrors GChunk
+    behaviour via is_chunk_non_empty.
+    """
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    litellm._custom_providers.append("my-custom-provider")
+
+    chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(content=""),
+                finish_reason="stop",
+            )
+        ],
+    )
+
+    result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+    assert result is None
+    assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
+
+
+def test_chunk_creator_stops_iteration_on_trailing_chunk(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+):
+    """
+    After received_finish_reason is set, any empty trailing chunk (e.g. provider
+    metadata flush) must raise StopIteration to end the stream cleanly.
+    """
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    initialized_custom_stream_wrapper.received_finish_reason = "stop"
+    litellm._custom_providers.append("my-custom-provider")
+
+    trailing_chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(content=None),
+                finish_reason="stop",
+            )
+        ],
+    )
+
+    with pytest.raises(StopIteration):
+        initialized_custom_stream_wrapper.chunk_creator(chunk=trailing_chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+
+def test_chunk_creator_strips_finish_reason_from_content_chunk(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+):
+    """
+    When content and finish_reason arrive in the same chunk, finish_reason must be
+    stripped so finish_reason_handler() emits it on the synthetic terminal chunk —
+    preventing two terminal chunks (double finish_reason bug).
+    """
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    litellm._custom_providers.append("my-custom-provider")
+
+    chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(content="Hello"),
+                finish_reason="stop",
+            )
+        ],
+    )
+
+    result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+    assert result is not None
+    assert (
+        result.choices[0].finish_reason is None
+    ), "finish_reason must be stripped from content chunks to avoid double terminal chunks"
+    assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
+
+
+def test_chunk_creator_tool_calls_not_dropped_on_finish(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+):
+    """
+    A terminal chunk with finish_reason="tool_calls" and delta.tool_calls must NOT
+    be silently dropped — tool_calls counts as content so the chunk is passed through
+    (with finish_reason stripped) rather than returning None.
+    """
+    from litellm.types.utils import ChatCompletionDeltaToolCall, Function
+
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    litellm._custom_providers.append("my-custom-provider")
+
+    chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(
+                    content=None,
+                    tool_calls=[
+                        ChatCompletionDeltaToolCall(
+                            id="call_abc",
+                            function=Function(name="get_weather", arguments='{"city":"NYC"}'),
+                            type="function",
+                            index=0,
+                        )
+                    ],
+                ),
+                finish_reason="tool_calls",
+            )
+        ],
+    )
+
+    result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+    assert result is not None, "tool_calls chunk must not be dropped"
+    assert result.choices[0].delta.tool_calls is not None
+    assert result.choices[0].finish_reason is None
+    assert initialized_custom_stream_wrapper.received_finish_reason == "tool_calls"

tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py

@@ -1,5 +1,7 @@
+import urllib.parse
 from unittest.mock import patch
 
+import litellm
 from litellm.llms.azure.image_edit.transformation import AzureImageEditConfig
 from litellm.types.router import GenericLiteLLMParams
 
@@ -138,3 +140,96 @@ def test_azure_finalize_image_edit_strips_model_after_openai_transform():
     assert data_out.get("prompt") == prompt
     assert data_out.get("n") == 1
     assert len(files) >= 1
+
+
+# ---------------------------------------------------------------------------
+# api_version fallback chain
+#
+# Pin the resolution order used by ``AzureImageEditConfig.get_complete_url``:
+#   litellm_params["api_version"]
+#     > litellm.api_version (module-global)
+#     > AZURE_API_VERSION env var
+#     > litellm.AZURE_DEFAULT_API_VERSION
+#
+# Before this fallback chain existed, image edit only read ``litellm_params``
+# and produced an unversioned URL when callers set api_version via the global
+# or the env var (Azure then 404s with "Resource not found"). The chat path
+# in ``litellm/llms/azure/common_utils.py`` already had this fallback.
+# ---------------------------------------------------------------------------
+
+
+_FALLBACK_API_BASE = "https://x.openai.azure.com"
+_FALLBACK_MODEL = "gpt-image-1"
+
+
+def _query_params(url: str) -> dict:
+    return dict(urllib.parse.parse_qsl(urllib.parse.urlparse(url).query))
+
+
+def test_api_version_uses_litellm_params_first(monkeypatch):
+    monkeypatch.setattr(litellm, "api_version", "from-global", raising=False)
+    monkeypatch.setenv("AZURE_API_VERSION", "from-env")
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=_FALLBACK_API_BASE,
+        litellm_params={"api_version": "from-params"},
+    )
+
+    assert _query_params(url) == {"api-version": "from-params"}
+
+
+def test_api_version_falls_back_to_litellm_global(monkeypatch):
+    monkeypatch.setattr(litellm, "api_version", "from-global", raising=False)
+    monkeypatch.setenv("AZURE_API_VERSION", "from-env")
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=_FALLBACK_API_BASE,
+        litellm_params={},
+    )
+
+    assert _query_params(url) == {"api-version": "from-global"}
+
+
+def test_api_version_falls_back_to_env_var(monkeypatch):
+    monkeypatch.setattr(litellm, "api_version", None, raising=False)
+    monkeypatch.setenv("AZURE_API_VERSION", "from-env")
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=_FALLBACK_API_BASE,
+        litellm_params={},
+    )
+
+    assert _query_params(url) == {"api-version": "from-env"}
+
+
+def test_api_version_falls_back_to_azure_default(monkeypatch):
+    monkeypatch.setattr(litellm, "api_version", None, raising=False)
+    monkeypatch.delenv("AZURE_API_VERSION", raising=False)
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=_FALLBACK_API_BASE,
+        litellm_params={},
+    )
+
+    assert _query_params(url) == {"api-version": litellm.AZURE_DEFAULT_API_VERSION}
+
+
+def test_api_version_in_api_base_query_is_preserved(monkeypatch):
+    """``api_base`` already carrying ``?api-version=...`` must not be overridden."""
+    monkeypatch.setattr(litellm, "api_version", None, raising=False)
+    monkeypatch.delenv("AZURE_API_VERSION", raising=False)
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=(
+            f"{_FALLBACK_API_BASE}/openai/deployments/{_FALLBACK_MODEL}"
+            "/images/edits?api-version=2024-05-01-preview"
+        ),
+        litellm_params={"api_version": "would-be-overridden"},
+    )
+
+    assert _query_params(url) == {"api-version": "2024-05-01-preview"}

tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py

@@ -0,0 +1,283 @@
+"""
+Unit tests for Amazon Bedrock Mantle Responses API configuration.
+
+Mantle's gpt-5.5 / gpt-5.4 are served ONLY on the non-standard
+`/openai/v1/responses` path. These tests lock the URL construction and
+Bearer auth that make that routing work.
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../../../../.."))
+
+import pytest
+
+import litellm
+from litellm.llms.bedrock_mantle.responses.transformation import (
+    BedrockMantleResponsesAPIConfig,
+)
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+
+class TestBedrockMantleResponsesURL:
+    def test_url_uses_region_from_env(self, monkeypatch):
+        monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2")
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(api_base=None, litellm_params={})
+        assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+
+    def test_url_normalizes_v1_suffix(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(
+            api_base="https://bedrock-mantle.us-east-2.api.aws/v1",
+            litellm_params={},
+        )
+        assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+        assert "/v1/openai/v1/responses" not in url
+        url_trailing = cfg.get_complete_url(
+            api_base="https://bedrock-mantle.us-east-2.api.aws/v1/",
+            litellm_params={},
+        )
+        assert (
+            url_trailing
+            == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+        )
+
+    def test_url_does_not_double_openai_v1(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(
+            api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1",
+            litellm_params={},
+        )
+        assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+
+    def test_url_full_endpoint_base_not_doubled(self, monkeypatch):
+        # AWS model card tells users to set OPENAI_BASE_URL to the full endpoint.
+        # If copied into api_base, it must not be doubled.
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(
+            api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
+            litellm_params={},
+        )
+        assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+        assert url.count("/responses") == 1
+
+    def test_url_region_fallback_to_aws_region(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        monkeypatch.setenv("AWS_REGION", "us-west-2")
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(api_base=None, litellm_params={})
+        assert url == "https://bedrock-mantle.us-west-2.api.aws/openai/v1/responses"
+
+    def test_url_region_default_us_east_1(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        monkeypatch.delenv("AWS_REGION", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(api_base=None, litellm_params={})
+        assert url == "https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses"
+
+
+class TestBedrockMantleResponsesAuth:
+    def test_config_api_key_takes_priority(self, monkeypatch):
+        monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key")
+        cfg = BedrockMantleResponsesAPIConfig()
+        headers = cfg.validate_environment(
+            headers={},
+            model="openai.gpt-5.5",
+            litellm_params=GenericLiteLLMParams(api_key="config-key"),
+        )
+        assert headers["Authorization"] == "Bearer config-key"
+
+    def test_env_key_fallback(self, monkeypatch):
+        monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key")
+        monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        headers = cfg.validate_environment(
+            headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams()
+        )
+        assert headers["Authorization"] == "Bearer env-key"
+
+    def test_bedrock_bearer_token_fallback(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
+        monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bearer-key")
+        cfg = BedrockMantleResponsesAPIConfig()
+        headers = cfg.validate_environment(
+            headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams()
+        )
+        assert headers["Authorization"] == "Bearer bearer-key"
+
+    def test_missing_key_raises(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
+        monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        with pytest.raises(ValueError, match="Bedrock Mantle API key"):
+            cfg.validate_environment(
+                headers={},
+                model="openai.gpt-5.5",
+                litellm_params=GenericLiteLLMParams(),
+            )
+
+    def test_custom_llm_provider(self):
+        cfg = BedrockMantleResponsesAPIConfig()
+        assert cfg.custom_llm_provider == LlmProviders.BEDROCK_MANTLE
+
+    def test_native_websocket_disabled(self):
+        # Mantle Responses has no realtime/websocket transport, so the config
+        # must opt out; otherwise realtime routing would try a socket Mantle
+        # does not serve.
+        cfg = BedrockMantleResponsesAPIConfig()
+        assert cfg.supports_native_websocket() is False
+
+    def test_file_search_routes_to_emulation(self):
+        # Mantle cannot reach OpenAI's vector stores, so a native file_search
+        # tool forwarded as-is gets a 400. The config must opt out of native
+        # file_search so LiteLLM's emulation handles it instead of forwarding.
+        from litellm.responses.file_search.emulated_handler import (
+            should_use_emulated_file_search,
+        )
+
+        cfg = BedrockMantleResponsesAPIConfig()
+        assert cfg.supports_native_file_search() is False
+        assert (
+            should_use_emulated_file_search(
+                tools=[{"type": "file_search", "vector_store_ids": ["vs_1"]}],
+                provider_config=cfg,
+            )
+            is True
+        )
+
+
+class TestBedrockMantleResponsesRegistry:
+    def test_registry_returns_config_for_gpt_5_5(self):
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model="openai.gpt-5.5",
+        )
+        assert isinstance(cfg, BedrockMantleResponsesAPIConfig)
+
+    def test_registry_returns_config_for_gpt_5_4_enum(self):
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider=LlmProviders.BEDROCK_MANTLE,
+            model="openai.gpt-5.4",
+        )
+        assert isinstance(cfg, BedrockMantleResponsesAPIConfig)
+
+    def test_registry_returns_none_for_gpt_oss(self):
+        # Regression guard: gpt-oss must NOT get the native Responses config; it
+        # keeps the chat-completions emulation path (responses/main.py ~line 1109).
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model="openai.gpt-oss-120b",
+        )
+        assert cfg is None
+
+    def test_registry_returns_none_for_gpt_oss_safeguard(self):
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model="openai.gpt-oss-safeguard-20b",
+        )
+        assert cfg is None
+
+    def test_registry_returns_config_for_future_frontier_model(self):
+        # Forward-compatibility: an unseen OpenAI gpt frontier model (e.g. gpt-6) must
+        # get the native Responses config without a code change. The gate allow-lists
+        # the openai.gpt- family (minus gpt-oss), so gpt-6 matches automatically.
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model="openai.gpt-6",
+        )
+        assert isinstance(cfg, BedrockMantleResponsesAPIConfig)
+
+    @pytest.mark.parametrize(
+        "model",
+        [
+            "nvidia.nemotron-nano-9b-v2",
+            "mistral.ministral-3-3b-instruct",
+            "google.gemma-3-27b-it",
+            "zai.glm-4.6",
+        ],
+    )
+    def test_registry_returns_none_for_non_openai_models(self, model):
+        # Regression for the chat-only families on Mantle. These models 400 on
+        # /openai/v1/responses and are served on /v1/chat/completions, so the
+        # registry must NOT hand them the Responses config; they fall through to
+        # None and keep the chat-completions emulation.
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model=model,
+        )
+        assert cfg is None
+
+    def test_registry_returns_none_when_model_is_none(self):
+        # By-id operations (delete/get/cancel) call with model=None; keep returning
+        # None so those paths are unchanged.
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model=None,
+        )
+        assert cfg is None
+
+
+@pytest.fixture
+def local_cost_map(monkeypatch):
+    """Force the bundled backup cost map and re-derive the provider model sets.
+
+    ``litellm.model_cost`` is populated once at import time (here, from the
+    network-fetched ``main`` copy, which lags this branch). ``add_known_models``
+    only re-buckets whatever is already in ``model_cost``, so the cost map must
+    first be reloaded from the local backup before the new keys appear.
+    """
+    original_model_cost = litellm.model_cost
+    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
+    litellm.model_cost = litellm.get_model_cost_map(url="")
+    litellm.get_model_info.cache_clear()
+    litellm.add_known_models()
+    try:
+        yield
+    finally:
+        litellm.model_cost = original_model_cost
+        litellm.get_model_info.cache_clear()
+
+
+class TestBedrockMantleResponsesPricing:
+    def test_gpt_5_5_pricing_and_mode(self, local_cost_map):
+        info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5")
+        assert info["mode"] == "responses"
+        assert info["input_cost_per_token"] == pytest.approx(5.5e-06)
+        assert info["output_cost_per_token"] == pytest.approx(3.3e-05)
+        assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07)
+        assert info["max_input_tokens"] == 272000
+
+    def test_gpt_5_4_pricing_and_mode(self, local_cost_map):
+        info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.4")
+        assert info["mode"] == "responses"
+        assert info["input_cost_per_token"] == pytest.approx(2.75e-06)
+        assert info["output_cost_per_token"] == pytest.approx(1.65e-05)
+        assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07)
+        assert info["max_input_tokens"] == 272000
+
+    def test_models_registered(self, local_cost_map):
+        assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
+        assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models

tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py

@@ -496,3 +496,59 @@ def test_transform_tools_skips_non_function_tools():
         "type": "object",
         "properties": {"id": {"type": "string"}},
     }
+
+
+def test_map_response_format_passes_json_schema_through_unchanged():
+    """
+    json_schema response_format must reach Fireworks unchanged.
+
+    Regression guard for the prior downgrade to {type: json_object, schema: ...}
+    which silently dropped `strict` and `name` and disabled grammar-guided
+    decoding on the Fireworks side.
+    """
+    config = FireworksAIConfig()
+    response_format = {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "priority_classification",
+            "strict": True,
+            "schema": {
+                "type": "object",
+                "properties": {
+                    "priority": {
+                        "type": "string",
+                        "enum": ["high", "medium", "low"],
+                    }
+                },
+                "required": ["priority"],
+                "additionalProperties": False,
+            },
+        },
+    }
+
+    result = config.map_openai_params(
+        {"response_format": response_format},
+        {},
+        "fireworks_ai/accounts/fireworks/models/qwen3-32b",
+        drop_params=False,
+    )
+
+    rf = result["response_format"]
+    assert rf["type"] == "json_schema"
+    assert rf["json_schema"]["name"] == "priority_classification"
+    assert rf["json_schema"]["strict"] is True
+    assert rf["json_schema"]["schema"] == response_format["json_schema"]["schema"]
+
+
+def test_map_response_format_json_object_unchanged():
+    """
+    The plain json_object form keeps working as before.
+    """
+    config = FireworksAIConfig()
+    result = config.map_openai_params(
+        {"response_format": {"type": "json_object"}},
+        {},
+        "fireworks_ai/accounts/fireworks/models/qwen3-32b",
+        drop_params=False,
+    )
+    assert result == {"response_format": {"type": "json_object"}}

tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py

@@ -2,6 +2,7 @@
 Tests for Gemini (Veo) video generation transformation.
 """
 
+import io
 import json
 import os
 from unittest.mock import MagicMock, Mock, patch
@@ -132,6 +133,87 @@ class TestGeminiVideoConfig:
         assert data["parameters"]["durationSeconds"] == 8
         assert data["parameters"]["resolution"] == "1080p"
 
+    def test_transform_video_create_request_image_goes_to_instance(self):
+        """Image belongs in instances[0], not in parameters (per Veo API)."""
+        prompt = "Animate this still"
+        api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning"
+        image_dict = {"bytesBase64Encoded": "aGVsbG8=", "mimeType": "image/jpeg"}
+
+        data, _, _ = self.config.transform_video_create_request(
+            model="veo-3.0-generate-preview",
+            prompt=prompt,
+            api_base=api_base,
+            video_create_optional_request_params={
+                "image": image_dict,
+                "aspectRatio": "16:9",
+                "durationSeconds": 4,
+            },
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+
+        assert data["instances"][0]["prompt"] == prompt
+        assert data["instances"][0]["image"] == image_dict
+        assert "image" not in data.get("parameters", {})
+        assert data["parameters"]["aspectRatio"] == "16:9"
+        assert data["parameters"]["durationSeconds"] == 4
+
+    def test_transform_video_create_request_image_filelike_goes_to_instance(self):
+        """File-like image (BytesIO) gets base64-encoded into instances[0]['image']."""
+        prompt = "Animate this still"
+        api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning"
+        # 1x1 PNG (8 bytes after magic + minimal IHDR is not legal — but the
+        # transformer only cares that ImageEditRequestUtils can sniff a MIME and
+        # that .read() returns bytes; an explicit name="image.jpeg" hands the
+        # MIME sniffer a clean answer regardless of payload).
+        image_bytes = b"\xff\xd8\xff\xe0fake-jpeg-bytes"
+        image_file = io.BytesIO(image_bytes)
+        image_file.name = "still.jpeg"
+
+        data, _, _ = self.config.transform_video_create_request(
+            model="veo-3.0-generate-preview",
+            prompt=prompt,
+            api_base=api_base,
+            video_create_optional_request_params={
+                "image": image_file,
+                "aspectRatio": "16:9",
+            },
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+
+        # File-like took the _convert_image_to_gemini_format branch and landed
+        # in instances[0]["image"], not in parameters.
+        instance_image = data["instances"][0]["image"]
+        assert isinstance(instance_image, dict)
+        assert instance_image["mimeType"].startswith("image/")
+        assert instance_image["bytesBase64Encoded"]
+        # Round-trip the base64 — should equal the original bytes.
+        import base64
+
+        assert base64.b64decode(instance_image["bytesBase64Encoded"]) == image_bytes
+        assert "image" not in data.get("parameters", {})
+
+    def test_transform_video_create_request_image_none_is_dropped(self):
+        """Explicit image=None is popped and never reaches parameters."""
+        prompt = "no image at all"
+        api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning"
+
+        data, _, _ = self.config.transform_video_create_request(
+            model="veo-3.0-generate-preview",
+            prompt=prompt,
+            api_base=api_base,
+            video_create_optional_request_params={
+                "image": None,
+                "aspectRatio": "16:9",
+            },
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+
+        assert "image" not in data["instances"][0]
+        assert "image" not in data.get("parameters", {})
+
     def test_map_openai_params(self):
         """Test parameter mapping from OpenAI format to Veo format."""
         openai_params = {

tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py

@@ -104,6 +104,23 @@ class TestHuggingFaceEmbedding:
         assert "source_sentence" not in str(request_data)
         assert "sentences" not in str(request_data)
 
+    def test_embedding_allows_special_token_looking_input(self):
+        input_text = ["hello <|fim_prefix|> world"]
+
+        response = litellm.embedding(
+            model=self.model,
+            input=input_text,
+            input_type="embed",
+        )
+
+        self.mock_http.assert_called_once()
+        post_call_args = self.mock_http.call_args
+        request_data = json.loads(post_call_args[1]["data"])
+
+        assert request_data["inputs"] == input_text
+        assert response.usage.prompt_tokens > 0
+        assert response.usage.total_tokens == response.usage.prompt_tokens
+
     def test_embedding_with_sentence_similarity_task(self):
         """Test embedding when task type is sentence-similarity (requires 2+ sentences)"""
 

tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py

@@ -201,9 +201,12 @@ class TestContextCachingEndpoints:
         assert returned_params == optional_params
         assert returned_cache == "existing_cache_name"
 
-        # Verify cache key was generated with tools and model
+        # Verify cache key was generated with tools, tool_choice and model
         mock_cache_obj.get_cache_key.assert_called_once_with(
-            messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro"
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=None,
+            model="gemini-1.5-pro",
         )
 
     @pytest.mark.parametrize(
@@ -474,9 +477,12 @@ class TestContextCachingEndpoints:
         assert returned_params == optional_params
         assert returned_cache == "existing_cache_name"
 
-        # Verify cache key was generated with tools and model
+        # Verify cache key was generated with tools, tool_choice and model
         mock_cache_obj.get_cache_key.assert_called_once_with(
-            messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro"
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=None,
+            model="gemini-1.5-pro",
         )
 
     @pytest.mark.asyncio
@@ -800,6 +806,546 @@ class TestContextCachingEndpoints:
             # But original tools should still be available for comparison
             assert original_tools == self.sample_tools
 
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    def test_check_and_create_cache_tool_choice_popped_from_optional_params(
+        self, custom_llm_provider
+    ):
+        """tool_choice is popped from optional_params when cached messages exist."""
+        with patch(
+            "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+        ) as mock_separate:
+            cached_messages = [self.sample_messages[0]]
+            non_cached_messages = [self.sample_messages[1]]
+            mock_separate.return_value = (cached_messages, non_cached_messages)
+
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}}
+
+            with patch.object(
+                self.context_caching, "check_cache", return_value="existing_cache"
+            ):
+                self.context_caching.check_and_create_cache(
+                    messages=self.sample_messages,
+                    optional_params=optional_params,
+                    api_key="test_key",
+                    api_base=None,
+                    model="gemini-1.5-pro",
+                    client=self.mock_client,
+                    timeout=30.0,
+                    logging_obj=self.mock_logging,
+                    custom_llm_provider=custom_llm_provider,
+                    vertex_project="test_project",
+                    vertex_location="test_location",
+                    vertex_auth_header="vertext_test_token",
+                )
+
+            assert "tool_choice" not in optional_params
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(
+        self, custom_llm_provider
+    ):
+        """tool_choice is NOT popped when there are no cached messages (early return)."""
+        with patch(
+            "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+        ) as mock_separate:
+            mock_separate.return_value = ([], self.sample_messages)
+
+            tool_choice = {"functionCallingConfig": {"mode": "AUTO"}}
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = tool_choice
+
+            self.context_caching.check_and_create_cache(
+                messages=self.sample_messages,
+                optional_params=optional_params,
+                api_key="test_key",
+                api_base=None,
+                model="gemini-1.5-pro",
+                client=self.mock_client,
+                timeout=30.0,
+                logging_obj=self.mock_logging,
+                custom_llm_provider=custom_llm_provider,
+                vertex_project="test_project",
+                vertex_location="test_location",
+                vertex_auth_header="vertext_test_token",
+            )
+
+            assert optional_params.get("tool_choice") == tool_choice
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params(
+        self, custom_llm_provider
+    ):
+        """Async equivalent of test_check_and_create_cache_tool_choice_popped_from_optional_params."""
+        with patch(
+            "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+        ) as mock_separate:
+            cached_messages = [self.sample_messages[0]]
+            non_cached_messages = [self.sample_messages[1]]
+            mock_separate.return_value = (cached_messages, non_cached_messages)
+
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}}
+
+            with patch.object(
+                self.context_caching, "async_check_cache", return_value="existing_cache"
+            ):
+                await self.context_caching.async_check_and_create_cache(
+                    messages=self.sample_messages,
+                    optional_params=optional_params,
+                    api_key="test_key",
+                    api_base=None,
+                    model="gemini-1.5-pro",
+                    client=self.mock_async_client,
+                    timeout=30.0,
+                    logging_obj=self.mock_logging,
+                    custom_llm_provider=custom_llm_provider,
+                    vertex_project="test_project",
+                    vertex_location="test_location",
+                    vertex_auth_header="vertext_test_token",
+                )
+
+            assert "tool_choice" not in optional_params
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(
+        self, custom_llm_provider
+    ):
+        """Async equivalent of test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages."""
+        with patch(
+            "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+        ) as mock_separate:
+            mock_separate.return_value = ([], self.sample_messages)
+
+            tool_choice = {"functionCallingConfig": {"mode": "AUTO"}}
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = tool_choice
+
+            await self.context_caching.async_check_and_create_cache(
+                messages=self.sample_messages,
+                optional_params=optional_params,
+                api_key="test_key",
+                api_base=None,
+                model="gemini-1.5-pro",
+                client=self.mock_async_client,
+                timeout=30.0,
+                logging_obj=self.mock_logging,
+                custom_llm_provider=custom_llm_provider,
+                vertex_project="test_project",
+                vertex_location="test_location",
+                vertex_auth_header="vertext_test_token",
+            )
+
+            assert optional_params.get("tool_choice") == tool_choice
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    def test_check_and_create_cache_tool_choice_in_request_body(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """End-to-end: tool_choice ends up as `toolConfig` on the cache-creation HTTP POST body."""
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None  # cache miss -> create new
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_client.post.return_value = mock_response
+
+        tool_choice = {"functionCallingConfig": {"mode": "ANY"}}
+        optional_params = self.sample_optional_params.copy()
+        optional_params["tool_choice"] = tool_choice
+
+        self.context_caching.check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        self.mock_client.post.assert_called_once()
+        call_args = self.mock_client.post.call_args
+        assert call_args.kwargs["json"]["tools"] == self.sample_tools
+        assert call_args.kwargs["json"]["toolConfig"] == tool_choice
+        mock_cache_obj.get_cache_key.assert_called_once_with(
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=tool_choice,
+            model="gemini-1.5-pro",
+        )
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "async_check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    async def test_async_check_and_create_cache_tool_choice_in_request_body(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """Async equivalent of test_check_and_create_cache_tool_choice_in_request_body."""
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_async_client.post = AsyncMock(return_value=mock_response)
+
+        tool_choice = {"functionCallingConfig": {"mode": "ANY"}}
+        optional_params = self.sample_optional_params.copy()
+        optional_params["tool_choice"] = tool_choice
+
+        await self.context_caching.async_check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_async_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        call_args = self.mock_async_client.post.call_args
+        assert call_args.kwargs["json"]["tools"] == self.sample_tools
+        assert call_args.kwargs["json"]["toolConfig"] == tool_choice
+        mock_cache_obj.get_cache_key.assert_called_once_with(
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=tool_choice,
+            model="gemini-1.5-pro",
+        )
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    def test_check_and_create_cache_omits_tool_config_when_tool_choice_unset(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """When the caller didn't pass tool_choice, toolConfig must NOT appear in the cache body."""
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_client.post.return_value = mock_response
+
+        optional_params = self.sample_optional_params.copy()
+
+        self.context_caching.check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        call_args = self.mock_client.post.call_args
+        assert "tools" in call_args.kwargs["json"]
+        assert "toolConfig" not in call_args.kwargs["json"]
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    def test_check_and_create_cache_tool_choice_function_pin(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """tool_choice as a function-pin dict survives the cache body intact."""
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_client.post.return_value = mock_response
+
+        function_pin = {
+            "functionCallingConfig": {
+                "mode": "ANY",
+                "allowed_function_names": ["get_current_weather"],
+            }
+        }
+        optional_params = self.sample_optional_params.copy()
+        optional_params["tool_choice"] = function_pin
+
+        self.context_caching.check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        call_args = self.mock_client.post.call_args
+        assert call_args.kwargs["json"]["toolConfig"] == function_pin
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    def test_check_and_create_cache_tool_choice_typed_constructor(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """Exercise the actual ToolConfig(FunctionCallingConfig(...)) constructor that map_tool_choice_values produces.
+
+        ToolConfig / FunctionCallingConfig are TypedDicts (litellm/types/llms/vertex_ai.py:158, 277)
+        so this is functionally identical to the dict-literal tests above at
+        runtime — but exercising the typed constructor pins the test to the
+        same call shape map_tool_choice_values uses and auto-follows if
+        either type ever migrates to a Pydantic model upstream.
+        """
+        from litellm.types.llms.vertex_ai import (
+            FunctionCallingConfig,
+            ToolConfig,
+        )
+
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_client.post.return_value = mock_response
+
+        tool_choice = ToolConfig(
+            functionCallingConfig=FunctionCallingConfig(mode="ANY")
+        )
+        optional_params = self.sample_optional_params.copy()
+        optional_params["tool_choice"] = tool_choice
+
+        self.context_caching.check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        call_args = self.mock_client.post.call_args
+        assert call_args.kwargs["json"]["toolConfig"] == tool_choice
+        assert call_args.kwargs["json"]["toolConfig"] == {
+            "functionCallingConfig": {"mode": "ANY"}
+        }
+        mock_cache_obj.get_cache_key.assert_called_once_with(
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=tool_choice,
+            model="gemini-1.5-pro",
+        )
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    def test_check_and_create_cache_distinct_tool_choices_use_distinct_keys(
+        self,
+        mock_check_cache,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """Two requests with different tool_choice values must produce different cache keys.
+
+        Runs the real local_cache_obj.get_cache_key to verify the hashed
+        output actually differs — mocking it would only prove that distinct
+        arguments are forwarded, not that they produce distinct keys.
+        """
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_check_cache.return_value = "existing_cache"
+
+        auto_tool_choice = {"functionCallingConfig": {"mode": "AUTO"}}
+        any_tool_choice = {"functionCallingConfig": {"mode": "ANY"}}
+        for choice in (auto_tool_choice, any_tool_choice):
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = choice
+            self.context_caching.check_and_create_cache(
+                messages=self.sample_messages,
+                optional_params=optional_params,
+                api_key="test_key",
+                api_base=None,
+                model="gemini-1.5-pro",
+                client=self.mock_client,
+                timeout=30.0,
+                logging_obj=self.mock_logging,
+                custom_llm_provider=custom_llm_provider,
+                vertex_project="test_project",
+                vertex_location="test_location",
+                vertex_auth_header="vertext_test_token",
+            )
+
+        check_cache_calls = mock_check_cache.call_args_list
+        assert len(check_cache_calls) == 2
+        first_cache_key = check_cache_calls[0].kwargs["cache_key"]
+        second_cache_key = check_cache_calls[1].kwargs["cache_key"]
+        assert first_cache_key != second_cache_key
+
     @pytest.mark.parametrize(
         "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
     )

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py

@@ -0,0 +1,211 @@
+"""
+Tests for the MCP elicitation handler.
+
+Covers the gateway-mode relay logic (`elicitation/create` requests from an
+upstream MCP server being forwarded to the connected downstream client) as
+well as the decline paths used in tool-bridge mode or when the downstream
+client lacks the requested elicitation capability.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from mcp.types import (
+    ElicitRequestFormParams,
+    ElicitRequestURLParams,
+    ElicitResult,
+    ErrorData,
+)
+
+from litellm.proxy._experimental.mcp_server import elicitation_handler
+from litellm.proxy._experimental.mcp_server.elicitation_handler import (
+    _relay_elicitation_to_downstream,
+    handle_elicitation_request,
+)
+
+
+def _form_params(message: str = "fill the form") -> ElicitRequestFormParams:
+    return ElicitRequestFormParams(
+        mode="form",
+        message=message,
+        requestedSchema={"type": "object", "properties": {}},
+    )
+
+
+def _url_params(message: str = "please authorize") -> ElicitRequestURLParams:
+    return ElicitRequestURLParams(
+        mode="url",
+        message=message,
+        url="https://example.com/oauth",
+        elicitationId="elc-1",
+    )
+
+
+def _caps(*, url=True, form=True) -> SimpleNamespace:
+    elicit = SimpleNamespace(
+        url=object() if url else None,
+        form=object() if form else None,
+    )
+    return SimpleNamespace(elicitation=elicit)
+
+
+class TestHandleElicitationRequest:
+    async def test_should_decline_when_no_downstream_session(self):
+        result = await handle_elicitation_request(
+            context=SimpleNamespace(),
+            params=_form_params(),
+            downstream_session=None,
+        )
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+
+    async def test_should_relay_to_downstream_when_session_present(self):
+        accepted = ElicitResult(action="accept", content={"name": "ada"})
+        session = SimpleNamespace(elicit_form=AsyncMock(return_value=accepted))
+
+        result = await handle_elicitation_request(
+            context=SimpleNamespace(),
+            params=_form_params(),
+            downstream_session=session,
+            downstream_capabilities=None,
+        )
+
+        assert result is accepted
+        session.elicit_form.assert_awaited_once()
+
+    async def test_should_return_error_data_when_unavailable(self, monkeypatch):
+        monkeypatch.setattr(elicitation_handler, "MCP_ELICITATION_AVAILABLE", False)
+        result = await handle_elicitation_request(
+            context=SimpleNamespace(),
+            params=_form_params(),
+            downstream_session=SimpleNamespace(),
+        )
+        assert isinstance(result, ErrorData)
+        assert "not available" in result.message
+
+    async def test_should_return_error_data_on_unexpected_failure(self):
+        class _ExplodingParams:
+            mode = "form"
+
+            @property
+            def message(self):
+                raise RuntimeError("boom")
+
+        result = await handle_elicitation_request(
+            context=SimpleNamespace(),
+            params=_ExplodingParams(),
+            downstream_session=None,
+        )
+        assert isinstance(result, ErrorData)
+        assert "boom" in result.message
+
+
+class TestRelayElicitationToDownstream:
+    async def test_should_relay_form_mode(self):
+        accepted = ElicitResult(action="accept", content={"name": "ada"})
+        session = SimpleNamespace(elicit_form=AsyncMock(return_value=accepted))
+
+        params = _form_params("collect name")
+        result = await _relay_elicitation_to_downstream(
+            params=params,
+            downstream_session=session,
+            downstream_capabilities=_caps(form=True),
+        )
+
+        assert result is accepted
+        session.elicit_form.assert_awaited_once()
+        _, kwargs = session.elicit_form.call_args
+        assert kwargs["message"] == "collect name"
+        assert kwargs["requestedSchema"] == params.requestedSchema
+
+    async def test_should_relay_url_mode(self):
+        accepted = ElicitResult(action="accept")
+        session = SimpleNamespace(elicit_url=AsyncMock(return_value=accepted))
+
+        result = await _relay_elicitation_to_downstream(
+            params=_url_params(),
+            downstream_session=session,
+            downstream_capabilities=_caps(url=True),
+        )
+
+        assert result is accepted
+        session.elicit_url.assert_awaited_once()
+        _, kwargs = session.elicit_url.call_args
+        assert kwargs["url"] == "https://example.com/oauth"
+        assert kwargs["elicitation_id"] == "elc-1"
+
+    async def test_should_use_generic_elicit_for_unknown_param_type(self):
+        accepted = ElicitResult(action="accept")
+        session = SimpleNamespace(elicit=AsyncMock(return_value=accepted))
+
+        # A bare params object that is neither Form nor URL params triggers
+        # the generic fallback path.
+        params = SimpleNamespace(mode="form", message="hi", requestedSchema={})
+        result = await _relay_elicitation_to_downstream(
+            params=params,
+            downstream_session=session,
+            downstream_capabilities=None,
+        )
+
+        assert result is accepted
+        session.elicit.assert_awaited_once()
+
+    async def test_should_decline_when_elicitation_unsupported(self):
+        session = SimpleNamespace(elicit_form=AsyncMock())
+        caps = SimpleNamespace(elicitation=None)
+
+        result = await _relay_elicitation_to_downstream(
+            params=_form_params(),
+            downstream_session=session,
+            downstream_capabilities=caps,
+        )
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+        session.elicit_form.assert_not_awaited()
+
+    async def test_should_decline_url_mode_when_url_unsupported(self):
+        session = SimpleNamespace(elicit_url=AsyncMock())
+
+        result = await _relay_elicitation_to_downstream(
+            params=_url_params(),
+            downstream_session=session,
+            downstream_capabilities=_caps(url=False, form=True),
+        )
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+        session.elicit_url.assert_not_awaited()
+
+    async def test_should_decline_form_mode_when_form_unsupported(self):
+        session = SimpleNamespace(elicit_form=AsyncMock())
+
+        result = await _relay_elicitation_to_downstream(
+            params=_form_params(),
+            downstream_session=session,
+            downstream_capabilities=_caps(url=True, form=False),
+        )
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+        session.elicit_form.assert_not_awaited()
+
+    async def test_should_decline_when_downstream_relay_raises(self):
+        session = SimpleNamespace(
+            elicit_form=AsyncMock(side_effect=RuntimeError("transport closed"))
+        )
+
+        result = await _relay_elicitation_to_downstream(
+            params=_form_params(),
+            downstream_session=session,
+            downstream_capabilities=_caps(form=True),
+        )
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py

@@ -16,7 +16,6 @@ from typing import Any, Dict, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from fastapi import HTTPException
 
 from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
 from litellm.proxy._types import UserAPIKeyAuth
@@ -549,11 +548,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()
@@ -593,11 +588,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()
@@ -643,11 +634,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()
@@ -703,11 +690,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()
@@ -755,11 +738,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py

@@ -0,0 +1,254 @@
+"""
+Tests for the MCP sampling completion pipeline.
+
+Covers building the internal `acompletion` kwargs from MCP request params
+(messages, sampling options, tools, tool choice, metadata), routing the call
+through the proxy router / guardrails, and the end-to-end
+`handle_sampling_create_message` success and error-propagation behaviour.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from mcp.types import CreateMessageResult, ErrorData
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _build_completion_kwargs,
+    _run_guardrails_and_call_llm,
+    handle_sampling_create_message,
+)
+
+
+def _params(**overrides):
+    base = dict(
+        messages=[
+            SimpleNamespace(
+                role="user", content=SimpleNamespace(type="text", text="hi")
+            )
+        ],
+        systemPrompt="be concise",
+        maxTokens=128,
+        temperature=None,
+        stopSequences=None,
+        tools=None,
+        toolChoice=None,
+        metadata=None,
+        modelPreferences=None,
+    )
+    base.update(overrides)
+    return SimpleNamespace(**base)
+
+
+def _passthrough_add_data():
+    async def _add(data, **kwargs):
+        return data
+
+    return _add
+
+
+class TestBuildCompletionKwargs:
+    async def test_should_include_sampling_options_and_tools(self):
+        params = _params(
+            temperature=0.3,
+            stopSequences=["STOP"],
+            tools=[
+                SimpleNamespace(
+                    name="search", description="d", inputSchema={"type": "object"}
+                )
+            ],
+            toolChoice=SimpleNamespace(mode="required"),
+            metadata={"trace": "abc"},
+        )
+        with patch(
+            "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
+            side_effect=_passthrough_add_data(),
+        ):
+            kwargs = await _build_completion_kwargs(
+                params=params,
+                model="gpt-4o",
+                user_api_key_auth=SimpleNamespace(user_id="u1"),
+                raw_headers=None,
+                client_ip=None,
+            )
+
+        assert kwargs["model"] == "gpt-4o"
+        assert kwargs["max_tokens"] == 128
+        assert kwargs["temperature"] == 0.3
+        assert kwargs["stop"] == ["STOP"]
+        assert kwargs["tools"][0]["function"]["name"] == "search"
+        assert kwargs["tool_choice"] == "required"
+        assert kwargs["metadata"]["mcp_metadata"] == {"trace": "abc"}
+        assert kwargs["user"] == "u1"
+        assert kwargs["messages"][0] == {"role": "system", "content": "be concise"}
+
+    async def test_should_omit_optional_fields_when_unset(self):
+        with patch(
+            "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
+            side_effect=_passthrough_add_data(),
+        ):
+            kwargs = await _build_completion_kwargs(
+                params=_params(),
+                model="gpt-4o",
+                user_api_key_auth=SimpleNamespace(user_id=None),
+                raw_headers=None,
+                client_ip=None,
+            )
+
+        assert "temperature" not in kwargs
+        assert "stop" not in kwargs
+        assert "tools" not in kwargs
+        assert "tool_choice" not in kwargs
+        assert kwargs["metadata"] == {}
+
+
+class TestRunGuardrailsAndCallLlm:
+    async def test_should_route_through_llm_router_when_available(self):
+        router = MagicMock()
+        router.acompletion = AsyncMock(return_value="router-response")
+        with (
+            patch("litellm.proxy.proxy_server.proxy_logging_obj", None),
+            patch("litellm.proxy.proxy_server.llm_router", router),
+        ):
+            result = await _run_guardrails_and_call_llm(
+                completion_kwargs={"model": "gpt-4o", "messages": []},
+                user_api_key_auth=SimpleNamespace(),
+            )
+
+        assert result == "router-response"
+        router.acompletion.assert_awaited_once()
+
+    async def test_should_propagate_guardrail_rejection(self):
+        plo = MagicMock()
+        plo.pre_call_hook = AsyncMock(side_effect=ValueError("blocked by guardrail"))
+        with patch("litellm.proxy.proxy_server.proxy_logging_obj", plo):
+            with pytest.raises(ValueError, match="blocked by guardrail"):
+                await _run_guardrails_and_call_llm(
+                    completion_kwargs={"model": "gpt-4o", "messages": []},
+                    user_api_key_auth=SimpleNamespace(),
+                )
+
+
+class TestHandleSamplingCreateMessagePipeline:
+    async def test_should_return_message_result_on_success(self):
+        auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok")
+        response = SimpleNamespace(
+            choices=[
+                SimpleNamespace(
+                    message=SimpleNamespace(
+                        content="the answer is 42", tool_calls=None
+                    ),
+                    finish_reason="stop",
+                )
+            ],
+            model="gpt-4o",
+        )
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs",
+                new_callable=AsyncMock,
+                return_value={"model": "gpt-4o", "messages": []},
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_guardrails_and_call_llm",
+                new_callable=AsyncMock,
+                return_value=response,
+            ),
+        ):
+            result = await handle_sampling_create_message(
+                context=MagicMock(),
+                params=_params(),
+                default_model="gpt-4o",
+                user_api_key_auth=auth,
+            )
+
+        assert isinstance(result, CreateMessageResult)
+        assert result.content.text == "the answer is 42"
+        assert result.stopReason == "endTurn"
+
+    async def test_should_reraise_known_proxy_exceptions(self):
+        from litellm.exceptions import RateLimitError
+
+        auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok")
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs",
+                new_callable=AsyncMock,
+                side_effect=RateLimitError(
+                    "rate limited", llm_provider="openai", model="gpt-4o"
+                ),
+            ),
+        ):
+            with pytest.raises(RateLimitError):
+                await handle_sampling_create_message(
+                    context=MagicMock(),
+                    params=_params(),
+                    default_model="gpt-4o",
+                    user_api_key_auth=auth,
+                )
+
+    async def test_should_return_error_data_on_unexpected_failure(self):
+        auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok")
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs",
+                new_callable=AsyncMock,
+                side_effect=RuntimeError("kaboom"),
+            ),
+        ):
+            result = await handle_sampling_create_message(
+                context=MagicMock(),
+                params=_params(),
+                default_model="gpt-4o",
+                user_api_key_auth=auth,
+            )
+
+        assert isinstance(result, ErrorData)
+        assert "kaboom" in result.message
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py

@@ -0,0 +1,327 @@
+"""
+Tests for MCP sampling handler model-access enforcement.
+
+Verifies that handle_sampling_create_message and _check_model_access
+enforce the same model-permission checks as regular /chat/completions
+calls, preventing a malicious upstream MCP server from requesting
+inference on models the caller's API key is not authorized to use.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _check_model_access,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_user_api_key_auth(
+    *,
+    models=None,
+    team_id=None,
+    team_model_aliases=None,
+    api_key="sk-test-key",
+    token=None,
+    user_role=None,
+):
+    """Build a minimal UserAPIKeyAuth-like object for tests."""
+    auth = MagicMock()
+    auth.models = models or []
+    auth.team_id = team_id
+    auth.team_model_aliases = team_model_aliases or {}
+    auth.access_group_ids = []
+    auth.api_key = api_key
+    auth.token = token
+    auth.user_role = user_role
+    return auth
+
+
+# ---------------------------------------------------------------------------
+# _check_model_access
+# ---------------------------------------------------------------------------
+
+
+class TestCheckModelAccess:
+    """Tests for the _check_model_access helper that gates sampling requests."""
+
+    @pytest.mark.asyncio
+    async def test_should_return_none_when_no_auth_context(self):
+        """No auth context means no restriction — pass through."""
+        result = await _check_model_access("gpt-4o", user_api_key_auth=None)
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_should_allow_model_when_key_has_access(self):
+        """Key with explicit model access should be allowed."""
+        auth = _make_user_api_key_auth(models=["gpt-4o", "gpt-3.5-turbo"])
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            return_value=True,
+        ) as mock_check:
+            result = await _check_model_access("gpt-4o", user_api_key_auth=auth)
+
+        assert result is None
+        mock_check.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    async def test_should_deny_model_when_key_lacks_access(self):
+        """Key without model access should be denied with ErrorData."""
+        from litellm.proxy._types import ProxyException
+
+        auth = _make_user_api_key_auth(models=["gpt-3.5-turbo"])
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            side_effect=ProxyException(
+                message="key not allowed to access model",
+                type="key_model_access_denied",
+                param="model",
+                code=401,
+            ),
+        ):
+            result = await _check_model_access("gpt-4o", user_api_key_auth=auth)
+
+        # Should return ErrorData, not raise
+        assert result is not None
+        assert result.code == -1
+        assert "Model access denied" in result.message
+        assert "gpt-4o" in result.message
+
+    @pytest.mark.asyncio
+    async def test_should_allow_wildcard_model_access(self):
+        """Key with wildcard model access should allow any model."""
+        auth = _make_user_api_key_auth(models=["*"])
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            return_value=True,
+        ):
+            result = await _check_model_access(
+                "claude-3-opus-20240229", user_api_key_auth=auth
+            )
+
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_should_deny_expensive_model_requested_by_malicious_server(self):
+        """Simulates the attack: malicious MCP server hints at an expensive model
+        the caller's key is restricted from using."""
+        from litellm.proxy._types import ProxyException
+
+        # Key only has access to cheap models
+        auth = _make_user_api_key_auth(models=["gpt-3.5-turbo"])
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            side_effect=ProxyException(
+                message="key not allowed to access model. This key can only access models=['gpt-3.5-turbo']. Tried to access claude-3-opus-20240229",
+                type="key_model_access_denied",
+                param="model",
+                code=401,
+            ),
+        ):
+            result = await _check_model_access(
+                "claude-3-opus-20240229", user_api_key_auth=auth
+            )
+
+        assert result is not None
+        assert result.code == -1
+        assert "claude-3-opus-20240229" in result.message
+
+    @pytest.mark.asyncio
+    async def test_should_deny_empty_oauth_passthrough_placeholder(self):
+        """Regression: process_mcp_request() returns an empty UserAPIKeyAuth()
+        for OAuth2 upstream-token passthrough.  The None check alone is not
+        sufficient — the empty placeholder is truthy but has no api_key, no
+        token, and an empty models list.  can_key_call_model() would treat
+        that as all-model access, letting an OAuth-only user trigger sampling
+        calls on any proxy model without a LiteLLM key or budget."""
+        # Simulate the empty placeholder from process_mcp_request()
+        auth = _make_user_api_key_auth(
+            models=[],
+            api_key=None,
+            token=None,
+            user_role=None,
+        )
+
+        result = await _check_model_access("gpt-4o", user_api_key_auth=auth)
+
+        # Must be denied — not passed through to can_key_call_model
+        assert result is not None
+        assert result.code == -1
+        assert "sampling requires a valid LiteLLM" in result.message
+
+    @pytest.mark.asyncio
+    async def test_should_allow_proxy_admin_even_without_api_key(self):
+        """Proxy admins may not have a traditional api_key but should still
+        be allowed to use sampling."""
+        auth = _make_user_api_key_auth(
+            models=[],
+            api_key=None,
+            token=None,
+            user_role="proxy_admin",
+        )
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            return_value=True,
+        ):
+            result = await _check_model_access("gpt-4o", user_api_key_auth=auth)
+
+        assert result is None
+
+
+# ---------------------------------------------------------------------------
+# handle_sampling_create_message — auth + budget gating
+# ---------------------------------------------------------------------------
+
+
+class TestSamplingAuthAndBudgetGating:
+
+    @pytest.mark.asyncio
+    async def test_should_deny_when_no_auth_context(self):
+        """Sampling must reject calls with no user_api_key_auth."""
+        from litellm.proxy._experimental.mcp_server.sampling_handler import (
+            handle_sampling_create_message,
+        )
+
+        params = MagicMock()
+        params.modelPreferences = None
+        params.messages = []
+        params.systemPrompt = None
+        params.maxTokens = 100
+        params.temperature = None
+        params.stopSequences = None
+        params.tools = None
+        params.toolChoice = None
+        params.metadata = None
+
+        result = await handle_sampling_create_message(
+            context=MagicMock(),
+            params=params,
+            default_model="gpt-4o",
+            user_api_key_auth=None,
+        )
+
+        assert result is not None
+        assert result.code == -1
+        assert "authenticated" in result.message.lower()
+
+    @pytest.mark.asyncio
+    async def test_should_run_budget_checks(self):
+        """Sampling must call _run_budget_checks after model access check."""
+        from litellm.proxy._experimental.mcp_server.sampling_handler import (
+            handle_sampling_create_message,
+        )
+
+        auth = _make_user_api_key_auth(models=["gpt-4o"])
+        params = MagicMock()
+        params.modelPreferences = None
+        params.messages = []
+        params.systemPrompt = None
+        params.maxTokens = 100
+        params.temperature = None
+        params.stopSequences = None
+        params.tools = None
+        params.toolChoice = None
+        params.metadata = None
+
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=None,
+            ) as mock_budget,
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+            patch(
+                "litellm.proxy.proxy_server.llm_router",
+                new=None,
+            ),
+            patch(
+                "litellm.acompletion",
+                new_callable=AsyncMock,
+                return_value=MagicMock(
+                    choices=[
+                        MagicMock(
+                            message=MagicMock(content="hi", tool_calls=None),
+                            finish_reason="stop",
+                        )
+                    ],
+                    model="gpt-4o",
+                ),
+            ),
+        ):
+            await handle_sampling_create_message(
+                context=MagicMock(),
+                params=params,
+                default_model="gpt-4o",
+                user_api_key_auth=auth,
+            )
+
+        mock_budget.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    async def test_should_deny_over_budget_caller(self):
+        """When _run_budget_checks returns ErrorData, sampling must return it."""
+        from mcp.types import ErrorData
+        from litellm.proxy._experimental.mcp_server.sampling_handler import (
+            handle_sampling_create_message,
+        )
+
+        auth = _make_user_api_key_auth(models=["gpt-4o"])
+        params = MagicMock()
+        params.modelPreferences = None
+        params.messages = []
+        params.systemPrompt = None
+        params.maxTokens = 100
+        params.temperature = None
+        params.stopSequences = None
+        params.tools = None
+        params.toolChoice = None
+        params.metadata = None
+
+        budget_error = ErrorData(code=-1, message="ExceededBudget: over limit")
+
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=budget_error,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+        ):
+            result = await handle_sampling_create_message(
+                context=MagicMock(),
+                params=params,
+                default_model="gpt-4o",
+                user_api_key_auth=auth,
+            )
+
+        assert result is budget_error
+        assert "ExceededBudget" in result.message

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_resolution.py

@@ -0,0 +1,91 @@
+"""
+Tests for MCP sampling model resolution (hint matching and fallback chain).
+
+`_resolve_model_from_preferences` first tries to match upstream model hints
+against the proxy's available models (direct then substring), then priority
+scoring, then the caller default, the first available model, and finally the
+configured `default_mcp_sampling_model` before raising.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _resolve_model_from_preferences,
+)
+
+
+def _prefs(*, hints=None, cost=None, speed=None, intelligence=None):
+    return SimpleNamespace(
+        hints=hints or [],
+        costPriority=cost,
+        speedPriority=speed,
+        intelligencePriority=intelligence,
+    )
+
+
+class TestHintMatching:
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", [{"model_name": "gpt-4o"}, {"model_name": "claude-3"}])
+    def test_should_match_hint_as_substring(self):
+        prefs = _prefs(hints=[SimpleNamespace(name="gpt-4")])
+        assert _resolve_model_from_preferences(prefs) == "gpt-4o"
+
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", ["gpt-4o", "claude-3"])
+    def test_should_match_hint_against_string_model_list_entries(self):
+        prefs = _prefs(hints=[SimpleNamespace(name="claude-3")])
+        assert _resolve_model_from_preferences(prefs) == "claude-3"
+
+    @patch("litellm.model_list", None)
+    def test_should_use_router_model_names(self):
+        router = MagicMock()
+        router.get_model_names.return_value = ["router-gpt", "router-claude"]
+        with patch("litellm.proxy.proxy_server.llm_router", router):
+            prefs = _prefs(hints=[SimpleNamespace(name="router-claude")])
+            assert _resolve_model_from_preferences(prefs) == "router-claude"
+
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", [{"model_name": "gpt-4o"}])
+    def test_should_skip_hint_without_name(self):
+        prefs = _prefs(hints=[SimpleNamespace()])  # hint has no `.name`
+        assert (
+            _resolve_model_from_preferences(prefs, default_model="gpt-4o") == "gpt-4o"
+        )
+
+
+class TestFallbackChain:
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch(
+        "litellm.model_list", [{"model_name": "first-model"}, {"model_name": "second"}]
+    )
+    def test_should_fall_back_to_first_available_when_no_default(self):
+        prefs = _prefs(hints=[SimpleNamespace(name="no-such")])
+        assert _resolve_model_from_preferences(prefs) == "first-model"
+
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", [])
+    def test_should_use_configured_default_sampling_model(self, monkeypatch):
+        import litellm
+
+        monkeypatch.setattr(
+            litellm, "default_mcp_sampling_model", "fallback-model", raising=False
+        )
+        prefs = _prefs()
+        assert _resolve_model_from_preferences(prefs) == "fallback-model"
+
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", [])
+    def test_should_raise_when_nothing_resolvable(self, monkeypatch):
+        import litellm
+
+        monkeypatch.setattr(litellm, "default_mcp_sampling_model", None, raising=False)
+        prefs = _prefs()
+        with pytest.raises(ValueError, match="No model could be resolved"):
+            _resolve_model_from_preferences(prefs)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_priority_selection.py

@@ -0,0 +1,248 @@
+"""
+Tests for MCP sampling handler priority-based model selection.
+
+Verifies that _resolve_model_from_preferences honours costPriority,
+speedPriority, and intelligencePriority when hints don't match,
+per the MCP spec.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _has_priorities,
+    _resolve_model_from_preferences,
+    _select_model_by_priority,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _prefs(*, hints=None, cost=None, speed=None, intelligence=None):
+    """Build a minimal ModelPreferences-like object."""
+    return SimpleNamespace(
+        hints=hints or [],
+        costPriority=cost,
+        speedPriority=speed,
+        intelligencePriority=intelligence,
+    )
+
+
+# Model info stubs keyed by model name
+_MODEL_INFO = {
+    "gpt-3.5-turbo": {
+        "input_cost_per_token": 0.0000005,
+        "output_cost_per_token": 0.0000015,
+        "max_output_tokens": 4096,
+        "max_tokens": 4096,
+        "output_tokens_per_second": 50.0,
+    },
+    "gpt-4o": {
+        "input_cost_per_token": 0.0000025,
+        "output_cost_per_token": 0.0000100,
+        "max_output_tokens": 16384,
+        "max_tokens": 128000,
+        "output_tokens_per_second": 60.0,
+    },
+    "claude-3-opus": {
+        "input_cost_per_token": 0.0000150,
+        "output_cost_per_token": 0.0000750,
+        "max_output_tokens": 4096,
+        "max_tokens": 200000,
+        "output_tokens_per_second": 20.0,
+    },
+    "gpt-4o-mini": {
+        "input_cost_per_token": 0.00000015,
+        "output_cost_per_token": 0.0000006,
+        "max_output_tokens": 16384,
+        "max_tokens": 128000,
+        "output_tokens_per_second": 100.0,
+    },
+}
+
+
+def _mock_get_model_info(model, **kwargs):
+    """Mock litellm.get_model_info using our test data."""
+    if model in _MODEL_INFO:
+        return _MODEL_INFO[model]
+    raise Exception(f"Unknown model: {model}")
+
+
+# ---------------------------------------------------------------------------
+# _has_priorities
+# ---------------------------------------------------------------------------
+
+
+class TestHasPriorities:
+    def test_should_return_false_when_no_priorities_set(self):
+        prefs = _prefs()
+        assert _has_priorities(prefs) is False
+
+    def test_should_return_false_when_all_zero(self):
+        prefs = _prefs(cost=0, speed=0, intelligence=0)
+        assert _has_priorities(prefs) is False
+
+    def test_should_return_true_when_cost_set(self):
+        prefs = _prefs(cost=0.8)
+        assert _has_priorities(prefs) is True
+
+    def test_should_return_true_when_intelligence_set(self):
+        prefs = _prefs(intelligence=0.5)
+        assert _has_priorities(prefs) is True
+
+
+# ---------------------------------------------------------------------------
+# _select_model_by_priority
+# ---------------------------------------------------------------------------
+
+
+class TestSelectModelByPriority:
+    """Tests for the priority-based scoring logic."""
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_prefer_cheapest_when_cost_priority_high(self, _mock):
+        """High costPriority should select the cheapest model."""
+        prefs = _prefs(cost=1.0, speed=0, intelligence=0)
+        models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"]
+        result = _select_model_by_priority(models, prefs)
+        # gpt-4o-mini has the lowest combined cost
+        assert result == "gpt-4o-mini"
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_prefer_smartest_when_intelligence_priority_high(self, _mock):
+        """High intelligencePriority should select the model with highest max_output_tokens."""
+        prefs = _prefs(cost=0, speed=0, intelligence=1.0)
+        models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"]
+        result = _select_model_by_priority(models, prefs)
+        # gpt-4o and gpt-4o-mini both have 16384 max_output_tokens (tied)
+        # Either is acceptable
+        assert result in ("gpt-4o", "gpt-4o-mini")
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_balance_cost_and_intelligence(self, _mock):
+        """Balanced priorities should pick a middle-ground model."""
+        prefs = _prefs(cost=0.5, speed=0, intelligence=0.5)
+        models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"]
+        result = _select_model_by_priority(models, prefs)
+        # gpt-4o-mini is cheap AND has high max_output_tokens → best balance
+        assert result == "gpt-4o-mini"
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_prefer_fastest_when_speed_priority_high(self, _mock):
+        """High speedPriority should prefer cheaper (faster proxy) models."""
+        prefs = _prefs(cost=0, speed=1.0, intelligence=0)
+        models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"]
+        result = _select_model_by_priority(models, prefs)
+        # gpt-4o-mini has lowest cost → fastest proxy
+        assert result == "gpt-4o-mini"
+
+    @patch(
+        "litellm.get_model_info",
+        side_effect=lambda m, **kw: (_ for _ in ()).throw(Exception("no info")),
+    )
+    def test_should_return_none_when_no_model_info(self, _mock):
+        """If get_model_info fails for all models, return None."""
+        prefs = _prefs(cost=1.0)
+        models = ["unknown-model-1", "unknown-model-2"]
+        result = _select_model_by_priority(models, prefs)
+        assert result is None
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_handle_single_model(self, _mock):
+        """Single model should always be returned regardless of priorities."""
+        prefs = _prefs(cost=1.0, intelligence=1.0)
+        result = _select_model_by_priority(["gpt-4o"], prefs)
+        assert result == "gpt-4o"
+
+    def test_speed_priority_is_neutral_when_no_tps_data(self):
+        """When no candidate exposes output_tokens_per_second, speedPriority
+        must not fall back to context-window size as a latency proxy: that
+        biased selection toward the smallest-context model regardless of real
+        speed. With a neutral score the tie resolves to the first candidate,
+        so the larger-context model listed first is kept."""
+        no_tps_info = {
+            "big-ctx": {
+                "input_cost_per_token": 0.0,
+                "output_cost_per_token": 0.0,
+                "max_output_tokens": 100000,
+                "max_tokens": 100000,
+            },
+            "small-ctx": {
+                "input_cost_per_token": 0.0,
+                "output_cost_per_token": 0.0,
+                "max_output_tokens": 1000,
+                "max_tokens": 1000,
+            },
+        }
+
+        def info(model, **kwargs):
+            return no_tps_info[model]
+
+        with patch("litellm.get_model_info", side_effect=info):
+            prefs = _prefs(speed=1.0)
+            # The inverse-max_output proxy would pick "small-ctx" here; a
+            # neutral score keeps the first candidate.
+            assert _select_model_by_priority(["big-ctx", "small-ctx"], prefs) == (
+                "big-ctx"
+            )
+
+
+# ---------------------------------------------------------------------------
+# _resolve_model_from_preferences — priority integration
+# ---------------------------------------------------------------------------
+
+
+class TestResolveModelPriorityIntegration:
+    """End-to-end tests for priority selection within _resolve_model_from_preferences."""
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch(
+        "litellm.model_list",
+        [
+            {"model_name": "gpt-3.5-turbo"},
+            {"model_name": "gpt-4o"},
+            {"model_name": "gpt-4o-mini"},
+        ],
+    )
+    def test_should_use_priority_when_hints_empty(self, _mock_info):
+        """With no hints but priorities set, should use priority-based selection."""
+        prefs = _prefs(cost=1.0, speed=0, intelligence=0)
+        result = _resolve_model_from_preferences(prefs, default_model="gpt-4o")
+        # Should pick cheapest, NOT fall through to default_model
+        assert result == "gpt-4o-mini"
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch(
+        "litellm.model_list",
+        [
+            {"model_name": "gpt-3.5-turbo"},
+            {"model_name": "gpt-4o"},
+            {"model_name": "gpt-4o-mini"},
+        ],
+    )
+    def test_should_skip_priority_when_no_priorities_set(self, _mock_info):
+        """With no priorities set, should fall through to default_model."""
+        prefs = _prefs()  # no priorities
+        result = _resolve_model_from_preferences(prefs, default_model="gpt-4o")
+        assert result == "gpt-4o"
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch(
+        "litellm.model_list",
+        [
+            {"model_name": "gpt-3.5-turbo"},
+            {"model_name": "gpt-4o"},
+            {"model_name": "gpt-4o-mini"},
+        ],
+    )
+    def test_should_prefer_hint_over_priority(self, _mock_info):
+        """Hints should take precedence over priority-based selection."""
+        hints = [SimpleNamespace(name="gpt-4o")]
+        prefs = _prefs(hints=hints, cost=1.0)  # cost says cheap, but hint says gpt-4o
+        result = _resolve_model_from_preferences(prefs, default_model="gpt-3.5-turbo")
+        assert result == "gpt-4o"

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_request_builder.py

@@ -0,0 +1,147 @@
+"""
+Tests for _build_sampling_request header forwarding.
+
+Verifies that the synthetic FastAPI Request built for sampling sub-calls
+correctly propagates the original MCP connection's headers and client IP
+so that header-dependent guardrails, routing hooks, and trace correlation
+function correctly.
+"""
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _build_sampling_request,
+)
+
+
+class TestBuildSamplingRequest:
+    """Tests for the _build_sampling_request helper."""
+
+    def test_should_include_content_type_by_default(self):
+        """Even with no raw headers, content-type must be present."""
+        req = _build_sampling_request()
+        headers = dict(req.headers)
+        assert headers.get("content-type") == "application/json"
+
+    def test_should_forward_raw_headers(self):
+        """Headers from the original MCP connection should be forwarded."""
+        raw = {
+            "x-litellm-tags": "tag1,tag2",
+            "x-litellm-trace-id": "trace-abc-123",
+            "user-agent": "MCP-Client/1.0",
+            "authorization": "Bearer sk-test",
+        }
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+
+        assert headers.get("x-litellm-tags") == "tag1,tag2"
+        assert headers.get("x-litellm-trace-id") == "trace-abc-123"
+        assert headers.get("user-agent") == "MCP-Client/1.0"
+        assert headers.get("authorization") == "Bearer sk-test"
+
+    def test_should_skip_hop_by_hop_headers(self):
+        """content-length and transfer-encoding should not be forwarded."""
+        raw = {
+            "content-length": "42",
+            "transfer-encoding": "chunked",
+            "x-custom": "keep-me",
+        }
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+
+        assert "content-length" not in headers
+        assert "transfer-encoding" not in headers
+        assert headers.get("x-custom") == "keep-me"
+
+    def test_should_not_duplicate_content_type(self):
+        """If raw_headers includes content-type, don't add it twice."""
+        raw = {"content-type": "text/plain"}
+        req = _build_sampling_request(raw_headers=raw)
+        # Count how many content-type headers are present
+        ct_count = sum(1 for k, _ in req.scope["headers"] if k == b"content-type")
+        assert ct_count == 1
+
+    def test_should_inject_client_ip_as_x_forwarded_for(self):
+        """client_ip should be injected as x-forwarded-for."""
+        req = _build_sampling_request(client_ip="10.0.0.42")
+        headers = dict(req.headers)
+        assert headers.get("x-forwarded-for") == "10.0.0.42"
+
+    def test_should_not_override_existing_x_forwarded_for(self):
+        """Caller-supplied x-forwarded-for is stripped; resolved client_ip wins."""
+        raw = {"x-forwarded-for": "192.168.1.1"}
+        req = _build_sampling_request(raw_headers=raw, client_ip="10.0.0.42")
+        headers = dict(req.headers)
+        assert headers.get("x-forwarded-for") == "10.0.0.42"
+
+    def test_should_set_correct_path(self):
+        """The synthetic request should have the sampling path."""
+        req = _build_sampling_request()
+        assert req.scope["path"] == "/mcp/sampling/createMessage"
+
+    def test_server_should_default_to_litellm_port(self):
+        """Server tuple should use port 4000 (LiteLLM default), not 0."""
+        req = _build_sampling_request()
+        _host, _port = req.scope["server"]
+        assert _port == 4000, f"Expected default LiteLLM port 4000, got {_port}"
+
+    def test_should_populate_client_tuple_from_client_ip(self):
+        """request.client.host must return the real client IP for
+        IP-based routing and guardrails."""
+        req = _build_sampling_request(client_ip="10.0.0.42")
+        assert req.scope.get("client") is not None
+        assert req.scope["client"][0] == "10.0.0.42"
+        # Verify request.client.host works (Starlette Address)
+        assert req.client is not None
+        assert req.client.host == "10.0.0.42"
+
+    def test_should_not_set_client_when_no_ip(self):
+        """If no client_ip is provided, client should not be in scope."""
+        req = _build_sampling_request()
+        assert "client" not in req.scope
+
+    def test_should_skip_all_hop_by_hop_headers(self):
+        """All hop-by-hop headers must be filtered, not just content-length
+        and transfer-encoding."""
+        raw = {
+            "content-length": "42",
+            "transfer-encoding": "chunked",
+            "connection": "keep-alive",
+            "keep-alive": "timeout=5",
+            "upgrade": "websocket",
+            "te": "trailers",
+            "trailer": "Expires",
+            "x-custom": "keep-me",
+        }
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+
+        for hop_header in [
+            "content-length",
+            "transfer-encoding",
+            "connection",
+            "keep-alive",
+            "upgrade",
+            "te",
+            "trailer",
+        ]:
+            assert (
+                hop_header not in headers
+            ), f"Hop-by-hop header '{hop_header}' should be filtered"
+        assert headers.get("x-custom") == "keep-me"
+
+    def test_should_forward_traceparent_header(self):
+        """traceparent header must be forwarded for trace correlation."""
+        raw = {
+            "traceparent": "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01",
+        }
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+        assert headers.get("traceparent") == (
+            "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01"
+        )
+
+    def test_should_forward_x_litellm_api_key(self):
+        """x-litellm-api-key header must be forwarded for auth."""
+        raw = {"x-litellm-api-key": "sk-proxy-key-123"}
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+        assert headers.get("x-litellm-api-key") == "sk-proxy-key-123"

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py

@@ -0,0 +1,180 @@
+"""
+Tests for MCP sampling handler response/tool conversion.
+
+Covers the translation of a LiteLLM completion response back into MCP
+`CreateMessageResult` / `CreateMessageResultWithTools`, plus the helpers that
+convert MCP tool definitions, tool-choice modes, and image/audio content into
+OpenAI request format.
+"""
+
+import json
+from types import SimpleNamespace
+
+from mcp.types import (
+    CreateMessageResult,
+    CreateMessageResultWithTools,
+    ErrorData,
+    TextContent,
+    ToolUseContent,
+)
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _convert_mcp_content_to_openai,
+    _convert_mcp_tool_choice_to_openai,
+    _convert_mcp_tools_to_openai,
+    _convert_openai_response_to_mcp_result,
+    _convert_single_content,
+)
+
+
+def _tool_call(*, call_id: str, name: str, arguments):
+    return SimpleNamespace(
+        id=call_id, function=SimpleNamespace(name=name, arguments=arguments)
+    )
+
+
+def _response(*, content=None, tool_calls=None, finish_reason="stop", model="gpt-4o"):
+    message = SimpleNamespace(content=content, tool_calls=tool_calls)
+    choice = SimpleNamespace(message=message, finish_reason=finish_reason)
+    return SimpleNamespace(choices=[choice], model=model)
+
+
+class TestConvertOpenAIResponseToMcpResult:
+    def test_should_return_error_data_when_no_choices(self):
+        response = SimpleNamespace(choices=[], model="gpt-4o")
+        result = _convert_openai_response_to_mcp_result(response, "gpt-4o")
+        assert isinstance(result, ErrorData)
+        assert "no choices" in result.message.lower()
+
+    def test_should_convert_plain_text_response(self):
+        result = _convert_openai_response_to_mcp_result(
+            _response(content="hello world"), "gpt-4o"
+        )
+        assert isinstance(result, CreateMessageResult)
+        assert isinstance(result.content, TextContent)
+        assert result.content.text == "hello world"
+        assert result.role == "assistant"
+        assert result.stopReason == "endTurn"
+
+    def test_should_map_length_finish_reason_to_max_tokens(self):
+        result = _convert_openai_response_to_mcp_result(
+            _response(content="truncated", finish_reason="length"), "gpt-4o"
+        )
+        assert result.stopReason == "maxTokens"
+
+    def test_should_prefer_actual_model_from_response(self):
+        result = _convert_openai_response_to_mcp_result(
+            _response(content="hi", model="gpt-4o-2024-08-06"), "gpt-4o"
+        )
+        assert result.model == "gpt-4o-2024-08-06"
+
+    def test_should_convert_tool_calls_response(self):
+        tc = _tool_call(
+            call_id="call_1",
+            name="get_weather",
+            arguments=json.dumps({"city": "NYC"}),
+        )
+        result = _convert_openai_response_to_mcp_result(
+            _response(content=None, tool_calls=[tc], finish_reason="tool_calls"),
+            "gpt-4o",
+        )
+        assert isinstance(result, CreateMessageResultWithTools)
+        assert result.stopReason == "toolUse"
+        tool_uses = [c for c in result.content if isinstance(c, ToolUseContent)]
+        assert len(tool_uses) == 1
+        assert tool_uses[0].name == "get_weather"
+        assert tool_uses[0].id == "call_1"
+        assert tool_uses[0].input == {"city": "NYC"}
+
+    def test_should_keep_text_alongside_tool_calls(self):
+        tc = _tool_call(call_id="call_1", name="search", arguments="{}")
+        result = _convert_openai_response_to_mcp_result(
+            _response(
+                content="let me check", tool_calls=[tc], finish_reason="tool_calls"
+            ),
+            "gpt-4o",
+        )
+        texts = [c for c in result.content if isinstance(c, TextContent)]
+        assert texts and texts[0].text == "let me check"
+
+    def test_should_wrap_unparsable_tool_arguments_as_raw(self):
+        tc = _tool_call(call_id="call_1", name="bad", arguments="not-json{")
+        result = _convert_openai_response_to_mcp_result(
+            _response(tool_calls=[tc], finish_reason="tool_calls"), "gpt-4o"
+        )
+        tool_uses = [c for c in result.content if isinstance(c, ToolUseContent)]
+        assert tool_uses[0].input == {"raw": "not-json{"}
+
+
+class TestConvertMcpToolsToOpenAI:
+    def test_should_return_none_when_no_tools(self):
+        assert _convert_mcp_tools_to_openai(None) is None
+
+    def test_should_convert_tool_with_schema(self):
+        schema = {"type": "object", "properties": {"q": {"type": "string"}}}
+        tool = SimpleNamespace(
+            name="search", description="search the web", inputSchema=schema
+        )
+        result = _convert_mcp_tools_to_openai([tool])
+        assert result == [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search",
+                    "description": "search the web",
+                    "parameters": schema,
+                },
+            }
+        ]
+
+    def test_should_default_description_and_parameters(self):
+        tool = SimpleNamespace(name="noop", description=None, inputSchema=None)
+        result = _convert_mcp_tools_to_openai([tool])
+        fn = result[0]["function"]
+        assert fn["description"] == ""
+        assert fn["parameters"] == {"type": "object", "properties": {}}
+
+
+class TestConvertMcpToolChoiceToOpenAI:
+    def test_should_return_none_when_no_choice(self):
+        assert _convert_mcp_tool_choice_to_openai(None) is None
+
+    def test_should_map_known_modes(self):
+        for mode in ("auto", "required", "none"):
+            choice = SimpleNamespace(mode=mode)
+            assert _convert_mcp_tool_choice_to_openai(choice) == mode
+
+    def test_should_default_unknown_mode_to_auto(self):
+        choice = SimpleNamespace(mode="banana")
+        assert _convert_mcp_tool_choice_to_openai(choice) == "auto"
+
+
+class TestConvertImageAndAudioContent:
+    def test_should_convert_image_to_data_uri(self):
+        content = SimpleNamespace(type="image", data="aGVsbG8=", mimeType="image/jpeg")
+        result = _convert_single_content(content)
+        assert result == {
+            "type": "image_url",
+            "image_url": {"url": "data:image/jpeg;base64,aGVsbG8="},
+        }
+
+    def test_should_map_audio_mime_to_format(self):
+        content = SimpleNamespace(type="audio", data="Zm9v", mimeType="audio/mp3")
+        result = _convert_single_content(content)
+        assert result["type"] == "input_audio"
+        assert result["input_audio"] == {"data": "Zm9v", "format": "mp3"}
+
+    def test_should_default_unknown_audio_mime_to_wav(self):
+        content = SimpleNamespace(type="audio", data="Zm9v", mimeType="audio/weird")
+        result = _convert_single_content(content)
+        assert result["input_audio"]["format"] == "wav"
+
+    def test_should_flatten_list_content(self):
+        items = [
+            SimpleNamespace(type="text", text="a"),
+            SimpleNamespace(type="image", data="x", mimeType="image/png"),
+        ]
+        result = _convert_mcp_content_to_openai(items)
+        assert isinstance(result, list)
+        assert result[0] == {"type": "text", "text": "a"}
+        assert result[1]["type"] == "image_url"

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py

@@ -0,0 +1,312 @@
+"""
+Tests for MCP sampling handler tool_use / tool_result content conversion.
+
+Verifies that multi-turn tool-calling conversations from upstream MCP
+servers are faithfully converted to OpenAI format instead of being
+reduced to lossy plain-text stubs.
+"""
+
+import json
+from types import SimpleNamespace
+from typing import Any, Dict
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _convert_mcp_messages_to_openai,
+    _convert_single_content,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helpers — lightweight MCP type stand-ins
+# ---------------------------------------------------------------------------
+
+
+def _text(text: str) -> SimpleNamespace:
+    return SimpleNamespace(type="text", text=text)
+
+
+def _tool_use(*, name: str, tool_id: str, input_data: Dict[str, Any]) -> SimpleNamespace:
+    return SimpleNamespace(type="tool_use", name=name, id=tool_id, input=input_data)
+
+
+def _tool_result(
+    *, tool_use_id: str, content: Any = None, is_error: bool = False
+) -> SimpleNamespace:
+    if content is None:
+        content = []
+    return SimpleNamespace(
+        type="tool_result", toolUseId=tool_use_id, content=content, isError=is_error
+    )
+
+
+def _sampling_msg(role: str, content: Any) -> SimpleNamespace:
+    return SimpleNamespace(role=role, content=content)
+
+
+# ---------------------------------------------------------------------------
+# _convert_single_content — tool_use
+# ---------------------------------------------------------------------------
+
+
+class TestConvertSingleContentToolUse:
+    """Tests for the tool_use branch of _convert_single_content."""
+
+    def test_should_produce_function_call_dict(self):
+        """tool_use must produce a proper function-call dict, not a text stub."""
+        tu = _tool_use(name="get_weather", tool_id="call_123", input_data={"city": "NYC"})
+        result = _convert_single_content(tu)
+
+        assert result["_marker_type"] == "tool_use"
+        assert result["type"] == "function"
+        assert result["id"] == "call_123"
+        assert result["function"]["name"] == "get_weather"
+        assert json.loads(result["function"]["arguments"]) == {"city": "NYC"}
+
+    def test_should_not_produce_text_stub(self):
+        """Regression: the old code produced '[Tool call: get_weather]'."""
+        tu = _tool_use(name="get_weather", tool_id="call_1", input_data={})
+        result = _convert_single_content(tu)
+
+        # Must NOT be a text content part
+        assert result.get("type") != "text"
+        assert "Tool call" not in str(result)
+
+    def test_should_handle_empty_input(self):
+        tu = _tool_use(name="no_args_tool", tool_id="call_2", input_data={})
+        result = _convert_single_content(tu)
+
+        assert json.loads(result["function"]["arguments"]) == {}
+
+
+# ---------------------------------------------------------------------------
+# _convert_single_content — tool_result
+# ---------------------------------------------------------------------------
+
+
+class TestConvertSingleContentToolResult:
+    """Tests for the tool_result branch of _convert_single_content."""
+
+    def test_should_produce_tool_role_message(self):
+        """tool_result must produce a tool-role dict, not a text content part."""
+        tr = _tool_result(
+            tool_use_id="call_123",
+            content=[_text("Temperature: 72°F")],
+        )
+        result = _convert_single_content(tr)
+
+        assert result["_marker_type"] == "tool_result"
+        assert result["role"] == "tool"
+        assert result["tool_call_id"] == "call_123"
+        assert "72°F" in result["content"]
+
+    def test_should_handle_empty_content(self):
+        tr = _tool_result(tool_use_id="call_456", content=[])
+        result = _convert_single_content(tr)
+
+        assert result["role"] == "tool"
+        assert result["tool_call_id"] == "call_456"
+        assert result["content"] == ""
+
+    def test_should_concatenate_multiple_text_parts(self):
+        tr = _tool_result(
+            tool_use_id="call_789",
+            content=[_text("Line 1"), _text("Line 2")],
+        )
+        result = _convert_single_content(tr)
+        assert "Line 1" in result["content"]
+        assert "Line 2" in result["content"]
+
+
+# ---------------------------------------------------------------------------
+# _convert_mcp_messages_to_openai — multi-turn tool calling
+# ---------------------------------------------------------------------------
+
+
+class TestConvertMcpMessagesMultiTurnTools:
+    """End-to-end tests for multi-turn tool-calling message sequences."""
+
+    def test_should_convert_assistant_tool_use_to_tool_calls_array(self):
+        """An assistant message with tool_use content should produce
+        a proper tool_calls array, not a text stub."""
+        messages = [
+            _sampling_msg("assistant", _tool_use(
+                name="search", tool_id="call_1", input_data={"query": "LiteLLM"}
+            )),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert msg["role"] == "assistant"
+        assert "tool_calls" in msg
+        assert len(msg["tool_calls"]) == 1
+        tc = msg["tool_calls"][0]
+        assert tc["function"]["name"] == "search"
+        assert tc["id"] == "call_1"
+
+    def test_should_convert_user_tool_result_to_tool_role_message(self):
+        """A user message with tool_result content should produce
+        a separate role='tool' message."""
+        messages = [
+            _sampling_msg("user", _tool_result(
+                tool_use_id="call_1",
+                content=[_text("Found 42 results")],
+            )),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert msg["role"] == "tool"
+        assert msg["tool_call_id"] == "call_1"
+        assert "42 results" in msg["content"]
+
+    def test_should_handle_full_tool_calling_round_trip(self):
+        """Simulate a complete tool-calling conversation:
+        user → assistant(tool_use) → user(tool_result) → assistant(text)
+        """
+        messages = [
+            _sampling_msg("user", _text("What's the weather in NYC?")),
+            _sampling_msg("assistant", _tool_use(
+                name="get_weather", tool_id="call_w1",
+                input_data={"city": "NYC"},
+            )),
+            _sampling_msg("user", _tool_result(
+                tool_use_id="call_w1",
+                content=[_text("72°F, sunny")],
+            )),
+            _sampling_msg("assistant", _text("It's 72°F and sunny in NYC!")),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 4
+
+        # 1. User message
+        assert result[0]["role"] == "user"
+
+        # 2. Assistant with tool_calls
+        assert result[1]["role"] == "assistant"
+        assert "tool_calls" in result[1]
+        assert result[1]["tool_calls"][0]["function"]["name"] == "get_weather"
+
+        # 3. Tool result
+        assert result[2]["role"] == "tool"
+        assert result[2]["tool_call_id"] == "call_w1"
+
+        # 4. Final assistant text
+        assert result[3]["role"] == "assistant"
+        assert "72°F" in str(result[3]["content"])
+
+    def test_should_handle_mixed_text_and_tool_use_in_assistant(self):
+        """An assistant message with both text and tool_use content."""
+        messages = [
+            _sampling_msg("assistant", [
+                _text("Let me check that for you."),
+                _tool_use(name="lookup", tool_id="call_lu1", input_data={"id": 42}),
+            ]),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert msg["role"] == "assistant"
+        assert "tool_calls" in msg
+        assert msg["tool_calls"][0]["function"]["name"] == "lookup"
+        # Text content should also be present
+        assert msg.get("content") is not None
+
+    def test_should_handle_multiple_tool_uses_in_single_message(self):
+        """Multiple tool_use items in a single assistant message → multiple tool_calls."""
+        messages = [
+            _sampling_msg("assistant", [
+                _tool_use(name="tool_a", tool_id="call_a", input_data={}),
+                _tool_use(name="tool_b", tool_id="call_b", input_data={"x": 1}),
+            ]),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert len(msg["tool_calls"]) == 2
+        names = {tc["function"]["name"] for tc in msg["tool_calls"]}
+        assert names == {"tool_a", "tool_b"}
+
+    def test_should_handle_multiple_tool_results_in_single_message(self):
+        """Multiple tool_result items in a single user message → multiple tool messages."""
+        messages = [
+            _sampling_msg("user", [
+                _tool_result(tool_use_id="call_a", content=[_text("Result A")]),
+                _tool_result(tool_use_id="call_b", content=[_text("Result B")]),
+            ]),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 2
+        assert all(m["role"] == "tool" for m in result)
+        ids = {m["tool_call_id"] for m in result}
+        assert ids == {"call_a", "call_b"}
+
+    def test_should_preserve_system_prompt(self):
+        """System prompt should still be emitted first."""
+        messages = [_sampling_msg("user", _text("Hi"))]
+        result = _convert_mcp_messages_to_openai(
+            messages, system_prompt="You are helpful."
+        )
+
+        assert result[0]["role"] == "system"
+        assert result[0]["content"] == "You are helpful."
+
+
+# ---------------------------------------------------------------------------
+# _convert_mcp_messages_to_openai — marker hoisting on unexpected roles
+# ---------------------------------------------------------------------------
+
+
+class TestConvertMcpMessagesMarkerHoisting:
+    """The role-matched fast paths only fire for assistant/tool_use and
+    user/tool_result. Content that arrives on an unexpected role must still
+    be hoisted to the correct message position by the generic fallback,
+    not silently dropped or embedded inline as a content part."""
+
+    def test_should_hoist_tool_use_arriving_on_user_role(self):
+        messages = [
+            _sampling_msg("user", _tool_use(
+                name="search", tool_id="call_1", input_data={"q": "x"}
+            )),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        assert result[0]["role"] == "assistant"
+        assert result[0]["tool_calls"][0]["function"]["name"] == "search"
+
+    def test_should_hoist_tool_result_arriving_on_assistant_role(self):
+        messages = [
+            _sampling_msg("assistant", _tool_result(
+                tool_use_id="call_1", content=[_text("done")]
+            )),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        assert result[0]["role"] == "tool"
+        assert result[0]["tool_call_id"] == "call_1"
+        assert "done" in result[0]["content"]
+
+    def test_should_keep_text_when_hoisting_tool_use_on_user_role(self):
+        messages = [
+            _sampling_msg("user", [
+                _text("here you go"),
+                _tool_use(name="lookup", tool_id="call_2", input_data={}),
+            ]),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert msg["role"] == "assistant"
+        assert msg["tool_calls"][0]["function"]["name"] == "lookup"
+        assert any(
+            isinstance(p, dict) and p.get("text") == "here you go"
+            for p in msg["content"]
+        )

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py

@@ -1,5 +1,4 @@
 import asyncio
-import contextlib
 import contextvars
 from datetime import datetime, timedelta
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -894,7 +893,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
         extra_headers=None,
         add_prefix=True,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         if server.name == "working_server":
             # Working server returns tools
@@ -1000,7 +999,7 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
         extra_headers=None,
         add_prefix=True,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # All servers fail
         raise Exception(f"Server {server.name} connection failed")
@@ -1122,8 +1121,8 @@ async def test_concurrent_initialize_session_managers():
     # Reset state before test
     original_initialized = mcp_server._SESSION_MANAGERS_INITIALIZED
     original_session_cm = mcp_server._session_manager_cm
-    original_session_stateful_cm = mcp_server._session_manager_stateful_cm
-    original_sse_session_cm = mcp_server._sse_session_manager_cm
+    original_stateful_cm = mcp_server._session_manager_stateful_cm
+    original_sse_cm = mcp_server._sse_session_manager_cm
     original_cleanup_task = mcp_server._stateful_auth_context_cleanup_task
 
     try:
@@ -1131,30 +1130,38 @@ async def test_concurrent_initialize_session_managers():
         mcp_server._session_manager_cm = None
         mcp_server._session_manager_stateful_cm = None
         mcp_server._sse_session_manager_cm = None
-        mcp_server._stateful_auth_context_cleanup_task = None
 
-        # Mock the session managers to avoid actual MCP initialization
+        # Create mock context managers for all three session managers
+        mock_cm_stateless = AsyncMock()
+        mock_cm_stateless.__aenter__ = AsyncMock()
+        mock_cm_stateless.__aexit__ = AsyncMock()
+
+        mock_cm_stateful = AsyncMock()
+        mock_cm_stateful.__aenter__ = AsyncMock()
+        mock_cm_stateful.__aexit__ = AsyncMock()
+
+        mock_cm_sse = AsyncMock()
+        mock_cm_sse.__aenter__ = AsyncMock()
+        mock_cm_sse.__aexit__ = AsyncMock()
+
         with (
-            patch(
-                "litellm.proxy._experimental.mcp_server.server.session_manager_stateless"
-            ) as mock_session_manager_stateless,
-            patch(
-                "litellm.proxy._experimental.mcp_server.server.session_manager_stateful"
-            ) as mock_session_manager_stateful,
-            patch(
-                "litellm.proxy._experimental.mcp_server.server.sse_session_manager"
-            ) as mock_sse_session_manager,
+            patch.object(
+                mcp_server.session_manager_stateless,
+                "run",
+                return_value=mock_cm_stateless,
+            ) as mock_stateless_run,
+            patch.object(
+                mcp_server.session_manager_stateful,
+                "run",
+                return_value=mock_cm_stateful,
+            ) as mock_stateful_run,
+            patch.object(
+                mcp_server.sse_session_manager,
+                "run",
+                return_value=mock_cm_sse,
+            ) as mock_sse_run,
             patch("litellm.proxy._experimental.mcp_server.server.verbose_logger"),
         ):
-            # Mock the run() method to return a mock context manager
-            mock_cm = AsyncMock()
-            mock_cm.__aenter__ = AsyncMock()
-            mock_cm.__aexit__ = AsyncMock()
-
-            mock_session_manager_stateless.run.return_value = mock_cm
-            mock_session_manager_stateful.run.return_value = mock_cm
-            mock_sse_session_manager.run.return_value = mock_cm
-
             # Create multiple concurrent tasks that call initialize_session_managers
             async def init_task():
                 await initialize_session_managers()
@@ -1171,19 +1178,25 @@ async def test_concurrent_initialize_session_managers():
 
             # Each session manager.run() should only be called once due to the lock
             assert (
-                mock_session_manager_stateless.run.call_count == 1
-            ), f"Expected 1 call to session_manager_stateless.run(), got {mock_session_manager_stateless.run.call_count}"
+                mock_stateless_run.call_count == 1
+            ), f"Expected 1 call to session_manager_stateless.run(), got {mock_stateless_run.call_count}"
             assert (
-                mock_session_manager_stateful.run.call_count == 1
-            ), f"Expected 1 call to session_manager_stateful.run(), got {mock_session_manager_stateful.run.call_count}"
+                mock_stateful_run.call_count == 1
+            ), f"Expected 1 call to session_manager_stateful.run(), got {mock_stateful_run.call_count}"
             assert (
-                mock_sse_session_manager.run.call_count == 1
-            ), f"Expected 1 call to sse_session_manager.run(), got {mock_sse_session_manager.run.call_count}"
+                mock_sse_run.call_count == 1
+            ), f"Expected 1 call to sse_session_manager.run(), got {mock_sse_run.call_count}"
 
-            # The context managers should only be entered once each (3 managers)
+            # The context managers should only be entered once each
             assert (
-                mock_cm.__aenter__.call_count == 3
-            ), f"Expected 3 calls to __aenter__ (one per session manager), got {mock_cm.__aenter__.call_count}"
+                mock_cm_stateless.__aenter__.call_count == 1
+            ), f"Expected 1 call to stateless __aenter__, got {mock_cm_stateless.__aenter__.call_count}"
+            assert (
+                mock_cm_stateful.__aenter__.call_count == 1
+            ), f"Expected 1 call to stateful __aenter__, got {mock_cm_stateful.__aenter__.call_count}"
+            assert (
+                mock_cm_sse.__aenter__.call_count == 1
+            ), f"Expected 1 call to sse __aenter__, got {mock_cm_sse.__aenter__.call_count}"
 
             # State should be properly set
             assert mcp_server._SESSION_MANAGERS_INITIALIZED is True
@@ -1195,14 +1208,12 @@ async def test_concurrent_initialize_session_managers():
         leaked_task = mcp_server._stateful_auth_context_cleanup_task
         if leaked_task is not None and leaked_task is not original_cleanup_task:
             leaked_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError, Exception):
-                await leaked_task
 
         # Restore original state
         mcp_server._SESSION_MANAGERS_INITIALIZED = original_initialized
         mcp_server._session_manager_cm = original_session_cm
-        mcp_server._session_manager_stateful_cm = original_session_stateful_cm
-        mcp_server._sse_session_manager_cm = original_sse_session_cm
+        mcp_server._session_manager_stateful_cm = original_stateful_cm
+        mcp_server._sse_session_manager_cm = original_sse_cm
         mcp_server._stateful_auth_context_cleanup_task = original_cleanup_task
 
 
@@ -1637,10 +1648,7 @@ async def test_mcp_routing_initialize_rejected_when_owner_at_session_cap():
     active = {f"s{i}": 1 for i in range(cap)}  # all in flight -> cannot evict
     contexts = {f"s{i}": MagicMock() for i in range(cap)}
 
-    init_body = (
-        b'{"jsonrpc":"2.0","id":1,"method":"initialize",'
-        b'"params":{"protocolVersion":"2024-11-05"}}'
-    )
+    init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}'
     scope = {
         "type": "http",
         "method": "POST",
@@ -2587,6 +2595,134 @@ async def test_stateful_mcp_get_stream_does_not_block_post():
         mcp_server._stateful_session_locks.pop(session_id, None)
 
 
+def test_jsonrpc_text_has_top_level_method_ignores_nested_method():
+    """The top-level-key scan must not be fooled by a ``method`` field nested
+    inside a JSON-RPC response's ``result`` payload — a flat substring search
+    would, and that misread is what deadlocks the session lock."""
+    from litellm.proxy._experimental.mcp_server.server import (
+        _jsonrpc_text_has_top_level_method,
+    )
+
+    request = '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}'
+    assert _jsonrpc_text_has_top_level_method(request) is True
+
+    # method key out of order (after params) is still top-level
+    reordered = '{"jsonrpc":"2.0","params":{"x":1},"method":"foo"}'
+    assert _jsonrpc_text_has_top_level_method(reordered) is True
+
+    # response whose result nests a "method" key (and arrays of them)
+    response = (
+        '{"jsonrpc":"2.0","id":1,"result":{"toolResult":{"method":"GET"},'
+        '"steps":[{"method":"x"}]}}'
+    )
+    assert _jsonrpc_text_has_top_level_method(response) is False
+
+    # truncated response: result value never closes, no top-level method seen
+    truncated = '{"jsonrpc":"2.0","id":1,"result":{"text":"' + "q" * 5000
+    assert _jsonrpc_text_has_top_level_method(truncated) is False
+
+
+@pytest.mark.asyncio
+async def test_truncated_jsonrpc_response_with_nested_method_skips_lock():
+    """Regression: a large JSON-RPC *response* POST whose ``result`` payload
+    nests a ``method`` key must skip the per-session lock so it does not
+    deadlock behind the in-flight request POST that is holding the lock while
+    it awaits this very response (e.g. sampling/createMessage)."""
+    try:
+        from litellm.proxy._experimental.mcp_server import server as mcp_server
+        from litellm.proxy._experimental.mcp_server.server import (
+            handle_streamable_http_mcp,
+            session_manager_stateful,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    session_id = "nested-method-response-session"
+    owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner")
+    mcp_server._stateful_session_auth_contexts[session_id] = (
+        mcp_server.MCPAuthenticatedUser(user_api_key_auth=owner_auth)
+    )
+    mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for(
+        owner_auth
+    )
+
+    gate = asyncio.Event()
+    request_in_handle = asyncio.Event()
+    response_handled = asyncio.Event()
+
+    async def handle(s, r, se):
+        msg = await r()
+        body = msg.get("body", b"") or b""
+        if b'"result"' in body:
+            response_handled.set()
+        else:
+            request_in_handle.set()
+            await gate.wait()
+
+    async def call(body: bytes):
+        scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp",
+            "headers": [(b"mcp-session-id", session_id.encode())],
+        }
+        receive = AsyncMock(
+            return_value={
+                "type": "http.request",
+                "body": body,
+                "more_body": False,
+            }
+        )
+        await handle_streamable_http_mcp(scope, receive, AsyncMock())
+
+    # The in-flight request POST holds the session lock while blocked.
+    request_body = b'{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}'
+    # A JSON-RPC response larger than the routing peek cap so it can't be fully
+    # parsed, with a nested "method" key in the first bytes to trip a flat
+    # substring heuristic.
+    response_body = (
+        '{"jsonrpc":"2.0","id":99,"result":{"toolResult":'
+        '{"method":"GET","payload":"' + ("x" * 5000) + '"}}}'
+    ).encode()
+
+    try:
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+                new_callable=AsyncMock,
+                return_value=(owner_auth, None, None, None, None, None),
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+                True,
+            ),
+            patch.object(
+                session_manager_stateful, "handle_request", side_effect=handle
+            ),
+            patch.object(
+                session_manager_stateful,
+                "_server_instances",
+                {session_id: MagicMock()},
+            ),
+        ):
+            req_task = asyncio.create_task(call(request_body))
+            await asyncio.wait_for(request_in_handle.wait(), timeout=1.0)
+
+            resp_task = asyncio.create_task(call(response_body))
+            # Under a flat substring heuristic the response would acquire the
+            # lock held by req_task and this wait would time out (deadlock).
+            await asyncio.wait_for(response_handled.wait(), timeout=1.0)
+
+            gate.set()
+            await asyncio.gather(req_task, resp_task)
+    finally:
+        gate.set()
+        mcp_server._stateful_session_auth_contexts.pop(session_id, None)
+        mcp_server._stateful_session_owners.pop(session_id, None)
+        mcp_server._stateful_session_locks.pop(session_id, None)
+        mcp_server._stateful_session_active_request_counts.pop(session_id, None)
+
+
 @pytest.mark.asyncio
 @pytest.mark.no_parallel
 async def test_mcp_routing_with_conflicting_alias_and_group_name():
@@ -2729,7 +2865,7 @@ async def test_oauth2_headers_passed_to_mcp_client():
         mcp_auth_header=None,
         extra_headers=None,
         stdio_env=None,
-        subject_token=None,
+        **kwargs,
     ):
         # Capture the arguments for verification
         captured_client_args.update(
@@ -2738,7 +2874,7 @@ async def test_oauth2_headers_passed_to_mcp_client():
                 "mcp_auth_header": mcp_auth_header,
                 "extra_headers": extra_headers,
                 "stdio_env": stdio_env,
-                "subject_token": subject_token,
+                "kwargs": kwargs,
             }
         )
         # Return a mock client that doesn't actually connect
@@ -2764,6 +2900,16 @@ async def test_oauth2_headers_passed_to_mcp_client():
             "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
             AsyncMock(return_value=[oauth2_server]),
         ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server._prefetch_oauth_creds_for_user",
+            new_callable=AsyncMock,
+            return_value={},
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
+            new_callable=AsyncMock,
+            return_value=None,
+        ),
     ):
         # Call _get_tools_from_mcp_servers which should eventually call _create_mcp_client
         await _get_tools_from_mcp_servers(
@@ -2840,7 +2986,7 @@ async def test_list_tools_single_server_unprefixed_names():
         extra_headers=None,
         add_prefix=False,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         tool = MagicMock()
         tool.name = f"{server.alias}-toolA" if add_prefix else "toolA"
@@ -2922,7 +3068,7 @@ async def test_list_tools_multiple_servers_prefixed_names():
         extra_headers=None,
         add_prefix=True,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         tool = MagicMock()
         # When multiple servers, add_prefix should be True -> prefixed names
@@ -3189,7 +3335,7 @@ async def test_list_tools_filters_by_key_team_permissions():
         extra_headers=None,
         add_prefix=False,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # Return 4 tools, but only 2 should be allowed
         tool1 = MagicMock()
@@ -3299,7 +3445,7 @@ async def test_list_tools_with_team_tool_permissions_inheritance():
         extra_headers=None,
         add_prefix=False,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # Return 4 tools
         tool1 = MagicMock()
@@ -3395,7 +3541,7 @@ async def test_list_tools_with_no_tool_permissions_shows_all():
         extra_headers=None,
         add_prefix=False,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # Return 3 tools
         tool1 = MagicMock()
@@ -3494,7 +3640,7 @@ async def test_list_tools_strips_prefix_when_matching_permissions():
         extra_headers=None,
         add_prefix=True,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # Return tools WITH prefix (as they come from MCP server)
         tool1 = MagicMock()
@@ -5178,3 +5324,42 @@ def test_get_forwarded_auth_from_scope_skips_when_no_litellm_key_header():
         ]
     }
     assert _get_forwarded_auth_from_scope(scope) is None
+
+
+@pytest.mark.asyncio
+async def test_create_mcp_client_sampling_disabled_by_default():
+    """Sampling callback must be None when allow_sampling is not set (default False)."""
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+        MCPServerManager,
+    )
+
+    manager = MCPServerManager()
+    server = MCPServer(
+        server_id="no-sampling",
+        name="no-sampling",
+        url="https://example.com/mcp",
+        transport=MCPTransport.http,
+    )
+
+    client = await manager._create_mcp_client(server=server)
+    assert client._sampling_callback is None
+
+
+@pytest.mark.asyncio
+async def test_create_mcp_client_sampling_enabled():
+    """Sampling callback must be set when allow_sampling=True."""
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+        MCPServerManager,
+    )
+
+    manager = MCPServerManager()
+    server = MCPServer(
+        server_id="with-sampling",
+        name="with-sampling",
+        url="https://example.com/mcp",
+        transport=MCPTransport.http,
+        allow_sampling=True,
+    )
+
+    client = await manager._create_mcp_client(server=server)
+    assert client._sampling_callback is not None

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -320,9 +320,7 @@ class TestMCPServerManager:
         async def mock_get_tools_from_server(
             server,
             mcp_auth_header=None,
-            mcp_protocol_version=None,
-            raw_headers=None,
-            user_api_key_auth=None,
+            **kwargs,
         ):
             if server.name == "github":
                 tool1 = MagicMock()
@@ -375,9 +373,7 @@ class TestMCPServerManager:
         async def mock_get_tools_from_server(
             server,
             mcp_auth_header=None,
-            mcp_protocol_version=None,
-            raw_headers=None,
-            user_api_key_auth=None,
+            **kwargs,
         ):
             assert mcp_auth_header == "legacy-token"  # Should use legacy header
             tool = MagicMock()
@@ -414,9 +410,7 @@ class TestMCPServerManager:
         async def mock_get_tools_from_server(
             server,
             mcp_auth_header=None,
-            mcp_protocol_version=None,
-            raw_headers=None,
-            user_api_key_auth=None,
+            **kwargs,
         ):
             assert (
                 mcp_auth_header == "server-specific-token"
@@ -457,7 +451,7 @@ class TestMCPServerManager:
         captured_extra_headers = None
 
         async def capture_create_mcp_client(
-            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
+            server, mcp_auth_header, extra_headers, stdio_env, **kwargs
         ):  # pragma: no cover - helper
             nonlocal captured_extra_headers
             captured_extra_headers = extra_headers
@@ -507,7 +501,7 @@ class TestMCPServerManager:
         captured_extra_headers = None
 
         async def capture_create_mcp_client(
-            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
+            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs
         ):  # pragma: no cover - helper
             nonlocal captured_extra_headers
             captured_extra_headers = extra_headers
@@ -560,7 +554,7 @@ class TestMCPServerManager:
         captured_extra_headers = None
 
         async def capture_create_mcp_client(
-            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
+            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs
         ):  # pragma: no cover - helper
             nonlocal captured_extra_headers
             captured_extra_headers = extra_headers
@@ -616,7 +610,7 @@ class TestMCPServerManager:
         captured_extra_headers = None
 
         async def capture_create_mcp_client(
-            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
+            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs
         ):  # pragma: no cover - helper
             nonlocal captured_extra_headers
             captured_extra_headers = extra_headers
@@ -1166,9 +1160,7 @@ class TestMCPServerManager:
         async def mock_get_tools_from_server(
             server,
             mcp_auth_header=None,
-            mcp_protocol_version=None,
-            raw_headers=None,
-            user_api_key_auth=None,
+            **kwargs,
         ):
             assert (
                 mcp_auth_header == "server-specific-token"

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py

@@ -9,8 +9,6 @@ they may send a stale `mcp-session-id` header. This test verifies that:
 
 import asyncio
 from unittest.mock import AsyncMock, MagicMock, patch
-
-from fastapi import HTTPException
 from litellm.types.mcp import MCPAuth
 import pytest
 
@@ -600,6 +598,8 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
     Per-user OAuth server with no stored token should fail fast with 401 +
     WWW-Authenticate so PKCE can start.
     """
+    from fastapi import HTTPException
+
     try:
         from litellm.proxy._experimental.mcp_server.server import (
             handle_streamable_http_mcp,
@@ -612,8 +612,13 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
         "type": "http",
         "method": "POST",
         "path": "/mcp",
+        "scheme": "http",
+        "query_string": b"",
+        "root_path": "",
+        "server": ("localhost", 8000),
         "headers": [
             (b"content-type", b"application/json"),
+            (b"host", b"localhost:8000"),
         ],
     }
     receive = AsyncMock()
@@ -660,11 +665,12 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
         with pytest.raises(HTTPException) as exc_info:
             await handle_streamable_http_mcp(scope, receive, send)
 
-    exc = exc_info.value
-    assert exc.status_code == 401
-    assert "www-authenticate" in exc.headers
+    # Verify a 401 was raised
     assert mock_get_stored_token.await_count == 1
     assert mock_handle_request.await_count == 0
+    assert exc_info.value.status_code == 401
+    assert "www-authenticate" in exc_info.value.headers
+    assert "Bearer authorization_uri=" in exc_info.value.headers["www-authenticate"]
 
 
 @pytest.mark.asyncio
@@ -685,11 +691,22 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
         "type": "http",
         "method": "POST",
         "path": "/mcp",
+        "scheme": "http",
+        "query_string": b"",
+        "root_path": "",
+        "server": ("localhost", 8000),
         "headers": [
             (b"content-type", b"application/json"),
+            (b"host", b"localhost:8000"),
         ],
     }
-    receive = AsyncMock()
+    receive = AsyncMock(
+        return_value={
+            "type": "http.request",
+            "body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}',
+            "more_body": False,
+        }
+    )
     send = AsyncMock()
     user_auth = MagicMock()
     user_auth.user_id = "test-user-id"
@@ -729,6 +746,11 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
             "handle_request",
             new_callable=AsyncMock,
         ) as mock_handle_request,
+        patch.object(
+            session_manager_stateless,
+            "_server_instances",
+            {},
+        ),
     ):
         await handle_streamable_http_mcp(scope, receive, send)
 

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py

@@ -2,6 +2,7 @@
 Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
 """
 
+import json
 import os
 import re
 import sys
@@ -20,7 +21,7 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.tool_permission import (
     ToolPermissionGuardrail,
 )
-from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
 from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
     PermissionError,
 )
@@ -894,3 +895,153 @@ class TestToolPermissionGuardrailIntegration:
         is_allowed, rule_id, _ = guardrail._check_tool_permission("Read")
         assert is_allowed is False
         assert rule_id == "deny_read"
+
+
+class TestToolPermissionGuardrailInMemoryUpdate:
+    """Regression: an in-memory params update (PUT /guardrails path) must rebuild
+    the compiled rule maps, not just self.rules, so the new rules are enforced
+    without reinitializing the guardrail."""
+
+    def _bash(self, command):
+        return ChatCompletionMessageToolCall(
+            function={"name": "Bash", "arguments": json.dumps({"command": command})},
+            type="function",
+        )
+
+    def test_update_in_memory_recompiles_added_param_pattern(self):
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="tp",
+            rules=[{"id": "native-bash", "tool_name": r"^Bash$", "decision": "allow"}],
+            default_action="deny",
+            on_disallowed_action="block",
+        )
+        # No pattern yet: any Bash command is allowed.
+        assert (
+            guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0]
+            is True
+        )
+
+        guardrail.update_in_memory_litellm_params(
+            LitellmParams(
+                guardrail="tool_permission",
+                mode=["pre_call", "post_call"],
+                default_action="deny",
+                on_disallowed_action="block",
+                rules=[
+                    {
+                        "id": "native-bash",
+                        "tool_name": r"^Bash$",
+                        "decision": "allow",
+                        "allowed_param_patterns": {
+                            "command": r"^(?!(echo blockme)$).*$"
+                        },
+                    }
+                ],
+            )
+        )
+
+        # The compiled map must be rebuilt, and enforcement must reflect it.
+        assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {})
+        assert (
+            guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0]
+            is False
+        )
+        assert (
+            guardrail._get_permission_for_tool_call(self._bash("echo hello"))[0] is True
+        )
+
+    def test_update_in_memory_recompiles_tool_name_target(self):
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="tp",
+            rules=[],
+            default_action="allow",
+            on_disallowed_action="block",
+        )
+        # No rules: default_action allow lets Bash through.
+        assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is True
+
+        guardrail.update_in_memory_litellm_params(
+            LitellmParams(
+                guardrail="tool_permission",
+                mode=["pre_call", "post_call"],
+                default_action="allow",
+                on_disallowed_action="block",
+                rules=[{"id": "deny-bash", "tool_name": r"^Bash$", "decision": "deny"}],
+            )
+        )
+
+        # A newly added deny rule (new id) must match -> its compiled target was rebuilt.
+        assert "deny-bash" in guardrail._compiled_rule_targets
+        assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is False
+
+    def test_update_in_memory_preserves_rules_when_rules_absent(self):
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="tp",
+            rules=[
+                {
+                    "id": "native-bash",
+                    "tool_name": r"^Bash$",
+                    "decision": "allow",
+                    "allowed_param_patterns": {"command": r"^(?!(echo blockme)$).*$"},
+                }
+            ],
+            default_action="deny",
+            on_disallowed_action="block",
+        )
+        assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {})
+
+        # A partial update that does not carry `rules` must NOT wipe the existing
+        # ruleset / compiled maps.
+        guardrail.update_in_memory_litellm_params(
+            LitellmParams(
+                guardrail="tool_permission",
+                mode=["pre_call", "post_call"],
+                default_action="deny",
+                on_disallowed_action="block",
+            )
+        )
+
+        assert len(guardrail.rules) == 1
+        assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {})
+        assert (
+            guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0]
+            is False
+        )
+
+    def test_update_in_memory_rejects_invalid_regex_and_keeps_previous_rules(self):
+        """Regression: a live update whose rules contain an invalid regex must be
+        rejected atomically. The bad rule must not leak in as a compiled-target
+        wildcard (match-all), and the previously enforced ruleset must survive."""
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="tp",
+            rules=[{"id": "deny-secret", "tool_name": r"^Secret$", "decision": "deny"}],
+            default_action="allow",
+            on_disallowed_action="block",
+        )
+        # Baseline: only "Secret" is denied; any other tool is allowed.
+        assert guardrail._check_tool_permission("Secret")[0] is False
+        assert guardrail._check_tool_permission("Other")[0] is True
+
+        with pytest.raises(ValueError):
+            guardrail.update_in_memory_litellm_params(
+                LitellmParams(
+                    guardrail="tool_permission",
+                    mode=["pre_call", "post_call"],
+                    default_action="allow",
+                    on_disallowed_action="block",
+                    rules=[
+                        {
+                            "id": "deny-secret",
+                            "tool_name": r"^Secret$",
+                            "decision": "deny",
+                        },
+                        {"id": "bad", "tool_name": "[unclosed", "decision": "deny"},
+                    ],
+                )
+            )
+
+        # The bad rule must not have leaked in, and the prior ruleset must hold.
+        assert "bad" not in guardrail._compiled_rule_targets
+        assert all(rule.id != "bad" for rule in guardrail.rules)
+        assert guardrail._check_tool_permission("Other")[0] is True
+        assert guardrail._check_tool_permission("Secret")[0] is False

tests/test_litellm/proxy/test_caching_routes.py

@@ -123,33 +123,100 @@ def test_cache_ping_failure(mock_redis_failure):
     assert "message" in error_details
     assert "litellm_cache_params" in error_details
     assert "health_check_cache_params" in error_details
-    assert "traceback" in error_details
 
-    # Verify specific error message
-    assert "invalid username-password pair" in error_details["message"]
+    # Verify generic static message (exception text must not leak to clients)
+    assert error_details["message"] == "Service Unhealthy"
 
 
-def test_cache_ping_no_cache_initialized():
-    """Test cache ping when no cache is initialized"""
-    # Set cache to None
-    original_cache = litellm.cache
-    litellm.cache = None
-
+def test_cache_ping_failure_does_not_expose_traceback(mock_redis_failure):
+    """CWE-209: Stack trace and exception text must not appear in the HTTP 503 response body."""
     response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
     assert response.status_code == 503
 
     data = response.json()
-    print("response data=", json.dumps(data, indent=4))
-    assert "error" in data
-    error = data["error"]
+    error = data.get("error", {})
+    raw_body = json.dumps(data)
 
-    # Verify error contains all expected fields
-    assert "message" in error
+    # The word "traceback" (case-insensitive) must not appear anywhere in the response
+    assert (
+        "traceback" not in raw_body.lower()
+    ), "CWE-209: Python traceback exposed in HTTP 503 response body"
+    # Internal frame paths should not leak either
+    assert (
+        'File "' not in raw_body
+    ), "CWE-209: Python stack frame paths exposed in HTTP 503 response body"
+    # Exception text (e.g. Redis hostnames/IPs) must not leak either
+    assert (
+        "invalid username-password pair" not in raw_body
+    ), "CWE-209: Exception message text exposed in HTTP 503 response body"
+
+    # The error message should be the safe static string
     error_details = json.loads(error["message"])
-    assert "Cache not initialized. litellm.cache is None" in error_details["message"]
+    assert error_details["message"] == "Service Unhealthy"
 
-    # Restore original cache
-    litellm.cache = original_cache
+
+def test_cache_ping_no_cache_initialized():
+    """Test cache ping when no cache is initialized returns 503 with ProxyException envelope.
+
+    Verifies the exact response structure so that regressions in the error format
+    (e.g. message moving to a different field, or extra internal details leaking)
+    are caught immediately.
+    """
+    original_cache = litellm.cache
+    litellm.cache = None
+
+    try:
+        response = client.get(
+            "/cache/ping", headers={"Authorization": "Bearer sk-1234"}
+        )
+        assert response.status_code == 503
+
+        data = response.json()
+        print("response data=", json.dumps(data, indent=4))
+        # ProxyException is serialised as {"error": {"message": "...", "type": ..., ...}}
+        assert "error" in data
+        error_details = json.loads(data["error"]["message"])
+        assert (
+            error_details["message"] == "Cache not initialized. litellm.cache is None"
+        )
+    finally:
+        litellm.cache = original_cache
+
+
+def test_cache_ping_no_cache_does_not_expose_internals():
+    """CWE-209: No-cache 503 must use the ProxyException envelope with no internal details.
+
+    The null-cache path raises ProxyException directly (not HTTPException), so the
+    response is {"error": {"message": "...", ...}} — same envelope as other 503s from
+    this endpoint — with no tracebacks, source paths, or extra fields leaking.
+    """
+    original_cache = litellm.cache
+    litellm.cache = None
+
+    try:
+        response = client.get(
+            "/cache/ping", headers={"Authorization": "Bearer sk-1234"}
+        )
+        assert response.status_code == 503
+
+        raw_body = response.text
+        # No Python traceback or source-file paths must appear in the response
+        assert "traceback" not in raw_body.lower(), (
+            "CWE-209: Python traceback exposed in /cache/ping no-cache response"
+        )
+        assert 'File "' not in raw_body, (
+            "CWE-209: Python stack frame paths exposed in /cache/ping no-cache response"
+        )
+
+        data = response.json()
+        # Response must use the ProxyException envelope
+        assert "error" in data, f"Expected ProxyException envelope, got: {data}"
+        error_details = json.loads(data["error"]["message"])
+        assert (
+            error_details["message"] == "Cache not initialized. litellm.cache is None"
+        )
+    finally:
+        litellm.cache = original_cache
 
 
 def test_cache_ping_health_check_includes_only_cache_attributes(mock_redis_success):

tests/test_litellm/proxy/test_dynamic_mcp_route.py

@@ -486,3 +486,57 @@ async def test_dynamic_mcp_route_empty_access_group_returns_404():
             await dynamic_mcp_route("empty_group", request)
 
     assert exc_info.value.status_code == 404
+
+
+# ---------------------------------------------------------------------------
+# 6. Unexpected exception → 500 without leaking stack trace (CWE-209)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_dynamic_mcp_route_unexpected_exception_returns_500_without_traceback():
+    """CWE-209: an unexpected exception must return 500 with a generic message,
+    never leaking str(e) or a Python traceback to the caller."""
+    from litellm.proxy.proxy_server import dynamic_mcp_route
+
+    request = _make_request("/boom/mcp")
+
+    fake_mgr = MagicMock()
+    fake_mgr.get_mcp_server_by_name = MagicMock(
+        side_effect=RuntimeError("internal host: redis://10.0.0.1:6379")
+    )
+
+    with patch(_MCP_MANAGER, fake_mgr):
+        with pytest.raises(HTTPException) as exc_info:
+            await dynamic_mcp_route("boom", request)
+
+    assert exc_info.value.status_code == 500
+    assert exc_info.value.detail == "Internal server error"
+    assert "10.0.0.1" not in str(exc_info.value.detail)
+    assert "traceback" not in str(exc_info.value.detail).lower()
+
+
+@pytest.mark.asyncio
+async def test_toolset_mcp_route_unexpected_exception_returns_500_without_traceback():
+    """CWE-209: toolset_mcp_route must return 500 with a generic message on
+    unexpected errors, never leaking exception text to the caller."""
+    from litellm.proxy.proxy_server import toolset_mcp_route
+
+    request = _make_request("/toolset/broken_toolset/mcp")
+
+    fake_mgr = MagicMock()
+    fake_mgr.get_toolset_by_name_cached = AsyncMock(
+        side_effect=RuntimeError("connection to db-host:5432 refused")
+    )
+
+    with (
+        patch(_MCP_MANAGER, fake_mgr),
+        patch(_PRISMA, new=MagicMock()),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await toolset_mcp_route("broken_toolset", request)
+
+    assert exc_info.value.status_code == 500
+    assert exc_info.value.detail == "Internal server error"
+    assert "db-host" not in str(exc_info.value.detail)
+    assert "traceback" not in str(exc_info.value.detail).lower()

tests/test_litellm/test_main.py

@@ -849,12 +849,13 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict(
 
 
 @pytest.mark.parametrize(
-    "model, model_info, expected_model_param",
+    "model, model_info, expected_model_param, expected_base_model_param",
     [
-        ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro"),
+        ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro", None),
         (
             "gemini/gemini-3.1-pro",
             {"base_model": "gemini-3.1-pro-preview"},
+            "gemini-3.1-pro",
             "gemini-3.1-pro-preview",
         ),
     ],
@@ -863,7 +864,13 @@ def test_completion_optional_params_base_model(
     model: str,
     model_info: dict | None,
     expected_model_param: str,
+    expected_base_model_param: str | None,
 ):
+    """``model_info.base_model`` must reach ``get_optional_params`` as ``base_model``
+    (an additive capability hint), without overwriting ``model`` with the label.
+
+    Regression for #29618: overwriting ``model`` with a friendly ``base_model``
+    label made Bedrock drop ``tools``/``tool_choice`` under ``drop_params``."""
     with patch("litellm.main.get_optional_params") as mock_get_optional_params:
         mock_get_optional_params.return_value = MagicMock()
 
@@ -881,10 +888,9 @@ def test_completion_optional_params_base_model(
         litellm.completion(**kwargs)
 
         assert mock_get_optional_params.called is True
-        get_optional_params_model_param = mock_get_optional_params.call_args.kwargs[
-            "model"
-        ]
-        assert get_optional_params_model_param == expected_model_param
+        call_kwargs = mock_get_optional_params.call_args.kwargs
+        assert call_kwargs["model"] == expected_model_param
+        assert call_kwargs["base_model"] == expected_base_model_param
 
 
 @patch("litellm.completion_extras.responses_api_bridge.completion")

tests/test_litellm/test_utils.py

@@ -4144,3 +4144,51 @@ class TestValidateAndFixThinkingParam:
         validate_and_fix_thinking_param(thinking=thinking)
         assert "budgetTokens" in thinking
         assert "budget_tokens" not in thinking
+
+
+class TestBedrockBaseModelLabelKeepsTools:
+    """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly
+    label must not silently drop ``tools``/``tool_choice`` under ``drop_params``."""
+
+    TOOLS = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                },
+            },
+        }
+    ]
+
+    def test_base_model_label_keeps_tools_with_drop_params(self):
+        from litellm.utils import get_optional_params
+
+        result = get_optional_params(
+            model="eu.anthropic.claude-haiku-4-5-20251001-v1:0",
+            custom_llm_provider="bedrock",
+            base_model="claude-haiku-4-5",
+            tools=self.TOOLS,
+            tool_choice="auto",
+            drop_params=True,
+        )
+
+        assert "tools" in result
+        assert "tool_choice" in result
+
+    def test_base_model_label_alone_drops_tools(self):
+        """Without the real model id the label resolves to no tool support, so passing
+        the label as ``model`` is exactly what dropped tools before the fix."""
+        from litellm.utils import get_optional_params
+
+        result = get_optional_params(
+            model="claude-haiku-4-5",
+            custom_llm_provider="bedrock",
+            tools=self.TOOLS,
+            tool_choice="auto",
+            drop_params=True,
+        )
+
+        assert "tools" not in result