From 668e0142664ac6fa6ab6a6c53925c79cf20b1b6e Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 9 Sep 2026 00:11:05 +0000 Subject: [PATCH] fix(proxy): stop reservation lease renewal once the request task is gone A streaming /v1/messages client disconnect skips reconciliation, so the lease kept the orphaned reservation alive until request_timeout instead of the plain 60s counter TTL the base branch had. Stop renewing when the request task that took the reservation is done. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/budget_reservation.py | 10 ++++++-- .../proxy/test_budget_reservation.py | 23 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index bd9eb154839..c0af6685600 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -116,6 +116,7 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c budget_reservation=budget_reservation, counter_keys=counter_keys, interval=spend_counter_cache.redis_cache.default_ttl / 2, + request_task=asyncio.current_task(), ) ) _lease_renewals.add(task) @@ -123,14 +124,19 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c async def _renew_reservation_lease( - budget_reservation: Mapping[str, object], counter_keys: frozenset[str], interval: float + budget_reservation: Mapping[str, object], + counter_keys: frozenset[str], + interval: float, + request_task: asyncio.Task[object] | None, ) -> None: + """Stops on finalization or once the request task that took the reservation is gone, so a + disconnect path that skipped reconciliation falls back to the plain counter TTL.""" from litellm.proxy.proxy_server import refresh_spend_counter_ttl deadline: Final = time.monotonic() + litellm.request_timeout while time.monotonic() < deadline: await asyncio.sleep(interval) - if budget_reservation.get("finalized") is True: + if budget_reservation.get("finalized") is True or (request_task is not None and request_task.done()): return for counter_key in counter_keys: await refresh_spend_counter_ttl(counter_key=counter_key) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 3c33ca11c91..626c82ea397 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2304,6 +2304,29 @@ async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_reservation_lease_stops_when_request_task_ends_without_reconciling( + spend_counter_state, +): + """A request whose task ends without reconciling (client disconnect path that skips the + cost callbacks) must not keep renewing: the counter falls back to its plain TTL instead of + pinning the reservation until the request timeout.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-orphan", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease-orphan" + + reservation = await asyncio.create_task(_reserve(valid_token, 0.6, key_cache, proxy_logging_obj)) + assert reservation is not None + assert reservation["finalized"] is False + + await asyncio.sleep(0.5) + assert redis_cache.refresh_count == 0 + assert await redis_cache.async_get_cache(key=counter_key) is None + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state,