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,