mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
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:
parent
4ecc5fd13a
commit
83d5136f67
1 changed files with 91 additions and 0 deletions
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue