From e8e2fad99dbe440ed0d1c42f07800bf652fa1022 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 3 Nov 2026 00:06:07 +0000 Subject: [PATCH] fix(usage-export): only populate key and user metadata for admin view and strip credential and SCIM slots --- litellm/constants.py | 3 ++ litellm/proxy/common_utils/callback_utils.py | 3 +- .../daily_activity_routes.py | 8 +++- .../internal_user_endpoints.py | 8 +--- .../repositories/daily_activity_repository.py | 14 +++++- litellm/repositories/daily_activity_sql.py | 46 ++++++++++++------- .../proxy/management_endpoints/scim_v2.py | 3 ++ .../coverage_registry/quota_management.yaml | 2 +- .../test_usage_export_metadata_e2e.py | 22 ++++++++- .../spend/_daily_activity_fixtures.py | 19 +++++++- .../spend/test_daily_activity_repository.py | 26 +++++++++-- .../test_daily_activity_routes.py | 43 +++++++++++++++-- .../repositories/test_daily_activity_sql.py | 46 +++++++++++++------ 13 files changed, 190 insertions(+), 53 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index cd7c0176428..34e07d83b35 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2312,6 +2312,9 @@ USAGE_MODEL_TOP_KEYS_MAX: Final[int] = 100 USAGE_CACHE_LEAKAGE_KEYS_DEFAULT: Final[int] = 20 USAGE_CACHE_LEAKAGE_KEYS_MAX: Final[int] = 100 USAGE_EXPORT_BATCH_SIZE: Final[int] = 1000 +CALLBACK_CONFIG_METADATA_SLOTS: Final[frozenset[str]] = frozenset( + {"logging", "callback_settings", "secret_manager_settings"} +) # Furthest back the catch-up pass looks for unpriced PTU days when a deployment # declares no ptu_effective_from, bounding the scan for an open-ended window. PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90 diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index d1787f3ff8b..29037859ea1 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -14,6 +14,7 @@ from litellm import get_secret from litellm._logging import verbose_proxy_logger from litellm.constants import ( BAGGAGE_MAPPED_TAGS_METADATA_KEY, + CALLBACK_CONFIG_METADATA_SLOTS, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH, @@ -52,7 +53,7 @@ _CALLBACK_VAR_ENCRYPTED_PREFIX: Final = CALLBACK_VAR_ENCRYPTED_PREFIX # Metadata slots that hold operator-configured callback and secret-manager setup # (and therefore integration credentials). Resolved from UserAPIKeyAuth during # pre-call setup, never read back off the copies stamped into request metadata. -_CALLBACK_CONFIG_SLOTS: Final = frozenset({"logging", "callback_settings", "secret_manager_settings"}) +_CALLBACK_CONFIG_SLOTS: Final = CALLBACK_CONFIG_METADATA_SLOTS blue_color_code: Final = "\033[94m" reset_color_code: Final = "\033[0m" diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 1abda1aae33..7a1cb7db5f0 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -12,7 +12,7 @@ from fastapi.responses import StreamingResponse from litellm import constants from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth, user_api_key_has_admin_view from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import ( InvalidDateRange, @@ -497,7 +497,11 @@ def _register_export_route(router: APIRouter, resolver: EntityScopeResolver, pre prisma_client, user_aggregated=False, ) - rows: Final = repository.export_rows(resolved.scope, export_type=export_type) + rows: Final = repository.export_rows( + resolved.scope, + export_type=export_type, + include_metadata=user_api_key_has_admin_view(user_api_key_dict), + ) first_row: Final = await anext(rows, None) return StreamingResponse( _stream_export_rows(first_row, rows, file_format), diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 54a2d4062d2..405794eecb4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -106,9 +106,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( UserUpdateResult, ) from litellm.types.proxy.management_endpoints.scim_v2 import ( - SCIM_ENTERPRISE_METADATA_KEY, - SCIM_ENTITLEMENTS_METADATA_KEY, - SCIM_ROLES_METADATA_KEY, + SCIM_DIRECTORY_METADATA_KEYS, ) from litellm.types.utils import BudgetConfig @@ -882,9 +880,7 @@ async def _get_user_info_teams( return team_list, teams_1 -_SCIM_DIRECTORY_METADATA_KEYS: Final = frozenset( - {SCIM_ENTERPRISE_METADATA_KEY, SCIM_ENTITLEMENTS_METADATA_KEY, SCIM_ROLES_METADATA_KEY} -) +_SCIM_DIRECTORY_METADATA_KEYS: Final = SCIM_DIRECTORY_METADATA_KEYS def _redact_scim_enterprise_metadata( diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 321468c4976..e60e3cf2a8e 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -258,12 +258,22 @@ class DailyActivityRepository: query: Final = build_cache_leakage_keys_sql(scope, limit=limit) return _KEY_SPEND_ADAPTER.validate_python(await self._query(query)) - async def export_rows(self, scope: DailyActivityScope, *, export_type: ExportType) -> AsyncIterator[ExportRow]: + async def export_rows( + self, scope: DailyActivityScope, *, export_type: ExportType, include_metadata: bool = False + ) -> AsyncIterator[ExportRow]: batch_size: Final = constants.USAGE_EXPORT_BATCH_SIZE cursor: ExportCursor | None = None # rebind-ok: each page advances the export keyset cursor while True: batch: tuple[ExportRow, ...] = _EXPORT_ADAPTER.validate_python( - await self._query(build_export_sql(scope, export_type=export_type, after=cursor, batch_size=batch_size)) + await self._query( + build_export_sql( + scope, + export_type=export_type, + after=cursor, + batch_size=batch_size, + include_metadata=include_metadata, + ) + ) ) for row in batch: yield row diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index 5393fc42100..756bba85278 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -8,8 +8,9 @@ from typing import Final from typing_extensions import assert_never from litellm import constants -from litellm.constants import PTU_SENTINEL_API_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL +from litellm.constants import CALLBACK_CONFIG_METADATA_SLOTS, PTU_SENTINEL_API_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL from litellm.repositories.daily_activity_buckets import UNATTRIBUTED_API_KEY, UNKNOWN_MODEL_GROUP +from litellm.types.proxy.management_endpoints.scim_v2 import SCIM_DIRECTORY_METADATA_KEYS from litellm.types.repositories.daily_activity import DailyActivityScope, DailyActivityTable, ExportType _API_KEY_ROLLED_UP_BIT: Final = 32 @@ -447,11 +448,16 @@ def build_cache_leakage_keys_sql(scope: DailyActivityScope, *, limit: int) -> Sq def build_export_sql( - scope: DailyActivityScope, *, export_type: ExportType, after: ExportCursor | None, batch_size: int + scope: DailyActivityScope, + *, + export_type: ExportType, + after: ExportCursor | None, + batch_size: int, + include_metadata: bool = False, ) -> SqlQuery: _bounded_limit(batch_size) where_clause, where_params = build_where_clause(scope) - group_key, output_key, user_fields, type_joins = _export_grouping(export_type) + group_key, output_key, user_fields, type_joins = _export_grouping(export_type, include_metadata=include_metadata) grouping_keys: Final = ( f"scoped.date, COALESCE(scoped.\"{scope.entity_id_field}\", '')", *((group_key,) if export_type is not ExportType.DAILY else ()), @@ -520,18 +526,26 @@ def build_export_sql( ) -_EMPTY_METADATA_FIELDS: Final = ", '{}'::jsonb AS key_metadata, '{}'::jsonb AS user_metadata" -_USER_METADATA_FIELD: Final = ( - "COALESCE(MIN(CASE WHEN jsonb_typeof(u.metadata) = 'object' THEN u.metadata::text END), '{}')::jsonb " - "AS user_metadata" -) -_KEY_METADATA_FIELD: Final = ( - "COALESCE(MIN(CASE WHEN jsonb_typeof(COALESCE(vt.metadata, dvt.metadata)) = 'object' " - "THEN COALESCE(vt.metadata, dvt.metadata)::text END), '{}')::jsonb AS key_metadata" +def _metadata_object_field(expression: str, redacted_keys: frozenset[str], alias: str) -> str: + redacted: Final = " - ".join((expression, *(f"'{key}'" for key in sorted(redacted_keys)))) + return ( + f"COALESCE(MIN(CASE WHEN jsonb_typeof({expression}) = 'object' THEN ({redacted})::text END), '{{}}')::jsonb " + f"AS {alias}" + ) + + +_EMPTY_KEY_METADATA_FIELD: Final = "'{}'::jsonb AS key_metadata" +_EMPTY_USER_METADATA_FIELD: Final = "'{}'::jsonb AS user_metadata" +_EMPTY_METADATA_FIELDS: Final = f", {_EMPTY_KEY_METADATA_FIELD}, {_EMPTY_USER_METADATA_FIELD}" +_USER_METADATA_FIELD: Final = _metadata_object_field("u.metadata", SCIM_DIRECTORY_METADATA_KEYS, "user_metadata") +_KEY_METADATA_FIELD: Final = _metadata_object_field( + "COALESCE(vt.metadata, dvt.metadata)", CALLBACK_CONFIG_METADATA_SLOTS, "key_metadata" ) -def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, ...]]: +def _export_grouping(export_type: ExportType, *, include_metadata: bool) -> tuple[str, str, str, tuple[str, ...]]: + key_metadata: Final = _KEY_METADATA_FIELD if include_metadata else _EMPTY_KEY_METADATA_FIELD + user_metadata: Final = _USER_METADATA_FIELD if include_metadata else _EMPTY_USER_METADATA_FIELD if export_type is ExportType.DAILY: return ( "''", @@ -545,10 +559,8 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, _API_KEY_EXPR, "MAX(CASE WHEN scoped.api_key = '' THEN 'Unattributed' " "ELSE COALESCE(vt.key_alias, dvt.key_alias) END) AS key_alias, " - "MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email, " - + _KEY_METADATA_FIELD - + ", " - + _USER_METADATA_FIELD, + f"MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email, " + f"{key_metadata}, {user_metadata}", ( "LEFT JOIN \"LiteLLM_VerificationToken\" vt ON vt.token = scoped.api_key AND scoped.api_key <> ''", """LEFT JOIN LATERAL ( @@ -573,7 +585,7 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, "COALESCE(vt.user_id, dvt.user_id, '')", "NULL::text", "NULL::text AS key_alias, MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, " - "MAX(u.user_email) AS user_email, '{}'::jsonb AS key_metadata, " + _USER_METADATA_FIELD, + f"MAX(u.user_email) AS user_email, {_EMPTY_KEY_METADATA_FIELD}, {user_metadata}", ( 'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key', """LEFT JOIN LATERAL ( diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index fadb309e0cd..dd9df4f9a3e 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -18,6 +18,9 @@ SCIM_ENTERPRISE_USER_SCHEMA: Final = "urn:ietf:params:scim:schemas:extension:ent SCIM_ENTERPRISE_METADATA_KEY: Final = "scim_enterprise" SCIM_ENTITLEMENTS_METADATA_KEY: Final = "scim_entitlements" SCIM_ROLES_METADATA_KEY: Final = "scim_roles" +SCIM_DIRECTORY_METADATA_KEYS: Final[frozenset[str]] = frozenset( + {SCIM_ENTERPRISE_METADATA_KEY, SCIM_ENTITLEMENTS_METADATA_KEY, SCIM_ROLES_METADATA_KEY} +) SCIM_MANAGED_TEAM_METADATA_KEY: Final = "scim_managed" SCIM_TEAM_DATA_METADATA_KEY: Final = "scim_data" diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 75dd883742f..56f5e7a66ca 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -40,7 +40,7 @@ - {id: quota_management.budget.fallback.routes_to_fallback, module: quota_management, tier: P1, behavior: budget, variant: fallback, assertions: [routes_to_fallback], exercised_on: [messages], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "budget_fallbacks reroute to the fallback model once the primary's budget is exhausted"} - {id: quota_management.budget.spend_counter.reseed_matches_db, module: quota_management, tier: P2, behavior: budget, variant: spend_counter, assertions: [reseed_matches_db], exercised_on: [chat_completions], source: "proxy/spend_tracking/budget_reservation.py", rationale: "Concurrent cold-counter reseeds keep the enforcement counter equal to DB spend (#26829)"} - {id: quota_management.spend_tracking.chat_completions.logs_cost, module: quota_management, tier: P0, behavior: spend_tracking, variant: chat_completions, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A paid chat call writes a nonzero spend row"} -- {id: quota_management.spend_tracking.usage_export.carries_key_and_user_metadata, module: quota_management, tier: P1, behavior: spend_tracking, variant: usage_export, assertions: [carries_key_and_user_metadata], exercised_on: [chat_completions], source: "repositories/daily_activity_sql.py / LIT-8524", rationale: "The daily_with_keys usage export row carries the key's and the owning user's stored metadata objects"} +- {id: quota_management.spend_tracking.usage_export.carries_key_and_user_metadata, module: quota_management, tier: P1, behavior: spend_tracking, variant: usage_export, assertions: [carries_key_and_user_metadata, empty_metadata_without_admin_view], exercised_on: [chat_completions], source: "repositories/daily_activity_sql.py / LIT-8524", rationale: "The daily_with_keys usage export row carries the key's and the owning user's stored metadata objects for callers with admin view and {} for everyone else"} - {id: quota_management.spend_tracking.stream.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: stream, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Streaming responses aggregate token counts into a spend row"} - {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/messages/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"} - {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"} diff --git a/tests/e2e/quota_management/spend_tracking/test_usage_export_metadata_e2e.py b/tests/e2e/quota_management/spend_tracking/test_usage_export_metadata_e2e.py index e70eba87689..fee0c97d623 100644 --- a/tests/e2e/quota_management/spend_tracking/test_usage_export_metadata_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_usage_export_metadata_e2e.py @@ -4,7 +4,8 @@ A customer charging usage back by cost center keeps that cost center in the virtual key's metadata and the owner's user metadata. Before LIT-8524 the Usage page export only carried the key alias, user id and email, so every exported row needed a /key/info and /user/info round trip to recover the -metadata. The export now joins both objects onto each daily_with_keys row. +metadata. The export now joins both objects onto each daily_with_keys row for +callers with admin view, and returns {} for everyone else. """ from __future__ import annotations @@ -38,6 +39,7 @@ class UsageExportParams(BaseModel): start_date: str end_date: str api_key: str + user_id: str | None = None export_type: str = "daily_with_keys" format: str = "json" @@ -150,3 +152,21 @@ class TestUsageExportMetadata: assert row.user_metadata.department == f"research-{marker}", ( f"user metadata not exported: {row.user_metadata}" ) + + own_view: Final = client.proxy.transport.probe( + EXPORT_ROUTE, + params=UsageExportParams( + start_date=(started - timedelta(days=1)).strftime("%Y-%m-%d"), + end_date=end.strftime("%Y-%m-%d"), + api_key=token, + user_id=user_id, + ), + headers=client.proxy.transport.bearer(record.key), + ) + assert own_view.status_code == 200, f"{EXPORT_ROUTE} answered {own_view.status_code}: {own_view.body}" + own_rows: Final = _EXPORT_ROWS.validate_json(own_view.body) + assert own_rows, f"{EXPORT_ROUTE} returned no rows for the key's own user: {own_view.body}" + for row in own_rows: + assert (row.key_metadata, row.user_metadata) == (UsageExportKeyMetadata(), UsageExportUserMetadata()), ( + f"a caller without admin view got metadata objects: {row}" + ) diff --git a/tests/integration/spend/_daily_activity_fixtures.py b/tests/integration/spend/_daily_activity_fixtures.py index 15574d88e82..a7e4c0df918 100644 --- a/tests/integration/spend/_daily_activity_fixtures.py +++ b/tests/integration/spend/_daily_activity_fixtures.py @@ -171,7 +171,16 @@ def seed_daily_activity_fixture(connection: psycopg.Connection, *, schema: str, "INSERT INTO {} (token, key_alias, team_id, user_id, metadata, models) VALUES (%s, %s, %s, %s, %s, %s)" ).format(verification_token_table), ( - ("key-a", "alias-a", "team-1", "user-1", '{"tags": ["blue", "gold"]}', []), + ( + "key-a", + "alias-a", + "team-1", + "user-1", + '{"tags": ["blue", "gold"], "cost_center": "cc-1", ' + '"logging": [{"callback_name": "langfuse", "callback_vars": {"langfuse_secret_key": "sk"}}], ' + '"callback_settings": {"otel": {"headers": "x"}}, "secret_manager_settings": {"vault": "v"}}', + [], + ), ("key-b", "alias-b", "team-1", "user-1", '{"tags": []}', []), ("key-c", "alias-c", "team-1", "user-1", '{"tags": []}', []), ("key-cache", "alias-cache", "team-1", "user-1", '{"tags": []}', []), @@ -201,7 +210,13 @@ def seed_daily_activity_fixture(connection: psycopg.Connection, *, schema: str, sql.SQL("INSERT INTO {} (user_id, user_email, models, metadata) VALUES (%s, %s, %s, %s)").format( user_table ), - ("user-1", "user@example.com", [], '{"department": "research"}'), + ( + "user-1", + "user@example.com", + [], + '{"department": "research", "scim_enterprise": {"employeeNumber": "E-1"}, ' + '"scim_entitlements": ["ent"], "scim_roles": ["role"]}', + ), ) cursor.execute( sql.SQL("INSERT INTO {} (team_id, team_alias, admins, members, models) VALUES (%s, %s, %s, %s, %s)").format( diff --git a/tests/integration/spend/test_daily_activity_repository.py b/tests/integration/spend/test_daily_activity_repository.py index 4c16c5c951b..6229bbc13be 100644 --- a/tests/integration/spend/test_daily_activity_repository.py +++ b/tests/integration/spend/test_daily_activity_repository.py @@ -219,7 +219,12 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch: assert tuple(row.api_key for row in leakage_keys) == ("key-cache", "key-c") exports: Final = tuple( - [row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_KEYS)] + [ + row + async for row in repository.export_rows( + team_scope, export_type=ExportType.DAILY_WITH_KEYS, include_metadata=True + ) + ] ) assert tuple(row.api_key for row in exports) == ( "key-a", @@ -239,10 +244,20 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch: assert deleted_key_export.key_metadata == {"tags": ["archived"]} assert deleted_key_export.user_metadata == {"department": "research"} active_key_export: Final = next(row for row in exports if row.api_key == "key-a") - assert active_key_export.key_metadata == {"tags": ["blue", "gold"]} + assert active_key_export.key_metadata == {"tags": ["blue", "gold"], "cost_center": "cc-1"} assert active_key_export.user_metadata == {"department": "research"} + unprivileged_exports: Final = tuple( + [row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_KEYS)] + ) + assert tuple(row.api_key for row in unprivileged_exports) == tuple(row.api_key for row in exports) + assert all((row.key_metadata, row.user_metadata) == ({}, {}) for row in unprivileged_exports) user_exports: Final = tuple( - [row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_USERS)] + [ + row + async for row in repository.export_rows( + team_scope, export_type=ExportType.DAILY_WITH_USERS, include_metadata=True + ) + ] ) assert len(user_exports) == 1 assert (user_exports[0].user_id, user_exports[0].user_email, user_exports[0].spend) == ( @@ -253,7 +268,10 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch: assert user_exports[0].user_metadata == {"department": "research"} assert user_exports[0].key_metadata == {} daily_export: Final = tuple( - [row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY)] + [ + row + async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY, include_metadata=True) + ] ) assert len(daily_export) == 1 assert daily_export[0].spend == 1273.0 diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index f39dedb3dcb..bb6cf8196c6 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -260,6 +260,7 @@ class _FakeRepository: self.user_page_call: tuple[DailyActivityScope, int, int] | None = None self.user_metadata = AsyncMock(side_effect=self._user_metadata) self.export_rows_error: Exception | None = None + self.export_rows_include_metadata: bool | None = None def _matching_rows(self, scope: DailyActivityScope) -> tuple[_Activity, ...]: return tuple( @@ -369,7 +370,10 @@ class _FakeRepository: for key in api_keys } - async def export_rows(self, scope: DailyActivityScope, *, export_type: ExportType) -> AsyncIterator[ExportRow]: + async def export_rows( + self, scope: DailyActivityScope, *, export_type: ExportType, include_metadata: bool = False + ) -> AsyncIterator[ExportRow]: + self.export_rows_include_metadata = include_metadata if self.export_rows_error is not None: raise self.export_rows_error for row in self._matching_rows(scope): @@ -391,8 +395,8 @@ class _FakeRepository: failed_requests=row.failed_requests, cache_read_input_tokens=row.cache_read_input_tokens, cache_creation_input_tokens=row.cache_creation_input_tokens, - key_metadata={"cost_center": "=finance", "tier": 2}, - user_metadata={"department": "research"}, + key_metadata={"cost_center": "=finance", "tier": 2} if include_metadata else {}, + user_metadata={"department": "research"} if include_metadata else {}, ) @@ -1190,6 +1194,39 @@ def test_export_routes_stream_json_arrays( assert records[0]["user_metadata"] == {"department": "research"} +@pytest.mark.parametrize( + ("role", "expects_metadata"), + [ + (LitellmUserRoles.PROXY_ADMIN, True), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, True), + (LitellmUserRoles.INTERNAL_USER, False), + ], +) +def test_export_metadata_objects_only_for_callers_with_admin_view( + daily_activity_client: tuple[TestClient, _FakeRepository], + role: LitellmUserRoles, + expects_metadata: bool, +) -> None: + client, repository = daily_activity_client + response: Final = client.get( + "/user/daily/activity/export", + params={ + **_entity_params("user_id", "user-a"), + "export_type": ExportType.DAILY_WITH_KEYS.value, + "format": "json", + }, + headers={"x-user-role": role.value, "x-user-id": "user-a"}, + ) + assert response.status_code == 200, response.text + assert repository.export_rows_include_metadata is expects_metadata + records: Final = response.json() + assert records, response.text + expected_key: Final = {"cost_center": "=finance", "tier": 2} if expects_metadata else {} + expected_user: Final = {"department": "research"} if expects_metadata else {} + assert all(record["key_metadata"] == expected_key for record in records), response.text + assert all(record["user_metadata"] == expected_user for record in records), response.text + + def test_export_first_row_error_returns_json_error_before_streaming( daily_activity_client: tuple[TestClient, _FakeRepository], ) -> None: diff --git a/tests/unit/repositories/test_daily_activity_sql.py b/tests/unit/repositories/test_daily_activity_sql.py index aed120c85c5..d107d4e7e09 100644 --- a/tests/unit/repositories/test_daily_activity_sql.py +++ b/tests/unit/repositories/test_daily_activity_sql.py @@ -451,27 +451,33 @@ def test_export_uses_latest_deleted_key_metadata(export_type: ExportType) -> Non assert "COALESCE(vt.user_id, dvt.user_id)" in query.sql +_REDACTED_KEY_METADATA: Final = ( + "THEN (COALESCE(vt.metadata, dvt.metadata) - 'callback_settings' - 'logging' - 'secret_manager_settings')" + "::text END), '{}')::jsonb AS key_metadata" +) +_REDACTED_USER_METADATA: Final = ( + "THEN (u.metadata - 'scim_enterprise' - 'scim_entitlements' - 'scim_roles')::text END), '{}')::jsonb " + "AS user_metadata" +) +_EMPTY_KEY_METADATA: Final = "'{}'::jsonb AS key_metadata" +_EMPTY_USER_METADATA: Final = "'{}'::jsonb AS user_metadata" + + @pytest.mark.parametrize( ("export_type", "key_metadata", "user_metadata"), [ - ( - ExportType.DAILY_WITH_KEYS, - "THEN COALESCE(vt.metadata, dvt.metadata)::text END), '{}')::jsonb AS key_metadata", - "THEN u.metadata::text END), '{}')::jsonb AS user_metadata", - ), - ( - ExportType.DAILY_WITH_USERS, - "'{}'::jsonb AS key_metadata", - "THEN u.metadata::text END), '{}')::jsonb AS user_metadata", - ), - (ExportType.DAILY, "'{}'::jsonb AS key_metadata", "'{}'::jsonb AS user_metadata"), - (ExportType.DAILY_WITH_MODELS, "'{}'::jsonb AS key_metadata", "'{}'::jsonb AS user_metadata"), + (ExportType.DAILY_WITH_KEYS, _REDACTED_KEY_METADATA, _REDACTED_USER_METADATA), + (ExportType.DAILY_WITH_USERS, _EMPTY_KEY_METADATA, _REDACTED_USER_METADATA), + (ExportType.DAILY, _EMPTY_KEY_METADATA, _EMPTY_USER_METADATA), + (ExportType.DAILY_WITH_MODELS, _EMPTY_KEY_METADATA, _EMPTY_USER_METADATA), ], ) -def test_export_selects_stored_key_and_user_metadata_objects( +def test_export_selects_redacted_key_and_user_metadata_objects_for_admin_view( export_type: ExportType, key_metadata: str, user_metadata: str ) -> None: - query = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2) + query = build_export_sql( + _scope(entity_ids=None), export_type=export_type, after=None, batch_size=2, include_metadata=True + ) assert key_metadata in query.sql assert user_metadata in query.sql @@ -479,6 +485,18 @@ def test_export_selects_stored_key_and_user_metadata_objects( assert "SELECT key_alias, user_id, metadata" in query.sql +@pytest.mark.parametrize("export_type", tuple(ExportType)) +def test_export_without_metadata_selects_empty_objects_and_never_reads_metadata_columns( + export_type: ExportType, +) -> None: + query = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2) + + assert _EMPTY_KEY_METADATA in query.sql + assert _EMPTY_USER_METADATA in query.sql + assert "u.metadata" not in query.sql + assert "vt.metadata" not in query.sql + + @pytest.mark.parametrize("export_type", tuple(ExportType)) def test_export_without_cursor_omits_cursor_predicate_and_parameters(export_type: ExportType) -> None: query = build_export_sql(