fix: settle memory round spend before budget rechecks

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,

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

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

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():

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]

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"<html>private provider error</html>",
+        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

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()