From 84623c9ea4d4f8a9792e0ebb61835bfbbb691157 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 18:14:02 -0800 Subject: [PATCH] fix(proxy): keep the Responses WebSocket release method name and the failure hook's log line ProxyLogging.arelease_responses_websocket_slot shipped in v1.107.0, so it keeps its name. The failure hook logs its old "Error releasing TPM reservation on post-call failure" line again, and a failure hook whose refund fails still frees its slot without raising --- .../hooks/parallel_request_limiter_v3.py | 21 +++++++--- .../pass_through_endpoints.py | 2 +- .../proxy/response_api_endpoints/endpoints.py | 2 +- litellm/proxy/utils.py | 2 +- .../hooks/test_parallel_request_limiter.py | 2 +- .../hooks/test_parallel_request_limiter_v3.py | 40 ++++++++++++++++--- .../test_pass_through_endpoints.py | 12 +++--- .../response_api_endpoints/test_endpoints.py | 18 ++++----- 8 files changed, 69 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 14b39343d43..d0376d1133b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5222,28 +5222,39 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): relayed here with recovered partial usage settles the reservation at that usage instead of refunding it. """ - await self._release_stashed_admission( + release_error: Final = await self._release_stashed_admission( request_data, # pyright: ignore[reportUnknownArgumentType] # request_data is a bare dict user_api_key_dict.parent_otel_span, ) + if release_error is not None: + verbose_proxy_logger.error( + "Error releasing TPM reservation on post-call failure: %s", release_error, exc_info=release_error + ) @with_service_target("rate_limits") async def async_release_responses_websocket_admission(self, user_api_key_dict: UserAPIKeyAuth) -> None: """Free the slot and the TPM reservation a Responses WebSocket took at admission, once it closes. Its turns log under their own call ids and charge their own usage, so the reservation is refunded in full.""" - await self._release_stashed_admission(MappingProxyType({}), user_api_key_dict.parent_otel_span) + release_error: Final = await self._release_stashed_admission( + MappingProxyType({}), user_api_key_dict.parent_otel_span + ) + if release_error is not None: + verbose_proxy_logger.error( + "Error releasing a closed Responses WebSocket's admission: %s", release_error, exc_info=release_error + ) async def _release_stashed_admission( self, request_data: Mapping[str, object], parent_otel_span: Span | None - ) -> None: + ) -> Exception | None: try: stash: Final = get_request_stash() if stash is None: - return + return None await self._release_stashed_parallel_slot(stash, parent_otel_span) await self._settle_stashed_reservations(stash, request_data, parent_otel_span) except Exception as e: - verbose_proxy_logger.exception("Error releasing the request's rate limit admission: %s", e) + return e + return None async def _settle_stashed_reservations( self, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 022e95389b2..cdcb42ab761 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3598,7 +3598,7 @@ async def websocket_passthrough_request( if _client_socket_is_open(websocket): await websocket.close(code=1011, reason="WebSocket passthrough error") finally: - await proxy_logging_obj.arelease_responses_websocket_admission(user_api_key_dict) + await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict) if _client_socket_is_open(websocket): await websocket.close() diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index b29e8f82563..26d8d0c3e1b 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2712,7 +2712,7 @@ async def _responses_websocket_session( try: failure_row: Final = await relay_and_bill_failure() finally: - await proxy_logging_obj.arelease_responses_websocket_admission(user_api_key_dict) + await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict) if failure_row is not None: with contextlib.suppress(Exception): await asyncio.shield(failure_row) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index da2dead53cf..ea7d7979d6d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5293,7 +5293,7 @@ class ProxyLogging: arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests - async def arelease_responses_websocket_admission(self, user_api_key_dict: UserAPIKeyAuth) -> None: + async def arelease_responses_websocket_slot(self, user_api_key_dict: UserAPIKeyAuth) -> None: """Release what a Responses WebSocket took at admission. Only call this once that admission succeeded: the legacy limiter keeps no record of its acquisitions and decrements unconditionally.""" match self.get_proxy_hook("parallel_request_limiter"): diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 4acf450b713..0f5d672b778 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -441,6 +441,6 @@ async def test_a_rejected_realtime_attempt_does_not_free_a_slot_it_never_took(): after_rejection = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) assert after_rejection["current_requests"] == 1, after_rejection - await proxy_logging_obj.arelease_responses_websocket_admission(auth) + await proxy_logging_obj.arelease_responses_websocket_slot(auth) after_socket_close = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) assert after_socket_close["current_requests"] == 0, after_socket_close 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 18d6aee362c..99ae2b4244e 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5914,7 +5914,7 @@ async def test_a_closed_responses_websocket_gives_back_its_slot_and_tpm_reservat assert 2 * reserved > 100 assert handler._gauge_in_flight_from_cache_value(_local(handler, parallel_key)) == 1 - await proxy_logging_obj.arelease_responses_websocket_admission(auth) + await proxy_logging_obj.arelease_responses_websocket_slot(auth) assert int(_local(handler, tokens_key) or 0) == 0 assert handler._gauge_in_flight_from_cache_value(_local(handler, parallel_key)) == 0 @@ -5936,7 +5936,7 @@ async def test_a_responses_websocket_admission_reservation_is_refunded_exactly_o assert int(_local(handler, tokens_key) or 0) > 0 def close() -> Awaitable[None]: - return proxy_logging_obj.arelease_responses_websocket_admission(auth) + return proxy_logging_obj.arelease_responses_websocket_slot(auth) def failure_row() -> Awaitable[None]: return handler.async_post_call_failure_hook( @@ -5977,7 +5977,7 @@ async def test_a_closed_responses_websocket_gives_back_its_project_token_reserva await _admit_responses_websocket(handler, cache, auth, "project-socket") assert int(_local(handler, project_keys[1]) or 0) > 0 - await proxy_logging_obj.arelease_responses_websocket_admission(auth) + await proxy_logging_obj.arelease_responses_websocket_slot(auth) assert [int(_local(handler, key) or 0) for key in project_keys] == [0, 0] @@ -5992,6 +5992,11 @@ class _RefundFailingRateLimits(FakeRateLimits): raise ConnectionError("rate limit store unreachable") await super().increment_many(increments) + def set_local(self, entry: RateLimitEntry, value: LocalValue, ttl_seconds: float | None) -> None: + if self.refunds_fail and isinstance(entry, RateLimitCounter) and entry.kind == "tokens": + raise ConnectionError("rate limit store unreachable") + super().set_local(entry, value, ttl_seconds) + @pytest.mark.asyncio async def test_a_responses_websocket_close_whose_refund_fails_still_frees_its_slot() -> None: @@ -6009,7 +6014,30 @@ async def test_a_responses_websocket_close_whose_refund_fails_still_frees_its_sl await _admit_responses_websocket(handler, cache, auth, "socket") fakes.refunds_fail = True - await proxy_logging_obj.arelease_responses_websocket_admission(auth) + await proxy_logging_obj.arelease_responses_websocket_slot(auth) + + assert handler._gauge_in_flight_from_cache_value(_local(handler, RateLimitScope("api_key", api_key).gauge)) == 0 + + +@pytest.mark.asyncio +async def test_a_proxy_side_rejection_whose_refund_fails_still_frees_its_slot() -> None: + cache: Final = DualCache() + fakes: Final = _RefundFailingRateLimits() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), rate_limit_store=fakes.store() + ) + api_key: Final = hash_token("sk-proxy-rejection-refund-fails") + auth: Final = UserAPIKeyAuth( + api_key=api_key, tpm_limit=1_000, max_parallel_requests=1, metadata={"default_estimated_output_tokens": 60} + ) + await _admit_responses_websocket(handler, cache, auth, "socket") + fakes.refunds_fail = True + + await handler.async_post_call_failure_hook( + request_data={"model": "gpt-4o-mini", "litellm_call_id": "socket"}, + original_exception=Exception("guardrail rejected the request"), + user_api_key_dict=auth, + ) assert handler._gauge_in_flight_from_cache_value(_local(handler, RateLimitScope("api_key", api_key).gauge)) == 0 @@ -6038,7 +6066,7 @@ async def test_a_responses_websocket_that_closes_in_a_later_minute_leaves_that_m increments=[{"tokens": 20}], ) - await proxy_logging_obj.arelease_responses_websocket_admission(auth) + await proxy_logging_obj.arelease_responses_websocket_slot(auth) assert int(_local(handler, _counter("api_key", api_key)) or 0) == 20 @@ -6055,7 +6083,7 @@ async def test_a_closed_responses_websocket_refunds_a_project_reservation_stashe stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset({scope}) - await proxy_logging_obj.arelease_responses_websocket_admission( + await proxy_logging_obj.arelease_responses_websocket_slot( UserAPIKeyAuth(api_key=hash_token("sk-responses-websocket-windowless")) ) 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 fde486bf945..07898208286 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 @@ -7398,7 +7398,7 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.arelease_responses_websocket_admission = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) @@ -7451,7 +7451,7 @@ async def test_websocket_passthrough_propagates_active_trace_context( mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.arelease_responses_websocket_admission = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) @@ -7518,7 +7518,7 @@ async def test_websocket_passthrough_never_forwards_caller_credentials_upstream( mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.arelease_responses_websocket_admission = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) @@ -7612,7 +7612,7 @@ def _patched_websocket_passthrough_environment(upstream_ws): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.arelease_responses_websocket_admission = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) @@ -8639,7 +8639,7 @@ async def test_websocket_passthrough_frees_one_legacy_parallel_slot_only_when_th ) ) proxy_logging.pre_call_hook = AsyncMock(side_effect=admit) - proxy_logging.arelease_responses_websocket_admission = proxy_logging_obj.arelease_responses_websocket_admission + proxy_logging.arelease_responses_websocket_slot = proxy_logging_obj.arelease_responses_websocket_slot worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine)) ) @@ -8718,7 +8718,7 @@ async def test_websocket_passthrough_gives_back_its_tpm_reservation_when_the_ses patch(SUCCESS_HANDLER, new_callable=AsyncMock, side_effect=log_row), ): proxy_logging.pre_call_hook = AsyncMock(side_effect=admit) - proxy_logging.arelease_responses_websocket_admission = proxy_logging_obj.arelease_responses_websocket_admission + proxy_logging.arelease_responses_websocket_slot = proxy_logging_obj.arelease_responses_websocket_slot worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine)) ) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 9a64041931e..24ad068e6df 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3564,7 +3564,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_responses_websocket_admission = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -3654,7 +3654,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=lambda **kwargs: row_written.set()) - proxy_logging_obj.arelease_responses_websocket_admission = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) with ( patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above @@ -3683,7 +3683,7 @@ class TestResponsesWSFirstFrameModelAuth: await asyncio.wait_for(row_written.wait(), timeout=1) proxy_logging_obj.post_call_failure_hook.assert_awaited_once() - proxy_logging_obj.arelease_responses_websocket_admission.assert_awaited_once() + proxy_logging_obj.arelease_responses_websocket_slot.assert_awaited_once() booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs assert booked["original_exception"] is (crash if ending == "crash" else refusal) assert booked["request_data"]["litellm_metadata"] == { @@ -3749,7 +3749,7 @@ class TestResponsesWSFirstFrameModelAuth: async def release_over_the_network(user_api_key_dict): await asyncio.sleep(0.01) - proxy_logging_obj.arelease_responses_websocket_admission = AsyncMock(side_effect=release_over_the_network) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(side_effect=release_over_the_network) with ( patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above @@ -3783,7 +3783,7 @@ class TestResponsesWSFirstFrameModelAuth: crash_close = call(code=1011, reason="Internal server error") crashed = relay_outcome == "crashes_despite_the_cancel" assert ws.close.await_args_list == ([limit_close, crash_close] if crashed else [limit_close]) - proxy_logging_obj.arelease_responses_websocket_admission.assert_awaited_once() + proxy_logging_obj.arelease_responses_websocket_slot.assert_awaited_once() assert not row_written.is_set() row_may_finish.set() await asyncio.wait_for(row_written.wait(), timeout=1) @@ -3836,7 +3836,7 @@ class TestResponsesWSFirstFrameModelAuth: side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=slow_failure_hook) - proxy_logging_obj.arelease_responses_websocket_admission = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) with ( patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above @@ -3861,7 +3861,7 @@ class TestResponsesWSFirstFrameModelAuth: responses_websocket_endpoint(websocket=ws, model=None, user_api_key_dict=MagicMock()) ) await asyncio.wait_for(row_started.wait(), timeout=1) - proxy_logging_obj.arelease_responses_websocket_admission.assert_awaited_once() + proxy_logging_obj.arelease_responses_websocket_slot.assert_awaited_once() finished_early, _ = await asyncio.wait({endpoint}, timeout=0.05) assert endpoint not in finished_early row_may_finish.set() @@ -3913,7 +3913,7 @@ class TestResponsesWSFirstFrameModelAuth: } ) proxy_logging_obj.post_call_failure_hook = AsyncMock() - proxy_logging_obj.arelease_responses_websocket_admission = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) with ( patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above @@ -3968,7 +3968,7 @@ class TestResponsesWSFirstFrameModelAuth: ) proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_responses_websocket_admission = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with (