From 5f1d9a3accd12a97c91aa4986d089d46979efa75 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 15 Sep 2026 12:37:40 -0700 Subject: [PATCH] fix: settle memory round spend before budget rechecks --- litellm/proxy/common_request_processing.py | 16 ++++++- .../proxy/hooks/proxy_track_cost_callback.py | 10 ++++ litellm/proxy/memory/gateway.py | 4 +- litellm/proxy/memory/transport.py | 24 +++++++++- .../hooks/test_proxy_track_cost_callback.py | 47 +++++++++++++++++++ .../proxy/memory/test_memory_v2_boundaries.py | 31 ++++++++++++ .../proxy/memory/test_transport.py | 35 ++++++++++++++ 7 files changed, 164 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d0ff248d3de..ab7cac212f3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2361,11 +2361,24 @@ class ProxyBaseLLMRequestProcessing: ) -> Response: from litellm.proxy.auth.user_api_key_auth import ( _run_centralized_common_checks, # pyright: ignore[reportPrivateUsage] # Reuse the authenticated admission and budget checks. + _run_post_custom_auth_checks, # pyright: ignore[reportPrivateUsage] # Reuse expiry and model-budget checks on the already authenticated identity. + _should_skip_budget_checks, # pyright: ignore[reportPrivateUsage] # Preserve free-model budget exemptions. ) processor: Final = ProxyBaseLLMRequestProcessing(data=body) headers: Final = Response() try: + from litellm.proxy.auth.auth_checks import ( + _virtual_key_max_budget_check, # pyright: ignore[reportPrivateUsage] # Reuse key-budget enforcement for every billed round. + ) + + await _run_post_custom_auth_checks( + auth, inner_request, body, inner_request.url.path, auth.parent_otel_span + ) + if not _should_skip_budget_checks( + body, inner_request.url.path, inner_request, llm_router, auth.team_id + ): + await _virtual_key_max_budget_check(auth, proxy_logging_obj) await _run_centralized_common_checks(auth, inner_request, body, inner_request.url.path) result: Final = await processor._process_llm_request( request=inner_request, @@ -2431,9 +2444,10 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, ) - from litellm.proxy.memory.transport import in_gateway_round + from litellm.proxy.memory.transport import begin_gateway_accounting, in_gateway_round if in_gateway_round() and route_type in ("acompletion", "aresponses", "anthropic_messages"): + begin_gateway_accounting(logging_obj.litellm_call_id) self.data["caching"] = False self.data["cache"] = { # mutable-ok: The existing inference pipeline consumes native cache controls. "no-cache": True, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3a61f773001..17a679c36ad 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -105,6 +105,16 @@ class _ProxyDBLogger(CustomLogger): async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime ) -> None: + from litellm.proxy.memory.transport import gateway_accounting + + accounting: Final = gateway_accounting(str(kwargs.get("litellm_call_id"))) + if accounting is not None: + try: + await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) + finally: + if not accounting.done(): + accounting.set_result(None) + return 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 diff --git a/litellm/proxy/memory/gateway.py b/litellm/proxy/memory/gateway.py index 55cc870e00e..6812755293c 100644 --- a/litellm/proxy/memory/gateway.py +++ b/litellm/proxy/memory/gateway.py @@ -196,7 +196,9 @@ class GatewayMemoryLoop: if status >= 400: raise HTTPException( status_code=status, - detail=str(_OBJECT.validate_json(await call.read()).get("error", "Gateway model call failed")), + detail="The gateway model request was rate limited" + if status == 429 + else "The gateway model request failed", headers={ # mutable-ok: FastAPI's HTTPException accepts a native header dictionary. name.decode("latin-1"): value.decode("latin-1") for name, value in start.headers diff --git a/litellm/proxy/memory/transport.py b/litellm/proxy/memory/transport.py index 62bfc52b82b..1dab8907fcd 100644 --- a/litellm/proxy/memory/transport.py +++ b/litellm/proxy/memory/transport.py @@ -8,7 +8,7 @@ from typing import Final, TypeAlias from uuid import uuid4 import anyio -from fastapi import Request +from fastapi import HTTPException, Request from pydantic import BaseModel, ConfigDict, TypeAdapter from starlette.responses import Response, StreamingResponse @@ -16,6 +16,9 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import wait_for_request_parallel_release _GATEWAY_ROUND: Final[ContextVar[int | None]] = ContextVar("litellm_gateway_memory_round", default=None) +_ROUND_ACCOUNTING: Final[ContextVar[tuple[str, asyncio.Future[None]] | None]] = ContextVar( + "litellm_memory_round_accounting", default=None +) _OBJECT: Final = TypeAdapter(dict[str, object]) _ROUND_HEADERS: Final = frozenset(("idempotency-key", "x-request-id", "x-litellm-call-id")) RoundExecutor: TypeAlias = Callable[ @@ -23,6 +26,17 @@ RoundExecutor: TypeAlias = Callable[ ] # mutable-ok: The processor mutates its fresh request copy. +def begin_gateway_accounting(call_id: str) -> None: + _ROUND_ACCOUNTING.set((call_id, asyncio.get_running_loop().create_future())) + + +def gateway_accounting(call_id: str | None = None) -> asyncio.Future[None] | None: + accounting: Final = _ROUND_ACCOUNTING.get() + if accounting is None or (call_id is not None and accounting[0] != call_id): + return None + return accounting[1] + + def in_gateway_round() -> bool: return _GATEWAY_ROUND.get() is not None @@ -110,6 +124,14 @@ class GatewayRound: await self.writer.send(bytes(response.body)) if response.background is not None: await response.background() + accounting: Final = gateway_accounting() + if accounting is not None: + try: + await asyncio.wait_for(asyncio.shield(accounting), timeout=15) + except TimeoutError as exc: + raise HTTPException( + status_code=503, detail="Memory could not confirm model spend; retry later" + ) from exc await wait_for_request_parallel_release() except BaseException as exc: if not self.started.done(): diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 395ce68ec54..e55434c9321 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -2393,3 +2393,50 @@ async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_ == "/chat/completions: Invalid model name passed in. Call `/v1/models` to view available models for your key." ) assert error_information["error_class"] == "ProxyModelNotFoundError" + + +@pytest.mark.asyncio +async def test_memory_accounting_finishes_counters_before_a_sidecar_can_defer_them(tmp_path): + from starlette.requests import Request + from starlette.responses import Response + + from litellm.proxy.memory.transport import begin_gateway_accounting, gateway_accounting, gateway_round + + fallback = AsyncMock() + producer = SpendEventProducer( + address=UnixAddress(path=str(tmp_path / "absent.sock")), + on_unavailable="fallback", + buffer_size=10, + connect_timeout=1.0, + fallback=fallback, + ) + logger = _ProxyDBLogger(producer) + + async def execute(request, body, auth): + begin_gateway_accounting("call-1") + pending = gateway_accounting() + await logger.async_log_success_event( + {**_offload_kwargs(), "litellm_call_id": "nested-call"}, + _offload_response(), + datetime.now(), + datetime.now(), + ) + assert not pending.done() + await logger.async_log_success_event(_offload_kwargs(), _offload_response(), datetime.now(), datetime.now()) + return Response(b"done") + + async def run(): + async with gateway_round( + execute, + Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}), + {}, + UserAPIKeyAuth(), + ) as call: + assert await call.read() == b"done" + + row, counters, _ = await _spend_row_written_by(run) + assert row["spend"] == 0.0125 + assert counters["response_cost"] == 0.0125 + await producer.close(drain_timeout=1) + fallback.assert_awaited_once() + assert b"nested-call" in fallback.call_args.args[0] diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py index cd8f961b151..58e32e2d01f 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py @@ -956,3 +956,34 @@ async def test_previous_response_uses_owned_upstream_and_pending_tool_outputs(pr assert body["previous_response_id"] == "native-last" assert pending in body["input"] assert any(item.get("call_id") == "client-call" for item in body["input"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + ( + b"", + b"private provider error", + b'data: {"error":"private provider error"}\n\n', + b'{"error":"private provider error"}', + ), +) +@pytest.mark.parametrize("status", (429, 502)) +async def test_error_response_retains_status_and_retry_after_without_parsing_provider_body( + prisma_edge: MagicMock, body: bytes, status: int +) -> None: + execute = AsyncMock(return_value=Response(body, status_code=status, headers={"retry-after": "7"})) + loop = GatewayMemoryLoop( + execute, + request(), + {"messages": [{"role": "user", "content": "hi"}]}, + "acompletion", + store(prisma_edge), + UserAPIKeyAuth(), + ) + with pytest.raises(HTTPException) as error: + async for _ in loop.run(): + pass + assert error.value.status_code == status + assert error.value.headers["retry-after"] == "7" + assert "private provider" not in error.value.detail diff --git a/tests/test_litellm/proxy/memory/test_transport.py b/tests/test_litellm/proxy/memory/test_transport.py index 38b610e3101..dc69c61cd7d 100644 --- a/tests/test_litellm/proxy/memory/test_transport.py +++ b/tests/test_litellm/proxy/memory/test_transport.py @@ -106,3 +106,38 @@ async def test_memory_rounds_share_one_rpm_admission_but_keep_token_limits(): pass assert token_limited.value.status_code == 429 assert "tokens" in token_limited.value.detail + + +@pytest.mark.asyncio +async def test_round_waits_for_accounting_before_it_can_continue(): + from starlette.responses import Response + + from litellm.proxy.memory.transport import begin_gateway_accounting, gateway_accounting + + delivered = asyncio.Event() + accounting = [] + + async def execute(inner, data, auth): + begin_gateway_accounting(data["litellm_call_id"]) + accounting.append(gateway_accounting()) + return Response(b"answer") + + async def read(): + async with gateway_round( + execute, + Request({"type": "http", "method": "POST", "path": "/v1/messages", "headers": []}), + {}, + UserAPIKeyAuth(), + ) as call: + async for chunk in call.chunks(): + assert chunk == b"answer" + delivered.set() + + task = asyncio.create_task(read()) + try: + await asyncio.wait_for(delivered.wait(), timeout=1) + assert not task.done() + accounting[0].set_result(None) + await asyncio.wait_for(task, timeout=1) + finally: + task.cancel()