tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-09 02:24:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #25647 from BerriAI/litellm_yj_apr_11

Browse Source 
[Infra] Merge dev branch with main
This commit is contained in:
yuneng-jiang
2026-04-13 17:28:38 -07:00
committed by GitHub
parent c25b4b2ce8 df75e79615
commit 8427534f13
13 changed files with 521 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/oidc.md
+21 -2
View File
@@ -57,12 +57,31 @@ oidc/config_name_here/
#### Unofficial Providers (not recommended)
For the unofficial `file` provider, you can use the following format:
For the unofficial `file` provider, you can use the following format
(note the double slash — the path after `oidc/file/` must be absolute):
```
oidc/file/home/user/dave/this_is_a_file_with_a_token.txt
oidc/file//var/run/secrets/my-token
```
For safety, the resolved path must live inside an allowed credential
directory. By default the following directories are allowed:
- `/var/run/secrets`
- `/run/secrets`
If your deployment mounts credentials elsewhere, set the
`LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS` environment variable to a
comma-separated list of absolute directories. The value replaces the
default list, so include the defaults if you still need them:
```bash
export LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS="/var/run/secrets,/etc/litellm/creds"
```
Paths that resolve (after following symlinks and `..`) outside the
allowlist are rejected.
For the unofficial `env`, use the following format, where `SECRET_TOKEN` is the name of the environment variable that contains the token:
```
docs/my-website/docs/proxy/config_settings.md
+1
View File
@@ -830,6 +830,7 @@ router_settings:
| LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS | TTL in seconds for the distributed lock used by the key rotation job. Default is 600 (10 minutes).
| LITELLM_LICENSE | License key for LiteLLM usage
| LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS | Set to `True` to use the local bundled Anthropic beta headers config only, disabling remote fetching. Default is `False`
| LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS | Comma-separated list of absolute directories from which the `oidc/file/` provider is permitted to read token files. Defaults to `/var/run/secrets,/run/secrets`.
| LITELLM_LOCAL_BLOG_POSTS | When set to `True`, uses the local bundled blog posts only, disabling remote fetching from GitHub. Default is `False`
| LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM
| LITELLM_LOCAL_POLICY_TEMPLATES | When set to "true", uses local backup policy templates instead of fetching from GitHub. Policy templates are fetched from https://raw.githubusercontent.com/BerriAI/litellm/main/policy_templates.json by default, with automatic fallback to local backup on failure
litellm/litellm_core_utils/initialize_dynamic_callback_params.py
+26 -13
View File
@@ -1,7 +1,28 @@
from typing import Dict, Optional
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import StandardCallbackDynamicParams
def _is_env_reference(value: object) -> bool:
return isinstance(value, str) and "os.environ/" in value
def _raise_env_reference_error(param: str, *, source: str) -> None:
raise ValueError(
f"Callback param '{param}' (from {source}) contains an 'os.environ/' "
"reference. Environment references in request-supplied parameters are "
"no longer resolved server-side for security reasons.\n"
"To resolve:\n"
" 1. Remove the 'os.environ/' reference from your request body / "
"metadata.\n"
" 2. Either (a) configure this callback value in your proxy "
"config.yaml under 'litellm_settings' / 'general_settings', or "
"(b) pass the resolved secret value directly in the request.\n"
"See https://docs.litellm.ai/docs/proxy/logging for server-side "
"callback configuration."
)
# Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict
_supported_callback_params = [
"langfuse_public_key",
@@ -46,12 +67,8 @@ def initialize_standard_callback_dynamic_params(
for param in _supported_callback_params:
if param in kwargs:
_param_value = kwargs.get(param)
if (
_param_value is not None
and isinstance(_param_value, str)
and "os.environ/" in _param_value
):
_param_value = get_secret_str(secret_name=_param_value)
if _is_env_reference(_param_value):
_raise_env_reference_error(param, source="request body")
standard_callback_dynamic_params[param] = _param_value # type: ignore
# 2. Fallback: check "metadata" or "litellm_params" -> "metadata"
@@ -64,12 +81,8 @@ def initialize_standard_callback_dynamic_params(
for param in _supported_callback_params:
if param not in standard_callback_dynamic_params and param in metadata:
_param_value = metadata.get(param)
if (
_param_value is not None
and isinstance(_param_value, str)
and "os.environ/" in _param_value
):
_param_value = get_secret_str(secret_name=_param_value)
if _is_env_reference(_param_value):
_raise_env_reference_error(param, source="metadata")
standard_callback_dynamic_params[param] = _param_value # type: ignore
return standard_callback_dynamic_params
litellm/proxy/health_endpoints/_health_endpoints.py
+31 -67
View File
@@ -5,7 +5,7 @@ import os
import time
import traceback
from datetime import datetime, timedelta
from typing import Any, Dict, Literal, Optional, Union, cast
from typing import Any, Dict, Iterable, Literal, Optional, Union, cast
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
@@ -36,79 +36,43 @@ from litellm.proxy.health_check import (
from litellm.proxy.middleware.in_flight_requests_middleware import (
get_in_flight_requests,
)
from litellm.secret_managers.main import get_secret
#### Health ENDPOINTS ####
def _resolve_os_environ_variables(params: dict) -> dict:
def _reject_os_environ_references(params: dict) -> None:
"""
Resolve ``os.environ/`` environment variables in ``litellm_params``.
This walks the input dict/list structure iteratively (no Python recursion) to
avoid unbounded recursion / stack overflows on deeply nested inputs.
Validate that the provided params do not contain any ``os.environ/``
references. Values with that prefix are expected to come only from
server-side configuration (already resolved before reaching here). If a
request-supplied value still carries the prefix, raise ``HTTPException``.
"""
if not isinstance(params, dict):
return params
return
# Use an explicit stack to avoid recursion and handle nested dicts/lists.
# We also keep a `seen` set to guard against accidental cycles.
resolved_root: dict = {}
stack: list[tuple[object, object]] = [(params, resolved_root)]
stack: list[object] = [params]
seen: set[int] = {id(params)}
while stack:
src, dst = stack.pop()
src = stack.pop()
if isinstance(src, dict):
values: Iterable[object] = src.values()
elif isinstance(src, list):
values = src
else:
continue
if isinstance(src, dict) and isinstance(dst, dict):
for key, value in src.items():
# Direct string replacement for os.environ/ references
if isinstance(value, str) and value.startswith("os.environ/"):
dst[key] = get_secret(value)
elif isinstance(value, dict):
if id(value) in seen:
# Cycle detected – keep a shallow copy reference to prevent infinite loops
dst[key] = {}
continue
seen.add(id(value))
new_dict: dict = {}
dst[key] = new_dict
stack.append((value, new_dict))
elif isinstance(value, list):
if id(value) in seen:
dst[key] = []
continue
seen.add(id(value))
new_list: list = []
dst[key] = new_list
stack.append((value, new_list))
else:
dst[key] = value
elif isinstance(src, list) and isinstance(dst, list):
for item in src:
if isinstance(item, str) and item.startswith("os.environ/"):
dst.append(get_secret(item))
elif isinstance(item, dict):
if id(item) in seen:
dst.append({})
continue
seen.add(id(item))
new_dict = {}
dst.append(new_dict)
stack.append((item, new_dict))
elif isinstance(item, list):
if id(item) in seen:
dst.append([])
continue
seen.add(id(item))
new_list = []
dst.append(new_list)
stack.append((item, new_list))
else:
dst.append(item)
return resolved_root
for value in values:
if isinstance(value, str) and value.startswith("os.environ/"):
raise HTTPException(
status_code=400,
detail={
"error": "Environment variable references are not permitted in request parameters."
},
)
if isinstance(value, (dict, list)) and id(value) not in seen:
seen.add(id(value))
stack.append(value)
def get_callback_identifier(callback):
@@ -1510,6 +1474,10 @@ async def test_model_connection(
# Get model name from litellm_params
request_litellm_params = litellm_params or {}
# Reject request-supplied os.environ/ references. Config values are
# already resolved before reaching this endpoint; any remaining
# reference must have come from the request body.
_reject_os_environ_references(request_litellm_params)
model_name = request_litellm_params.get("model")
# Look up model configuration from router if model name is provided
@@ -1546,11 +1514,7 @@ async def test_model_connection(
# Merge: config params (from proxy config) as base, request params override
# This allows users to override specific params while using config for credentials
merged_litellm_params = {**config_litellm_params, **request_litellm_params}
# Resolve os.environ/ environment variables in any remaining request params
# This handles cases where user explicitly passes os.environ/ values to override config
litellm_params = _resolve_os_environ_variables(merged_litellm_params)
litellm_params = {**config_litellm_params, **request_litellm_params}
## Auth check
await ModelManagementAuthChecks.can_user_make_model_call(
litellm/proxy/management_endpoints/internal_user_endpoints.py
+33
View File
@@ -1131,6 +1131,16 @@ async def _update_single_user_helper(
if prisma_client is None:
raise Exception("Not connected to DB!")
# Only proxy admins can modify user_role
if (
user_request.user_role is not None
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
):
raise HTTPException(
status_code=403,
detail="Only proxy admins can modify user roles.",
)
# Validate user identifier
if not user_request.user_id and not user_request.user_email:
raise ValueError("Either user_id or user_email must be provided")
@@ -1508,12 +1518,35 @@ async def bulk_user_update(
detail={"error": "Database not connected"},
)
# Only proxy admins can modify user_role in bulk updates
_bulk_role = (
getattr(data.user_updates, "user_role", None) if data.user_updates else None
)
if _bulk_role is None and data.users:
_bulk_role = next(
(u.user_role for u in data.users if u.user_role is not None), None
)
if (
_bulk_role is not None
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
):
raise HTTPException(
status_code=403,
detail="Only proxy admins can modify user roles.",
)
# Determine the list of users to update
users_to_update: Union[
List[UpdateUserRequest], List[UpdateUserRequestNoUserIDorEmail]
] = []
if data.all_users and data.user_updates:
# Only proxy admins can update all users at once
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
raise HTTPException(
status_code=403,
detail="Only proxy admins can update all users at once.",
)
# Optimized path for updating all users directly in database
all_users_in_db = await prisma_client.db.litellm_usertable.find_many(
order={"created_at": "desc"}
litellm/proxy/management_endpoints/key_management_endpoints.py
+27 -22
View File
@@ -768,9 +768,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
request_type="key", **data_json, table_name="key"
)
response["soft_budget"] = (
data.soft_budget
) # include the user-input soft budget in the response
response[
"soft_budget"
] = data.soft_budget # include the user-input soft budget in the response
response = GenerateKeyResponse(**response)
@@ -1961,8 +1961,13 @@ async def _validate_update_key_data(
user_api_key_cache=user_api_key_cache,
)
# Admin-only: only proxy admins, team admins, or org admins can modify max_budget
if data.max_budget is not None and data.max_budget != existing_key_row.max_budget:
# Admin-only: only proxy admins, team admins, or org admins can modify max_budget or spend
if (
data.max_budget is not None and data.max_budget != existing_key_row.max_budget
) or (
data.spend is not None
and data.spend != getattr(existing_key_row, "spend", None)
):
if prisma_client is not None:
hashed_key = existing_key_row.token
await _check_key_admin_access(
@@ -1970,7 +1975,7 @@ async def _validate_update_key_data(
hashed_token=hashed_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
route="/key/update (max_budget)",
route="/key/update (max_budget/spend)",
)
# Check team limits if key has a team_id (from request or existing key)
@@ -3272,10 +3277,10 @@ async def delete_verification_tokens(
try:
if prisma_client:
tokens = [_hash_token_if_needed(token=key) for key in tokens]
_keys_being_deleted: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
)
_keys_being_deleted: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
)
if len(_keys_being_deleted) == 0:
@@ -3475,9 +3480,9 @@ async def _rotate_master_key( # noqa: PLR0915
from litellm.proxy.proxy_server import proxy_config
try:
models: Optional[List] = (
await prisma_client.db.litellm_proxymodeltable.find_many()
)
models: Optional[
List
] = await prisma_client.db.litellm_proxymodeltable.find_many()
except Exception:
models = None
# 2. process model table
@@ -4117,11 +4122,11 @@ async def validate_key_list_check(
param="user_id",
code=status.HTTP_403_FORBIDDEN,
)
complete_user_info_db_obj: Optional[BaseModel] = (
await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
)
complete_user_info_db_obj: Optional[
BaseModel
] = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
)
if complete_user_info_db_obj is None:
@@ -4204,10 +4209,10 @@ async def _fetch_user_team_objects(
if complete_user_info is None or not complete_user_info.teams:
return []
teams: Optional[List[BaseModel]] = (
await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
)
teams: Optional[
List[BaseModel]
] = await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
)
if teams is None:
return []
litellm/proxy/spend_tracking/spend_management_endpoints.py
+28 -44
View File
@@ -2297,10 +2297,13 @@ async def view_spend_logs( # noqa: PLR0915
}
if api_key is not None and isinstance(api_key, str):
filter_query["api_key"] = api_key # type: ignore
elif request_id is not None and isinstance(request_id, str):
if api_key.startswith("sk-"):
filter_query["api_key"] = prisma_client.hash_token(token=api_key) # type: ignore
else:
filter_query["api_key"] = api_key # type: ignore
if request_id is not None and isinstance(request_id, str):
filter_query["request_id"] = request_id # type: ignore
elif user_id is not None and isinstance(user_id, str):
if user_id is not None and isinstance(user_id, str):
filter_query["user"] = user_id # type: ignore
# Check if user wants unsummarized data
@@ -2375,49 +2378,30 @@ async def view_spend_logs( # noqa: PLR0915
return response
elif api_key is not None and isinstance(api_key, str):
if api_key.startswith("sk-"):
hashed_token = prisma_client.hash_token(token=api_key)
else:
hashed_token = api_key
spend_log = await prisma_client.get_data(
table_name="spend",
query_type="find_all",
key_val={"key": "api_key", "value": hashed_token},
)
if spend_log is None:
return []
if isinstance(spend_log, list):
return spend_log
else:
return [spend_log]
elif request_id is not None:
spend_log = await prisma_client.get_data(
table_name="spend",
query_type="find_unique",
key_val={"key": "request_id", "value": request_id},
)
if spend_log is None:
return []
return [spend_log]
elif user_id is not None:
spend_log = await prisma_client.get_data(
table_name="spend",
query_type="find_all",
key_val={"key": "user", "value": user_id},
)
if spend_log is None:
return []
if isinstance(spend_log, list):
return spend_log
else:
return [spend_log]
else:
spend_logs = await prisma_client.get_data(
table_name="spend", query_type="find_all"
)
scoped_filter: Dict[str, Any] = {}
if api_key is not None and isinstance(api_key, str):
if api_key.startswith("sk-"):
hashed_token = prisma_client.hash_token(token=api_key)
else:
hashed_token = api_key
scoped_filter["api_key"] = hashed_token
if request_id is not None and isinstance(request_id, str):
scoped_filter["request_id"] = request_id
if user_id is not None and isinstance(user_id, str):
scoped_filter["user"] = user_id
return spend_logs
if not scoped_filter:
spend_logs = await prisma_client.get_data(
table_name="spend", query_type="find_all"
)
return spend_logs
data = await prisma_client.db.litellm_spendlogs.find_many(
where=scoped_filter, # type: ignore
order={"startTime": "desc"},
)
return data
return None
litellm/secret_managers/main.py
+49 -2
View File
@@ -16,6 +16,52 @@ from litellm.secret_managers.secret_manager_handler import get_secret_from_manag
oidc_cache = DualCache()
_DEFAULT_OIDC_ALLOWED_CREDENTIAL_DIRS = ("/var/run/secrets", "/run/secrets")
def _get_oidc_allowed_credential_dirs() -> list[str]:
"""
Return the absolute, normalized list of directories from which
``oidc/file/`` is permitted to read token files.
Defaults to standard container credential mount points. Operators can
override via the ``LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS`` environment
variable (comma-separated list of absolute paths).
"""
override = os.getenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS")
raw_dirs = (
[d.strip() for d in override.split(",") if d.strip()]
if override
else list(_DEFAULT_OIDC_ALLOWED_CREDENTIAL_DIRS)
)
return [os.path.realpath(d) for d in raw_dirs]
def _resolve_oidc_file_path(requested_path: str) -> str:
"""
Resolve ``requested_path`` and verify it falls within one of the allowed
credential directories. Raises ``ValueError`` otherwise.
"""
if not os.path.isabs(requested_path):
raise ValueError(
"oidc/file path must be absolute. Use the format "
"'oidc/file//var/run/secrets/<name>' (note the leading slash "
"after 'oidc/file/')."
)
resolved = os.path.realpath(requested_path)
for allowed in _get_oidc_allowed_credential_dirs():
try:
if os.path.commonpath([resolved, allowed]) == allowed:
return resolved
except ValueError:
# commonpath raises when paths are on different drives (Windows);
# treat as not-matching and continue.
continue
raise ValueError(
"oidc/file path is outside the allowed credential directories. "
"Set LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS to extend the allowlist."
)
def _get_oidc_http_handler(timeout: Optional[httpx.Timeout] = None) -> HTTPHandler:
"""
@@ -196,8 +242,9 @@ def get_secret( # noqa: PLR0915
oidc_token = f.read()
return oidc_token
elif oidc_provider == "file":
# Load token from a file
with open(oidc_aud, "r") as f:
# Load token from a file within an allowed credential directory.
safe_path = _resolve_oidc_file_path(oidc_aud)
with open(safe_path, "r") as f:
oidc_token = f.read()
return oidc_token
elif oidc_provider == "env":
tests/litellm_utils_tests/test_secret_manager.py
+7 -6
View File
@@ -183,13 +183,14 @@ def test_oidc_env_variable():
del os.environ[env_var_name]
def test_oidc_file():
# Create a temporary file
with tempfile.NamedTemporaryFile(mode="w+") as temp_file:
def test_oidc_file(monkeypatch):
# Create a temporary file inside a directory added to the allowlist.
with tempfile.TemporaryDirectory() as temp_dir:
monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", temp_dir)
temp_file_path = os.path.join(temp_dir, "token.txt")
secret_value = "secret-" + uuid4().hex
temp_file.write(secret_value)
temp_file.flush()
temp_file_path = temp_file.name
with open(temp_file_path, "w") as temp_file:
temp_file.write(secret_value)
secret_val = get_secret(f"oidc/file/{temp_file_path}")
tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py
+99
View File
@@ -0,0 +1,99 @@
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
initialize_standard_callback_dynamic_params,
)
def test_resolves_plain_values_at_top_level():
kwargs = {
"langfuse_public_key": "pk-test",
"langfuse_secret_key": "sk-test",
}
params = initialize_standard_callback_dynamic_params(kwargs)
assert params.get("langfuse_public_key") == "pk-test"
assert params.get("langfuse_secret_key") == "sk-test"
def test_resolves_plain_values_from_metadata():
kwargs = {
"metadata": {
"langfuse_public_key": "pk-meta",
"langfuse_host": "https://test.langfuse.com",
}
}
params = initialize_standard_callback_dynamic_params(kwargs)
assert params.get("langfuse_public_key") == "pk-meta"
assert params.get("langfuse_host") == "https://test.langfuse.com"
def test_env_reference_at_top_level_raises_with_guidance():
kwargs = {"langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY"}
with pytest.raises(ValueError) as exc_info:
initialize_standard_callback_dynamic_params(kwargs)
message = str(exc_info.value)
assert "langfuse_public_key" in message
assert "request body" in message
assert "os.environ/" in message
assert "config.yaml" in message
def test_env_reference_in_metadata_raises_with_guidance():
kwargs = {
"metadata": {
"langsmith_api_key": "os.environ/LANGSMITH_API_KEY",
}
}
with pytest.raises(ValueError) as exc_info:
initialize_standard_callback_dynamic_params(kwargs)
message = str(exc_info.value)
assert "langsmith_api_key" in message
assert "metadata" in message
def test_env_reference_in_litellm_params_metadata_raises():
kwargs = {
"litellm_params": {
"metadata": {
"gcs_bucket_name": "os.environ/GCS_BUCKET",
}
}
}
with pytest.raises(ValueError) as exc_info:
initialize_standard_callback_dynamic_params(kwargs)
assert "gcs_bucket_name" in str(exc_info.value)
def test_non_string_values_are_not_flagged():
kwargs = {
"langsmith_sampling_rate": 0.5,
"turn_off_message_logging": True,
}
params = initialize_standard_callback_dynamic_params(kwargs)
assert params.get("langsmith_sampling_rate") == 0.5
assert params.get("turn_off_message_logging") is True
def test_empty_kwargs_returns_empty_params():
params = initialize_standard_callback_dynamic_params(None)
assert dict(params) == {}
params = initialize_standard_callback_dynamic_params({})
assert dict(params) == {}
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
+6 -6
View File
@@ -450,10 +450,10 @@ async def test_test_model_connection_loads_config_from_router():
params["messages"] = [{"role": "user", "content": "test"}]
return params
# Mock _resolve_os_environ_variables
def mock_resolve_os_environ(params):
return params
# Mock _reject_os_environ_references
def mock_reject_os_environ(params):
return None
with patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
@@ -476,8 +476,8 @@ async def test_test_model_connection_loads_config_from_router():
"litellm.proxy.health_endpoints._health_endpoints._update_litellm_params_for_health_check",
mock_update_params,
), patch(
"litellm.proxy.health_endpoints._health_endpoints._resolve_os_environ_variables",
mock_resolve_os_environ,
"litellm.proxy.health_endpoints._health_endpoints._reject_os_environ_references",
mock_reject_os_environ,
):
# Call the endpoint with only model name (no credentials)
result = await health_test_model_connection(
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
+173
View File
@@ -2761,3 +2761,176 @@ async def test_ui_view_spend_logs_team_member_no_permission_blocked(
assert response.status_code == 403
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
class _CaptureFilterDB:
"""Mock DB that records the `where` filter passed to find_many."""
def __init__(self):
self.litellm_spendlogs = self
self.captured_where = None
async def find_many(self, *args, **kwargs):
self.captured_where = kwargs.get("where")
return []
async def group_by(self, *args, **kwargs):
self.captured_where = kwargs.get("where")
return []
class _CapturePrismaClient:
def __init__(self):
self.db = _CaptureFilterDB()
def hash_token(self, token):
return "hashed::" + token
@pytest.mark.asyncio
async def test_view_spend_logs_internal_user_combines_user_with_api_key(
client, monkeypatch
):
"""Internal users must have their user filter applied alongside api_key."""
mock_client = _CapturePrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_client)
start_date = "2024-01-01"
end_date = "2024-12-31"
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="internal-user-1",
)
try:
response = client.get(
"/spend/logs",
params={
"start_date": start_date,
"end_date": end_date,
"summarize": "false",
"api_key": "sk-some-raw-token",
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
where = mock_client.db.captured_where
assert where is not None
assert where["user"] == "internal-user-1"
assert where["api_key"] == "hashed::sk-some-raw-token"
assert "startTime" in where
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_view_spend_logs_internal_user_combines_user_with_request_id(
client, monkeypatch
):
"""Internal users must have their user filter applied alongside request_id."""
mock_client = _CapturePrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_client)
start_date = "2024-01-01"
end_date = "2024-12-31"
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="internal-user-2",
)
try:
response = client.get(
"/spend/logs",
params={
"start_date": start_date,
"end_date": end_date,
"summarize": "false",
"request_id": "req-abc",
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
where = mock_client.db.captured_where
assert where is not None
assert where["user"] == "internal-user-2"
assert where["request_id"] == "req-abc"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_view_spend_logs_non_date_range_combines_user_with_request_id(
client, monkeypatch
):
"""Non-date-range path must also combine user + request_id filters."""
mock_client = _CapturePrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_client)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="internal-user-3",
)
try:
response = client.get(
"/spend/logs",
params={"request_id": "req-xyz"},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
where = mock_client.db.captured_where
assert where is not None
assert where["user"] == "internal-user-3"
assert where["request_id"] == "req-xyz"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_view_spend_logs_non_date_range_hashes_sk_api_key(client, monkeypatch):
"""Non-date-range path must hash sk- prefixed api_keys before filtering."""
mock_client = _CapturePrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_client)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
)
try:
response = client.get(
"/spend/logs",
params={"api_key": "sk-raw-admin-token"},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
where = mock_client.db.captured_where
assert where is not None
assert where["api_key"] == "hashed::sk-raw-admin-token"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_view_spend_logs_date_range_hashes_sk_api_key(client, monkeypatch):
"""Date-range path must hash sk- prefixed api_keys before filtering."""
mock_client = _CapturePrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_client)
start_date = "2024-01-01"
end_date = "2024-12-31"
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
)
try:
response = client.get(
"/spend/logs",
params={
"start_date": start_date,
"end_date": end_date,
"summarize": "false",
"api_key": "sk-raw-admin-token",
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
where = mock_client.db.captured_where
assert where is not None
assert where["api_key"] == "hashed::sk-raw-admin-token"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
tests/test_litellm/secret_managers/test_secret_managers_main.py
+20 -1
View File
@@ -199,9 +199,10 @@ def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider, monkeypat
mock_token_provider.assert_called_once_with()
def test_oidc_file_success(tmp_path):
def test_oidc_file_success(tmp_path, monkeypatch):
token_file = tmp_path / "token.txt"
token_file.write_text("file_token")
monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path))
secret_name = f"oidc/file/{token_file}"
result = get_secret(secret_name)
@@ -209,6 +210,24 @@ def test_oidc_file_success(tmp_path):
assert result == "file_token"
def test_oidc_file_rejects_path_outside_allowlist(tmp_path, monkeypatch):
outside_file = tmp_path / "outside.txt"
outside_file.write_text("should_not_read")
# Allowlist a different directory.
allowed_dir = tmp_path / "allowed"
allowed_dir.mkdir()
monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(allowed_dir))
with pytest.raises(ValueError, match="outside the allowed credential directories"):
get_secret(f"oidc/file/{outside_file}")
def test_oidc_file_rejects_relative_path(tmp_path, monkeypatch):
monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path))
with pytest.raises(ValueError, match="must be absolute"):
get_secret("oidc/file/relative/path/token")
def test_oidc_env_success(mock_env):
mock_env["CUSTOM_TOKEN"] = "env_token"