From 8c0fe877b9bbfa0a6934d8ec9373fcaf029a2082 Mon Sep 17 00:00:00 2001 From: Tanishq <74068132+tanishqbaweja@users.noreply.github.com> Date: Mon, 26 Oct 2026 03:41:44 +0530 Subject: [PATCH] fix(proxy): use current spend for virtual key soft budget alerts --- litellm/proxy/auth/auth_checks.py | 15 ++- .../spend/test_key_budget_lockout.py | 73 ++++++++++++++- tests/unit/proxy/auth/test_auth_checks.py | 91 +++++++++++++++++++ 3 files changed, 174 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 79e385c376a..007edef8470 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6801,16 +6801,25 @@ async def virtual_key_soft_budget_check( """ - if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: + if not valid_token.soft_budget: + return + + from litellm.proxy.proxy_server import get_current_spend + + spend: Final = await get_current_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + ) + if spend >= valid_token.soft_budget: verbose_proxy_logger.debug( "Crossed Soft Budget for token %s, spend %s, soft_budget %s", valid_token.token, - valid_token.spend, + spend, valid_token.soft_budget, ) call_info: Final = CallInfo( token=valid_token.token, - spend=valid_token.spend, + spend=spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, diff --git a/tests/integration/spend/test_key_budget_lockout.py b/tests/integration/spend/test_key_budget_lockout.py index a3d1d3d0a78..f18b67ee7fa 100644 --- a/tests/integration/spend/test_key_budget_lockout.py +++ b/tests/integration/spend/test_key_budget_lockout.py @@ -1,5 +1,6 @@ import json import os +import re import uuid from hashlib import sha256 from pathlib import Path @@ -7,9 +8,10 @@ from typing import Final import httpx import pytest -from integration._support.client import Gateway, eventually, object_value +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import OwnedState, owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, JsonValue, TypeAdapter @@ -201,6 +203,73 @@ def test_an_exhausted_key_is_refused_inference_but_can_still_read_its_own_info(g assert own["max_budget"] == 0.06 +def test_soft_budget_alert_reads_spend_from_another_proxy(gateway: Gateway, tmp_path: Path) -> None: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "general_settings": { + **object_value(base["general_settings"]), + "alerting": ["slack"], + "alert_types": ["budget_alerts"], + }, + } + path: Final = tmp_path / "soft-budget-alerting.yaml" + path.write_text(yaml.safe_dump(config)) + email: Final = "budget-owner@example.test" + alias: Final = f"soft-budget-{uuid.uuid4().hex}" + with gateway.scenario() as scenario, wire_server(lambda _: Reply(body=b"ok", content_type="text/plain")) as alerts: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + user_id: Final = scenario.user(user_email=email) + key: Final = scenario.key(user_id=user_id, models=[model], key_alias=alias, soft_budget=0.1) + digest: Final = sha256(key.encode()).hexdigest() + overrides: Final = { + "SLACK_WEBHOOK_URL": f"{alerts.url}/alerts", + "PROXY_BASE_URL": "http://proxy.integration.test", + } + with owned_proxy(gateway, tmp_path, overrides, config=path, state=OwnedState.SHARED_WITH_WORKER) as candidate: + first: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "warm the alerting proxy"}]}, + key=key, + ) + assert first.status_code == 200, first.text + assert float(first.headers["x-litellm-response-cost"]) == pytest.approx(20 * 0.001 + 20 * 0.002) + eventually(lambda: _spend(gateway, digest), lambda spend: spend == pytest.approx(0.06), seconds=70) + second: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "cross the budget on the other proxy"}]}, + key=key, + ) + assert second.status_code == 200, second.text + eventually(lambda: _spend(gateway, digest), lambda spend: spend == pytest.approx(0.12), seconds=70) + assert alerts.received.qsize() == 0 + third: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "check the shared spend"}]}, + key=key, + ) + assert third.status_code == 200, third.text + eventually(alerts.received.qsize, lambda count: count >= 1, seconds=40) + requests: Final = alerts.drain() + assert len(requests) == 1, requests + assert (requests[0].method, requests[0].target) == ("POST", "/alerts") + payload: Final = JSON_OBJECT.validate_json(requests[0].body) + message: Final = string_value(payload["text"]) + normalized: Final = re.sub(r"Timestamp: `\d{2}:\d{2}:\d{2}`", "Timestamp: ``", message) + assert {**payload, "text": normalized} == { + "text": ( + "Alert type: `budget_alerts`\nLevel: `High`\nTimestamp: ``\n\n" + "Message: Soft Budget Crossed: Total Soft Budget:`0.1`\n\n" + f"*spend:* `0.12`\n*soft_budget:* `0.1`\n*user_id:* `{user_id}`\n" + f"*user_email:* `{email}`\n*key_alias:* `{alias}`\n*event_group:* `key`\n" + "\n\nProxy URL: `http://proxy.integration.test`" + ) + } + + def _bounded_chat(gateway: Gateway, model: str, key: str, content: str | None = None) -> httpx.Response: return gateway.request( "POST", diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e1afa81d372..d3ed39b498d 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -48,6 +48,8 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.utils import ProxyLogging from litellm.proxy.utils import CallInfo +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER @@ -1063,6 +1065,95 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert): ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" +@pytest.mark.parametrize( + "cached_spend, counter_spend, soft_budget, expected_spend", + [ + (0.0, 100.0, 50.0, 100.0), + (0.0, 50.0, 50.0, 50.0), + (100.0, 25.0, 50.0, None), + (100.0, 0.0, 50.0, None), + (100.0, None, 50.0, 100.0), + (0.0, 100.0, None, None), + (0.0, 100.0, 0.0, None), + ], +) +@pytest.mark.asyncio +async def test_soft_budget_alert_uses_current_spend_without_blocking_requests( + monkeypatch: pytest.MonkeyPatch, + cached_spend: float, + counter_spend: float | None, + soft_budget: float | None, + expected_spend: float | None, +) -> None: + from litellm.proxy import proxy_server + + counter: Final = DualCache() + monkeypatch.setattr(proxy_server, "spend_counter_cache", counter) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setenv("WEBHOOK_URL", "https://alerts.example.test/budget") + if counter_spend is not None: + counter.in_memory_cache.set_cache(key="spend:key:soft-budget-key", value=counter_spend) + received: Final[list[httpx.Request]] = [] + + def receive(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response(200, text="ok") + + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(receive)) + logging: Final = ProxyLogging(UserApiKeyCache()) + logging.alerting = ["webhook"] + logging.slack_alerting_instance = SlackAlerting(alerting=["webhook"], async_http_handler=handler) + token: Final = UserAPIKeyAuth( + token="soft-budget-key", + spend=cached_spend, + soft_budget=soft_budget, + user_id="budget-user", + org_id="budget-org", + key_alias="budget-key", + ) + tasks_before: Final = asyncio.all_tasks() + try: + await virtual_key_soft_budget_check( + token, logging, LiteLLM_UserTable(user_id="budget-user", user_email="owner@example.test") + ) + await asyncio.gather(*(asyncio.all_tasks() - tasks_before)) + expected: Final = ( + [] + if expected_spend is None + else [ + { + "spend": expected_spend, + "max_budget": None, + "soft_budget": soft_budget, + "token": "soft-budget-key", + "customer_id": None, + "user_id": "budget-user", + "team_id": None, + "team_alias": None, + "organization_id": "budget-org", + "user_email": "owner@example.test", + "key_alias": "budget-key", + "projected_exceeded_date": None, + "projected_spend": None, + "event_group": "key", + "alert_emails": None, + "max_budget_alert_emails": None, + "event": "soft_budget_crossed", + "event_message": f"Soft Budget Crossed: Total Soft Budget:`{soft_budget}`", + } + ] + ) + assert [json.loads(request.content) for request in received] == expected + assert token.spend == cached_spend + finally: + pending: Final = asyncio.all_tasks() - tasks_before + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + await handler.close() + + @pytest.mark.parametrize( "spend, soft_budget, expect_alert, metadata, expected_alert_emails", [