Merge pull request #26859 from stuxf/chore/audit-log-team-callback-mutations
chore(team): audit-log team-callback admin mutations
This commit is contained in:
yuneng-jiang
GitHub
commit
d51d96f405
@ -4,15 +4,22 @@ Endpoints to control callbacks per team
|
||||
Use this when each team should control its own callbacks
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import traceback
|
||||
from typing import List, Optional
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy._types import (
|
||||
AddTeamCallback,
|
||||
LiteLLM_AuditLogs,
|
||||
LitellmTableNames,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
TeamCallbackMetadata,
|
||||
@ -24,6 +31,106 @@ from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
_CALLBACK_VARS_REDACTED = "***REDACTED***"
|
||||
|
||||
|
||||
def _redact_callback_secrets(metadata: Any) -> Any:
|
||||
"""Strip secret values out of a team-metadata snapshot before audit logging.
|
||||
|
||||
Both ``team_metadata["logging"]`` (list of ``AddTeamCallback`` dicts) and
|
||||
``team_metadata["callback_settings"]["callback_vars"]`` carry provider
|
||||
credentials such as ``langfuse_secret_key``, ``langsmith_api_key``, and
|
||||
``gcs_path_service_account``. Persisting them verbatim into
|
||||
``LiteLLM_AuditLogs`` would let anyone with read access to the audit
|
||||
table harvest team callback credentials, so we replace each value with
|
||||
a fixed marker. The keys themselves are kept so the audit reader can
|
||||
still see *which* fields changed.
|
||||
"""
|
||||
if not isinstance(metadata, dict):
|
||||
return metadata
|
||||
redacted = copy.deepcopy(metadata)
|
||||
logging_entries = redacted.get("logging")
|
||||
if isinstance(logging_entries, list):
|
||||
for entry in logging_entries:
|
||||
if isinstance(entry, dict) and isinstance(entry.get("callback_vars"), dict):
|
||||
entry["callback_vars"] = {
|
||||
k: _CALLBACK_VARS_REDACTED for k in entry["callback_vars"]
|
||||
}
|
||||
callback_settings = redacted.get("callback_settings")
|
||||
if isinstance(callback_settings, dict) and isinstance(
|
||||
callback_settings.get("callback_vars"), dict
|
||||
):
|
||||
callback_settings["callback_vars"] = {
|
||||
k: _CALLBACK_VARS_REDACTED for k in callback_settings["callback_vars"]
|
||||
}
|
||||
return redacted
|
||||
|
||||
|
||||
def _log_audit_task_exception(task: "asyncio.Task[None]") -> None:
|
||||
"""Surface a fire-and-forget audit-log task failure.
|
||||
|
||||
``asyncio.create_task`` swallows exceptions silently — if the audit
|
||||
write fails (transient DB error etc.) we'd otherwise lose the row
|
||||
without any signal. Log at warning level so the operator sees there's
|
||||
a gap in the audit trail.
|
||||
"""
|
||||
if task.cancelled():
|
||||
return
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
verbose_proxy_logger.warning("Failed to write team-callback audit log: %s", exc)
|
||||
|
||||
|
||||
async def _emit_team_callback_audit_log(
|
||||
*,
|
||||
team_id: str,
|
||||
before_metadata: Any,
|
||||
after_metadata: Any,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: Optional[str],
|
||||
) -> None:
|
||||
"""Emit an audit-log row for a team-callback mutation.
|
||||
|
||||
Mirrors the ``store_audit_logs``-gated pattern used in
|
||||
``team_endpoints.py``: the call is async-fire-and-forget and is a no-op
|
||||
when audit logging is not enabled on the proxy. Captured under
|
||||
``LitellmTableNames.TEAM_TABLE_NAME`` so the row co-locates with other
|
||||
team mutations in the audit table.
|
||||
|
||||
Callback secrets are redacted before serialization so the audit table
|
||||
cannot itself become a credential-harvest sink.
|
||||
"""
|
||||
if litellm.store_audit_logs is not True:
|
||||
return
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
redacted_before = _redact_callback_secrets(before_metadata)
|
||||
redacted_after = _redact_callback_secrets(after_metadata)
|
||||
|
||||
task = asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
object_id=team_id,
|
||||
action="updated",
|
||||
updated_values=json.dumps({"metadata": redacted_after}, default=str),
|
||||
before_value=json.dumps({"metadata": redacted_before}, default=str),
|
||||
)
|
||||
)
|
||||
)
|
||||
task.add_done_callback(_log_audit_task_exception)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/team/{team_id:path}/callback",
|
||||
tags=["team management"],
|
||||
@ -123,6 +230,7 @@ async def add_team_callbacks(
|
||||
param="callback_name",
|
||||
)
|
||||
|
||||
before_metadata = copy.deepcopy(team_metadata)
|
||||
team_callback_settings.append(data.model_dump())
|
||||
|
||||
team_metadata["logging"] = team_callback_settings
|
||||
@ -132,6 +240,14 @@ async def add_team_callbacks(
|
||||
where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore
|
||||
)
|
||||
|
||||
await _emit_team_callback_audit_log(
|
||||
team_id=team_id,
|
||||
before_metadata=before_metadata,
|
||||
after_metadata=team_metadata,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"data": new_team_row,
|
||||
@ -165,6 +281,10 @@ async def disable_team_logging(
|
||||
http_request: Request,
|
||||
team_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
litellm_changed_by: Optional[str] = Header(
|
||||
None,
|
||||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||||
),
|
||||
):
|
||||
"""
|
||||
Disable all logging callbacks for a team
|
||||
@ -198,6 +318,7 @@ async def disable_team_logging(
|
||||
|
||||
# Update team metadata to disable logging
|
||||
team_metadata = _existing_team.metadata
|
||||
before_metadata = copy.deepcopy(team_metadata)
|
||||
team_callback_settings = team_metadata.get("callback_settings", {})
|
||||
team_callback_settings_obj = TeamCallbackMetadata(**team_callback_settings)
|
||||
|
||||
@ -222,6 +343,17 @@ async def disable_team_logging(
|
||||
},
|
||||
)
|
||||
|
||||
# Disabling a team's logging callbacks is itself a logging-control
|
||||
# action — emit an audit-log row so the action remains traceable
|
||||
# even though the team's own observability is now off.
|
||||
await _emit_team_callback_audit_log(
|
||||
team_id=team_id,
|
||||
before_metadata=before_metadata,
|
||||
after_metadata=team_metadata,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Logging disabled for team {team_id}",
|
||||
|
||||
@ -0,0 +1,294 @@
|
||||
"""
|
||||
Audit-log emission for the team-callback admin endpoints.
|
||||
|
||||
The endpoints in ``team_callback_endpoints.py`` mutate a team's logging
|
||||
callbacks (``add_team_callbacks``) or zero them out entirely
|
||||
(``disable_team_logging``). Both are admin-only mutations, and the
|
||||
disable variant is itself a logging-control action, so when the operator
|
||||
has Enterprise audit logging enabled (``litellm.store_audit_logs = True``)
|
||||
each call must emit a row that captures who did it and what the metadata
|
||||
looked like before/after.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import Request
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import (
|
||||
AddTeamCallback,
|
||||
LitellmTableNames,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_callback_endpoints import (
|
||||
add_team_callbacks,
|
||||
disable_team_logging,
|
||||
)
|
||||
|
||||
|
||||
def _admin_auth() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="hashed",
|
||||
user_id="admin-user",
|
||||
user_role="proxy_admin",
|
||||
)
|
||||
|
||||
|
||||
def _existing_team_row(metadata: dict) -> MagicMock:
|
||||
row = MagicMock()
|
||||
row.team_id = "team-1"
|
||||
row.metadata = metadata
|
||||
return row
|
||||
|
||||
|
||||
def _patch_prisma(existing_metadata: dict):
|
||||
"""Build a context-manager that patches the proxy's ``prisma_client``
|
||||
to return ``existing_metadata`` from ``get_data`` and a stub team row
|
||||
from ``litellm_teamtable.update``."""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.get_data = AsyncMock(return_value=_existing_team_row(existing_metadata))
|
||||
|
||||
updated_row = MagicMock()
|
||||
updated_row.team_id = "team-1"
|
||||
mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_row)
|
||||
return mock_prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disable_team_logging_emits_audit_log_when_enabled(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "store_audit_logs", True)
|
||||
mock_prisma = _patch_prisma(
|
||||
{
|
||||
"callback_settings": {
|
||||
"success_callback": ["langfuse"],
|
||||
"failure_callback": [],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
audit_calls = []
|
||||
|
||||
async def capture(request_data):
|
||||
audit_calls.append(request_data)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
patch(
|
||||
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
|
||||
new=capture,
|
||||
),
|
||||
):
|
||||
await disable_team_logging(
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id="team-1",
|
||||
user_api_key_dict=_admin_auth(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
# asyncio.create_task fires the coroutine eagerly; await one tick to let
|
||||
# the audit-log emit run before the test exits.
|
||||
import asyncio
|
||||
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert len(audit_calls) == 1
|
||||
log = audit_calls[0]
|
||||
assert log.table_name == LitellmTableNames.TEAM_TABLE_NAME
|
||||
assert log.object_id == "team-1"
|
||||
assert log.action == "updated"
|
||||
assert log.changed_by == "admin-user"
|
||||
|
||||
before = json.loads(log.before_value)
|
||||
after = json.loads(log.updated_values)
|
||||
# Before: the team's pre-existing success_callback survives in the snapshot.
|
||||
assert before["metadata"]["callback_settings"]["success_callback"] == ["langfuse"]
|
||||
# After: callbacks zeroed out by the endpoint.
|
||||
assert after["metadata"]["callback_settings"]["success_callback"] == []
|
||||
assert after["metadata"]["callback_settings"]["failure_callback"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disable_team_logging_no_audit_when_disabled(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "store_audit_logs", False)
|
||||
mock_prisma = _patch_prisma(
|
||||
{
|
||||
"callback_settings": {
|
||||
"success_callback": ["langfuse"],
|
||||
"failure_callback": [],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
audit_calls = []
|
||||
|
||||
async def capture(request_data):
|
||||
audit_calls.append(request_data)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
|
||||
patch(
|
||||
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
|
||||
new=capture,
|
||||
),
|
||||
):
|
||||
await disable_team_logging(
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id="team-1",
|
||||
user_api_key_dict=_admin_auth(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert audit_calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_team_callbacks_emits_audit_log_when_enabled(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "store_audit_logs", True)
|
||||
mock_prisma = _patch_prisma({"logging": []})
|
||||
|
||||
audit_calls = []
|
||||
|
||||
async def capture(request_data):
|
||||
audit_calls.append(request_data)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
patch(
|
||||
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
|
||||
new=capture,
|
||||
),
|
||||
):
|
||||
await add_team_callbacks(
|
||||
data=AddTeamCallback(
|
||||
callback_name="langfuse",
|
||||
callback_type="success",
|
||||
callback_vars={
|
||||
"langfuse_public_key": "pk",
|
||||
"langfuse_secret_key": "sk",
|
||||
},
|
||||
),
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id="team-1",
|
||||
user_api_key_dict=_admin_auth(),
|
||||
litellm_changed_by="ops-on-call",
|
||||
)
|
||||
import asyncio
|
||||
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert len(audit_calls) == 1
|
||||
log = audit_calls[0]
|
||||
assert log.table_name == LitellmTableNames.TEAM_TABLE_NAME
|
||||
assert log.object_id == "team-1"
|
||||
assert log.action == "updated"
|
||||
# ``litellm_changed_by`` header takes precedence over the auth user_id.
|
||||
assert log.changed_by == "ops-on-call"
|
||||
|
||||
before = json.loads(log.before_value)
|
||||
after = json.loads(log.updated_values)
|
||||
assert before["metadata"]["logging"] == []
|
||||
assert len(after["metadata"]["logging"]) == 1
|
||||
assert after["metadata"]["logging"][0]["callback_name"] == "langfuse"
|
||||
|
||||
# Callback secrets MUST NOT leak into the audit log payload.
|
||||
callback_vars = after["metadata"]["logging"][0]["callback_vars"]
|
||||
assert callback_vars["langfuse_public_key"] != "pk"
|
||||
assert callback_vars["langfuse_secret_key"] != "sk"
|
||||
# Key names are preserved so the auditor can see which fields changed.
|
||||
assert "langfuse_public_key" in callback_vars
|
||||
assert "langfuse_secret_key" in callback_vars
|
||||
# And no plaintext secret should appear anywhere in the serialized row.
|
||||
assert "sk" not in log.updated_values.replace("sk-", "") # crude leak check
|
||||
assert "pk" not in (log.updated_values.replace("pk-", "").replace("public_key", ""))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disable_team_logging_redacts_existing_callback_secrets(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "store_audit_logs", True)
|
||||
# Existing team has populated callback_vars containing secrets — redaction
|
||||
# must apply to the BEFORE snapshot too.
|
||||
mock_prisma = _patch_prisma(
|
||||
{
|
||||
"callback_settings": {
|
||||
"success_callback": ["langfuse"],
|
||||
"failure_callback": [],
|
||||
"callback_vars": {
|
||||
"langfuse_public_key": "pk-real",
|
||||
"langfuse_secret_key": "sk-real-secret",
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
audit_calls = []
|
||||
|
||||
async def capture(request_data):
|
||||
audit_calls.append(request_data)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
patch(
|
||||
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
|
||||
new=capture,
|
||||
),
|
||||
):
|
||||
await disable_team_logging(
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id="team-1",
|
||||
user_api_key_dict=_admin_auth(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
import asyncio
|
||||
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert len(audit_calls) == 1
|
||||
log = audit_calls[0]
|
||||
# The pre-existing secret_key value must NOT appear in the serialized
|
||||
# before_value or updated_values.
|
||||
assert "sk-real-secret" not in log.before_value
|
||||
assert "sk-real-secret" not in log.updated_values
|
||||
assert "pk-real" not in log.before_value
|
||||
assert "pk-real" not in log.updated_values
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_team_callbacks_no_audit_when_disabled(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "store_audit_logs", False)
|
||||
mock_prisma = _patch_prisma({"logging": []})
|
||||
|
||||
audit_calls = []
|
||||
|
||||
async def capture(request_data):
|
||||
audit_calls.append(request_data)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
|
||||
patch(
|
||||
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
|
||||
new=capture,
|
||||
),
|
||||
):
|
||||
await add_team_callbacks(
|
||||
data=AddTeamCallback(
|
||||
callback_name="langfuse",
|
||||
callback_type="success",
|
||||
callback_vars={
|
||||
"langfuse_public_key": "pk",
|
||||
"langfuse_secret_key": "sk",
|
||||
},
|
||||
),
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id="team-1",
|
||||
user_api_key_dict=_admin_auth(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert audit_calls == []
|
||||
Loading…
Reference in New Issue
Block a user