diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 27186800e2b..6b7aa4d15d2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2101,6 +2101,11 @@ class BlockModelRequest(LiteLLMPydanticObjectBase): model_id: str # required +CALLBACK_VAR_NAMES: Final = frozenset( + {*StandardCallbackDynamicParams.__annotations__, CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR} +) + + class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str callback_type: Literal["success", "failure", "success_and_failure"] | None = "success_and_failure" @@ -2110,14 +2115,9 @@ class AddTeamCallback(LiteLLMPydanticObjectBase): @classmethod def validate_callback_vars(cls, values): callback_vars: Final = values.get("callback_vars", {}) - valid_keys: Final = { - *StandardCallbackDynamicParams.__annotations__, - CAPTURE_MESSAGE_CONTENT_VAR, - EXCLUDED_SERVICES_VAR, - } for key, value in callback_vars.items(): - if key not in valid_keys: - raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}") + if key not in CALLBACK_VAR_NAMES: + raise ValueError(f"Invalid callback variable: {key}. Must be one of {set(CALLBACK_VAR_NAMES)}") if key in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) and not isinstance(value, str): raise ValueError(f"{key} must be a string, got {type(value).__name__}") callback_vars[key] = str(value) diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 302978d17a4..a486d6665fc 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -7,9 +7,13 @@ per-integration checks here. import math from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType -from typing import Final, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings +from typing import Final, Literal, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings +from pydantic import TypeAdapter, ValidationError + +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR _NEWRELIC_CALLBACK: Final = "newrelic" @@ -276,7 +280,7 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str """Validate every ``logging`` entry of a team/key metadata payload.""" if not metadata: return None - entries: Final = _object_sequence(metadata.get("logging")) + entries: Final = object_sequence(metadata.get("logging")) if entries is None: return None entry_error: Final = next( @@ -302,11 +306,11 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str def stored_callback_entries(logging_entries: object) -> tuple[tuple[str | None, Mapping[str, str]], ...]: """The callback name and vars of every stored ``logging`` entry.""" - entries: Final = _object_sequence(logging_entries) or () + entries: Final = object_sequence(logging_entries) or () return tuple((_entry_callback_name(entry), _entry_callback_vars(entry)) for entry in entries) -def _object_sequence(value: object) -> Sequence[object] | None: +def object_sequence(value: object) -> Sequence[object] | None: return value if isinstance(value, Sequence) and not isinstance(value, (str, bytes)) else None @@ -445,3 +449,72 @@ def _newrelic_config_error(callback_vars: Mapping[str, str]) -> str | None: if region is not None and (not api_key or api_key == "None"): return "newrelic_region requires newrelic_api_key; the region rides the team's own key." return None + + +@dataclass(frozen=True, slots=True) +class CredentialReferenceRejection: + status_code: Literal[400, 403] + message: str + + +def credential_is_defined(credential_name: str) -> bool: + return CredentialAccessor.find_credential(credential_name) is not None + + +_STRING_KEYED: Final = TypeAdapter(Mapping[str, object]) + + +def string_keyed(value: object) -> Mapping[str, object]: + try: + return _STRING_KEYED.validate_python(value) + except ValidationError: + return MappingProxyType({}) + + +def _logging_entries(metadata: object) -> object: + return string_keyed(metadata).get("logging") + + +def _entry_credential_name(entry: object) -> str | None: + match entry: + case {"litellm_credential_name": str() as credential_name}: + return credential_name + case _: + return None + + +def logging_credential_names(logging_entries: object) -> tuple[str, ...]: + return tuple( + name for entry in object_sequence(logging_entries) or () if (name := _entry_credential_name(entry)) is not None + ) + + +def credential_reference_rejection( + requested_metadata: object, + stored_logging: object, + *, + is_proxy_admin: bool, + is_defined: Callable[[str], bool] = credential_is_defined, +) -> CredentialReferenceRejection | None: + """Gate a key metadata write whose logging entries reference a proxy credential. + + ``requested_metadata`` is None when the write leaves the key's metadata untouched. The admin check runs before + the existence check so a non-admin cannot probe which credential names are defined. + """ + if requested_metadata is None: + return None + requested_logging: Final = _logging_entries(requested_metadata) + requested_names: Final = logging_credential_names(requested_logging) + stored_names: Final = logging_credential_names(stored_logging) + references_involved: Final = bool(requested_names or stored_names) + if references_involved and not is_proxy_admin and requested_logging != stored_logging: + return CredentialReferenceRejection( + 403, + "Only proxy admins can change the logging entries of a key that uses litellm_credential_name", + ) + undefined: Final = sorted({name for name in set(requested_names) - set(stored_names) if not is_defined(name)}) + if undefined: + return CredentialReferenceRejection( + 400, f"metadata.logging references undefined credential(s): {', '.join(undefined)}" + ) + return None diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c195779aa0a..0aa8d2291ef 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -26,8 +26,10 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams +from litellm.proxy._types import CALLBACK_VAR_NAMES, CommonProxyErrors, LiteLLMPromptInjectionParams +from litellm.proxy.common_utils.callback_config_validation import object_sequence, string_keyed from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -736,6 +738,47 @@ def decrypt_callback_vars(metadata: object) -> Any: return _transform_callback_vars(metadata, _decrypt_or_passthrough) +def decrypted_callback_metadata(metadata: object) -> object: + return _transform_callback_vars(metadata, _decrypt_or_passthrough) + + +def decrypted_logging_entries(metadata: object) -> object: + return string_keyed(decrypted_callback_metadata(metadata)).get("logging") + + +def key_logging_entries(metadata: object) -> tuple[Mapping[str, object], ...] | None: + entries: Final = decrypted_logging_entries(metadata) + if entries is None: + return None + return tuple(resolved for entry in _logging_objects(entries) if (resolved := _with_credential(entry)) is not None) + + +def _logging_objects(entries: object) -> tuple[Mapping[str, object], ...]: + items: Final = () if entries in ("", {}) else object_sequence(entries) + if items is None or not all(isinstance(item, Mapping) for item in items): + raise TypeError(f"key metadata.logging must be a list of objects, got {type(entries).__name__}") + return tuple(string_keyed(item) for item in items) + + +def _with_credential(entry: Mapping[str, object]) -> Mapping[str, object] | None: + credential_name: Final = entry.get("litellm_credential_name") + if not isinstance(credential_name, str): + return entry + credential_vars: Final = { + name: str(value) + for name, value in CredentialAccessor.get_credential_values(credential_name).items() + if name in CALLBACK_VAR_NAMES and value is not None + } + if not credential_vars: + verbose_proxy_logger.error( + "Ignoring key logging entry for %s: credential %r is not defined or holds no logging values", + entry.get("callback_name"), + credential_name, + ) + return None + return {**entry, "callback_vars": {**credential_vars, **string_keyed(entry.get("callback_vars"))}} + + def _transform_callback_vars(metadata: object, transform: Callable[[str, object], object]) -> object: if not isinstance(metadata, dict): return metadata diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9cc18294cd1..0bb4a803609 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -88,6 +88,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, get_metadata_variable_name_from_kwargs, + key_logging_entries, strip_callback_config, ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports @@ -1246,9 +1247,11 @@ def convert_key_logging_metadata_to_callback( return team_callback_settings_obj -def get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallback | None: +def get_validated_callback_metadata(item: object, *, source: str) -> AddTeamCallback | None: + if not isinstance(item, Mapping): + raise TypeError(f"{source} metadata.logging must be a list of objects, got an item of {type(item).__name__}") try: - return AddTeamCallback(**item) + return AddTeamCallback.model_validate(item) except (PydanticValidationError, ValueError) as e: verbose_proxy_logger.warning( "Ignoring invalid %s callback metadata: %s", @@ -1267,9 +1270,9 @@ class KeyAndTeamLoggingSettings: """ @staticmethod - def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): + def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth) -> tuple[Mapping[str, object], ...] | None: if user_api_key_dict.metadata is not None and "logging" in user_api_key_dict.metadata: - return decrypt_callback_vars(user_api_key_dict.metadata).get("logging") + return key_logging_entries(user_api_key_dict.metadata) return None @staticmethod @@ -1283,9 +1286,7 @@ def get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> TeamCallbackMetadata | None: callback_settings_obj: TeamCallbackMetadata | None = None - key_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( - user_api_key_dict - ) + key_dynamic_logging_settings: Final = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) team_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( user_api_key_dict ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e35a65d8ba8..f88ac445a0b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -79,6 +79,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( # noq check_permissions_caller_permission, generate_key_helper_fn, prepare_metadata_fields, + raise_on_logging_credential_change, ) from litellm.proxy.management_helpers.bulk_user_deletion import delete_users_from_current_teams from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports @@ -598,6 +599,7 @@ async def new_user( if data.auto_create_key and isinstance(user_api_key_dict, UserAPIKeyAuth): enforce_batch_limits_are_admin_only(data, None, user_api_key_dict, "key") enforce_vector_store_index_grants_are_admin_only(data, None, None, user_api_key_dict, "key") + raise_on_logging_credential_change(data.metadata, None, user_api_key_dict) check_permissions_caller_permission( data=data, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 50517eeefcf..038c1bdf5a6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, c import fastapi import yaml from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -68,9 +68,14 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( evict_and_broadcast, ) -from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error +from litellm.proxy.common_utils.callback_config_validation import ( + credential_reference_rejection, + logging_metadata_config_error, +) from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, + decrypted_callback_metadata, + decrypted_logging_entries, encrypt_callback_vars, ) from litellm.proxy.common_utils.config_sync_pubsub import ( @@ -266,6 +271,12 @@ class _StoredKeyRouterSettings(LiteLLMBaseModel): router_settings: Mapping[str, object] | None = None +class _StoredKeyMetadata(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + metadata: object = None + + class _KeyRowWhere(TypedDict): token: ReadOnly[str] @@ -934,6 +945,23 @@ def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) - raise HTTPException(status_code=400, detail={"error": error}) +def raise_on_logging_credential_change( + requested_metadata: object, existing_key_row: LiteLLM_VerificationToken | None, user_api_key_dict: UserAPIKeyAuth +) -> None: + stored_metadata: Final = ( + None + if existing_key_row is None + else _StoredKeyMetadata.model_validate(existing_key_row, from_attributes=True).metadata + ) + rejection: Final = credential_reference_rejection( + decrypted_callback_metadata(requested_metadata), + decrypted_logging_entries(stored_metadata), + is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, + ) + if rejection is not None: + raise HTTPException(status_code=rejection.status_code, detail={"error": rejection.message}) + + def common_key_access_checks( user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest | UpdateKeyRequest, @@ -1485,6 +1513,7 @@ async def _common_key_generation_helper( validate_budget_duration(data.budget_duration) raise_on_invalid_key_logging_config(data.metadata) + raise_on_logging_credential_change(data.metadata, None, user_api_key_dict) if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( @@ -3430,6 +3459,7 @@ async def _process_single_key_update( ) # Prepare update data + raise_on_logging_credential_change(key_request.metadata, existing_key_row, user_api_key_dict) non_default_values = await prepare_key_update_data( data=key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router ) @@ -4176,6 +4206,7 @@ async def update_key_fn( # Enforce upperbound key params on update (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) + raise_on_logging_credential_change(data.metadata, existing_key_row, user_api_key_dict) non_default_values: Final = await prepare_key_update_data( data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router ) @@ -6363,6 +6394,7 @@ async def _execute_virtual_key_regeneration( await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request) # Enforce upperbound key params on regenerate (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) + raise_on_logging_credential_change(data.metadata, key_in_db, user_api_key_dict) non_default_values = await prepare_key_update_data( data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router ) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index e900941f67b..77e37cb11cb 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -34,6 +34,7 @@ from litellm.proxy.auth.auth_utils import ( enforce_vector_store_index_grants_are_admin_only, ) from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.common_utils.callback_config_validation import credential_reference_rejection from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks @@ -225,7 +226,14 @@ def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str enforce_vector_store_index_grants_are_admin_only(item, None, None, user_api_key_dict, "key") except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only return _error_message(exc) - return None + rejection: Final = ( + credential_reference_rejection( + item.metadata, None, is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ) + if item.auto_create_key + else None + ) + return None if rejection is None else rejection.message def _normalized_email(email: str | None) -> str | None: diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index 58785c89033..3fc4b2e7628 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -672,10 +672,12 @@ def _write_config( callback_settings: Mapping[str, JsonValue] | None = None, litellm_settings: Mapping[str, JsonValue] | None = None, general_settings: Mapping[str, JsonValue] | None = None, + extra_config: Mapping[str, JsonValue] | None = None, ) -> Path: loaded: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config: Final = { **loaded, + **(extra_config or {}), "litellm_settings": { **loaded["litellm_settings"], "callbacks": list(callbacks), @@ -752,6 +754,7 @@ def _rig( callback_settings: Mapping[str, JsonValue] | None = None, litellm_settings: Mapping[str, JsonValue] | None = None, general_settings: Mapping[str, JsonValue] | None = None, + extra_config: Mapping[str, JsonValue] | None = None, environment: Mapping[str, str] | None = None, remove_environment: tuple[str, ...] = (), disabled_environment: tuple[str, ...] = (), @@ -782,6 +785,7 @@ def _rig( callback_settings=resolved_settings, litellm_settings=litellm_settings, general_settings=general_settings, + extra_config=extra_config, ) preset_environment: Final = { **({"PHOENIX_COLLECTOR_ENDPOINT": destination.url + "/v1/traces"} if "arize_phoenix" in callbacks else {}), diff --git a/tests/integration/observability/test_arize_key_credential_reference.py b/tests/integration/observability/test_arize_key_credential_reference.py new file mode 100644 index 00000000000..b1087c4f4c5 --- /dev/null +++ b/tests/integration/observability/test_arize_key_credential_reference.py @@ -0,0 +1,428 @@ +from __future__ import annotations + +import operator +import time +import uuid +from collections.abc import Callable, Iterator +from itertools import accumulate +from pathlib import Path +from typing import Final + +import httpx +import pytest +from _openinference_support import ( + _anthropic_response, + _chat_request_marker, + _chat_response, + _chat_stream_response, + _chat_tool_call, + _matching_marker_spans, + _responses_response, + _rig, +) +from integration._support.client import Gateway, eventually +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +CREDENTIAL_NAME: Final = "arize-space-a" +SECRET_API_KEY: Final = "space-a-arize-key-" + uuid.uuid4().hex +SECRET_SPACE_ID: Final = "space-a-id-" + uuid.uuid4().hex +AMBIENT_ARIZE_ENVIRONMENT: Final = ("ARIZE_API_KEY", "ARIZE_SPACE_ID") +CREDENTIAL_CONFIG: Final[dict[str, JsonValue]] = { + "credential_list": [ + { + "credential_name": CREDENTIAL_NAME, + "credential_values": { + "arize_api_key": "os.environ/SPACE_A_ARIZE_API_KEY", + "arize_space_id": "os.environ/SPACE_A_ARIZE_SPACE_ID", + }, + "credential_info": {}, + }, + { + "credential_name": "model-only", + "credential_values": {"api_key": "sk-not-a-logging-value"}, + "credential_info": {}, + }, + ] +} +SECRET_ENVIRONMENT: Final = {"SPACE_A_ARIZE_API_KEY": SECRET_API_KEY, "SPACE_A_ARIZE_SPACE_ID": SECRET_SPACE_ID} + + +def _reference_entry(credential_name: str = CREDENTIAL_NAME) -> dict[str, JsonValue]: + return { + "callback_name": "arize", + "callback_type": "success", + "litellm_credential_name": credential_name, + "callback_vars": {}, + } + + +def _echo_marker(request: Request) -> Reply: + return _chat_response(_chat_request_marker(request)) + + +def _send_chat(proxy: Gateway, model: str, key: str, marker: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker}], + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + key=key, + ) + + +def _chat(proxy: Gateway, model: str, key: str, marker: str) -> None: + response: Final = _send_chat(proxy, model, key, marker) + assert response.status_code == 200, response.text + + +def _drains(destination: Wire) -> Iterator[tuple[Request, ...]]: + while True: + yield destination.drain() + time.sleep(0.1) + + +def _carries(request: Request, marker: str) -> bool: + return any(True for _ in _matching_marker_spans((request,), marker)) + + +def _received_until(destination: Wire, marker: str, seconds: float = 30) -> tuple[Request, ...]: + deadline: Final = time.monotonic() + seconds + received: Final = next( + batch + for batch in accumulate(_drains(destination), operator.add) + if any(_carries(request, marker) for request in batch) or time.monotonic() >= deadline + ) + assert any(_carries(request, marker) for request in received), f"no span for {marker} reached the destination" + return received + + +def test_key_referencing_a_credential_exports_with_the_secret_manager_values(gateway: Gateway, tmp_path: Path) -> None: + unreferenced_marker: Final = "plain-" + uuid.uuid4().hex + referenced_marker: Final = "ref-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + team: Final = scenario.team() + referencing_key: Final = scenario.key(team_id=team, metadata={"logging": [_reference_entry()]}) + unreferenced_key: Final = scenario.key(team_id=team) + + _chat(rig.proxy, rig.model, unreferenced_key, unreferenced_marker) + _chat(rig.proxy, rig.model, referencing_key, referenced_marker) + received: Final = _received_until(rig.destination, referenced_marker) + + exports: Final = tuple(request for request in received if _carries(request, referenced_marker)) + for export in exports: + headers = {name.lower(): value for name, value in export.headers.items()} + assert headers.get("api_key") == SECRET_API_KEY, headers + assert headers.get("arize-space-id") == SECRET_SPACE_ID, headers + assert not any(_carries(request, unreferenced_marker) for request in received) + + info: Final = rig.proxy.request("GET", "/key/info", params={"key": referencing_key}) + assert info.status_code == 200, info.text + assert CREDENTIAL_NAME in info.text + assert SECRET_API_KEY not in info.text + assert SECRET_SPACE_ID not in info.text + + +def test_only_proxy_admins_change_logging_on_a_key_that_references_a_credential( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + team: Final = scenario.team() + team_admin: Final = scenario.member(team, role="admin") + team_admin_key: Final = scenario.key(user_id=team_admin, team_id=team) + langfuse_entry: Final[dict[str, JsonValue]] = { + "callback_name": "langfuse", + "callback_type": "success", + "callback_vars": {"langfuse_public_key": "pk-lf-integration", "langfuse_secret_key": "sk-lf-integration"}, + } + referencing_key: Final = scenario.key(team_id=team, metadata={"logging": [_reference_entry(), langfuse_entry]}) + + def as_team_admin(path: str, body: dict[str, JsonValue]) -> httpx.Response: + return rig.proxy.request("POST", path, body, key=team_admin_key) + + def assert_gated(response: httpx.Response) -> None: + assert response.status_code == 403, response.text + assert "litellm_credential_name" in response.text, response.text + + assert_gated(as_team_admin("/key/generate", {"team_id": team, "metadata": {"logging": [_reference_entry()]}})) + + sibling: Final = as_team_admin( + "/key/update", + { + "key": referencing_key, + "metadata": { + "logging": [ + _reference_entry(), + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {"arize_endpoint": "http://127.0.0.1:9/v1/traces"}, + }, + ] + }, + }, + ) + assert_gated(sibling) + assert_gated(as_team_admin("/key/update", {"key": referencing_key, "metadata": {"team": "unrelated"}})) + assert_gated(as_team_admin("/key/regenerate", {"key": referencing_key, "metadata": {"team": "unrelated"}})) + + bulk: Final = as_team_admin( + "/team/key/bulk_update", + {"team_id": team, "key_ids": [referencing_key], "update_fields": {"metadata": {"team": "unrelated"}}}, + ) + assert bulk.status_code == 200, bulk.text + assert bulk.json()["successful_updates"] == [], bulk.text + assert "litellm_credential_name" in bulk.json()["failed_updates"][0]["failed_reason"], bulk.text + + renamed: Final = as_team_admin( + "/key/update", {"key": referencing_key, "key_alias": "alias-" + uuid.uuid4().hex} + ) + assert renamed.status_code == 200, renamed.text + + read_back: Final = rig.proxy.request("GET", "/key/info", params={"key": referencing_key}, key=team_admin_key) + assert read_back.status_code == 200, read_back.text + resaved: Final = as_team_admin( + "/key/update", + {"key": referencing_key, "metadata": {**read_back.json()["info"]["metadata"], "note": "unrelated edit"}}, + ) + assert resaved.status_code == 200, resaved.text + + stored: Final = rig.proxy.get("/key/info", params={"key": referencing_key}) + logging_entries: Final = stored["info"]["metadata"]["logging"] + assert [entry["callback_name"] for entry in logging_entries] == ["arize", "langfuse"], stored + assert logging_entries[0] == _reference_entry(), stored + + undefined: Final = rig.proxy.request( + "POST", "/key/generate", {"metadata": {"logging": [_reference_entry("no-such-credential")]}} + ) + assert undefined.status_code == 400, undefined.text + assert "no-such-credential" in undefined.text, undefined.text + + +def test_only_proxy_admins_create_a_user_whose_key_references_a_credential(gateway: Gateway, tmp_path: Path) -> None: + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + organization: Final = scenario.organization() + org_admin_key: Final = scenario.key(user_id=scenario.org_member(organization, "org_admin")) + referencing_metadata: Final = {"logging": [_reference_entry()]} + + single: Final = rig.proxy.request( + "POST", + "/user/new", + {"user_id": "user-" + uuid.uuid4().hex, "organizations": [organization], "metadata": referencing_metadata}, + key=org_admin_key, + ) + assert single.status_code == 403, single.text + assert "litellm_credential_name" in single.text, single.text + + delegate_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), allowed_routes=["/management/v1/users/bulk"] + ) + bulk: Final = rig.proxy.request( + "POST", + "/management/v1/users/bulk", + { + "users": [ + {"user_id": "user-" + uuid.uuid4().hex, "auto_create_key": True, "metadata": referencing_metadata} + ] + }, + key=delegate_key, + ) + assert bulk.status_code == 200, bulk.text + assert [row["success"] for row in bulk.json()["data"]] == [False], bulk.text + assert "litellm_credential_name" in bulk.json()["data"][0]["error"], bulk.text + + +def _carries_credential(request: Request) -> bool: + headers: Final = {name.lower(): value for name, value in request.headers.items()} + return headers.get("api_key") == SECRET_API_KEY and headers.get("arize-space-id") == SECRET_SPACE_ID + + +def _credential_exports(destination: Wire, seconds: float = 30) -> tuple[Request, ...]: + return eventually( + lambda: tuple(request for request in destination.drain() if _carries_credential(request)), + bool, + seconds=seconds, + ) + + +def _surface_reply(marker: str) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + if request.target.endswith("/responses"): + return _responses_response(marker) + if request.target.endswith("/messages"): + return _anthropic_response(marker) + return _chat_stream_response(marker, (_chat_tool_call(marker),)) + + return reply + + +_SURFACES: Final[dict[str, tuple[str, str, dict[str, JsonValue]]]] = { + "responses": ("/v1/responses", "openai/gpt-4o-mini", {"input": "weather in Paris?"}), + "messages": ( + "/v1/messages", + "anthropic/claude-opus-5-5", + {"max_tokens": 64, "messages": [{"role": "user", "content": "weather in Paris?"}]}, + ), + "chat-stream": ( + "/v1/chat/completions", + "openai/gpt-4o-mini", + {"stream": True, "messages": [{"role": "user", "content": "weather in Paris?"}]}, + ), +} + + +@pytest.mark.parametrize("surface", sorted(_SURFACES)) +def test_referencing_key_exports_on_each_surface(surface: str, gateway: Gateway, tmp_path: Path) -> None: + marker: Final = f"{surface}-" + uuid.uuid4().hex + path, model_name, body = _SURFACES[surface] + with ( + _rig( + gateway, + tmp_path, + _surface_reply(marker), + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + model_name=model_name, + api_base_suffix="" if surface == "messages" else "/v1", + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + response: Final = rig.proxy.request( + "POST", + path, + {"model": rig.model, **body, "metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + key=key, + ) + assert response.status_code == 200, response.text + received: Final = _received_until(rig.destination, marker) + exports: Final = tuple(request for request in received if _carries(request, marker)) + assert exports and all(_carries_credential(request) for request in exports), [ + dict(request.headers) for request in exports + ] + + +def test_referencing_key_exports_through_anthropic_pass_through(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "passthrough-" + uuid.uuid4().hex + with wire_server(lambda _request: _anthropic_response(marker)) as anthropic: + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment={ + **SECRET_ENVIRONMENT, + "ANTHROPIC_API_BASE": anthropic.url, + "ANTHROPIC_API_KEY": "sk-ant-synthetic", + }, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + response: Final = rig.proxy.request( + "POST", + "/anthropic/v1/messages", + {"model": "claude-opus-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers={"anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert response.json()["id"] == marker, response.text + assert anthropic.drain(), "the pass-through never reached the scripted Anthropic upstream" + assert _credential_exports(rig.destination) + + +def test_referencing_key_exports_on_otel_v1(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "otel-v1-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment={**SECRET_ENVIRONMENT, "LITELLM_OTEL_V2": "false"}, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + _chat(rig.proxy, rig.model, key, marker) + assert _credential_exports(rig.destination) + + +def test_reference_to_a_credential_without_logging_values_exports_nothing_and_is_unhealthy( + gateway: Gateway, tmp_path: Path +) -> None: + dropped_marker: Final = "model-only-" + uuid.uuid4().hex + sentinel_marker: Final = "sentinel-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + dropped_key: Final = scenario.key(metadata={"logging": [_reference_entry("model-only")]}) + sentinel_key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + + _chat(rig.proxy, rig.model, dropped_key, dropped_marker) + _chat(rig.proxy, rig.model, sentinel_key, sentinel_marker) + received: Final = _received_until(rig.destination, sentinel_marker) + assert not any(_carries(request, dropped_marker) for request in received) + + health: Final = rig.proxy.request("POST", "/key/health", {}, key=dropped_key) + assert health.status_code == 200, health.text + assert health.json()["key"] == "unhealthy", health.text + + +def test_malformed_key_logging_fails_requests_like_main(gateway: Gateway, tmp_path: Path) -> None: + stored_logging: Final[dict[str, JsonValue]] = { + "string": "abc", + "list-with-string": ["x"], + "object": {"callback_name": "arize"}, + "empty-list": [], + } + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + statuses: Final = { + label: _send_chat( + rig.proxy, rig.model, scenario.key(metadata={"logging": value}), "malformed-" + uuid.uuid4().hex + ).status_code + for label, value in stored_logging.items() + } + + assert statuses == {"string": 500, "list-with-string": 500, "object": 500, "empty-list": 200}, statuses diff --git a/tests/unit/proxy/common_utils/test_callback_config_validation.py b/tests/unit/proxy/common_utils/test_callback_config_validation.py index f42387d3e76..ac1fe2d785e 100644 --- a/tests/unit/proxy/common_utils/test_callback_config_validation.py +++ b/tests/unit/proxy/common_utils/test_callback_config_validation.py @@ -7,6 +7,7 @@ from litellm.proxy.common_utils.callback_config_validation import ( callback_config_error, conflicting_callback_var_error, conflicting_span_scope_error, + credential_reference_rejection, cross_entry_family_error, logging_metadata_config_error, ) @@ -499,3 +500,78 @@ def test_key_logging_entries_compare_excluded_services_by_parsed_value( def test_a_failure_only_entry_rejects_excluded_services() -> None: error: Final = callback_config_error("langfuse_otel", {"excluded_services": "redis"}, "failure") assert error is not None and "excluded_services" in error and "success_and_failure" in error + + +_REFERENCE: Final = {"callback_name": "arize", "callback_type": "success", "litellm_credential_name": "space-a"} +_INLINE: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} +_SIBLING: Final = {"callback_name": "arize", "callback_vars": {"arize_endpoint": "http://attacker.example/v1/traces"}} + + +@pytest.mark.parametrize( + ("requested", "stored_logging", "is_proxy_admin", "expected_status"), + [ + pytest.param({"logging": [_REFERENCE]}, None, False, 403, id="non-admin-adds-reference"), + pytest.param({"logging": [_INLINE, _REFERENCE]}, None, False, 403, id="non-admin-adds-reference-second"), + pytest.param({"team": "x"}, [_REFERENCE], False, 403, id="non-admin-removes-reference"), + pytest.param({"logging": [_REFERENCE, _SIBLING]}, [_REFERENCE], False, 403, id="non-admin-adds-sibling"), + pytest.param({"logging": [_REFERENCE]}, [_REFERENCE], False, None, id="non-admin-resends-same"), + pytest.param(None, [_REFERENCE], False, None, id="non-admin-leaves-metadata-alone"), + pytest.param({"logging": [_INLINE]}, [], False, None, id="non-admin-without-references"), + pytest.param({"logging": [_REFERENCE]}, [], True, None, id="admin-adds-defined-reference"), + pytest.param({"team": "x"}, [_REFERENCE], True, None, id="admin-removes-reference"), + ], +) +def test_credential_reference_rejection_lets_only_proxy_admins_change_referencing_logging( + requested: dict[str, object] | None, + stored_logging: list[object] | None, + is_proxy_admin: bool, + expected_status: int | None, +) -> None: + rejection: Final = credential_reference_rejection( + requested, stored_logging, is_proxy_admin=is_proxy_admin, is_defined=lambda name: name == "space-a" + ) + assert (None if rejection is None else rejection.status_code) == expected_status + + +def test_credential_reference_rejection_names_undefined_credentials_only_to_admins() -> None: + requested: Final = {"logging": [_INLINE, {**_REFERENCE, "litellm_credential_name": "missing"}]} + + admin: Final = credential_reference_rejection(requested, None, is_proxy_admin=True, is_defined=lambda _: False) + non_admin: Final = credential_reference_rejection(requested, None, is_proxy_admin=False, is_defined=lambda _: False) + + assert admin is not None and admin.status_code == 400 and "missing" in admin.message + assert non_admin is not None and non_admin.status_code == 403 and "missing" not in non_admin.message + + +@pytest.mark.parametrize("is_proxy_admin", [True, False], ids=["admin", "non-admin"]) +def test_credential_reference_rejection_lets_a_stored_reference_to_a_deleted_credential_stay( + is_proxy_admin: bool, +) -> None: + stale: Final = {**_REFERENCE, "litellm_credential_name": "deleted"} + + kept: Final = credential_reference_rejection( + {"logging": [stale], "note": "unrelated edit"}, + [stale], + is_proxy_admin=is_proxy_admin, + is_defined=lambda _: False, + ) + added: Final = credential_reference_rejection( + {"logging": [stale, {**_REFERENCE, "litellm_credential_name": "also-missing"}]}, + [stale], + is_proxy_admin=True, + is_defined=lambda _: False, + ) + + assert kept is None + assert added is not None and added.status_code == 400 + assert "also-missing" in added.message and "deleted" not in added.message + + +@pytest.mark.parametrize("name", [5, ["space-a"], None], ids=["int", "list", "null"]) +def test_a_non_string_credential_name_is_not_a_reference(name: object) -> None: + entry: Final = {**_REFERENCE, "litellm_credential_name": name} + + assert ( + credential_reference_rejection({"logging": [entry]}, None, is_proxy_admin=False, is_defined=lambda _: False) + is None + ) diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index c8d909ac2eb..c7739b9cc2f 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -1,5 +1,6 @@ import copy import json +import logging import sys from types import ModuleType, SimpleNamespace from typing import Final @@ -21,6 +22,7 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, initialize_callbacks_on_proxy, + key_logging_entries, normalize_callback_names, process_callback, sanitize_openai_provider_metadata, @@ -28,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import ( ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CredentialItem def test_get_remaining_tokens_and_requests_from_request_data(): @@ -775,3 +778,114 @@ def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_con assert len(litellm.callbacks) == 1 assert isinstance(litellm.callbacks[0], CustomLogger) + + +def _credential(name: str, values: dict[str, object]) -> CredentialItem: + return CredentialItem(credential_name=name, credential_values=values, credential_info={}) + + +def test_key_logging_entries_merges_each_referenced_credential_under_its_entry(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm, + "credential_list", + [ + _credential("space-a", {"arize_api_key": "key-a", "arize_space_id": "space-a", "api_base": "x"}), + _credential("space-b", {"arize_api_key": "key-b", "arize_space_id": "space-b"}), + ], + ) + entries: Final = [ + {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}, + {"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {"arize_space_id": "own"}}, + {"callback_name": "arize", "litellm_credential_name": "space-b", "callback_vars": {}}, + ] + + resolved: Final = key_logging_entries({"logging": entries}) + + assert resolved is not None + assert [entry["callback_vars"] for entry in resolved] == [ + {"langfuse_public_key": "pk"}, + {"arize_api_key": "key-a", "arize_space_id": "own"}, + {"arize_api_key": "key-b", "arize_space_id": "space-b"}, + ] + + +@pytest.mark.parametrize( + "credential_values", + [{"api_key": "sk-model-key"}, {"arize_api_key": None}], + ids=["model-credential", "null-values"], +) +def test_key_logging_entries_drops_an_entry_whose_credential_supplies_no_logging_values( + monkeypatch: pytest.MonkeyPatch, credential_values: dict[str, object] +) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("model-only", credential_values)]) + kept: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} + entries: Final = [ + kept, + {"callback_name": "arize", "litellm_credential_name": "model-only", "callback_vars": {}}, + ] + + assert key_logging_entries({"logging": entries}) == (kept,) + + +def test_a_dropped_credential_reference_is_logged_as_an_error_so_key_health_reports_it( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "credential_list", []) + entry: Final = {"callback_name": "arize", "litellm_credential_name": "missing", "callback_vars": {}} + + with caplog.at_level(logging.WARNING): + key_logging_entries({"logging": [entry]}) + + assert [record.levelno for record in caplog.records if "missing" in record.getMessage()] == [logging.ERROR] + + +def test_key_logging_entries_drops_an_entry_whose_credential_is_not_defined(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("space-a", {"arize_api_key": "key-a"})]) + kept: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} + entries: Final = [ + kept, + {"callback_name": "arize", "litellm_credential_name": "missing", "callback_vars": {}}, + ] + + assert key_logging_entries({"logging": entries}) == (kept,) + + +@pytest.mark.parametrize( + "logging_value", + [ + pytest.param("abc", id="string"), + pytest.param({"callback_name": "arize"}, id="object"), + pytest.param(5, id="number"), + pytest.param(["x"], id="list-with-string"), + pytest.param([{"callback_name": "langfuse", "callback_vars": {}}, None], id="list-with-null"), + ], +) +def test_key_logging_entries_fails_on_non_object_logging_like_the_request_path_always_has( + logging_value: object, +) -> None: + with pytest.raises(TypeError): + key_logging_entries({"logging": logging_value}) + + +@pytest.mark.parametrize( + ("logging_value", "expected"), + [ + pytest.param("", (), id="empty-string"), + pytest.param({}, (), id="empty-object"), + pytest.param([], (), id="empty-list"), + pytest.param([{"bad": 1}], ({"bad": 1},), id="object-with-unknown-fields"), + ], +) +def test_key_logging_entries_passes_empty_and_object_entries_through( + logging_value: object, expected: tuple[dict[str, object], ...] +) -> None: + assert key_logging_entries({"logging": logging_value}) == expected + + +def test_key_logging_entries_leaves_an_entry_with_a_non_string_credential_name_as_written( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("5", {"arize_api_key": "k"})]) + entry: Final = {"callback_name": "arize", "litellm_credential_name": 5, "callback_vars": {}} + + assert key_logging_entries({"logging": [entry]}) == (entry,) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 650bc89c917..2e18a0c1099 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber +from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, budget_windows_absent_cache_key, @@ -101,6 +102,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( persist_deleted_verification_tokens, prepare_key_update_data, prepare_metadata_fields, + raise_on_logging_credential_change, reset_key_spend_fn, update_key_fn, validate_key_list_check, @@ -112,6 +114,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyResponse, CustomKeyPolicyRequest, ) +from litellm.types.utils import CredentialItem from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters client = TestClient(app) @@ -24407,3 +24410,29 @@ async def test_key_aliases_answers_503_when_the_database_is_unreachable(monkeypa ) assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") + + +def test_non_admin_resaving_referencing_key_metadata_read_back_encrypted_is_not_a_logging_change(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-32-bytes-aaaaaaaaaaaaaa") + monkeypatch.setattr( + litellm, + "credential_list", + [CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})], + ) + stored = encrypt_callback_vars( + { + "logging": [ + {"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}}, + {"callback_name": "langfuse", "callback_vars": {"langfuse_secret_key": "sk-lf-secret"}}, + ] + } + ) + assert stored["logging"][1]["callback_vars"]["langfuse_secret_key"] != "sk-lf-secret" + key_row = LiteLLM_VerificationToken(token="hashed", metadata=stored) + internal_user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) + + raise_on_logging_credential_change({**stored, "note": "unrelated edit"}, key_row, internal_user) + with pytest.raises(HTTPException) as removed: + raise_on_logging_credential_change({"note": "unrelated edit"}, key_row, internal_user) + + assert removed.value.status_code == 403 diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py index 2b966a8ebeb..96c84578111 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py @@ -7,6 +7,7 @@ import pytest from prisma.errors import UniqueViolationError from pydantic import BaseModel, ConfigDict, ValidationError +import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth from litellm.proxy.list_api.common import ManagementProblem @@ -15,6 +16,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkNewUserItem, BulkNewUserRequest, ) +from litellm.types.utils import CredentialItem ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) @@ -434,6 +436,39 @@ async def test_only_proxy_admin_sets_batch_limits_on_a_created_key(): assert calls[-1]["metadata"] == limits +@pytest.mark.asyncio +async def test_only_proxy_admin_creates_a_key_whose_logging_references_a_credential( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm, + "credential_list", + [CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})], + ) + reference: Final = {"logging": [{"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}}]} + calls: list[dict[str, object]] = [] + + async def generate_key(**kwargs: object) -> dict[str, object]: + calls.append(kwargs) + return {"token": f"sk-{kwargs['user_id']}"} + + rows: Final = [ + {"user_id": "u1", "auto_create_key": True, "metadata": reference}, + {"user_id": "u2", "auto_create_key": True}, + ] + + response = await _run(_FakePrisma(), rows, caller=INTERNAL, generate_key=generate_key) + + assert [r.success for r in response.data] == [False, True] + assert "litellm_credential_name" in (response.data[0].error or "") + assert [call["user_id"] for call in calls] == ["u2"] + + admin_response = await _run(_FakePrisma(), rows[:1], caller=ADMIN, generate_key=generate_key) + + assert [r.key for r in admin_response.data] == ["sk-u1"] + assert calls[-1]["metadata"] == reference + + @pytest.mark.asyncio async def test_license_is_checked_once_against_the_whole_batch(): prisma = _FakePrisma() diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index fbd4de8ef5c..e4a560adca8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -26,6 +26,7 @@ from litellm.constants import ( SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, ) +from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, @@ -59,6 +60,7 @@ from litellm.proxy.litellm_pre_call_utils import ( get_dynamic_logging_metadata, get_metadata_variable_name, move_guardrails_to_metadata, + resolve_tenant_otel_destinations, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES @@ -2779,7 +2781,7 @@ def test_key_dynamic_logging_settings(): team_metadata={}, ) result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_arize) - assert result == [{"callback_name": "arize", "callback_type": "success"}] + assert result == ({"callback_name": "arize", "callback_type": "success"},) # Test with langfuse logging key_with_langfuse = UserAPIKeyAuth( @@ -2788,7 +2790,7 @@ def test_key_dynamic_logging_settings(): team_metadata={}, ) result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse) - assert result == [{"callback_name": "langfuse", "callback_type": "success"}] + assert result == ({"callback_name": "langfuse", "callback_type": "success"},) # Test with no logging metadata key_without_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) @@ -2912,6 +2914,22 @@ def test_get_dynamic_logging_metadata_with_arize_team_logging(): assert result.callback_vars["arize_space_id"] == "test_arize_space_id" +@pytest.mark.parametrize( + "logging_value", + [ + pytest.param({"callback_name": "langsmith"}, id="object"), + pytest.param("langsmith", id="string"), + pytest.param(["langsmith"], id="list-with-string"), + pytest.param([None], id="list-with-null"), + ], +) +def test_non_object_team_logging_metadata_fails_as_a_server_error_not_a_bad_request(logging_value: object) -> None: + user_api_key_dict: Final = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={"logging": logging_value}) + + with pytest.raises(TypeError): + get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=MagicMock()) + + def test_add_team_callback_rejects_env_reference(): with pytest.raises(PydanticValidationError) as exc_info: AddTeamCallback( @@ -10267,6 +10285,48 @@ def test_arize_otlp_protocol_on_a_key_logging_entry_reaches_the_destination(monk is_otel_v2_enabled.cache_clear() +def test_a_key_referencing_a_credential_routes_to_that_credential_with_its_own_excluded_services(monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "https://arize.internal.example/v1/traces") + monkeypatch.delenv("ARIZE_API_KEY", raising=False) + monkeypatch.delenv("ARIZE_SPACE_ID", raising=False) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="space-a", + credential_values={"arize_api_key": "k", "arize_space_id": "s"}, + credential_info={}, + ) + ], + ) + is_otel_v2_enabled.cache_clear() + try: + auth = UserAPIKeyAuth( + api_key="hashed-key", + metadata={ + "logging": [ + { + "callback_name": "arize", + "callback_type": "success", + "litellm_credential_name": "space-a", + "callback_vars": {"excluded_services": "redis,postgres"}, + } + ] + }, + team_metadata={}, + ) + + destinations = resolve_tenant_otel_destinations(auth) + + assert [(d.headers, d.excluded_db_systems) for d in destinations] == [ + ({"arize-space-id": "s", "api_key": "k"}, frozenset({"redis", "postgresql"})) + ] + finally: + is_otel_v2_enabled.cache_clear() + + class TestResolveUserProviderCredentials: """``_resolve_user_provider_credentials_for_request`` injects the calling user's stored provider tokens into secret_fields for deployments carrying a per_user_oauth credential."""