Merge ad971a7381 into b9e3587721

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -29,6 +29,8 @@ from litellm.proxy.auth.auth_checks import (
     get_key_object,
     get_team_object,
 )
+from litellm.proxy.auth.auth_utils import is_invalid_virtual_key_error
+from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
 from litellm.proxy.db.db_spend_update_writer import (
@@ -401,7 +403,9 @@ class _ProxyDBLogger(CustomLogger):
 
         request_route: Final = user_api_key_dict.request_route
         if (
-            ProxyDBLogger._should_track_errors_in_db() is False
+            isinstance(original_exception, KeyNotFoundError)
+            or is_invalid_virtual_key_error(original_exception)
+            or ProxyDBLogger._should_track_errors_in_db() is False
             or ProxyDBLogger._skips_pre_routing_rejections()
             and is_pre_routing_rejection(original_exception, _REQUEST_DATA.validate_python(request_data))
             or request_route is not None

tests/integration/spend/test_expired_ui_session_failure_rows.py

@@ -125,3 +125,46 @@ def test_an_expired_dashboard_session_key_still_logs_a_failed_llm_call(owned: Ow
         ).count("acompletion")
         == 1
     )
+
+
+@pytest.mark.parametrize("key_prefix", ["sk-", "wordlist-"])
+def test_unrecognized_keys_do_not_fill_spend_tables(gateway: Gateway, key_prefix: str) -> None:
+    attempts: Final = tuple(
+        (route, key_prefix + uuid.uuid4().hex)
+        for route in ("/v1/models", "/v1/chat/completions", "/v1/responses", "/v1/messages")
+    )
+    for route, key in attempts:
+        response: Final = gateway.request(
+            "GET" if route == "/v1/models" else "POST",
+            route,
+            {
+                "model": "integration-unrouted",
+                "messages": [{"role": "user", "content": "invalid key"}],
+                "input": "invalid key",
+                "max_tokens": 10,
+            },
+            key=key,
+        )
+        assert response.status_code == 401, response.text
+
+    with gateway.scenario() as scenario:
+        known_key: Final = scenario.key(duration="1s")
+        _rejected_as_expired(gateway, known_key, "/v1/models")
+        assert len(eventually(lambda: _failure_rows(known_key), lambda rows: len(rows) == 1, seconds=70)) == 1
+
+        keys: Final = [key for _, key in attempts] + [hash_token(key) for _, key in attempts]
+        empty: Final = [{"spend_logs": 0, "daily_user": 0, "daily_team": 0}]
+        assert (
+            held(
+                lambda: read_rows(
+                    'SELECT (SELECT count(*)::int FROM "LiteLLM_SpendLogs" WHERE api_key = ANY(%s)) AS spend_logs, '
+                    '(SELECT count(*)::int FROM "LiteLLM_DailyUserSpend" WHERE api_key = ANY(%s)) AS daily_user, '
+                    '(SELECT count(*)::int FROM "LiteLLM_DailyTeamSpend" WHERE api_key = ANY(%s)) AS daily_team',
+                    (keys, keys, keys),
+                ),
+                lambda rows: rows == empty,
+                holding=SPEND_LOG_FLUSH_SECONDS,
+                seconds=SPEND_LOG_FLUSH_SECONDS + 5,
+            )
+            == empty
+        )

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -15,6 +15,8 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
 from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth
+from litellm.proxy.auth.auth_utils import mark_invalid_virtual_key_error
+from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError
 from litellm.proxy.collector import SpendEventConsumer
 from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
 from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
@@ -41,6 +43,73 @@ from litellm.types.router import Deployment, LiteLLM_Params
 from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ["/v1/models", "/v1/chat/completions"])
+@pytest.mark.parametrize(
+    ("error", "expected_rows"),
+    [
+        pytest.param(KeyNotFoundError("unrecognized-key"), 0, id="unknown-key"),
+        pytest.param(
+            mark_invalid_virtual_key_error(
+                ProxyException(message="Malformed virtual key", type="auth_error", param="key", code=401), True
+            ),
+            0,
+            id="malformed-key",
+        ),
+        pytest.param(HTTPException(status_code=401, detail="Expired Key"), 1, id="expired-key"),
+        pytest.param(HTTPException(status_code=401, detail="Malformed virtual key"), 1, id="unmarked-auth-message"),
+        pytest.param(
+            litellm.AuthenticationError(message="Invalid upstream key", llm_provider="openai", model="unit-model"),
+            1,
+            id="provider-auth",
+        ),
+        pytest.param(litellm.BudgetExceededError(current_cost=2.0, max_budget=1.0), 1, id="exhausted-budget"),
+        pytest.param(HTTPException(status_code=403, detail="Model access denied"), 1, id="model-permission"),
+    ],
+)
+async def test_unrecognized_keys_write_no_spend_rows_but_recognized_key_failures_remain(
+    monkeypatch: pytest.MonkeyPatch, route: str, error: Exception, expected_rows: int
+) -> None:
+    from litellm.proxy import proxy_server
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+    from litellm.proxy.utils import PrismaClient, ProxyLogging
+
+    key_cache: Final = UserApiKeyCache()
+    logging: Final = ProxyLogging(user_api_key_cache=key_cache)
+    monkeypatch.setenv("DATABASE_URL", "postgresql://localhost/unit-tests")
+    client: Final = PrismaClient(database_url="postgresql://localhost/unit-tests", proxy_logging_obj=logging)
+    writer: Final = DBSpendUpdateWriter()
+    identity: Final = UserAPIKeyAuth(api_key="sk-unit-test-key", key_alias="unit-test-key", request_route=route)
+    monkeypatch.setattr(client, "spend_log_transactions", [])
+    monkeypatch.setattr(proxy_server, "prisma_client", client)
+    monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging)
+    monkeypatch.setattr(proxy_server, "user_api_key_cache", key_cache)
+    monkeypatch.setattr(proxy_server, "general_settings", {})
+    monkeypatch.setattr(proxy_server, "disable_spend_logs", False)
+    monkeypatch.setattr(proxy_server, "llm_router", None)
+    before: Final = frozenset(asyncio.all_tasks())
+
+    await ProxyDBLogger(spend_writer=lambda: writer).async_post_call_failure_hook(
+        request_data={
+            "model": "unit-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "start_time": datetime(2026, 1, 1, 12, tzinfo=timezone.utc),
+        },
+        original_exception=error,
+        user_api_key_dict=identity,
+    )
+    await asyncio.gather(*(asyncio.all_tasks() - before))
+
+    assert [(row["api_key"], row["status"], row["spend"]) for row in client.spend_log_transactions] == [
+        (identity.api_key, "failure", 0.0)
+    ] * expected_rows
+    for queue in (writer.daily_spend_update_queue, writer.daily_team_spend_update_queue):
+        daily: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions()
+        assert [(row["api_key"], row["failed_requests"], row["spend"]) for row in daily.values()] == [
+            (identity.api_key, 1, 0.0)
+        ] * expected_rows
+
+
 @pytest.mark.asyncio
 async def test_async_post_call_failure_hook():
     # Setup