scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 359

Merge a297a7d7ad into ff22f3de06

Browse source
This commit is contained in:
m6atlas 2026-10-29 06:55:28 +00:00 • committed by GitHub
commit a26cb1e1d1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 540 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
litellm/llms/anthropic/messages/native/mcp_handler.py
View file

@ -82,6 +82,7 @@ async def anthropic_messages_with_mcp(
from litellm.mcp_client.tools import (
transform_mcp_tool_to_anthropic_tool,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import request_slot_held_by_tool_loop
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
@ -145,52 +146,53 @@ async def anthropic_messages_with_mcp(
else base_call_args
)
working_messages: Sequence[Mapping[str, object]] = tuple(messages)
response: AnthropicMessagesResponse = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
messages=list(working_messages), stream=False, **loop_call_args
)
for _ in range(MAX_MCP_TOOL_USE_ITERATIONS):
if _get_stop_reason(response) != "tool_use":
break
tool_use_blocks = _extract_tool_use_blocks(response)
if not tool_use_blocks:
break
tool_results = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(
tool_server_map=tool_server_map,
served_tools=deduplicated_mcp_tools,
tool_calls=list(tool_use_blocks),
user_api_key_auth=context.user_api_key_auth,
mcp_auth_header=context.mcp_auth_header,
mcp_server_auth_headers=context.mcp_server_auth_headers,
oauth2_headers=context.oauth2_headers,
raw_headers=context.raw_headers,
litellm_call_id=context.litellm_call_id,
litellm_trace_id=context.litellm_trace_id,
request_tags=list(context.request_tags) if context.request_tags else None,
guardrail_context=context.guardrail_context,
)
# Every tool call was skipped, so there is nothing to feed back; a
# tool_result message with empty content is rejected by Anthropic.
if not tool_results:
break
working_messages = (
*working_messages,
{"role": "assistant", "content": list(_get_response_content(response))},
_build_tool_result_message(tool_results),
)
response = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
with request_slot_held_by_tool_loop(context.litellm_call_id):
response: AnthropicMessagesResponse = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
messages=list(working_messages), stream=False, **loop_call_args
)
else:
verbose_logger.warning(
"MCP tool loop hit its %s iteration cap for model %s; returning the last response",
MAX_MCP_TOOL_USE_ITERATIONS,
model,
)
for _ in range(MAX_MCP_TOOL_USE_ITERATIONS):
if _get_stop_reason(response) != "tool_use":
break
tool_use_blocks = _extract_tool_use_blocks(response)
if not tool_use_blocks:
break
tool_results = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(
tool_server_map=tool_server_map,
served_tools=deduplicated_mcp_tools,
tool_calls=list(tool_use_blocks),
user_api_key_auth=context.user_api_key_auth,
mcp_auth_header=context.mcp_auth_header,
mcp_server_auth_headers=context.mcp_server_auth_headers,
oauth2_headers=context.oauth2_headers,
raw_headers=context.raw_headers,
litellm_call_id=context.litellm_call_id,
litellm_trace_id=context.litellm_trace_id,
request_tags=list(context.request_tags) if context.request_tags else None,
guardrail_context=context.guardrail_context,
)
# Every tool call was skipped, so there is nothing to feed back; a
# tool_result message with empty content is rejected by Anthropic.
if not tool_results:
break
working_messages = (
*working_messages,
{"role": "assistant", "content": list(_get_response_content(response))},
_build_tool_result_message(tool_results),
)
response = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
messages=list(working_messages), stream=False, **loop_call_args
)
else:
verbose_logger.warning(
"MCP tool loop hit its %s iteration cap for model %s; returning the last response",
MAX_MCP_TOOL_USE_ITERATIONS,
model,
)
client_response: Final = (
await with_safeguard_results(response, request_safeguards_evaluator(kwargs, working_messages))

12
litellm/proxy/common_request_processing.py
View file

@ -2351,6 +2351,7 @@ class ProxyBaseLLMRequestProcessing:
rate_limited_model: str | None = None,
skip_guardrails: bool = False,
logging_auth: UserAPIKeyAuth | None = None,
litellm_call_id: str | None = None,
) -> tuple[dict, LiteLLMLoggingObj]:
start_time: Final = datetime.now() # start before calling guardrail hooks
@ -2487,10 +2488,10 @@ class ProxyBaseLLMRequestProcessing:
)
)
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
self.data["litellm_call_id"] = litellm_call_id
bind_telemetry_call_id(litellm_call_id)
DDSpanTagger.tag_call_id(litellm_call_id)
call_id: Final = litellm_call_id or resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
self.data["litellm_call_id"] = call_id
bind_telemetry_call_id(call_id)
DDSpanTagger.tag_call_id(call_id)
DDSpanTagger.tag_request(
user_api_key_dict=user_api_key_dict,
requested_model=aliased_model,
@ -2634,6 +2635,7 @@ class ProxyBaseLLMRequestProcessing:
else None
)
pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
try:
return await self.common_processing_pre_call_logic(
@ -2651,6 +2653,7 @@ class ProxyBaseLLMRequestProcessing:
model=model,
route_type=route_type,
llm_router=llm_router,
litellm_call_id=litellm_call_id,
)
except ProxyRateLimitError as original_exc:
rate_limited_data: Final = self.data
@ -2707,6 +2710,7 @@ class ProxyBaseLLMRequestProcessing:
route_type=route_type,
llm_router=llm_router,
rate_limited_model=original_model,
litellm_call_id=litellm_call_id,
)
except ProxyRateLimitError as fallback_exc:
if is_hard_per_model_limit(fallback_exc):

2
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -616,7 +616,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
if "model" not in data:
return None
claim_request_stash_for_data(data)
claim_request_stash_for_data(data, call_type)
model: Final = data["model"]
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)

43
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -11,8 +11,8 @@ import json
import logging
import os
import uuid
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set
from contextlib import asynccontextmanager
from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence, Set
from contextlib import asynccontextmanager, contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime, timezone
@ -713,6 +713,7 @@ class RequestRateLimiterStash:
rate_limit_response: RateLimitResponse | None = None
parallel_slot: ParallelSlotAcquisition | None = None
parallel_slot_reused_at: datetime | None = None
parallel_slot_held_by_tool_loop: bool = False
parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False)
reserved_tokens: int = 0
reserved_model: RateLimitedModel | None = None
@ -801,10 +802,28 @@ def get_or_create_request_stash() -> RequestRateLimiterStash:
return stash
def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash:
def hold_request_slot_for_tool_loop(litellm_call_id: str | None) -> RequestRateLimiterStash | None:
stash: Final = get_request_stash_for_call(litellm_call_id)
if stash is None or stash.parallel_slot is None:
return None
stash.parallel_slot_held_by_tool_loop = True
return stash
@contextmanager
def request_slot_held_by_tool_loop(litellm_call_id: str | None) -> Generator[None, None, None]:
held_stash: Final = hold_request_slot_for_tool_loop(litellm_call_id)
try:
yield
finally:
if held_stash is not None:
held_stash.parallel_slot_held_by_tool_loop = False
def claim_request_stash_for_data(data: dict, call_type: str | None) -> RequestRateLimiterStash:
stash: Final = get_or_create_request_stash()
owner_call_id: Final = data.get("litellm_call_id")
if isinstance(owner_call_id, str):
if isinstance(owner_call_id, str) and call_type != CallTypes.pass_through.value:
stash.owner_litellm_call_id = owner_call_id
return stash
@ -2117,7 +2136,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return
async with stash.parallel_slot_release_lock:
acquisition: Final = stash.parallel_slot
if acquisition is None or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at):
if (
acquisition is None
or (in_logging_callback and stash.parallel_slot_held_by_tool_loop)
or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at)
):
return
deferred: Final = in_logging_callback and await self._defer_parallel_slot_release(
acquisition, parent_otel_span
@ -4052,7 +4075,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""
if not isinstance(data, dict):
return
stash: Final = claim_request_stash_for_data(data)
stash: Final = claim_request_stash_for_data(data, call_type)
io_token_descriptors: Final = [
d for d in descriptors if d["key"] in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY)
]
@ -4281,7 +4304,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""
verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook")
stash: Final = claim_request_stash_for_data(data)
stash: Final = claim_request_stash_for_data(data, call_type)
reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None
if reuses_request_slot:
stash.parallel_slot_reused_at = self._get_current_time()
@ -5787,7 +5810,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
@with_service_target("rate_limits")
async def async_post_call_failure_hook(
self,
request_data: dict,
request_data: dict[str, object],
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: str | None = None,
@ -5810,7 +5833,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
that usage instead of refunding it.
"""
try:
stash: Final = get_request_stash()
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(request_data))
if stash is None:
return
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
@ -5851,7 +5874,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
ptu_ops: Final = (
self._build_ptu_failure_settlement_ops(
stash,
request_data, # pyright: ignore[reportUnknownArgumentType] # request_data is a bare dict
request_data,
tpm_actual,
)
if reserved_tokens > 0

22
litellm/responses/mcp/chat_completions_handler.py
View file

@ -215,6 +215,7 @@ async def acompletion_with_mcp(
# Create a custom async generator that collects chunks and handles tool execution
from litellm.main import stream_chunk_builder
from litellm.proxy.hooks.parallel_request_limiter_v3 import hold_request_slot_for_tool_loop
from litellm.types.utils import ModelResponseStream
class MCPStreamingIterator:
@ -258,6 +259,7 @@ async def acompletion_with_mcp(
self.follow_up_stream = None
self.follow_up_iterator = None
self.follow_up_exhausted = False
self.holds_request_slot = hold_request_slot_for_tool_loop(context.litellm_call_id) is not None
def __aiter__(self):
return self
@ -332,7 +334,21 @@ async def acompletion_with_mcp(
"Error draining inner MCP stream after final chunk; spend logging may be incomplete"
)
async def __anext__(self):
async def __anext__(self) -> ModelResponseStream:
try:
return await self._next_chunk()
except StopAsyncIteration:
await self.release_held_request_slot()
raise
async def release_held_request_slot(self) -> None:
if not self.holds_request_slot:
return
from litellm.proxy.proxy_server import proxy_logging_obj
await proxy_logging_obj.arelease_max_parallel_requests(self.user_api_key_auth)
async def _next_chunk(self) -> ModelResponseStream:
# Phase 1: Collect and yield initial stream chunks
if not self.stream_exhausted:
# Get the iterator from the stream wrapper
@ -548,6 +564,10 @@ async def acompletion_with_mcp(
def __aiter__(self):
return self._custom_iterator
async def aclose(self) -> None:
await self._custom_iterator.release_held_request_slot()
await super().aclose()
def __iter__(self):
# For synchronous iteration, create a sync wrapper
if self._sync_iterator is None:

10
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -742,6 +742,13 @@ class LiteLLM_Proxy_MCP_Handler:
tool_arguments,
tool_call_id,
) = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call)
tool_logging_call_id = (
f"{litellm_call_id}_{tool_call_id or uuid.uuid4()}" if litellm_call_id else str(uuid.uuid4())
)
logging_request_data = {
"call_type": CallTypes.call_mcp_tool.value,
"litellm_call_id": tool_logging_call_id,
}
if not tool_name:
verbose_logger.warning("Tool call missing name: %s", tool_call)
@ -772,9 +779,6 @@ class LiteLLM_Proxy_MCP_Handler:
},
}
]
tool_logging_call_id = (
f"{litellm_call_id}_{tool_call_id or uuid.uuid4()}" if litellm_call_id else str(uuid.uuid4())
)
logging_metadata: dict[str, object] = {
"tool_call_id": tool_call_id,
"tool_name": sanitized_tool_name,

81
tests/unit/llms/anthropic/messages/native/test_mcp_handler.py
View file

@ -454,3 +454,84 @@ async def test_anthropic_messages_with_mcp_loads_a_filtered_routes_beta_map_off_
@pytest.mark.asyncio
async def test_anthropic_messages_with_mcp_never_loads_the_beta_map_for_first_party_anthropic():
assert await _mcp_safeguards_turn_recording_beta_map_loads("anthropic", "claude-sonnet-4-6") == ()
@pytest.mark.asyncio
async def test_anthropic_messages_with_mcp_keeps_the_requests_parallel_slot_while_its_tools_run():
import asyncio
import contextvars
from datetime import datetime
from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.llms.anthropic.messages.native import mcp_handler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.utils import InternalUsageCache, hash_token
from litellm.responses.mcp.request_context import MCPRequestContext
from litellm.types.utils import CallTypes
local_cache: Final = DualCache()
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(local_cache))
key: Final = UserAPIKeyAuth(api_key=hash_token("sk-messages-tool-loop"), max_parallel_requests=1)
async def admit(data: dict[str, object]) -> None:
await limiter.async_pre_call_hook(
user_api_key_dict=key, cache=local_cache, data=data, call_type=CallTypes.anthropic_messages.value
)
async def second_request_status() -> int:
try:
await contextvars.Context().run(asyncio.create_task, admit({"model": "claude-x", "messages": []}))
except HTTPException as rejected:
return rejected.status_code
return 200
async def log_the_requests_success() -> None:
await limiter.async_log_success_event(
kwargs={"litellm_call_id": "messages-owner", "stream": False},
response_obj=None,
start_time=None,
end_time=datetime.now(),
)
async def run_tool_after_the_first_round_logs(**_: object) -> list[dict[str, str]]:
await log_the_requests_success()
assert await second_request_status() == 429, "A round logged while tools run must not free the request's slot"
return [{"tool_call_id": "toolu_1", "result": "ok", "name": "get_weather"}]
responses = [
{
"stop_reason": "tool_use",
"content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}],
},
{"stop_reason": "end_turn", "content": [{"type": "text", "text": "done"}]},
]
handler_class: Final = import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler
await admit(
{"model": "claude-x", "messages": [{"role": "user", "content": "hi"}], "litellm_call_id": "messages-owner"}
)
with (
patch.object(MCPRequestContext, "resolve", return_value=MCPRequestContext(user_api_key_auth="auth")),
patch.object(
handler_class,
"process_mcp_tools_without_openai_transform",
new=AsyncMock(return_value=([], {"get_weather": "weather"})),
),
patch.object(
handler_class, "execute_tool_calls", new=AsyncMock(side_effect=run_tool_after_the_first_round_logs)
),
patch("litellm.anthropic_messages", new=AsyncMock(side_effect=responses)),
):
result = await mcp_handler.anthropic_messages_with_mcp(
max_tokens=100,
messages=[{"role": "user", "content": "hi"}],
model="claude-sonnet-4-5",
tools=[MCP_REFERENCE],
)
assert result == responses[1]
await log_the_requests_success()
assert await second_request_status() == 200, "The request's own log after its tools finish must free its slot"

141
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -41,6 +41,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_request_stash,
get_or_create_request_stash,
get_request_stash,
request_slot_held_by_tool_loop,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler,
@ -4271,6 +4272,146 @@ async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot()
assert await in_flight() == 1
@pytest.mark.asyncio
async def test_failed_mcp_tool_call_keeps_its_requests_parallel_slot_and_reservation(monkeypatch) -> None:
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
api_key: Final = hash_token("sk-mcp-auto-execute-tool-failure")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1, tpm_limit=10_000)
parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests"
tokens_key: Final = handler.create_rate_limit_keys(key="api_key", value=api_key, rate_limit_type="tokens")
owner_call_id: Final = "mcp-tool-failure-owner"
owner_data: Final[dict[str, object]] = {
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "weather in Paris"}],
"max_tokens": 50,
"litellm_call_id": owner_call_id,
}
async def in_flight() -> int:
return handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key))
async def reserved_tokens() -> int:
return int(await local_cache.async_get_cache(key=tokens_key) or 0)
async def pre_call(data: dict[str, object], call_type: str) -> None:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=data,
call_type=call_type,
)
await pre_call(owner_data, CallTypes.acompletion.value)
await pre_call(
{"model": "MCP: get_weather", "messages": [{"role": "user", "content": "Tool: get_weather"}]},
CallTypes.call_mcp_tool.value,
)
reserved_before_tool_failure: Final = await reserved_tokens()
await handler.async_post_call_failure_hook(
request_data={"call_type": CallTypes.call_mcp_tool.value, "litellm_call_id": f"{owner_call_id}_call_1"},
original_exception=HTTPException(status_code=504, detail="MCP tool call timed out after 3.0s"),
user_api_key_dict=user_api_key_dict,
)
assert reserved_before_tool_failure > 0
assert await reserved_tokens() == reserved_before_tool_failure
with pytest.raises(HTTPException) as rejected:
await contextvars.Context().run(
asyncio.create_task, pre_call({"model": "gpt-4o-mini", "input": "next"}, CallTypes.aresponses.value)
)
assert rejected.value.status_code == 429
await handler.async_post_call_failure_hook(
request_data=owner_data,
original_exception=Exception("request failed"),
user_api_key_dict=user_api_key_dict,
)
assert await in_flight() == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True], ids=["non_stream_round", "stream_round"])
async def test_round_logged_while_a_tool_loop_holds_the_slot_keeps_it_until_the_request_ends(stream: bool) -> None:
now: Final = datetime(2026, 10, 28, 12, 0, 0)
api_key: Final = hash_token(f"sk-tool-loop-hold-{stream}")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache), time_provider=lambda: now
)
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
owner_call_id: Final = f"tool-loop-hold-owner-{stream}"
async def pre_call(data: dict[str, object]) -> None:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=data,
call_type=CallTypes.anthropic_messages.value,
)
async def next_request_status() -> int:
try:
await contextvars.Context().run(asyncio.create_task, pre_call({"model": "claude-x", "messages": []}))
except HTTPException as rejected:
return rejected.status_code
return 200
await pre_call(
{"model": "claude-x", "messages": [{"role": "user", "content": "hi"}], "litellm_call_id": owner_call_id}
)
with request_slot_held_by_tool_loop(owner_call_id):
with request_slot_held_by_tool_loop(f"{owner_call_id}-silent-shadow"):
pass
await handler.async_post_call_failure_hook(
request_data={"call_type": CallTypes.call_mcp_tool.value, "litellm_call_id": f"{owner_call_id}_toolu_1"},
original_exception=KeyError("lookup_order"),
user_api_key_dict=user_api_key_dict,
)
await handler.async_log_success_event(
kwargs={"litellm_call_id": owner_call_id, "stream": stream},
response_obj=None,
start_time=None,
end_time=now - timedelta(milliseconds=1),
)
assert await next_request_status() == 429
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
assert await next_request_status() == 200
@pytest.mark.asyncio
async def test_failed_pass_through_request_frees_its_slot_when_the_client_body_carries_a_litellm_call_id() -> None:
api_key: Final = hash_token("sk-pass-through-client-call-id")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests"
async def pre_call(data: dict[str, object]) -> None:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=data,
call_type=CallTypes.pass_through.value,
)
await pre_call({"model": "gpt-4o-mini", "max_completion_tokens": -5, "litellm_call_id": "client-chosen-id"})
await handler.async_post_call_failure_hook(
request_data={"model": "gpt-4o-mini", "litellm_call_id": "proxy-assigned-id"},
original_exception=HTTPException(status_code=400, detail="max_completion_tokens must be positive"),
user_api_key_dict=user_api_key_dict,
)
assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0
await contextvars.Context().run(asyncio.create_task, pre_call({"model": "gpt-4o-mini"}))
@pytest.mark.asyncio
@pytest.mark.parametrize(
("stream", "in_flight_after_first_round_logs"),

37
tests/unit/proxy/test_common_request_processing.py
View file

@ -8092,6 +8092,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=run_limiter)
proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=limiter.async_post_call_failure_hook)
guardrails_by_group = model_guardrails or {}
router = litellm.Router(
model_list=[
@ -8259,6 +8260,42 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
assert processor.data["litellm_logging_obj"].model == primary_model
assert processor.data["litellm_call_id"]
@pytest.mark.asyncio
async def test_failed_request_served_by_a_fallback_frees_its_parallel_slot(self, monkeypatch: pytest.MonkeyPatch):
"""/completions reports a failure with the endpoint's own request body, the one the
rate-limited first attempt ran on, so the fallback attempt that took the slot has to carry
the same call id or the slot never comes back."""
primary_model: Final = "gpt-4.1"
fallback_model: Final = "gpt-4.1-mini"
key: Final = ProxyUserAPIKeyAuth(api_key="hashed-key", max_parallel_requests=1)
rig: Final = self._v3_limiter_rig(
monkeypatch,
key,
[{primary_model: [fallback_model]}],
untagged_rate_limited_models=frozenset({primary_model}),
)
async def request_that_fails_after_falling_back() -> None:
client_data: Final[dict[str, object]] = {
"model": primary_model,
"messages": [{"role": "user", "content": "hi"}],
}
await self._pre_call(client_data, key, rig)
await rig[0].post_call_failure_hook(
user_api_key_dict=key,
original_exception=Exception("upstream failed"),
request_data=client_data,
)
await contextvars.Context().run(asyncio.create_task, request_that_fails_after_falling_back())
_, (next_data, _) = await contextvars.Context().run(
asyncio.create_task,
self._pre_call({"model": fallback_model, "messages": [{"role": "user", "content": "next"}]}, key, rig),
)
assert next_data["model"] == fallback_model
assert rig[3] == [primary_model, fallback_model, fallback_model]
HARD_PER_MODEL_LIMITS: Final = MappingProxyType({"disable_fallbacks_on_per_model_rate_limits": True})
@pytest.mark.parametrize("cap_owner", ["key", "team"])

143
tests/unit/responses/mcp/test_chat_completions_handler.py
View file

@ -26,6 +26,7 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import (
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.utils import ModelResponse, StandardLoggingPayload
from litellm.utils import CustomStreamWrapper
@pytest.mark.asyncio
@ -1712,3 +1713,145 @@ async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follo
assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19)
assert recorder.failed[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE))
assert recorder.failed[0]["response_cost"] > 0
async def _read_the_whole_stream(stream: CustomStreamWrapper) -> str:
return "".join([part.choices[0].delta.content or "" async for part in stream])
async def _hang_up_after_the_first_follow_up_chunk(stream: CustomStreamWrapper) -> str:
parts: Final = stream.__aiter__()
read: Final = (await parts.__anext__(), await parts.__anext__())
await stream.aclose()
return "".join(part.choices[0].delta.content or "" for part in read)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("end_stream", "expected_text"),
[(_read_the_whole_stream, "Sunny."), (_hang_up_after_the_first_follow_up_chunk, "Sunny")],
ids=["stream_ends", "client_hangs_up"],
)
@pytest.mark.parametrize(
("stream_call_id", "status_after_stream"),
[("chat-owner", 200), ("silent-shadow", 429)],
ids=["requests_own_stream", "silent_model_shadow_stream"],
)
async def test_streamed_mcp_auto_exec_keeps_the_requests_parallel_slot_until_the_stream_ends_or_closes(
monkeypatch, end_stream, expected_text, stream_call_id: str, status_after_stream: int
):
import contextvars
from unittest.mock import MagicMock
from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.utils import InternalUsageCache, hash_token
from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices
from litellm.utils import CustomStreamWrapper
local_cache: Final = DualCache()
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(local_cache))
key: Final = UserAPIKeyAuth(api_key=hash_token("sk-chat-stream-tool-loop"), max_parallel_requests=1)
tools: Final = [{"type": "function", "function": {"name": "get_weather"}}]
tool_call: Final = {"id": "call-1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}}
async def admit(data: dict[str, object]) -> None:
await limiter.async_pre_call_hook(
user_api_key_dict=key, cache=local_cache, data=data, call_type=CallTypes.acompletion.value
)
async def second_request_status() -> int:
try:
await contextvars.Context().run(asyncio.create_task, admit({"model": "gpt-x", "messages": []}))
except HTTPException as rejected:
return rejected.status_code
return 200
def chunk(content: str, finish_reason: str | None) -> ModelResponseStream:
return ModelResponseStream(
id="stream",
model="gpt-x",
created=0,
object="chat.completion.chunk",
choices=[
StreamingChoices(index=0, delta=Delta(content=content, role="assistant"), finish_reason=finish_reason)
],
)
logging_obj: Final = MagicMock()
logging_obj.model_call_details = {}
class FixedStream(CustomStreamWrapper):
def __init__(self, chunks: tuple[ModelResponseStream, ...]):
super().__init__(completion_stream=None, model="gpt-x", logging_obj=logging_obj)
self.pending = iter(chunks)
def __aiter__(self):
return self
async def __anext__(self):
try:
return next(self.pending)
except StopIteration:
raise StopAsyncIteration
async def run_tool_after_the_first_round_logs(**_):
await limiter.async_log_success_event(
kwargs={"litellm_call_id": stream_call_id, "stream": True},
response_obj=None,
start_time=None,
end_time=datetime.now(),
)
assert await second_request_status() == 429, "The first round's stream log must not free the request's slot"
return [{"tool_call_id": "call-1", "result": "sunny"}]
acompletion: Final = AsyncMock(
side_effect=[FixedStream((chunk("", "tool_calls"),)), FixedStream((chunk("Sunny", None), chunk(".", "stop")))]
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj",
types.SimpleNamespace(arelease_max_parallel_requests=limiter.async_release_max_parallel_requests_on_disconnect),
)
monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "should_use_litellm_mcp_gateway", staticmethod(lambda tools: True))
monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "_parse_mcp_tools", staticmethod(lambda tools: (tools, [])))
monkeypatch.setattr(
LiteLLM_Proxy_MCP_Handler,
"process_mcp_tools_without_openai_transform",
AsyncMock(return_value=(tools, {"get_weather": "weather"})),
)
monkeypatch.setattr(
LiteLLM_Proxy_MCP_Handler, "transform_mcp_tools_to_openai", staticmethod(lambda *_, **__: tools)
)
monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "should_auto_execute_tools", staticmethod(lambda **_: True))
monkeypatch.setattr(
LiteLLM_Proxy_MCP_Handler, "extract_tool_calls_from_chat_response", staticmethod(lambda **_: [tool_call])
)
monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", run_tool_after_the_first_round_logs)
monkeypatch.setattr(
LiteLLM_Proxy_MCP_Handler,
"create_follow_up_messages_for_chat",
staticmethod(lambda **_: [{"role": "tool", "tool_call_id": "call-1", "content": "sunny"}]),
)
monkeypatch.setattr(
ResponsesAPIRequestUtils, "extract_mcp_headers_from_request", staticmethod(lambda **_: (None, None, None, None))
)
await admit({"model": "gpt-x", "messages": [{"role": "user", "content": "hi"}], "litellm_call_id": "chat-owner"})
with patch("litellm.acompletion", acompletion):
stream = await acompletion_with_mcp(
model="gpt-x",
messages=[{"role": "user", "content": "hi"}],
tools=tools,
stream=True,
litellm_call_id=stream_call_id,
)
streamed_text: Final = await end_stream(stream)
assert (
await second_request_status() == status_after_stream
), "Only the request's own stream may free its slot, once that stream ends or closes"
assert streamed_text == expected_text

23
tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py
View file

@ -33,7 +33,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.responses.main import OutputFunctionToolCall
from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs, ModelResponse
class _DummyMCPResult:
@ -419,6 +419,27 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey
assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool"
@pytest.mark.asyncio
async def test_execute_tool_calls_reports_an_unmapped_tool_failure_under_the_tool_calls_own_id(monkeypatch):
call_tool_mock: Final = _setup_mcp_call_environment(monkeypatch)
post_call_failure_hook: Final = _setup_proxy_logging(monkeypatch)
results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(
tool_server_map={},
tool_calls=[{"id": "call-unmapped", "function": {"name": "unknown_tool", "arguments": "{}"}}],
user_api_key_auth=types.SimpleNamespace(api_key="test_key", user_id="test_user"),
litellm_call_id="cid",
)
assert results[0]["result"].startswith("Error executing tool")
call_tool_mock.assert_not_awaited()
post_call_failure_hook.assert_awaited_once()
assert post_call_failure_hook.await_args is not None
reported_request: Final = post_call_failure_hook.await_args.kwargs["request_data"]
assert reported_request["litellm_call_id"] == "cid_call-unmapped"
assert reported_request["call_type"] == CallTypes.call_mcp_tool.value
@pytest.mark.asyncio
async def test_execute_tool_calls_logs_each_tool_call_under_its_own_id_on_the_request_trace(
monkeypatch,