diff --git a/litellm/proxy/spend_tracking/spend_log_api_key_scope.py b/litellm/proxy/spend_tracking/spend_log_api_key_scope.py index 6507011a977..8b5c899553a 100644 --- a/litellm/proxy/spend_tracking/spend_log_api_key_scope.py +++ b/litellm/proxy/spend_tracking/spend_log_api_key_scope.py @@ -47,7 +47,9 @@ _API_KEY_ROWS: Final = TypeAdapter(tuple[_ApiKeyRow, ...]) def _narrowing_keys(rows: Sequence[Mapping[str, object]]) -> tuple[str, ...] | None: """None when the list cannot narrow the read: nothing known about the owner, or too many keys to help.""" keys: Final = tuple(row.api_key for row in _API_KEY_ROWS.validate_python(rows)) - return None if not keys or len(keys) > SPEND_LOG_API_KEY_SCOPE_MAX_KEYS else keys + if not keys or len(keys) > SPEND_LOG_API_KEY_SCOPE_MAX_KEYS: + return None + return keys if "" in keys else (*keys, "") def _daily_dates_around(start: date, end: date) -> tuple[str, str]: diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index ac46ccc1cfc..29e9be69776 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -17400,7 +17400,7 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m assert any("statement_timeout = 30000" in statement for statement in transaction.statements) mock_db_client.db.tx.assert_not_called() sql, *params = transaction.query_raw.call_args.args - assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta", ["sk-alpha", "sk-beta"]] + assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta", ["sk-alpha", "sk-beta", ""]] assert 'FROM "LiteLLM_SpendLogs" sl' in sql assert 'sl."startTime" >= $1::timestamp' in sql assert "sl.\"startTime\" < $2::timestamp + INTERVAL '1 day'" in sql @@ -17440,7 +17440,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien ) sql, *params = transaction.query_raw.call_args.args - assert params == ["2026-09-01", "2026-09-04", "team-alpha", ["sk-alpha"]] + assert params == ["2026-09-01", "2026-09-04", "team-alpha", ["sk-alpha", ""]] assert 'sl."user" = $' not in sql assert "sl.api_key = ANY($4::text[])" in sql @@ -17465,7 +17465,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli ) sql, *params = transaction.query_raw.call_args.args - assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob", ["sk-bob"]] + assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob", ["sk-bob", ""]] assert "sl.team_id IN ($3)" in sql assert 'AND sl."user" = $4' in sql assert "sl.api_key = ANY($5::text[])" in sql diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py index 60b263948b6..2e196d6a8c2 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py @@ -1,5 +1,3 @@ -import asyncio -import importlib import json from collections.abc import Iterator from typing import Final @@ -18,6 +16,9 @@ from litellm.router import Deployment, LiteLLM_Params from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +pytestmark: Final = pytest.mark.usefixtures("drained_logging_worker") + + @pytest.fixture def isolate_passthrough_endpoint_router_state(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -177,28 +178,6 @@ def _vcr_outcome_gate(request, vcr): record_vcr_outcome(request, vcr) -@pytest.fixture(scope="function", autouse=True) -def setup_and_teardown(): - """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. - """ - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - - asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) - importlib.reload(litellm) - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - importlib.reload(litellm.proxy.proxy_server) - except Exception as e: - print(f"Error reloading litellm.proxy.proxy_server: {e}") - loop = asyncio.get_event_loop_policy().new_event_loop() - asyncio.set_event_loop(loop) - print(litellm) - yield - loop.close() - asyncio.set_event_loop(None) - - @pytest.mark.asyncio async def test_gemini_pass_through_endpoint_forwards_key_and_body( respx_mock: "respx.MockRouter", httpx_transport: None, monkeypatch: pytest.MonkeyPatch 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 29e2f97d11c..337aa2939dd 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -944,10 +944,49 @@ async def test_ui_view_spend_logs_user_filter_reads_spend_logs_through_the_users assert '"LiteLLM_VerificationToken"' in lookup_sql assert lookup_params == ("alice", "2026-08-31", "2026-09-04", SPEND_LOG_API_KEY_SCOPE_MAX_KEYS + 1) page_sql, _, scoped_api_keys = _spend_logs_page_query(queries) - assert scoped_api_keys == ["sk-deleted", "sk-live"] + assert scoped_api_keys == ["sk-deleted", "sk-live", ""] assert '"user" = $' in page_sql +@pytest.mark.asyncio +@pytest.mark.parametrize("owner_filter", ["user_id", "team_id"]) +async def test_ui_view_spend_logs_owner_filter_preserves_empty_key_rows( + client: TestClient, monkeypatch: pytest.MonkeyPatch, owner_filter: str +) -> None: + logs: Final = [ + _owned_log("known", "sk-live", "alice", "team-alpha"), + _owned_log("unknown-key", "", "alice", "team-alpha"), + _owned_log("foreign", "", "bob", "team-beta"), + ] + monkeypatch.setattr( + ps, + "prisma_client", + make_ui_spend_logs_mock_prisma( + logs, + _owner_filtered_logs(logs), + owner_api_keys_fn=_owner_api_keys( + daily_rows={}, live_keys={"alice": {"sk-live"}, "team-alpha": {"sk-live"}} + ), + ), + ) + monkeypatch.setitem( + app.dependency_overrides, + ps.user_api_key_auth, + lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"), + ) + response: Final = client.get( + "/spend/logs/ui", + params={ + owner_filter: "alice" if owner_filter == "user_id" else "team-alpha", + "start_date": "2026-09-01 00:00:00", + "end_date": "2026-09-03 23:59:59", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert {row["request_id"] for row in response.json()["data"]} == {"known", "unknown-key"} + + @pytest.mark.asyncio async def test_ui_view_spend_logs_team_filter_reads_spend_logs_through_the_teams_api_keys(client, monkeypatch): logs = [ @@ -989,7 +1028,7 @@ async def test_ui_view_spend_logs_team_filter_reads_spend_logs_through_the_teams _, lookup_params = next((sql, params) for sql, params in queries if '"LiteLLM_DailyTeamSpend"' in sql) assert lookup_params[:3] == (["team-alpha"], "2026-08-31", "2026-09-04") page_sql, _, scoped_api_keys = _spend_logs_page_query(queries) - assert scoped_api_keys == ["sk-alpha-1", "sk-alpha-2"] + assert scoped_api_keys == ["sk-alpha-1", "sk-alpha-2", ""] assert "team_id = $" in page_sql