scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

feat(proxy): add a search param to key, memory, audit, and spend log listings

Browse source 
GET /key/list?search= matches the key hash (a raw sk- key is hashed
first) or a case-insensitive alias substring, and key_hash= now hashes a
raw sk- value too. GET /v1/memory?search= matches a key prefix or an
exact memory_id. GET /audit?search= matches id, object_id, changed_by,
or changed_by_api_key. GET /spend/logs/ui?search= matches request_id
across all time and api_key, team_id, user, end_user, session_id, or
model_id inside the date window; session grouping is skipped while a
search is active.

Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf
This commit is contained in:
ryan-crabbe-berri 2026-09-03 15:20:13 -07:00
parent 8699998c9e
commit 4811041048
9 changed files with 767 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
View file

@ -18,6 +18,7 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import (
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.utils import _hash_token_if_needed
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import AuditLogRepository
@ -48,6 +49,19 @@ def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, objec
}
def _build_search_condition(search: str) -> dict[str, object]:
"""Match any id column; a raw sk- key is hashed for the two columns that store key hashes."""
hashed: Final = _hash_token_if_needed(search)
return {
"OR": (
{"id": search},
{"changed_by": search},
{"object_id": hashed},
{"changed_by_api_key": hashed},
)
}
@router.get(
"/audit",
tags=["Audit Logging"],
@ -83,6 +97,13 @@ async def get_audit_logs(
None,
description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
),
search: str | None = Query(
None,
description=(
"Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value "
"(a raw sk- virtual key is hashed first)"
),
),
# Sorting parameters
sort_by: str | None = Query(
None,
@ -118,6 +139,11 @@ async def get_audit_logs(
*([_build_json_field_or_condition("token", object_key_hash)] if object_key_hash else []),
]
and_conditions: Final[tuple[dict[str, object], ...]] = (
*json_field_conditions,
*((_build_search_condition(search),) if search else ()),
)
# Build filter conditions
where_conditions: Final[dict[str, object]] = {
**({"changed_by": changed_by} if changed_by else {}),
@ -126,7 +152,7 @@ async def get_audit_logs(
**({"table_name": table_name} if table_name else {}),
**({"object_id": object_id} if object_id else {}),
**({"updated_at": date_filter} if start_date or end_date else {}),
**({"AND": json_field_conditions} if json_field_conditions else {}),
**({"AND": and_conditions} if and_conditions else {}),
}
order_by: Final[dict[str, str]] = (

30
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -147,6 +147,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyResponse,
BulkUpdateTeamKeysRequest,
FailedKeyUpdate,
KeySearchWhere,
SuccessfulKeyUpdate,
)
from litellm.types.router import Deployment
@ -5800,6 +5801,10 @@ async def list_keys(
None,
description="Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.",
),
search: str | None = Query(
None,
description="Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive).",
),
return_full_object: bool = Query(False, description="Return full key object"),
include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."),
include_created_by_keys: bool = Query(False, description="Include keys created by the user"),
@ -5862,13 +5867,17 @@ async def list_keys(
detail={"error": "Invalid expires value. Supported: 'active', 'expired'."},
)
hashed_key_hash: Final[str | None] = (
_hash_token_if_needed(token=key_hash) if isinstance(key_hash, str) else None
)
complete_user_info: Final = await validate_key_list_check(
user_api_key_dict=user_api_key_dict,
user_id=user_id,
team_id=team_id,
organization_id=organization_id,
key_alias=key_alias,
key_hash=key_hash,
key_hash=hashed_key_hash,
prisma_client=prisma_client,
)
@ -5928,7 +5937,7 @@ async def list_keys(
user_id=user_id,
team_id=team_id,
key_alias=key_alias,
key_hash=key_hash,
key_hash=hashed_key_hash,
return_full_object=return_full_object,
organization_id=organization_id,
admin_team_ids=admin_team_ids,
@ -5943,6 +5952,7 @@ async def list_keys(
agent_id=agent_id,
use_substring_matching=use_substring_matching,
expires_filter=expires if isinstance(expires, str) else None,
search=search,
)
verbose_proxy_logger.debug("Successfully prepared response")
@ -6162,6 +6172,16 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str,
return {"OR": [{"expires": None}, {"expires": {"gte": now}}]}
def _build_key_search_where(search: str) -> KeySearchWhere:
search_where: Final[KeySearchWhere] = {
"OR": (
{"token": _hash_token_if_needed(token=search)},
{"key_alias": {"contains": search, "mode": "insensitive"}},
)
}
return search_where
def _build_key_filter_conditions(
user_id: str | None,
team_id: str | None,
@ -6177,6 +6197,7 @@ def _build_key_filter_conditions(
agent_id: str | None = None,
use_substring_matching: bool = False,
expires_filter: str | None = None,
search: str | None = None,
) -> Mapping[str, object]:
"""Build filter conditions for key listing.
@ -6266,7 +6287,7 @@ def _build_key_filter_conditions(
# Apply team_id, project_id and access_group_id as global AND filters so they
# narrow results across all visibility conditions (own keys, team keys, etc.)
global_filters: Final[tuple[dict[str, object], ...]] = (
global_filters: Final[tuple[Mapping[str, object], ...]] = (
*(
(
{"key_alias": {"contains": key_alias, "mode": "insensitive"}}
@ -6277,6 +6298,7 @@ def _build_key_filter_conditions(
else ()
),
*(({"token": key_hash},) if key_hash and isinstance(key_hash, str) else ()),
*((_build_key_search_where(search),) if isinstance(search, str) and search else ()),
*(({"team_id": team_id},) if team_id and isinstance(team_id, str) else ()),
*(({"project_id": project_id},) if project_id else ()),
*(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()),
@ -6316,6 +6338,7 @@ async def _list_key_helper(
agent_id: str | None = None,
use_substring_matching: bool = False,
expires_filter: str | None = None,
search: str | None = None,
) -> KeyListResponseObject:
"""
Helper function to list keys
@ -6354,6 +6377,7 @@ async def _list_key_helper(
agent_id=agent_id,
use_substring_matching=use_substring_matching,
expires_filter=expires_filter,
search=search,
)
# Calculate skip for pagination

56
litellm/proxy/memory/memory_endpoints.py
View file

@ -22,6 +22,7 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Query
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
@ -91,6 +92,36 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object
return {"OR": ors}
class _StartsWith(TypedDict):
startsWith: ReadOnly[str]
class _MemoryKeyWhere(TypedDict):
key: ReadOnly[str | _StartsWith]
class _MemoryIdWhere(TypedDict):
memory_id: ReadOnly[str]
class _MemorySearchWhere(TypedDict):
OR: ReadOnly[tuple[_MemoryKeyWhere, _MemoryIdWhere]]
def _key_filter(search: str | None, key_prefix: str | None, key: str | None) -> Mapping[str, object] | None:
"""`search` matches a key prefix or an exact memory_id; otherwise `key_prefix` wins over `key`."""
if search is not None:
search_where: Final[_MemorySearchWhere] = {"OR": ({"key": {"startsWith": search}}, {"memory_id": search})}
return search_where
if key_prefix is not None:
prefix_where: Final[_MemoryKeyWhere] = {"key": {"startsWith": key_prefix}}
return prefix_where
if key is not None:
exact_where: Final[_MemoryKeyWhere] = {"key": key}
return exact_where
return None
def _row_to_model(row: "prisma_models.LiteLLM_MemoryTable") -> LiteLLM_MemoryRow:
return LiteLLM_MemoryRow(
memory_id=row.memory_id,
@ -326,6 +357,13 @@ async def list_memory(
"Mutually exclusive with `key`; if both are provided, `key_prefix` wins."
),
),
search: str | None = Query(
None,
description=(
"Match entries whose key starts with this value or whose memory_id equals it. "
"Takes precedence over `key_prefix` and `key` when provided."
),
),
page: int = Query(1, ge=1),
page_size: int = Query(50, ge=1, le=500),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -333,22 +371,16 @@ async def list_memory(
"""List memory entries visible to the caller."""
prisma_client: Final = _require_prisma()
# Build the key filter first (prefix wins if both `key` and `key_prefix`
# are passed). Then AND it with the visibility filter via an explicit
# top-level "AND" — safer than `dict.update` since future visibility
# filters could grow an "OR" key that would clobber this one if merged
# by key.
key_filter: Final[dict[str, object]] = {}
if key_prefix is not None:
key_filter["key"] = {"startsWith": key_prefix}
elif key is not None:
key_filter["key"] = key
# AND the key filter with the visibility filter via an explicit top-level
# "AND": both sides can carry an "OR" key (`search`, non-admin visibility),
# so merging them by key would let one clobber the other and leak rows.
key_filter: Final = _key_filter(search=search, key_prefix=key_prefix, key=key)
vis: Final = _visibility_filter(user_api_key_dict)
where: Mapping[str, object]
where: Mapping[str, object] | None
if vis is None:
where = key_filter
elif not key_filter:
elif key_filter is None:
where = vis
else:
where = {"AND": [key_filter, vis]}

62
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -2229,6 +2229,33 @@ async def calculate_spend(request: SpendCalculateRequest):
)
class _SpendLogSearchCondition(NamedTuple):
sql: str
params: tuple[object, ...]
def _build_spend_log_search_condition(
search: str,
start_date: datetime,
end_date: datetime,
next_param_index: int,
) -> _SpendLogSearchCondition:
"""request_id (indexed) matches across all time; the unindexed id columns only inside the window (sk- keys hashed)."""
raw: Final = f"${next_param_index}"
hashed: Final = f"${next_param_index + 1}"
window_start: Final = f"${next_param_index + 2}"
window_end: Final = f"${next_param_index + 3}"
sql: Final = (
f"(request_id = {raw} OR ("
f"\"startTime\" >= ({window_start}::timestamptz AT TIME ZONE 'UTC') "
f"AND \"startTime\" <= ({window_end}::timestamptz AT TIME ZONE 'UTC') "
f'AND (api_key = {hashed} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} '
f"OR session_id = {raw} OR model_id = {raw})))"
)
hashed_search: Final = hash_token(token=search) if search.startswith("sk-") else search
return _SpendLogSearchCondition(sql=sql, params=(search, hashed_search, start_date, end_date))
@router.get(
"/spend/logs/v2",
tags=["Budget & Spend Tracking"],
@ -2329,6 +2356,14 @@ async def ui_view_spend_logs(
"UI route only, honored when sorting by startTime"
),
),
search: str | None = fastapi.Query(
default=None,
description=(
"Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, "
"session_id, or model_id equals this value. request_id matches across all time; the other columns "
"match inside start_date/end_date, which stay required"
),
),
):
"""
View spend logs with pagination support.
@ -2392,8 +2427,10 @@ async def ui_view_spend_logs(
try:
is_admin_view: Final = _is_admin_view_safe(user_api_key_dict=user_api_key_dict)
is_request_id_lookup: Final = request_id is not None and not is_v2
is_search_lookup: Final = search is not None
search_owns_window: Final = is_search_lookup and not is_v2
if is_request_id_lookup:
if is_request_id_lookup and not is_search_lookup:
# request_id is the @id primary key: it identifies a single row, so a
# time window is meaningless. The dashboard always sends a default 24h
# window, which hid ids copied from an older page (LIT-3981). Drop the
@ -2576,7 +2613,7 @@ async def ui_view_spend_logs(
# Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison
# against the plain `timestamp` column does not depend on the DB session
# timezone (see #22529). Absent for a request_id-only lookup (see above).
if start_date_obj is not None and end_date_obj is not None:
if start_date_obj is not None and end_date_obj is not None and not search_owns_window:
sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')")
sql_params.append(start_date_obj)
p += 1
@ -2584,6 +2621,17 @@ async def ui_view_spend_logs(
sql_params.append(end_date_obj)
p += 1
if search is not None and start_date_obj is not None and end_date_obj is not None:
search_condition: Final = _build_spend_log_search_condition(
search=search,
start_date=start_date_obj,
end_date=end_date_obj,
next_param_index=p,
)
sql_conditions.append(search_condition.sql)
sql_params.extend(search_condition.params)
p += len(search_condition.params) # rebind-ok: advances the file's shared $N placeholder counter
# Equality filters - read effective values from where_conditions (post-authorization)
for sql_col, wc_key in [
("team_id", "team_id"),
@ -2662,7 +2710,13 @@ async def ui_view_spend_logs(
sql_params.append(f"%{error_message}%")
p += 1
if group_by_session is True and not is_v2 and not is_request_id_lookup and sort_by == "startTime":
if (
group_by_session is True
and not is_v2
and not is_request_id_lookup
and not is_search_lookup
and sort_by == "startTime"
):
return await _ui_session_grouped_spend_logs(
prisma_client=prisma_client,
sql_conditions=sql_conditions,
@ -2696,7 +2750,7 @@ async def ui_view_spend_logs(
_order_expr = order_column
joined_conditions: Final = " AND ".join(sql_conditions)
session_grouping: Final = group_by_session is True
session_grouping: Final = group_by_session is True and not is_search_lookup
count_group_clause: Final = f"GROUP BY {_SESSION_GROUP_KEY_SQL}" if session_grouping else ""
count_query: Final = f"""
SELECT COUNT(*) AS total_count

17
litellm/types/proxy/management_endpoints/key_management_endpoints.py
View file

@ -2,6 +2,23 @@ from datetime import datetime
from typing import Any, Final, Literal
from pydantic import BaseModel, ConfigDict, model_validator
from typing_extensions import ReadOnly, TypedDict
from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains
class KeyTokenWhere(TypedDict):
token: ReadOnly[str]
class KeyAliasContainsWhere(TypedDict):
key_alias: ReadOnly[InsensitiveContains]
class KeySearchWhere(TypedDict):
"""Prisma filter behind `/key/list?search=`: exact token (sk- keys hashed) or alias substring, case-insensitive."""
OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]]
class BulkUpdateKeyRequestItem(BaseModel):

66
tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py
View file

@ -1,5 +1,7 @@
import hashlib
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, patch
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import FastAPI
@ -8,10 +10,12 @@ from litellm_enterprise.proxy.audit_logging_endpoints import router as audit_rou
from litellm_enterprise.types.proxy.audit_logging_endpoints import AuditLogResponse
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
# Create an app with just the audit router for testing
app = FastAPI()
app.include_router(audit_router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin")
client = TestClient(app)
# Mock data for testing
@ -130,3 +134,63 @@ async def test_get_audit_log_by_id_not_found(mock_prisma_client):
data = response.json()
assert "message" in data["detail"]
assert "not found" in data["detail"]["message"].lower()
def _list_audit_logs_where(mock_prisma_client: MagicMock, query: str) -> dict[str, object]:
mock_prisma_client.db.litellm_auditlog.find_many.return_value = []
mock_prisma_client.db.litellm_auditlog.count.return_value = 0
response: Final = client.get(f"/audit?{query}")
assert response.status_code == 200, response.text
find_many_where: Final = mock_prisma_client.db.litellm_auditlog.find_many.call_args.kwargs["where"]
assert mock_prisma_client.db.litellm_auditlog.count.call_args.kwargs["where"] == find_many_where
return find_many_where
def test_search_matches_any_id_column_alongside_the_other_filters(mock_prisma_client):
where: Final = _list_audit_logs_where(mock_prisma_client, "search=abc-123&action=create&object_team_id=team-1")
assert where == {
"action": "create",
"AND": (
{
"OR": [
{"before_value": {"path": ["team_id"], "string_contains": "team-1"}},
{"updated_values": {"path": ["team_id"], "string_contains": "team-1"}},
]
},
{
"OR": (
{"id": "abc-123"},
{"changed_by": "abc-123"},
{"object_id": "abc-123"},
{"changed_by_api_key": "abc-123"},
)
},
),
}
def test_search_hashes_a_raw_virtual_key_for_the_hashed_columns(mock_prisma_client):
where: Final = _list_audit_logs_where(mock_prisma_client, "search=sk-raw")
hashed: Final = hashlib.sha256(b"sk-raw").hexdigest()
assert where == {
"AND": (
{
"OR": (
{"id": "sk-raw"},
{"changed_by": "sk-raw"},
{"object_id": hashed},
{"changed_by_api_key": hashed},
)
},
)
}
def test_an_empty_search_leaves_the_where_clause_unchanged(mock_prisma_client):
where: Final = _list_audit_logs_where(mock_prisma_client, "action=create&search=")
assert where == {"action": "create"}

152
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -6347,6 +6347,114 @@ def test_build_key_filter_conditions_key_hash_narrows_team_admin_visibility():
assert {"token": "hashed-token-123"} in where["AND"], f"key_hash not ANDed: {where}"
def _search_clause(search: str, token: str) -> dict:
return {"OR": [{"token": token}, {"key_alias": {"contains": search, "mode": "insensitive"}}]}
def test_build_key_filter_conditions_search_hashes_raw_key_and_ors_alias_contains():
"""
LIT-4741: `search` matches a key by its alias (case-insensitive contains) OR by
its ID. A pasted raw sk- key is hashed to its token first; an already-hashed
value is used verbatim.
"""
from litellm.proxy._types import hash_token
from litellm.proxy.management_endpoints.key_management_endpoints import (
_build_key_filter_conditions,
)
raw_where = json.loads(
json.dumps(
_build_key_filter_conditions(
user_id=None,
team_id=None,
organization_id=None,
key_alias=None,
key_hash=None,
exclude_team_id=None,
admin_team_ids=None,
search="sk-raw",
)
)
)
assert _search_clause("sk-raw", hash_token("sk-raw")) in raw_where["AND"], f"raw search not ANDed: {raw_where}"
hashed_where = json.loads(
json.dumps(
_build_key_filter_conditions(
user_id=None,
team_id=None,
organization_id=None,
key_alias=None,
key_hash=None,
exclude_team_id=None,
admin_team_ids=None,
search="already-hashed-token",
)
)
)
assert _search_clause("already-hashed-token", "already-hashed-token") in hashed_where["AND"], (
f"hashed search not used verbatim: {hashed_where}"
)
def test_build_key_filter_conditions_search_narrows_team_admin_visibility():
"""
LIT-4741, same class as LIT-3243: `search` must be a top-level AND so it
narrows a team admin's admin-team branch instead of being bypassed by it.
"""
from litellm.proxy._types import hash_token
from litellm.proxy.management_endpoints.key_management_endpoints import (
_build_key_filter_conditions,
)
where = json.loads(
json.dumps(
_build_key_filter_conditions(
user_id="team-admin-user",
team_id=None,
organization_id=None,
key_alias=None,
key_hash=None,
exclude_team_id=None,
admin_team_ids=["team-a"],
member_team_ids=["team-a"],
include_created_by_keys=False,
search="sk-member",
)
)
)
assert where.get("AND"), f"expected top-level AND, got: {where}"
assert _search_clause("sk-member", hash_token("sk-member")) in where["AND"], f"search not ANDed: {where}"
assert json.dumps({"team_id": {"in": ["team-a"]}}) in json.dumps(where)
@pytest.mark.asyncio
async def test_list_key_helper_applies_search_to_prisma_where():
"""LIT-4741: `search` given to _list_key_helper must reach the Prisma where clause."""
from litellm.proxy._types import hash_token
mock_prisma_client = AsyncMock()
mock_find_many = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many
mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)
await _list_key_helper(
prisma_client=mock_prisma_client,
page=1,
size=50,
user_id=None,
team_id=None,
organization_id=None,
key_alias=None,
key_hash=None,
search="sk-raw",
)
where = json.loads(json.dumps(mock_find_many.call_args.kwargs["where"]))
assert _search_clause("sk-raw", hash_token("sk-raw")) in where["AND"], f"search not in Prisma where: {where}"
@pytest.mark.asyncio
async def test_generate_key_negative_max_budget():
"""
@ -14870,6 +14978,50 @@ async def test_list_keys_non_admin_cannot_opt_into_substring():
assert kwargs["user_id"] == "alice"
@pytest.mark.asyncio
async def test_list_keys_hashes_raw_key_hash_before_validation():
"""LIT-4741: a raw sk- key pasted as key_hash is hashed before the ownership
check and the query, so a non-admin filtering by their own raw key gets the
row instead of the 'Key Hash not found.' 403."""
from litellm.proxy._types import hash_token
user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice")
validate = AsyncMock(
return_value=LiteLLM_UserTable(
user_id="alice", user_email="alice@example.com", teams=[], organization_memberships=[]
)
)
helper = AsyncMock(return_value={"keys": [], "total_count": 0, "current_page": 1, "total_pages": 0})
with (
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check",
validate,
),
patch("litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper", helper),
):
await list_keys(
request=MagicMock(),
user_api_key_dict=user,
status=None,
user_id=None,
key_hash="sk-raw",
)
assert validate.call_args.kwargs["key_hash"] == hash_token("sk-raw")
assert helper.call_args.kwargs["key_hash"] == hash_token("sk-raw")
@pytest.mark.asyncio
async def test_list_keys_search_is_honored_for_non_admin():
"""LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's
search reaches the helper while their own-user scoping stays in place."""
user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice")
kwargs = await _list_keys_capture_helper_kwargs(user, user_id=None, search="sk-raw")
assert kwargs["search"] == "sk-raw"
assert kwargs["user_id"] == "alice"
@pytest.mark.asyncio
async def test_cli_session_token_delegation_ceiling_blocked_by_team_budget():
team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0)

90
tests/test_litellm/proxy/memory/test_memory_endpoints.py
View file

@ -615,6 +615,96 @@ class TestMemoryEndpoints:
assert keys == {"user:profile"}
assert body["total"] == 1
def test_list_memory_search_matches_key_prefix_or_memory_id_within_scope(self):
"""
`search` matches a key prefix OR an exact memory_id, and stays ANDed
with the visibility filter so a pasted foreign id cannot leak a row.
"""
table = self.prisma.db.litellm_memorytable
table.rows.extend(
[
_make_row(memory_id="mem-own", key="user:profile", user_id="user-a", team_id=None),
_make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None),
_make_row(memory_id="mem-foreign", key="user:secret", user_id="user-b", team_id=None),
]
)
client = _make_client(_user_auth("user-a", "team-a"))
with _patch_prisma(self.prisma):
by_id = client.get("/v1/memory?search=mem-target")
by_prefix = client.get("/v1/memory?search=user:")
foreign_id = client.get("/v1/memory?search=mem-foreign")
assert by_id.status_code == 200, by_id.text
assert [m["memory_id"] for m in by_id.json()["memories"]] == ["mem-target"]
assert by_id.json()["total"] == 1
assert by_prefix.status_code == 200, by_prefix.text
assert {m["key"] for m in by_prefix.json()["memories"]} == {"user:profile"}
assert by_prefix.json()["total"] == 1
assert foreign_id.status_code == 200, foreign_id.text
assert foreign_id.json()["memories"] == []
assert foreign_id.json()["total"] == 0
def test_list_memory_search_by_memory_id_for_admin_sees_any_scope(self):
"""Admins have no visibility filter, so an id search returns the row whoever owns it."""
table = self.prisma.db.litellm_memorytable
table.rows.extend(
[
_make_row(memory_id="mem-a", key="a", user_id="user-a", team_id=None),
_make_row(memory_id="mem-b", key="b", user_id="user-b", team_id=None),
]
)
client = _make_client(_admin_auth())
with _patch_prisma(self.prisma):
resp = client.get("/v1/memory?search=mem-b")
assert resp.status_code == 200, resp.text
assert [m["memory_id"] for m in resp.json()["memories"]] == ["mem-b"]
assert resp.json()["total"] == 1
def test_list_memory_search_wins_over_key_prefix(self):
"""When both are sent, `search` decides the match and `key_prefix` is ignored."""
table = self.prisma.db.litellm_memorytable
table.rows.extend(
[
_make_row(memory_id="mem-own", key="user:profile", user_id="user-a", team_id=None),
_make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None),
]
)
client = _make_client(_user_auth("user-a", "team-a"))
with _patch_prisma(self.prisma):
resp = client.get("/v1/memory?search=mem-target&key_prefix=user:")
assert resp.status_code == 200, resp.text
assert [m["memory_id"] for m in resp.json()["memories"]] == ["mem-target"]
assert resp.json()["total"] == 1
def test_list_memory_key_prefix_never_matches_memory_id(self):
"""`key_prefix` stays a pure key-prefix match; only `search` consults memory_id."""
table = self.prisma.db.litellm_memorytable
table.rows.append(_make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None))
client = _make_client(_user_auth("user-a", "team-a"))
with _patch_prisma(self.prisma):
resp = client.get("/v1/memory?key_prefix=mem-target")
assert resp.status_code == 200, resp.text
assert resp.json()["memories"] == []
assert resp.json()["total"] == 0
def test_list_memory_key_exact_filter(self):
"""`key` is an exact match, never a prefix."""
table = self.prisma.db.litellm_memorytable
table.rows.extend(
[
_make_row(memory_id="m1", key="user:profile", user_id="user-a", team_id=None),
_make_row(memory_id="m2", key="user:profile:archived", user_id="user-a", team_id=None),
]
)
client = _make_client(_user_auth("user-a", "team-a"))
with _patch_prisma(self.prisma):
resp = client.get("/v1/memory?key=user:profile")
assert resp.status_code == 200, resp.text
assert [m["memory_id"] for m in resp.json()["memories"]] == ["m1"]
assert resp.json()["total"] == 1
def test_list_memory_admin_sees_all(self):
table = self.prisma.db.litellm_memorytable
table.rows.extend(

288
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -58,6 +58,25 @@ def _filter_logs_by_date_range(logs, where):
return filtered
_SEARCH_CLAUSE_RE = re.compile(
r'\(request_id = \$(\d+) OR \("startTime" >= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) '
r'AND "startTime" <= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) '
r'AND \(api_key = \$(\d+) OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 '
r"OR session_id = \$\1 OR model_id = \$\1\)\)\)"
)
def _matches_spend_log_search(log, search):
"""Mirror the search clause: request_id across all time, the other id columns inside the window."""
if log.get("request_id") == search["value"]:
return True
if not _filter_logs_by_date_range([log], {"startTime": {"gte": search["gte"], "lte": search["lte"]}}):
return False
if log.get("api_key") == search["api_key"]:
return True
return any(log.get(col) == search["value"] for col in ("team_id", "user", "end_user", "session_id", "model_id"))
def _reconstruct_ui_where_from_sql(sql_query, params):
"""
Rebuild the Prisma-style ``where`` dict the filter_fns below expect from the
@ -77,6 +96,17 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
def _iso(value):
return value.isoformat() if hasattr(value, "isoformat") else str(value)
search_clause = _SEARCH_CLAUSE_RE.search(clause.group(1))
if search_clause:
raw_index, start_index, end_index, hashed_index = (int(g) for g in search_clause.groups())
where["search"] = {
"value": params[raw_index - 1],
"api_key": params[hashed_index - 1],
"gte": _iso(params[start_index - 1]),
"lte": _iso(params[end_index - 1]),
}
remaining = clause.group(1) if search_clause is None else clause.group(1).replace(search_clause.group(0), "")
eq_cols = {
"team_id": "team_id",
'"user"': "user",
@ -89,7 +119,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
}
date_bounds: dict = {}
metadata_conds: list = []
for cond in (c.strip() for c in clause.group(1).split(" AND ")):
for cond in (c.strip() for c in remaining.split(" AND ")):
gte = re.search(r'"startTime" >= \(\$(\d+)', cond)
lte = re.search(r'"startTime" <= \(\$(\d+)', cond)
alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond)
@ -2352,6 +2382,219 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
app.dependency_overrides.pop(ps.user_api_key_auth, None)
def test_build_spend_log_search_condition_windows_every_branch_except_request_id():
"""LIT-4741: request_id matches across all time; the six other id columns only inside the window,
and a raw sk- key is hashed for the api_key branch alone."""
start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc)
end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc)
condition = spend_management_endpoints._build_spend_log_search_condition(
search="sk-raw-key", start_date=start, end_date=end, next_param_index=3
)
assert condition.sql == (
"(request_id = $3 OR (\"startTime\" >= ($5::timestamptz AT TIME ZONE 'UTC') "
"AND \"startTime\" <= ($6::timestamptz AT TIME ZONE 'UTC') "
'AND (api_key = $4 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))'
)
assert condition.params == ("sk-raw-key", hashlib.sha256(b"sk-raw-key").hexdigest(), start, end)
def test_build_spend_log_search_condition_leaves_non_key_values_unhashed():
start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc)
end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc)
condition = spend_management_endpoints._build_spend_log_search_condition(
search="sess-42", start_date=start, end_date=end, next_param_index=1
)
assert condition.params == ("sess-42", "sess-42", start, end)
def _search_fixture_logs(today):
recent = (today - datetime.timedelta(days=1)).isoformat()
old = (today - datetime.timedelta(days=90)).isoformat()
base = {
"api_key": "hashed-other",
"user": "user-x",
"team_id": "team-x",
"end_user": "cust-x",
"session_id": "sess-x",
"model_id": "mdl-x",
"spend": 0.01,
"model": "gpt-4",
}
return [
{**base, "request_id": "req-session", "session_id": "sess-42", "startTime": recent},
{**base, "request_id": "req-session-old", "session_id": "sess-42", "startTime": old},
{**base, "request_id": "req-key", "api_key": hashlib.sha256(b"sk-raw-key").hexdigest(), "startTime": recent},
{**base, "request_id": "req-team", "team_id": "team-7", "startTime": recent},
{**base, "request_id": "req-user", "user": "user-7", "startTime": recent},
{**base, "request_id": "req-end-user", "end_user": "cust-7", "startTime": recent},
{**base, "request_id": "req-model", "model_id": "mdl-7", "startTime": recent},
]
def _search_filter_fn(logs, captured):
def filter_fn(where):
captured["where"] = where
rows = _filter_logs_by_date_range(logs, where)
if "user" in where:
rows = [row for row in rows if row["user"] == where["user"]]
if "search" in where:
rows = [row for row in rows if _matches_spend_log_search(row, where["search"])]
return rows
return filter_fn
def _five_day_window(today):
return {
"start_date": (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S"),
"end_date": today.strftime("%Y-%m-%d %H:%M:%S"),
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"search,expected_request_ids",
[
("req-session-old", {"req-session-old"}),
("sess-42", {"req-session"}),
("sk-raw-key", {"req-key"}),
("team-7", {"req-team"}),
("user-7", {"req-user"}),
("cust-7", {"req-end-user"}),
("mdl-7", {"req-model"}),
("no-such-id", set()),
],
)
async def test_ui_view_spend_logs_search_matches_any_id(client, monkeypatch, search, expected_request_ids):
"""LIT-4741: one box matches any id column. A request_id is found across all time (the 5-day
window excludes the 90-day-old row), every other column only inside the window, and a raw
sk- key is hashed before it is compared with api_key. The window is not applied globally."""
today = datetime.datetime.now(timezone.utc)
logs = _search_fixture_logs(today)
captured = {}
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)),
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
try:
response = client.get(
"/spend/logs/ui",
params={"search": search, **_five_day_window(today)},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200, response.text
data = response.json()
assert {row["request_id"] for row in data["data"]} == expected_request_ids
assert data["total"] == len(expected_request_ids)
assert "startTime" not in captured["where"]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_spend_logs_v2_search_keeps_global_window(client, monkeypatch):
"""The public route keeps the caller's window on the whole query, so a search only finds rows
inside it even by request_id; the windowless request_id branch is a dashboard-only relaxation."""
today = datetime.datetime.now(timezone.utc)
logs = _search_fixture_logs(today)
captured = {}
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)),
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
try:
response = client.get(
"/spend/logs/v2",
params={"search": "req-session-old", **_five_day_window(today)},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200, response.text
data = response.json()
assert data["data"] == []
assert data["total"] == 0
assert "startTime" in captured["where"]
assert captured["where"]["search"]["value"] == "req-session-old"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"params",
[
{"search": "req-old"},
{"search": "req-old", "request_id": "req-old"},
],
)
async def test_ui_view_spend_logs_search_requires_dates(client, monkeypatch, params):
"""A search needs the window for its non-request_id branches, so it stays required even
alongside a request_id, which on its own may drop the window."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
make_ui_spend_logs_mock_prisma([], lambda where: []),
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
try:
response = client.get("/spend/logs/ui", params=params, headers={"Authorization": "Bearer sk-test"})
assert response.status_code == 400
assert "date" in response.text.lower()
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"search,expected_request_ids",
[("sess-9", {"req-own"}), ("req-foreign", set())],
)
async def test_ui_view_spend_logs_search_keeps_non_admin_scope(client, monkeypatch, search, expected_request_ids):
"""A search is scoped like any other listing: an internal user only sees their own rows even
when the id is on someone else's row, and the request_id ownership shortcut is not used."""
yesterday = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=1)).isoformat()
base = {"api_key": "hashed-key", "team_id": None, "spend": 0.01, "startTime": yesterday, "model": "gpt-4"}
logs = [
{**base, "request_id": "req-own", "user": "internal_user_1", "session_id": "sess-9"},
{**base, "request_id": "req-own-other", "user": "internal_user_1", "session_id": "sess-other"},
{**base, "request_id": "req-foreign", "user": "internal_user_2", "session_id": "sess-9"},
]
captured = {}
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)),
)
monkeypatch.setattr(
"litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs",
AsyncMock(return_value=[]),
)
ownership_check = AsyncMock()
monkeypatch.setattr(
"litellm.proxy.spend_tracking.spend_management_endpoints._assert_user_can_view_request_id",
ownership_check,
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1"
)
try:
start_date, end_date = _default_date_range()
response = client.get(
"/spend/logs/ui",
params={"search": search, "start_date": start_date, "end_date": end_date},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200, response.text
assert {row["request_id"] for row in response.json()["data"]} == expected_request_ids
assert captured["where"]["user"] == "internal_user_1"
ownership_check.assert_not_awaited()
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_unauthorized(client):
# Test without authorization header
@ -6351,3 +6594,46 @@ async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort
assert "OFFSET" in emitted_sql[1]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_session(client, monkeypatch):
"""The dashboard lists sessions by default; a search for an id lists every matching row instead,
so both calls of a session show up rather than one representative, and no session cursor is returned."""
rows = [_session_representative_row("req-1", "sess-1"), _session_representative_row("req-2", "sess-1")]
async def mock_query_raw(sql_query, *params):
if "mcp_tool_call_count" in sql_query:
return []
grouped = "DISTINCT ON" in sql_query or "GROUP BY" in sql_query
visible = rows[:1] if grouped else rows
if "COUNT(*)" in sql_query:
return [{"total_count": len(visible)}]
return visible
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
)
try:
start_date, end_date = _default_date_range()
response = client.get(
"/spend/logs/ui",
params={
"search": "sess-1",
"group_by_session": "true",
"start_date": start_date,
"end_date": end_date,
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200, response.text
data = response.json()
assert [row["request_id"] for row in data["data"]] == ["req-1", "req-2"]
assert data["total"] == 2
assert "next_session_cursor" not in data
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)