test(proxy): cover the max_parallel_requests slot a streamed /v1/messages MCP auto-execute keeps

A streamed /v1/messages request that auto-executes MCP tools replays its answer
as a stream after its model rounds already logged success. The stream path frees
the request's max_parallel_requests slot only on a client disconnect, so after
message_stop the slot stays taken and the key's next request answers 429

The test streams such a replay to the end and expects the slot to be free. It
fails until the fix lands
This commit is contained in:
mateo-berri 2026-10-20 21:39:42 -07:00
parent 4ecc5fd13a
commit 83d5136f67

View file

@ -4332,6 +4332,97 @@ async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request
assert await in_flight() == 0
@pytest.mark.asyncio
async def test_completed_mcp_replay_stream_frees_slot_its_logged_first_round_kept() -> None:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
reused_at: Final = datetime(2026, 1, 5, 12, 0, 0)
api_key: Final = hash_token("sk-mcp-replayed-stream")
local_cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache),
time_provider=lambda: reused_at,
)
proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler
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"
owner_call_id: Final = "mcp-replayed-stream-owner"
messages: Final = [{"role": "user", "content": "add 2 and 3"}]
async def in_flight() -> int:
return handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=parallel_key)
)
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(
{"model": "claude-haiku-5-5", "messages": messages, "stream": True, "litellm_call_id": owner_call_id},
CallTypes.anthropic_messages.value,
)
await pre_call(
{"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"},
CallTypes.call_mcp_tool.value,
)
logging_obj: Final = LiteLLMLoggingObj(
model="claude-haiku-5-5",
messages=messages,
stream=True,
call_type=CallTypes.anthropic_messages.value,
start_time=None,
litellm_call_id=owner_call_id,
function_id="mcp-replayed-stream",
)
assert logging_obj.claim_async_success_log()
await handler.async_log_success_event(
kwargs={"litellm_call_id": owner_call_id, "stream": False},
response_obj=None,
start_time=None,
end_time=reused_at - timedelta(seconds=1),
)
assert await in_flight() == 1
replay: Final = FakeAnthropicMessagesStreamIterator(
response={
"id": "msg_replayed",
"type": "message",
"role": "assistant",
"model": "claude-haiku-5-5",
"content": [{"type": "text", "text": "2 + 3 = 5"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 20, "output_tokens": 8},
}
)
with _override_litellm_callbacks([]):
streamed: Final = [
chunk
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
response=replay,
user_api_key_dict=user_api_key_dict,
request_data={
"model": "claude-haiku-5-5",
"litellm_call_id": owner_call_id,
"litellm_logging_obj": logging_obj,
},
proxy_logging_obj=proxy_logging_obj,
)
]
assert b"event: message_stop" in streamed[-1]
assert await in_flight() == 0
@pytest.mark.asyncio
async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None:
api_key: Final = hash_token("sk-mcp-reuse-headers")