From ba239b74f491691d9202c523a58df31cc6290780 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:43:47 -0700 Subject: [PATCH] fix(proxy): hold a websocket session's parallel slot until the socket closes Each finished turn logged its spend row under the session's own call id, so the first turn's success callback freed the key's max_parallel_requests slot while the socket was still open, and a key limited to one session could open several at once. The session now claims its slot under its own call id, every turn row logs under a call id of its own, and the slot is released when the session ends, as before Priced turn rows also take their row's call id as their id, so two turns that start in the same instant no longer share a spend log id and drop one row --- ...i_websocket_passthrough_logging_handler.py | 6 +- .../pass_through_endpoints.py | 10 ++- ...i_websocket_passthrough_logging_handler.py | 26 +++++- .../test_pass_through_endpoints.py | 88 ++++++++++++++++++- .../test_success_handler.py | 2 +- 5 files changed, 123 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py index eb2b4edf2fe..4ca71a6fc5e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -200,9 +200,9 @@ def _requested_model(kwargs: Mapping[str, object]) -> str | None: return model_group if isinstance(model_group, str) else None -def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse: +def _priced_result(model: str, billing: _Billing, start_time: datetime, call_id: str) -> ModelResponse: result: Final = ModelResponse( - id=f"openai-websocket-{start_time.timestamp()}", + id=call_id, object="chat.completion", created=int(start_time.timestamp()), model=model, @@ -246,7 +246,7 @@ class OpenAIWebsocketPassthroughLoggingHandler: logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost})) priced: Final[PassThroughEndpointLoggingTypedDict] = { - "result": _priced_result(model, billing, start_time), + "result": _priced_result(model, billing, start_time, logging_obj.litellm_call_id), "kwargs": { **kwargs, "model": model, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7b304331f28..61178caf35f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3124,14 +3124,16 @@ async def websocket_passthrough_request( def log_turn(lane: str | None, in_row: Callable[[Mapping[str, object]], bool]) -> None: end_time: Final = datetime.now() - if ledger.logged_turns == 0: + if upstream_frame_kind is None: log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time, lane, in_row) else: turn_call_id: Final = str(uuid.uuid4()) turn_logging: Final = new_call_logging(ledger.started, turn_call_id) turn_payload: Final = new_call_payload() turn_kwargs: Final = init_call_kwargs( - user_api_key_dict.model_copy(update={"budget_reservation": None}), + user_api_key_dict + if ledger.logged_turns == 0 + else user_api_key_dict.model_copy(update={"budget_reservation": None}), turn_payload, turn_call_id, turn_logging, @@ -3141,7 +3143,9 @@ async def websocket_passthrough_request( ledger.close_turn(lane, end_time) ### CALL HOOKS ### - modify incoming data / reject request before calling the model - websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True) + websocket_hook_data: Final = guardrail_request_data_with_streaming( + MappingProxyType({"litellm_call_id": litellm_call_id}), is_streaming=True + ) await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=websocket_hook_data, diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py index baf52af112b..35e12800ce3 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -1,3 +1,4 @@ +import uuid from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType @@ -15,6 +16,8 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket OpenAIWebsocketPassthroughLoggingHandler, openai_websocket_frame_kind, ) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id +from litellm.types.utils import CallTypes REALTIME_MODEL: Final = "gpt-realtime-2.1" RESPONSES_MODEL: Final = "gpt-5.4-mini" @@ -24,7 +27,7 @@ SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0) def _logging_obj() -> MagicMock: - logging_obj: Final = MagicMock(model_call_details={}) + logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id=str(uuid.uuid4())) return logging_obj @@ -290,6 +293,27 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row(): assert handled["kwargs"]["model"] == RESPONSES_MODEL +def test_turns_that_start_in_the_same_instant_each_keep_their_own_spend_row(): + """ + Spend logs key each row by the id its result carries and skip a repeated id + without a word, so two turns that start in the same instant, on one socket + or on two, each need an id of their own or one of them is never billed. + """ + turns: Final = [ + _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)], + url_route="/openai_passthrough/v1/responses", + ) + for _ in range(2) + ] + + spend_log_ids: Final = { + get_spend_logs_id(CallTypes.pass_through.value, turn["result"].model_dump(), turn["kwargs"]) for turn in turns + } + assert len(spend_log_ids) == 2 + assert None not in spend_log_ids + + def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): """ A client that closes before response.completed leaves no usage behind to diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 28c8689518d..06395bf72a9 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1,5 +1,6 @@ import asyncio import base64 +import contextvars import gzip import json import logging @@ -39,6 +40,8 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( @@ -77,7 +80,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes @@ -8204,6 +8207,89 @@ async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_th relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() +@pytest.mark.asyncio +async def test_websocket_passthrough_holds_the_keys_parallel_slot_until_the_session_ends(): + """ + A key allowed one parallel request holds that slot for as long as its + socket stays open. A turn that finishes and reaches spend mid-session must + not free it, or a second session gets in while the first is still running. + """ + limiter_cache: Final = DualCache() + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(limiter_cache)) + key: Final = UserAPIKeyAuth(api_key=hash_token("sk-one-socket"), max_parallel_requests=1) + logged_rows: Final[list[asyncio.Task[None]]] = [] + + async def admit(user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> None: + await limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=limiter_cache, data=data, call_type=call_type + ) + + async def release(data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object) -> None: + await limiter.async_post_call_success_hook(data=data, user_api_key_dict=user_api_key_dict, response=response) + + async def log_row(**row: object) -> None: + await limiter.async_log_success_event( + kwargs=row, response_obj=None, start_time=row["start_time"], end_time=row["end_time"] + ) + + async def opens_another_session() -> bool: + await asyncio.gather(*logged_rows) + try: + await asyncio.create_task( + admit(key, {"litellm_call_id": str(uuid.uuid4())}, "pass_through_endpoint"), + context=contextvars.Context(), + ) + except ProxyRateLimitError: + return False + return True + + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], answered_by='"response.completed"' + ) + take_frame: Final = client.receive + while_open: Final[list[bool]] = [] + + async def receive() -> Mapping[str, str | bytes]: + message: Final = await take_frame() + if message["type"] == "websocket.disconnect": + while_open.append(await opens_another_session()) + return message + + client.receive = receive + + with ( + _patched_websocket_passthrough_environment( + ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini")]) + ) as proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker, + patch(SUCCESS_HANDLER, new_callable=AsyncMock, side_effect=log_row), + ): + proxy_logging.pre_call_hook = AsyncMock(side_effect=admit) + proxy_logging.post_call_success_hook = AsyncMock(side_effect=release) + worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine)) + ) + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target="wss://api.openai.com/v1/responses", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=key, + forward_headers=False, + endpoint="/openai_passthrough/v1/responses", + accept_websocket=False, + client_frame_model_gate=RecordingFrameModelGate(), + upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=0.05, + ), + timeout=10, + ) + after_close: Final = await opens_another_session() + + assert len(logged_rows) == 1 + assert (while_open, after_close) == ([False], True) + + @pytest.mark.asyncio async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_asked_for_each(): """ diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index 2e2446b8ecf..302c696682d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -668,7 +668,7 @@ REALTIME_FRAMES: Final = [ def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict: - logging_obj: Final = MagicMock(model_call_details={}) + logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id="call-fake") return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( httpx_response=MagicMock(spec=httpx.Response), response_body=response_body,