diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 84278f4619d..2e839579c8e 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2575,7 +2575,7 @@ async def ui_view_spend_logs( description="Filter logs by model ID (litellm model deployment id)", ), model_group: str | None = fastapi.Query(default=None, description="Filter logs by model group"), - key_alias: str | None = fastapi.Query(default=None, description="Filter logs by key alias"), + key_alias: str | None = fastapi.Query(default=None, description="Filter logs by exact key alias match"), end_user: str | None = fastapi.Query(default=None, description="Filter logs by end user"), error_code: str | None = fastapi.Query(default=None, description="Filter logs by error code (e.g., '404', '500')"), error_message: str | None = fastapi.Query( @@ -2754,14 +2754,6 @@ async def ui_view_spend_logs( # Build metadata filters metadata_filters: Final = [] - if key_alias is not None: - metadata_filters.append( - { - "path": ["user_api_key_alias"], - "string_contains": key_alias, - } - ) - if error_code is not None: metadata_filters.append( { @@ -2783,6 +2775,27 @@ async def ui_view_spend_logs( where_conditions["metadata"] = metadata_filters[0] else: where_conditions["AND"] = [{"metadata": filter_cond} for filter_cond in metadata_filters] + if key_alias is not None: + active_key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many( + where={"key_alias": key_alias}, + ) + deleted_key_rows: Final = await prisma_client.db.litellm_deletedverificationtoken.find_many( + where={"key_alias": key_alias}, + ) + key_tokens: Final = list(dict.fromkeys(row.token for row in (*active_key_rows, *deleted_key_rows))) + key_alias_filter: Final = {"api_key": {"in": key_tokens}} + if where_conditions.get("api_key") is None: + where_conditions["api_key"] = key_alias_filter["api_key"] + else: + existing_and_conditions: Final = where_conditions.get("AND", []) + where_conditions["AND"] = [ + *( + existing_and_conditions + if isinstance(existing_and_conditions, list) + else [existing_and_conditions] + ), + key_alias_filter, + ] if end_user is not None: where_conditions["end_user"] = end_user @@ -2929,8 +2942,11 @@ async def ui_view_spend_logs( # Metadata JSON filters (PostgreSQL JSONB operators) if key_alias is not None: - sql_conditions.append(f"metadata->>'user_api_key_alias' LIKE ${p}") - sql_params.append(f"%{key_alias}%") + sql_conditions.append( + f'api_key IN (SELECT token FROM "LiteLLM_VerificationToken" WHERE key_alias = ${p} ' + f'UNION SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE key_alias = ${p})' + ) + sql_params.append(key_alias) p += 1 if error_code is not None: sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}") diff --git a/tests/integration/spend/test_spend_log_read_scope.py b/tests/integration/spend/test_spend_log_read_scope.py index 8a208d61aed..79b5ad03bd0 100644 --- a/tests/integration/spend/test_spend_log_read_scope.py +++ b/tests/integration/spend/test_spend_log_read_scope.py @@ -4,14 +4,17 @@ from collections.abc import AsyncIterator from dataclasses import dataclass from datetime import datetime, timedelta, timezone from hashlib import sha256 +from pathlib import Path from typing import Final from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit +import httpx import psycopg import pytest import pytest_asyncio -from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.client import Gateway, delete_key_if_present, eventually, object_value, string_value from integration._support.database import read_rows +from integration._support.process import owned_proxy_process from prisma import Prisma from psycopg import sql from psycopg.types.json import Jsonb @@ -77,6 +80,37 @@ def _seed_rows( ) +def _seed_alias_rows( + connection: psycopg.Connection, + rows: tuple[tuple[str, str, str | None, str], ...], + session_id: str, + started: datetime, +) -> None: + utc_timestamp: Final = started.astimezone(timezone.utc).replace(tzinfo=None) + with connection.cursor() as cursor: + cursor.executemany( + sql.SQL( + 'INSERT INTO {} (request_id, api_key, "user", session_id, "startTime", "endTime", ' + "metadata, messages, response, call_type) " + "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, 'acompletion')" + ).format(sql.Identifier("public", "LiteLLM_SpendLogs")), + tuple( + ( + request_id, + api_key, + user_id, + session_id, + utc_timestamp, + utc_timestamp, + Jsonb({"user_api_key_alias": key_alias}), + Jsonb([{"role": "user", "content": request_id + " payload"}]), + Jsonb({"id": request_id}), + ) + for request_id, api_key, user_id, key_alias in rows + ), + ) + + @pytest_asyncio.fixture(loop_scope="function") async def spend_database() -> AsyncIterator[Prisma]: schema: Final = f"integration_spend_scope_{uuid.uuid4().hex}" @@ -157,8 +191,8 @@ async def test_payload_sql_filters_foreign_collisions_and_prefers_exact_ids_for_ assert tuple(row.request_id for row in REQUEST_IDS.validate_python(result)) == expected -def _delete_session(session_id: str) -> None: - with psycopg.connect(os.environ["DATABASE_URL"]) as connection: +def _delete_session(session_id: str, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: connection.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE session_id = %s', (session_id,)) @@ -240,6 +274,99 @@ def _chat_request_id(gateway: Gateway, model: str, key: str) -> str: return str(body["id"]) +def test_spend_log_key_alias_is_exact_includes_deleted_keys_and_preserves_scope(tmp_path: Path) -> None: + with httpx.Client(base_url="http://127.0.0.1:1", timeout=15, trust_env=False) as source_client: + source_gateway: Final = Gateway(source_client, "sk-integration-master", "http://127.0.0.1:8190") + with owned_proxy_process(source_gateway, tmp_path, {}) as owned: + gateway: Final = owned.gateway + with gateway.scenario() as scenario: + assert gateway.database_url is not None + database_url: Final = gateway.database_url + owner: Final = scenario.user(user_role="internal_user") + other_owner: Final = scenario.user(user_role="internal_user") + key_alias: Final = f"prod-key-{uuid.uuid4().hex}" + overlapping_key_alias: Final = f"{key_alias}-2" + deleted_key_alias: Final = f"deleted-key-{uuid.uuid4().hex}" + key: Final = scenario.key(user_id=owner, key_alias=key_alias) + overlapping_key: Final = scenario.key( + user_id=other_owner, + key_alias=overlapping_key_alias, + ) + deleted_key_response: Final = gateway.post( + "/key/generate", + {"key_alias": deleted_key_alias}, + ) + deleted_key: Final = string_value(deleted_key_response["key"]) + scenario.cleanups.callback(delete_key_if_present, gateway, deleted_key) + request_ids: Final = tuple(f"alias-{uuid.uuid4().hex}" for _ in range(3)) + session_id: Final = f"key-alias-{uuid.uuid4().hex}" + started: Final = datetime.now(timezone.utc) - timedelta(hours=1) + scenario.cleanups.callback(_delete_session, session_id, database_url) + with psycopg.connect(database_url) as connection: + _seed_alias_rows( + connection, + ( + (request_ids[0], sha256(key.encode()).hexdigest(), owner, key_alias), + ( + request_ids[1], + sha256(overlapping_key.encode()).hexdigest(), + other_owner, + overlapping_key_alias, + ), + ( + request_ids[2], + sha256(deleted_key.encode()).hexdigest(), + other_owner, + deleted_key_alias, + ), + ), + session_id, + started, + ) + gateway.post("/key/delete", {"keys": [deleted_key]}) + date_filters: Final = { + "start_date": started.strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (started + timedelta(hours=2)).strftime("%Y-%m-%d %H:%M:%S"), + } + + for path in ("/spend/logs/ui", "/spend/logs/v2"): + exact: Final = gateway.request( + "GET", + path, + params={**date_filters, "key_alias": key_alias}, + ) + assert exact.status_code == 200, exact.text + assert exact.json()["total"] == 1, exact.text + assert [row["request_id"] for row in exact.json()["data"]] == [request_ids[0]], exact.text + + deleted: Final = gateway.request( + "GET", + path, + params={**date_filters, "key_alias": deleted_key_alias}, + ) + assert deleted.status_code == 200, deleted.text + assert deleted.json()["total"] == 1, deleted.text + assert [row["request_id"] for row in deleted.json()["data"]] == [request_ids[2]], deleted.text + + absent: Final = gateway.request( + "GET", + path, + params={**date_filters, "key_alias": f"absent-{uuid.uuid4().hex}"}, + ) + assert absent.status_code == 200, absent.text + assert absent.json()["total"] == 0, absent.text + assert absent.json()["data"] == [], absent.text + + foreign_alias: Final = gateway.request( + "GET", + "/spend/logs/ui", + params={**date_filters, "key_alias": overlapping_key_alias}, + key=key, + ) + assert foreign_alias.status_code == 200, foreign_alias.text + assert foreign_alias.json()["data"] == [], foreign_alias.text + + def _legacy_spend_rows(gateway: Gateway, key: str, params: dict[str, str]) -> list[dict[str, object]]: response: Final = gateway.request("GET", "/spend/logs", key=key, params=params) assert response.status_code == 200, response.text diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 23bbe521670..b8e8e408f99 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -4757,7 +4757,7 @@ async def test_ui_view_spend_logs_with_error_message(client): @pytest.mark.asyncio async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): - """Test merging error_code and key_alias filters with AND logic""" + """Test exact alias resolution while preserving API key and metadata filters""" mock_spend_logs = [ { "id": "log1", @@ -4773,7 +4773,7 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): { "id": "log2", "request_id": "req2", - "api_key": "sk-test-key", + "api_key": "sk-other-key", "user": "test_user_2", "team_id": "team1", "spend": 0.10, @@ -4795,34 +4795,50 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): ] def filter_by_error_code_and_key_alias(where): - if "AND" in where: - key_alias = error_code = None - for cond in where["AND"]: - if "metadata" in cond: - mf = cond["metadata"] - if mf.get("path") == ["user_api_key_alias"]: - key_alias = mf.get("string_contains") - elif mf.get("path") == ["error_information", "error_code"]: - error_code = str(mf.get("equals", "")).strip('"') - if key_alias == "test-key-1" and error_code == "500": - return [mock_spend_logs[2]] + metadata_filter = where.get("metadata", {}) + has_alias_substring_filter = any( + condition.get("metadata", {}).get("path") == ["user_api_key_alias"] + for condition in where.get("AND", []) + ) + if has_alias_substring_filter: + return mock_spend_logs + if ( + metadata_filter.get("path") != ["error_information", "error_code"] + or str(metadata_filter.get("equals", "")).strip('"') != "500" + ): + return mock_spend_logs + api_key_filter = where.get("api_key") + if api_key_filter == "sk-other-key": + return [] + if api_key_filter in ("sk-test-key", None) or api_key_filter == { + "in": ["sk-test-key", "sk-deleted-key"] + }: + return [mock_spend_logs[2]] return mock_spend_logs + query_observer = MagicMock() + prisma_mock = make_ui_spend_logs_mock_prisma( + mock_spend_logs, + filter_by_error_code_and_key_alias, + query_observer=query_observer, + ) + active_key_lookup = AsyncMock(return_value=[MagicMock(token="sk-test-key")]) + deleted_key_lookup = AsyncMock(return_value=[MagicMock(token="sk-deleted-key")]) + prisma_mock.db.litellm_verificationtoken = MagicMock(find_many=active_key_lookup) + prisma_mock.db.litellm_deletedverificationtoken = MagicMock(find_many=deleted_key_lookup) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" ) try: - with patch.object( - ps, - "prisma_client", - make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_error_code_and_key_alias), - ): + with patch.object(ps, "prisma_client", prisma_mock): start_date, end_date = _default_date_range() response = client.get( "/spend/logs/ui", params={ + "api_key": "sk-test-key", "error_code": "500", "key_alias": "test-key-1", "start_date": start_date, @@ -4842,6 +4858,36 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): assert metadata["user_api_key_alias"] == "test-key-1" assert "error_information" in metadata assert metadata["error_information"]["error_code"] == "500" + + mismatch_response = client.get( + "/spend/logs/ui", + params={ + "api_key": "sk-other-key", + "error_code": "500", + "key_alias": "test-key-1", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert mismatch_response.status_code == 200 + assert mismatch_response.json()["total"] == 0 + assert mismatch_response.json()["data"] == [] + page_query_call = next( + call + for call in query_observer.call_args_list + if 'FROM "LiteLLM_SpendLogs"' in call.args[0] and "ORDER BY" in call.args[0] + ) + page_query: Final = page_query_call.args[0] + page_params: Final = page_query_call.args[1] + alias_subquery: Final = re.search( + r'api_key IN \(SELECT token FROM "LiteLLM_VerificationToken" WHERE key_alias = \$(\d+) ' + r'UNION SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE key_alias = \$\1\)', + page_query, + ) + assert alias_subquery is not None + assert page_params[int(alias_subquery.group(1)) - 1] == "test-key-1" + assert page_params.count("test-key-1") == 1 finally: app.dependency_overrides.pop(ps.user_api_key_auth, None)