From 163117c4c9df462a7d30961daf2c79127c7fe4d1 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Mon, 9 Nov 2026 15:52:12 -0800
Subject: [PATCH] fix(proxy): drop a Responses WebSocket close refund whose
 rate-limit window already rolled over

---
 .../hooks/parallel_request_limiter_v3.py      | 13 ++++++---
 .../hooks/test_parallel_request_limiter_v3.py | 29 +++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py
index 5ff7a797125..14b39343d43 100644
--- a/litellm/proxy/hooks/parallel_request_limiter_v3.py
+++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py
@@ -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:
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 1020b45b629..18d6aee362c 100644
--- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
+++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
@@ -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()