scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(proxy): drop a Responses WebSocket close refund whose rate-limit window already rolled over

Browse source
This commit is contained in:
mateo-berri 2026-11-09 15:52:12 -08:00
parent e74b8b1a49
commit 163117c4c9
2 changed files with 38 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -499,6 +499,7 @@ class RequestRateLimiterStash:
reserved_tokens: int = 0
reserved_model: RateLimitedModel | None = None
reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset)
ptu_ceiling: PTUTeamCeiling | None = None
ptu_reserved_tokens: int = 0
ptu_reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset)
@ -4031,6 +4032,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
for identity in tpm_response.get("reservation_windows", frozenset())
if identity[0].scope.descriptor_key == PTU_TEAM_DESCRIPTOR_KEY
)
stash.reserved_window_identities = tpm_response.get("reservation_windows", frozenset())
stash.reserved_model = self._rate_limited_model(requested_model)
stash.reserved_scopes = frozenset(
(d["key"], d["value"])
@ -5271,11 +5273,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(request_data)
combined_ops: Final = (
self._build_reservation_aware_tpm_ops(
self._build_project_reservation_ops(
targets=tuple(scope for scope in stash.reserved_scopes if scope[0] != PTU_TEAM_DESCRIPTOR_KEY),
reserved_scopes=stash.reserved_scopes,
actual_tokens=tpm_actual,
reserved_tokens=reserved_tokens,
reservation_window_identities=stash.reserved_window_identities,
)
if reserved_tokens > 0
else ()
@ -5327,9 +5330,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
otpm_reserved,
)
stash.mark_reservation_released()
if combined_ops:
await self.rate_limits.increment_many(combined_ops)
for project_ops in (itpm_ops, otpm_ops, ptu_ops):
unguarded_tpm_ops: Final = [op for op in combined_ops if isinstance(op, CounterIncrement)]
guarded_tpm_ops: Final = tuple(op for op in combined_ops if isinstance(op, WindowGuardedIncrement))
if unguarded_tpm_ops:
await self.rate_limits.increment_many(unguarded_tpm_ops)
for project_ops in (guarded_tpm_ops, itpm_ops, otpm_ops, ptu_ops):
if isinstance(project_ops, list):
await self.rate_limits.increment_many(project_ops)
elif project_ops:

29
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -6014,6 +6014,35 @@ async def test_a_responses_websocket_close_whose_refund_fails_still_frees_its_sl
assert handler._gauge_in_flight_from_cache_value(_local(handler, RateLimitScope("api_key", api_key).gauge)) == 0
@pytest.mark.asyncio
async def test_a_responses_websocket_that_closes_in_a_later_minute_leaves_that_minutes_key_tpm_alone() -> None:
cache: Final = DualCache()
clock = [datetime(2026, 11, 9, 12, 0, 0)]
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), time_provider=lambda: clock[-1]
)
proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler
api_key: Final = hash_token("sk-responses-websocket-late-close")
auth: Final = UserAPIKeyAuth(api_key=api_key, tpm_limit=1_000, metadata={"default_estimated_output_tokens": 60})
await _admit_responses_websocket(handler, cache, auth, "socket")
clock.append(clock[-1] + timedelta(seconds=handler.window_size + 1))
await handler.atomic_check_and_increment_by_n(
descriptors=[
{
"key": "api_key",
"value": api_key,
"rate_limit": {"tokens_per_unit": 1_000, "window_size": handler.window_size},
}
],
increments=[{"tokens": 20}],
)
await proxy_logging_obj.arelease_responses_websocket_admission(auth)
assert int(_local(handler, _counter("api_key", api_key)) or 0) == 20
@pytest.mark.asyncio
async def test_a_closed_responses_websocket_refunds_a_project_reservation_stashed_without_its_window() -> None:
cache: Final = DualCache()