From eb6df4f6454062b893fc04500f77337cf8a0538f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 17 Aug 2026 10:22:04 -0700 Subject: [PATCH] refactor(proxy): audit alias follow-ups from re-review Actor fields now use plain or-fallthrough so a None from the auth context still resolves through the credential-confirmed lookup while the spoofing gate stays; the payload keys become NotRequired so external StandardAuditLogPayload constructors keep type-checking; the dead _serialized_blob identity helper is gone since mask_api_keys always re-serializes blobs to strings; the backfill runbook gains batching and VACUUM guidance and is cited from the writer and endpoint docstrings; ruff's Dict and Optional modernization applied to the touched enterprise files --- db_scripts/backfill_audit_log_aliases.sql | 7 ++++ .../proxy/audit_logging_endpoints.py | 36 ++++++++++--------- .../types/proxy/audit_logging_endpoints.py | 8 ++--- .../proxy/management_helpers/audit_logs.py | 30 ++++++---------- litellm/types/utils.py | 12 +++---- .../integrations/test_azure_sentinel.py | 23 ++++++++---- .../test_audit_log_callbacks.py | 11 +++--- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++- 8 files changed, 73 insertions(+), 58 deletions(-) diff --git a/db_scripts/backfill_audit_log_aliases.sql b/db_scripts/backfill_audit_log_aliases.sql index 03bff0c124a..e4610426885 100644 --- a/db_scripts/backfill_audit_log_aliases.sql +++ b/db_scripts/backfill_audit_log_aliases.sql @@ -10,6 +10,13 @@ -- Every statement only touches rows where the target column is NULL, so the -- script is idempotent and safe to re-run (including after a partial run). -- +-- On a large audit table, run this in batches instead of one shot: the first +-- object_alias statement and the object_team_id statement rewrite every +-- matching row, so wrap each UPDATE with an id-range or updated_at-range +-- predicate and loop until no rows change. Run VACUUM (ANALYZE) +-- "LiteLLM_AuditLog" afterward to reclaim the dead tuples the rewrites leave +-- behind. +-- -- Sources, matching what the writer produces: -- - object_alias comes from the before/updated JSON blobs captured at change -- time (updated_values wins over before_value; users prefer user_alias then diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index c96ab570de4..239a4fbf689 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -7,7 +7,7 @@ GET - /audit/{id} - Get audit log by id GET - /audit - Get all audit logs """ -from typing import Any, Dict, Optional +from typing import Any #### AUDIT LOGGING #### from fastapi import APIRouter, Depends, HTTPException, Query @@ -22,7 +22,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() -def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]: +def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, Any]: """ Build an OR condition that matches a value inside a JSON column at the given key, checking both before_value and updated_values. @@ -43,7 +43,7 @@ def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]: } -def _build_object_team_condition(object_team: str) -> Dict[str, Any]: +def _build_object_team_condition(object_team: str) -> dict[str, Any]: return { "OR": [ {"object_team_id": object_team}, @@ -62,14 +62,14 @@ async def get_audit_logs( page: int = Query(1, ge=1), page_size: int = Query(10, ge=1, le=100), # Filter parameters - changed_by: Optional[str] = Query(None, description="Filter by user or system that performed the action"), - changed_by_api_key: Optional[str] = Query(None, description="Filter by API key hash that performed the action"), - action: Optional[str] = Query(None, description="Filter by action type (create, update, delete)"), - table_name: Optional[str] = Query(None, description="Filter by table name that was modified"), - object_id: Optional[str] = Query(None, description="Filter by ID of the object that was modified"), - start_date: Optional[str] = Query(None, description="Filter logs after this date"), - end_date: Optional[str] = Query(None, description="Filter logs before this date"), - object_team_id: Optional[str] = Query( + changed_by: str | None = Query(None, description="Filter by user or system that performed the action"), + changed_by_api_key: str | None = Query(None, description="Filter by API key hash that performed the action"), + action: str | None = Query(None, description="Filter by action type (create, update, delete)"), + table_name: str | None = Query(None, description="Filter by table name that was modified"), + object_id: str | None = Query(None, description="Filter by ID of the object that was modified"), + start_date: str | None = Query(None, description="Filter logs after this date"), + end_date: str | None = Query(None, description="Filter logs before this date"), + object_team_id: str | None = Query( None, description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)", ), @@ -80,12 +80,12 @@ async def get_audit_logs( "or rows whose object_team_alias contains this value" ), ), - object_key_hash: Optional[str] = Query( + object_key_hash: str | None = Query( None, description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)", ), # Sorting parameters - sort_by: Optional[str] = Query( + sort_by: str | None = Query( None, description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')", ), @@ -98,7 +98,9 @@ async def get_audit_logs( Note: object_team_id and object_key_hash use Prisma JSON path filtering, which requires PostgreSQL. object_team filters on the denormalized - object_team_id and object_team_alias columns instead. + object_team_id and object_team_alias columns instead. Rows written before + those columns existed return NULL aliases until an operator runs the + db_scripts/backfill_audit_log_aliases.sql runbook. """ from litellm.proxy.proxy_server import prisma_client @@ -109,7 +111,7 @@ async def get_audit_logs( ) # Build filter conditions - where_conditions: Dict[str, Any] = {} + where_conditions: dict[str, Any] = {} if changed_by: where_conditions["changed_by"] = changed_by if changed_by_api_key: @@ -121,7 +123,7 @@ async def get_audit_logs( if object_id: where_conditions["object_id"] = object_id if start_date or end_date: - date_filter: Dict[str, Any] = {} + date_filter: dict[str, Any] = {} if start_date: date_filter["gte"] = start_date if end_date: @@ -142,7 +144,7 @@ async def get_audit_logs( where_conditions["AND"] = where_conditions.get("AND", []) + [_build_object_team_condition(object_team)] # Build sort conditions - order_by: Dict[str, Any] = {} + order_by: dict[str, Any] = {} if sort_by and isinstance(sort_by, str): order_by[sort_by] = sort_order else: diff --git a/enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py index 5a73655fb8a..1cf736c197b 100644 --- a/enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, Dict, List, Optional +from typing import Any from pydantic import BaseModel, Field @@ -14,8 +14,8 @@ class AuditLogResponse(BaseModel): action: str table_name: str object_id: str - before_value: Optional[Dict[str, Any]] = None - updated_values: Optional[Dict[str, Any]] = None + before_value: dict[str, Any] | None = None + updated_values: dict[str, Any] | None = None object_alias: str | None = None object_team_id: str | None = None object_team_alias: str | None = None @@ -26,7 +26,7 @@ class AuditLogResponse(BaseModel): class PaginatedAuditLogResponse(BaseModel): """Response model for paginated audit logs""" - audit_logs: List[AuditLogResponse] + audit_logs: list[AuditLogResponse] total: int = Field(..., description="Total number of audit logs matching the filters") page: int = Field(..., description="Current page number") page_size: int = Field(..., description="Number of items per page") diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index 65fa88359e1..c138d002ffa 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -1,9 +1,13 @@ """ Functions to create audit logs for LiteLLM Proxy + +New rows are stamped with denormalized alias columns (object_alias, object_team_id, +object_team_alias, changed_by_user_email, changed_by_key_alias) at write time. Rows +written before those columns existed keep NULLs until an operator runs the optional +db_scripts/backfill_audit_log_aliases.sql runbook. """ import asyncio -import json from datetime import datetime, timezone from typing import TYPE_CHECKING, Final, NamedTuple @@ -252,10 +256,6 @@ async def _lookup_key_alias(prisma_client: "PrismaClient", token: str) -> str | return actor_key.key_alias -def _serialized_blob(value: object) -> object: - return json.dumps(value) if isinstance(value, dict) else value - - async def _with_denormalized_aliases( request_data: LiteLLM_AuditLogs, prisma_client: "PrismaClient | None" ) -> LiteLLM_AuditLogs: @@ -300,24 +300,18 @@ async def _with_denormalized_aliases( need_actor_key: Final = ( prisma_client is not None and bool(request_data.changed_by_api_key) - and ("changed_by_key_alias" not in fields_set or "changed_by_user_email" not in fields_set) + and (not request_data.changed_by_key_alias or not request_data.changed_by_user_email) ) actor_key: Final = ( await _lookup_actor_key(prisma_client, request_data.changed_by_api_key) if need_actor_key and prisma_client is not None and request_data.changed_by_api_key else _ActorKey(key_alias=None, user_id=None) ) - changed_by_key_alias: Final = ( - request_data.changed_by_key_alias if "changed_by_key_alias" in fields_set else actor_key.key_alias - ) - changed_by_user_email: Final = ( - request_data.changed_by_user_email - if "changed_by_user_email" in fields_set - else ( - await _lookup_user_email(prisma_client, changed_by) - if prisma_client is not None and changed_by is not None and actor_key.user_id == changed_by - else None - ) + changed_by_key_alias: Final = request_data.changed_by_key_alias or actor_key.key_alias + changed_by_user_email: Final = request_data.changed_by_user_email or ( + await _lookup_user_email(prisma_client, changed_by) + if prisma_client is not None and changed_by is not None and actor_key.user_id == changed_by + else None ) return request_data.model_copy( update={ @@ -326,8 +320,6 @@ async def _with_denormalized_aliases( "object_team_alias": object_team_alias, "changed_by_user_email": changed_by_user_email, "changed_by_key_alias": changed_by_key_alias, - "updated_values": _serialized_blob(request_data.updated_values), - "before_value": _serialized_blob(request_data.before_value), } ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 69098716620..658a8e42d38 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -39,7 +39,7 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import ReadOnly, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -3127,11 +3127,11 @@ class StandardAuditLogPayload(TypedDict): object_id: str before_value: str | None updated_values: str | None - object_alias: ReadOnly[str | None] - object_team_id: ReadOnly[str | None] - object_team_alias: ReadOnly[str | None] - changed_by_user_email: ReadOnly[str | None] - changed_by_key_alias: ReadOnly[str | None] + object_alias: NotRequired[ReadOnly[str | None]] + object_team_id: NotRequired[ReadOnly[str | None]] + object_team_alias: NotRequired[ReadOnly[str | None]] + changed_by_user_email: NotRequired[ReadOnly[str | None]] + changed_by_key_alias: NotRequired[ReadOnly[str | None]] class StandardLoggingPayload(TypedDict): diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 7335316548d..ebe57b5f109 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -120,6 +120,11 @@ async def test_azure_sentinel_queues_audit_log_event(): object_id="team-1", before_value=None, updated_values='{"team_alias": "sentinel-demo"}', + object_alias="sentinel-demo", + object_team_id="team-1", + object_team_alias="sentinel-demo", + changed_by_user_email="user-1@example.com", + changed_by_key_alias="sentinel-key", ) await logger.async_log_audit_log_event(audit_log) @@ -155,6 +160,11 @@ async def test_azure_sentinel_sends_audit_log_payload_to_ingestion_api(): object_id="team-1", before_value=None, updated_values='{"team_alias": "sentinel-demo"}', + object_alias="sentinel-demo", + object_team_id="team-1", + object_team_alias="sentinel-demo", + changed_by_user_email="user-1@example.com", + changed_by_key_alias="sentinel-key", ) await logger.async_log_audit_log_event(audit_log) @@ -221,6 +231,11 @@ async def test_azure_sentinel_flushes_standard_and_audit_logs_separately(): object_id="team-1", before_value=None, updated_values='{"team_alias": "sentinel-demo"}', + object_alias="sentinel-demo", + object_team_id="team-1", + object_team_alias="sentinel-demo", + changed_by_user_email="user-1@example.com", + changed_by_key_alias="sentinel-key", ) logger.log_queue.append(standard_payload) @@ -250,17 +265,13 @@ async def test_azure_sentinel_flushes_standard_and_audit_logs_separately(): await logger.flush_queue() ingestion_calls = [ - call - for call in logger.async_httpx_client.post.call_args_list - if "dataCollectionRules" in call.kwargs["url"] + call for call in logger.async_httpx_client.post.call_args_list if "dataCollectionRules" in call.kwargs["url"] ] assert len(ingestion_calls) == 2 standard_call, audit_call = ingestion_calls assert "Custom-LiteLLM-Standard" in standard_call.kwargs["url"] - assert json.loads(standard_call.kwargs["data"].decode("utf-8")) == [ - standard_payload - ] + assert json.loads(standard_call.kwargs["data"].decode("utf-8")) == [standard_payload] assert "Custom-LiteLLM-Audit" in audit_call.kwargs["url"] assert json.loads(audit_call.kwargs["data"].decode("utf-8")) == [audit_log] diff --git a/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py index 89101fe34a8..6f9bbc7125f 100644 --- a/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py +++ b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py @@ -584,12 +584,13 @@ class TestAuditLogAliasDenormalization: assert db.litellm_teamtable.where_calls == [{"team_id": "team-9"}] @pytest.mark.asyncio - async def test_service_account_without_aliases_triggers_no_lookups(self): - """Explicitly supplied None actor fields (service keys with no alias or email) skip the - guaranteed-miss lookups instead of querying on every audit row.""" + async def test_none_actor_fields_fall_through_to_credential_confirmed_lookup(self): + """A None from the auth context still falls through to the lookup, but the email side only + resolves through the credential-confirmed gate: a service key with no owning user gets one + key lookup and never a user lookup.""" db = _FakeDb( user_row=SimpleNamespace(user_email="someone@example.com"), - key_row=SimpleNamespace(key_alias="some-key", user_id="someone"), + key_row=SimpleNamespace(key_alias=None, user_id=None), ) p1, p2, p3 = _gates(_FakePrismaClient(db)) with p1, p2, p3: @@ -611,8 +612,8 @@ class TestAuditLogAliasDenormalization: data = db.litellm_auditlog.created[0] assert "changed_by_user_email" not in data assert "changed_by_key_alias" not in data + assert db.litellm_verificationtoken.where_calls == [{"token": "hash-service"}] assert db.litellm_usertable.where_calls == [] - assert db.litellm_verificationtoken.where_calls == [] @pytest.mark.asyncio async def test_spoofed_changed_by_is_not_resolved_to_an_email(self): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2fbee15d638..047e8b0e54b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -721,7 +721,9 @@ export interface paths { * * Note: object_team_id and object_key_hash use Prisma JSON path filtering, * which requires PostgreSQL. object_team filters on the denormalized - * object_team_id and object_team_alias columns instead. + * object_team_id and object_team_alias columns instead. Rows written before + * those columns existed return NULL aliases until an operator runs the + * db_scripts/backfill_audit_log_aliases.sql runbook. */ get: operations["get_audit_logs_audit_get"]; put?: never;