diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8d0cccd54f4..1d180d9be51 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -541,6 +541,17 @@ _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata" _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) +def _proxy_server_request_with_body( + litellm_params: Mapping[str, object], request_body: Mapping[str, object] +) -> Mapping[str, object]: + proxy_server_request: Final = litellm_params.get("proxy_server_request") + return ( + {"proxy_server_request": {**_LITELLM_PARAMS.validate_python(proxy_server_request), "body": dict(request_body)}} + if isinstance(proxy_server_request, Mapping) + else {} + ) + + def _provider_response_id(source: object) -> str | None: candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None) return candidate if isinstance(candidate, str) and candidate else None @@ -767,7 +778,9 @@ class Logging(LiteLLMLoggingBaseClass): self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None - def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging": + def copy_for_new_call( + self, litellm_call_id: str, start_time: datetime.datetime, request_body: Mapping[str, object] + ) -> "Logging": connection_params: Final = _LITELLM_PARAMS.validate_python(self.litellm_params) litellm_params: Final[dict[str, object]] = { **connection_params, @@ -776,6 +789,7 @@ class Logging(LiteLLMLoggingBaseClass): for key, value in connection_params.items() if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) }, + **_proxy_server_request_with_body(connection_params, request_body), "litellm_call_id": litellm_call_id, } new_call: Final = copy.copy(self) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 70e67099832..0fb65b415b1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -742,15 +742,18 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def _request_owner_call_id(data: Mapping[str, object]) -> str | None: - from litellm.proxy.common_request_processing import request_litellm_call_id - - return request_litellm_call_id(data) +def _tpm_limited_tags(call_stash: RequestRateLimiterStash | None) -> frozenset[str]: + if call_stash is not None: + return call_stash.tpm_limited_tags + connection: Final = get_request_stash() + return connection.tpm_limited_tags if connection is not None and connection.websocket_connection else frozenset() def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: + from litellm.proxy.common_request_processing import request_litellm_call_id + stash: Final = get_or_create_request_stash() - owner_call_id: Final = _request_owner_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data)) + owner_call_id: Final = request_litellm_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data)) if owner_call_id is not None: stash.owner_litellm_call_id = owner_call_id return stash @@ -5067,7 +5070,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): targets: Final = self._collect_tpm_scope_targets( standard_logging_metadata=standard_logging_metadata, kwargs=kwargs, - tpm_limited_tags=stash.tpm_limited_tags if stash is not None else frozenset(), + tpm_limited_tags=_tpm_limited_tags(stash), model_group=reconcile_model.group if reconcile_model is not None else None, ) targets.extend( @@ -5458,12 +5461,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): both this hook and async_log_failure_event end up running in the same flow, only the first release/refund applies. A mid-stream failure relayed here with recovered partial usage settles the reservation at - that usage instead of refunding it. + that usage instead of refunding it. A native Responses WebSocket + connection keeps its slot and reservation through a failed turn, and + its disconnect releases them. """ try: - failed_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(request_data) - stash: Final = get_request_stash_for_call(_request_owner_call_id(failed_request)) - if stash is None: + stash: Final = get_request_stash() + if stash is None or stash.websocket_connection: return await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3c03f13e02d..d812b108403 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -139,6 +139,12 @@ class _ProxyDBLogger(CustomLogger): if self.spend_event_producer is None or not is_offloadable_success(response_obj): await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) return + _claim_budget_reservation( + budget_reservation=_get_budget_reservation_from_metadata( + metadata=get_litellm_metadata_from_kwargs(kwargs=kwargs) + ), + litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(kwargs)), + ) event: Final = build_spend_event( kwargs, response_obj, @@ -170,7 +176,7 @@ class _ProxyDBLogger(CustomLogger): ) -> None: budget_reservation: Final = _claim_budget_reservation( budget_reservation=user_api_key_dict.budget_reservation, - litellm_call_id=_request_call_id(_CALLBACK_DATA.validate_python(request_data)), + litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(request_data)), ) try: await _release_budget_reservation(budget_reservation=budget_reservation) @@ -353,7 +359,7 @@ class _ProxyDBLogger(CustomLogger): _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _claim_budget_reservation( budget_reservation=_get_budget_reservation_from_metadata(metadata=metadata), - litellm_call_id=_logged_call_id(_CALLBACK_DATA.validate_python(kwargs)), + litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(kwargs)), ) if ( isinstance(completion_response, LiteLLMBatch) @@ -709,14 +715,8 @@ def _get_budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dic _CALLBACK_DATA: Final = TypeAdapter(Mapping[str, object]) -def _request_call_id(request_data: Mapping[str, object]) -> str | None: - from litellm.proxy.common_request_processing import request_litellm_call_id - - return request_litellm_call_id(request_data) - - -def _logged_call_id(kwargs: Mapping[str, object]) -> str | None: - call_id: Final = kwargs.get("litellm_call_id") +def _string_call_id(callback_data: Mapping[str, object]) -> str | None: + call_id: Final = callback_data.get("litellm_call_id") return call_id if isinstance(call_id, str) else None diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 529959bc576..21b6fd8ee5e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2007,7 +2007,9 @@ class ResponsesWebSocketStreaming: def _start_turn(self, client_frame: str | None) -> None: turn_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes turn_logging_obj: Final = self.logging_obj.copy_for_new_call( - litellm_call_id=str(uuid.uuid4()), start_time=turn_start + litellm_call_id=str(uuid.uuid4()), + start_time=turn_start, + request_body=_json_object_or_none(client_frame) or {}, ) turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) turn_logging_obj.pre_call(input=client_frame, api_key="") @@ -2036,10 +2038,11 @@ class ResponsesWebSocketStreaming: def _turn_index_for_event(self, event_obj: Mapping[str, object]) -> int: """ An event naming a response id belongs to the turn whose ``response.created`` - carried it. The backend answers accepted turns strictly in order, so an - id-less event goes to the oldest turn still waiting for its ``response.created``, - except an id-less ``error``: that is the backend rejecting the newest pending - ``response.create`` at parse time while an earlier turn is still streaming. + carried it. The backend answers turns strictly in order, so an id-less event, + a rejected turn's ``error`` included, goes to the oldest turn still waiting for + its ``response.created``. The one exception is an id-less ``error`` while a turn + is streaming: the backend rejects a ``response.create`` it cannot parse as soon + as it reads it, so that error belongs to the newest pending turn. """ response_id: Final = _ws_event_response_id(event_obj) turn_ids: Final = tuple(turn.response_id for turn in self._open_turns) @@ -2048,7 +2051,7 @@ class ResponsesWebSocketStreaming: pending: Final = tuple(index for index, turn_id in enumerate(turn_ids) if turn_id is None) if not pending: return 0 - if response_id is None and event_obj.get("type") == "error": + if response_id is None and event_obj.get("type") == "error" and len(pending) < len(turn_ids): return pending[-1] return pending[0] diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 35b1ed16aa2..ccc4f00ed47 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -8267,10 +8267,12 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): reservation = {"reserved_cost": 0.5, "finalized": False} connection = _responses_ws_logging_obj() connection_start = connection.start_time + connection_request = {"url": "ws://proxy/v1/responses", "body": {"first_message": "turn 1 frame"}} connection.litellm_params = { "metadata": {"user_api_key_alias": "codex", "budget_reservation": reservation}, "litellm_metadata": {"model_info": {"id": "dep-1"}}, "litellm_call_id": connection.litellm_call_id, + "proxy_server_request": connection_request, } connection.model_call_details = { **connection.model_call_details, @@ -8279,7 +8281,8 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): } turn_start = datetime.datetime(2026, 10, 9, 17, 24, 47) - turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start) + turn_frame = {"type": "response.create", "input": "turn 2 frame"} + turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start, request_body=turn_frame) turn.litellm_params["metadata"]["status"] = "failure" turn.litellm_params["litellm_metadata"]["model_group"] = "gpt-4o" turn.model_call_details["standard_logging_object"] = {"id": "resp_1"} @@ -8290,6 +8293,14 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1-call" assert turn.model_call_details["first_api_call_start_time"] is None assert turn.litellm_params["metadata"]["budget_reservation"] is reservation + assert turn.model_call_details["litellm_params"]["proxy_server_request"] == { + "url": "ws://proxy/v1/responses", + "body": turn_frame, + } + assert connection.litellm_params["proxy_server_request"] == { + "url": "ws://proxy/v1/responses", + "body": {"first_message": "turn 1 frame"}, + } assert (connection.litellm_call_id, connection.start_time) == ("responses-ws-usage-test", connection_start) assert connection.litellm_params["litellm_call_id"] == "responses-ws-usage-test" assert "status" not in connection.litellm_params["metadata"] diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index d11b50243b5..8c5db21cb9a 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4383,9 +4383,11 @@ async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" - connection_data: Dict[str, Any] = {"model": "gpt-4o-mini", "litellm_call_id": "socket-call"} await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, cache=local_cache, data=connection_data, call_type="" + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4o-mini", "litellm_call_id": "socket-call"}, + call_type=CallTypes.aresponses_websocket.value, ) await handler.async_post_call_failure_hook( @@ -4393,14 +4395,10 @@ async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3() original_exception=Exception("previous_response_not_found"), user_api_key_dict=user_api_key_dict, ) - in_flight_after_turn_failure = handler._gauge_in_flight_from_cache_value( + in_flight_after_turn_failure: Final = handler._gauge_in_flight_from_cache_value( await local_cache.async_get_cache(key=counter_key) ) - await handler.async_post_call_failure_hook( - request_data=connection_data, - original_exception=Exception("socket closed"), - user_api_key_dict=user_api_key_dict, - ) + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) assert in_flight_after_turn_failure == 1 assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 @@ -4436,15 +4434,16 @@ async def _admit_pass_through_with_client_call_id( @pytest.mark.asyncio -async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3(): +async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3(monkeypatch: pytest.MonkeyPatch): _api_key = hash_token("sk-passthrough-client-id-failure") local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" logging_obj = await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict) + monkeypatch.setattr(litellm, "callbacks", [handler]) - await handler.async_post_call_failure_hook( + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_failure_hook( request_data={"litellm_call_id": "client-chosen-id", "litellm_logging_obj": logging_obj}, original_exception=Exception("rejected by a pre-call hook"), user_api_key_dict=user_api_key_dict, @@ -5224,6 +5223,41 @@ async def test_tag_429_names_the_tag_that_is_over_its_limit(monkeypatch): assert "cell-blocked" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_websocket_turns_count_toward_the_connection_tag_tpm_v3(monkeypatch): + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") + _request_stash.set(None) + resolver, _ = _static_tag_limits({"cell-1": TagRateLimit(rpm_limit=None, tpm_limit=1_000)}) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), tag_rate_limit_resolver=resolver + ) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-ws-tag-tpm")) + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4o-mini", "metadata": {"tags": ["cell-1"]}, "litellm_call_id": "socket-call"}, + call_type=CallTypes.aresponses_websocket.value, + ) + tokens_before_turns: Final = int(await local_cache.async_get_cache("{tag:cell-1}:tokens") or 0) + + for turn_call_id in ("turn-1-call", "turn-2-call"): + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": turn_call_id, + "standard_logging_object": {"metadata": {"user_api_key_hash": user_api_key_dict.api_key}}, + "model": "gpt-4o-mini", + }, + response_obj=ModelResponse(usage=Usage(prompt_tokens=60, completion_tokens=60, total_tokens=120)), + start_time=datetime.now(), + end_time=datetime.now(), + ) + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + + assert await local_cache.async_get_cache("{tag:cell-1}:tokens") == tokens_before_turns + 240 + + @pytest.mark.asyncio async def test_tag_object_tpm_limit_enforced_v3(monkeypatch): monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 1dd86c2c4a0..3ae4af4a4bc 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2718,6 +2718,44 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski assert event.response_cost == 0.0125 +@pytest.mark.asyncio +async def test_collector_events_let_only_the_first_turn_settle_a_shared_reservation(tmp_path): + from litellm.litellm_core_utils.core_helpers import ( + budget_reservation_from_metadata, + get_litellm_metadata_from_kwargs, + ) + from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation + from litellm.proxy.spend_tracking.spend_event import spend_event_callback_args + + handler = _RecordingHandler() + consumer = SpendEventConsumer(handler) + address = UnixAddress(path=str(tmp_path / "spend.sock")) + server = await consumer.serve(address) + producer = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback + ) + logger = ProxyDBLogger(producer) + connection_kwargs = _offload_kwargs() + + for turn_call_id in ("turn-1", "turn-2", "turn-3"): + await logger.async_log_success_event( + {**connection_kwargs, "litellm_call_id": turn_call_id}, _offload_response(), datetime.now(), datetime.now() + ) + await producer.close(drain_timeout=5.0) + server.close() + assert await consumer.drain(timeout=5.0) == 0 + + sidecar_kwargs = [spend_event_callback_args(decode_spend_event(line)).kwargs for line in handler.lines] + assert [ + claim_budget_reservation( + budget_reservation=budget_reservation_from_metadata(get_litellm_metadata_from_kwargs(kwargs)), + litellm_call_id=kwargs["litellm_call_id"], + ) + is not None + for kwargs in sidecar_kwargs + ] == [True, False, False] + + @pytest.mark.asyncio async def test_async_log_success_event_keeps_batch_retrieves_in_process(): producer = SpendEventProducer( diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 48bb4340fba..39bcf6710aa 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1907,12 +1907,13 @@ class _TurnLoggingRecorder: self.connection = MagicMock() self.connection.copy_for_new_call = MagicMock(side_effect=self._copy_for_new_call) - def _copy_for_new_call(self, litellm_call_id: str, start_time: object) -> MagicMock: + def _copy_for_new_call(self, litellm_call_id: str, start_time: object, request_body: object) -> MagicMock: from unittest.mock import AsyncMock turn = MagicMock() turn.litellm_call_id = litellm_call_id turn.start_time = start_time + turn.request_body = request_body turn.model_call_details = {} turn.dispatch_success_handlers = AsyncMock() turn.dispatch_failure_handlers = AsyncMock() @@ -3336,6 +3337,9 @@ class TestNativeWebSocketPerTurnLogging: [{"role": "user", "content": "two"}], [{"role": "user", "content": "three"}], ] + assert [turn.request_body for turn in recorder.turns] == [ + json.loads(_response_create(text)) for text in ("one", "two", "three") + ] for turn in recorder.turns: turn.dispatch_failure_handlers.assert_not_awaited() @@ -3402,6 +3406,47 @@ class TestNativeWebSocketPerTurnLogging: assert "max_output_tokens" in str(second.dispatch_failure_handlers.await_args[0][0]) assert second.record_partial_usage_for_failure.call_args[0][0].total_tokens == 0 + @pytest.mark.asyncio + async def test_a_rejected_turn_answered_in_order_fails_and_the_next_turn_keeps_its_response(self): + """OpenAI answers a response.create it parses but rejects (previous_response_not_found) in order, + before it starts the next frame's response. With both turns still waiting, the id-less error is the + older turn's failure and the next response.created belongs to the newer turn.""" + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + recorder = _TurnLoggingRecorder() + handler = _make_streaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=recorder.connection, + request_data={}, + authorized_model="gpt-5.4-mini", + custom_llm_provider="openai", + ) + created, _delta, completed = _completed_turn("resp_2", 40, 4) + rejection = { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "previous_response_not_found", + "message": "Previous response with id 'resp_missing' not found.", + }, + } + + handler._record_client_frame(_response_create("first")) + handler._record_client_frame(_response_create("second")) + for event in (rejection, created, completed): + handler._record_backend_event(json.dumps(event)) + assert await _until(lambda: recorder.logged_turns() == 2) + + first, second = recorder.turns + first.dispatch_success_handlers.assert_not_awaited() + assert "resp_missing" in str(first.dispatch_failure_handlers.await_args[0][0]) + second.dispatch_failure_handlers.assert_not_awaited() + assert [event["type"] for event in second.dispatch_success_handlers.await_args[0][0]] == [ + "response.created", + "response.completed", + ] + @pytest.mark.asyncio async def test_a_failed_turn_is_booked_alone_and_leaves_the_other_turns_successful(self): from unittest.mock import AsyncMock