From ad971a738117785fe3ae1efb87287f6993490227 Mon Sep 17 00:00:00 2001 From: Tanishq <74068132+tanishqbaweja@users.noreply.github.com> Date: Mon, 26 Oct 2026 04:38:45 +0530 Subject: [PATCH] fix(proxy): omit unrecognized keys from spend records --- .../proxy/hooks/proxy_track_cost_callback.py | 6 +- .../test_expired_ui_session_failure_rows.py | 43 ++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 69 +++++++++++++++++++ 3 files changed, 117 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c11b59a809a..4c4d2ec8ac4 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/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 diff --git a/tests/integration/spend/test_expired_ui_session_failure_rows.py b/tests/integration/spend/test_expired_ui_session_failure_rows.py index 7bd41eff8e9..39da316e729 100644 --- a/tests/integration/spend/test_expired_ui_session_failure_rows.py +++ b/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 + ) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 79193d37b04..53b4b1ca9a2 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/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