Merge 9f6aaaf334 into e4832c5d06

litellm/proxy/auth/auth_checks.py

@@ -6824,16 +6824,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,

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: `<clock>`", message)
+            assert {**payload, "text": normalized} == {
+                "text": (
+                    "Alert type: `budget_alerts`\nLevel: `High`\nTimestamp: `<clock>`\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",

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",
     [

tests/unit/proxy/auth/test_user_api_key_auth.py

@@ -642,80 +642,102 @@ async def test_auth_with_form_data_and_model():
     assert response.models == ["gpt-4"], "Model from virtual key should be preserved"
 
 
+@pytest.mark.parametrize("cached_spend", [0.0, 15.0])
 @pytest.mark.asyncio
-async def test_soft_budget_alert():
-    """
-    Test that when a token's spend exceeds soft_budget, it triggers a budget alert but allows the request
-    """
+async def test_soft_budget_alert(monkeypatch: pytest.MonkeyPatch, cached_spend: float) -> None:
     import asyncio
     import time
 
+    import httpx
     from fastapi import Request
-    from starlette.datastructures import URL
 
+    from litellm.caching.dual_cache import DualCache
+    from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
+    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+    from litellm.proxy import proxy_server
     from litellm.proxy._types import UserAPIKeyAuth
     from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
-    from litellm.proxy.proxy_server import hash_token, user_api_key_cache
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+    from litellm.proxy.utils import PrismaClient, ProxyLogging
 
-    # Setup
-    user_key = "sk-98765"
-    soft_budget = 10
-    current_spend = 15  # Spend exceeds soft budget
-
-    # Create a valid token with soft budget
-    valid_token = UserAPIKeyAuth(
-        token=hash_token(user_key),
-        soft_budget=soft_budget,
-        spend=current_spend,
-        last_refreshed_at=time.time(),
+    monkeypatch.setattr(time, "time", lambda: 1_700_000_000.0)
+    monkeypatch.setenv("WEBHOOK_URL", "https://alerts.example.test/budget")
+    user_key: Final = "sk-soft-budget-auth"
+    token_hash: Final = proxy_server.hash_token(user_key)
+    key_cache: Final = UserApiKeyCache()
+    counter: Final = DualCache()
+    key_cache.set_cache(
+        key=token_hash,
+        value=UserAPIKeyAuth(
+            token=token_hash,
+            soft_budget=10.0,
+            spend=cached_spend,
+            last_refreshed_at=1_700_000_000.0,
+        ),
     )
+    counter.in_memory_cache.set_cache(key=f"spend:key:{token_hash}", value=15.0)
+    received: Final[list[httpx.Request]] = []
 
-    # Store in cache
-    user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
+    def receive_alert(request: httpx.Request) -> httpx.Response:
+        received.append(request)
+        return httpx.Response(200, text="ok")
 
-    # Mock proxy server settings
-    setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
-    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
-    setattr(litellm.proxy.proxy_server, "prisma_client", AsyncMock())
+    async def receive_body() -> Message:
+        return {"type": "http.request", "body": b"{}", "more_body": False}
 
-    # Create request
-    request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []})
-    request._url = URL(url="/chat/completions")
-
-    # Track if budget_alerts was called
-    alert_called = False
-    original_budget_alerts = litellm.proxy.proxy_server.proxy_logging_obj.budget_alerts
-
-    async def mock_budget_alerts(*args, **kwargs):
-        nonlocal alert_called
-        if kwargs.get("type") == "soft_budget":
-            alert_called = True
-        return await original_budget_alerts(*args, **kwargs)
-
-    # Patch the budget_alerts method
-    setattr(
-        litellm.proxy.proxy_server.proxy_logging_obj,
-        "budget_alerts",
-        mock_budget_alerts,
+    handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(receive_alert))
+    proxy_logging: Final = ProxyLogging(key_cache)
+    proxy_logging.alerting = ["webhook"]
+    proxy_logging.slack_alerting_instance = SlackAlerting(alerting=["webhook"], async_http_handler=handler)
+    monkeypatch.setattr(proxy_server, "user_api_key_cache", key_cache)
+    monkeypatch.setattr(proxy_server, "spend_counter_cache", counter)
+    monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging)
+    monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY)
+    monkeypatch.setattr(proxy_server, "general_settings", {})
+    monkeypatch.setattr(
+        proxy_server,
+        "prisma_client",
+        PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging),
     )
-
+    request: Final = Request(
+        scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []},
+        receive=receive_body,
+    )
+    tasks_before: Final = asyncio.all_tasks()
+    proxy_logging.slack_alerting_instance.update_values(alerting=["webhook"])
+    background_tasks: Final = asyncio.all_tasks()
     try:
-        # Call user_api_key_auth
-        response = await user_api_key_auth(request=request, api_key="Bearer " + user_key)
-
-        # Assert the request was allowed (no exception raised)
-        assert response is not None
-        # Assert the alert was triggered
-        await asyncio.sleep(3)
-        assert alert_called == True, "Soft budget alert should have been triggered"
-
+        response: Final = await user_api_key_auth(request=request, api_key="Bearer " + user_key)
+        await asyncio.gather(*(asyncio.all_tasks() - background_tasks))
+        assert (response.token, response.spend, response.soft_budget) == (token_hash, cached_spend, 10.0)
+        assert [json.loads(alert.content) for alert in received] == [
+            {
+                "spend": 15.0,
+                "max_budget": None,
+                "soft_budget": 10.0,
+                "token": token_hash,
+                "customer_id": None,
+                "user_id": None,
+                "team_id": None,
+                "team_alias": None,
+                "organization_id": None,
+                "user_email": None,
+                "key_alias": None,
+                "projected_exceeded_date": None,
+                "projected_spend": None,
+                "event_group": "key",
+                "alert_emails": None,
+                "max_budget_alert_emails": None,
+                "event": "soft_budget_crossed",
+                "event_message": "Soft Budget Crossed: Total Soft Budget:`10.0`",
+            }
+        ]
     finally:
-        # Restore original budget_alerts
-        setattr(
-            litellm.proxy.proxy_server.proxy_logging_obj,
-            "budget_alerts",
-            original_budget_alerts,
-        )
+        pending: Final = asyncio.all_tasks() - tasks_before
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+        await handler.close()
 
 
 def test_is_allowed_route():