From e44e42dfe4cbd1460fdbfb9716e0a44d90ba3206 Mon Sep 17 00:00:00 2001 From: Jesus Velarde Date: Fri, 6 Nov 2026 09:35:11 -0800 Subject: [PATCH] feat(proxy): let a key reference a callback_settings definition by name A callback_settings entry with a callback_type other than generic_api is a named definition. A key's logging entry naming it runs that backend with the definition's values, which the config loader resolved through the secret manager at boot, so the key stores no credential. A reference entry carries the name and event type only, and is the only entry on the key for its backend, so no key-side value can merge with the definition's credentials. Only a proxy admin may point a key at a definition or change the logging list of a key that holds one, on /key/generate, /key/update, /key/{key}/regenerate and the metadata of /user/new. Boot fails on a definition named after a callback, with an unknown callback_type, a field no callback accepts, or a secret the secret manager cannot provide, naming the field and never its value --- .../callback_config_validation.py | 161 ++++++++++++++- litellm/proxy/litellm_pre_call_utils.py | 76 ++++++- .../internal_user_endpoints.py | 7 + .../key_management_endpoints.py | 55 +++++- litellm/proxy/proxy_server.py | 12 ++ .../test_callback_config_validation.py | 160 +++++++++++++++ .../test_internal_user_endpoints.py | 41 ++++ .../test_key_management_endpoints.py | 186 ++++++++++++++++++ .../proxy/proxy_server/test_proxy_config.py | 65 ++++++ .../unit/proxy/test_litellm_pre_call_utils.py | 108 ++++++++++ 10 files changed, 862 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 302978d17a4..f8926337390 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -6,10 +6,14 @@ per-integration checks here. """ import math -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, 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, TypeAlias, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings +from pydantic import TypeAdapter + +import litellm from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR _NEWRELIC_CALLBACK: Final = "newrelic" @@ -21,6 +25,148 @@ _ARIZE_OTLP_PROTOCOL_VAR: Final = "arize_otlp_protocol" _ARIZE_SAMPLING_RATE_VARS: Final[frozenset[str]] = frozenset( {"arize_success_sampling_rate", "arize_error_sampling_rate"} ) +_GENERIC_API_CALLBACK_TYPE: Final = "generic_api" +_CALLBACK_TYPE_FIELD: Final = "callback_type" +_SECRET_REFERENCE_PREFIX: Final = "os.environ/" +_STRING_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object]) + +LoggingEntryKey: TypeAlias = tuple[str | None, str | None, tuple[tuple[str, str], ...]] + + +@dataclass(frozen=True, slots=True) +class CallbackDefinition: + """A ``callback_settings`` entry a key's ``logging`` entry may reference by name. + + ``backend`` is the callback it runs on and ``callback_vars`` the values the config + loader resolved at boot, secrets included, so the key stores only the name. + """ + + name: str + backend: str + callback_vars: Mapping[str, str] + + +def callback_definition(callback_name: object) -> CallbackDefinition | None: + """The definition ``callback_name`` references, or None when it is a plain callback name.""" + if not isinstance(callback_name, str): + return None + return _parsed_definition(callback_name, litellm.callback_settings.get(callback_name)) + + +def _parsed_definition(name: str, raw: object) -> CallbackDefinition | None: + if not isinstance(raw, Mapping): + return None + definition: Final = _STRING_KEYED_MAPPING.validate_python(raw) + backend: Final = definition.get(_CALLBACK_TYPE_FIELD) + if not isinstance(backend, str) or backend == _GENERIC_API_CALLBACK_TYPE: + return None + return CallbackDefinition( + name=name, + backend=backend, + callback_vars=MappingProxyType( + { + key: "" if value is None else str(value) + for key, value in definition.items() + if key != _CALLBACK_TYPE_FIELD + } + ), + ) + + +def callback_backend_name(callback_name: str) -> str: + """The callback a ``logging`` entry runs on: its definition's backend for a reference, else its own name.""" + definition: Final = callback_definition(callback_name) + return callback_name if definition is None else definition.backend + + +def callback_definitions_error(callback_settings: Mapping[str, object]) -> str | None: + """Why a ``callback_settings`` definition could never serve a key, checked at config load. + + A typo in a field or secret name would otherwise surface only on the first request + of a key that already depends on the definition. Errors name the field, never its value. + """ + definitions: Final = ( + definition + for name, raw in callback_settings.items() + if (definition := _parsed_definition(name, raw)) is not None + ) + return next((error for definition in definitions if (error := _definition_error(definition)) is not None), None) + + +def _definition_error(definition: CallbackDefinition) -> str | None: + known_callbacks: Final = litellm._known_custom_logger_compatible_callbacks # pyright: ignore[reportPrivateUsage] # the registry of callback names has no public accessor + if definition.name in known_callbacks: + return ( + f"callback_settings.{definition.name} is named after the {definition.name} callback, " + "so a key could not tell them apart" + ) + if definition.backend not in known_callbacks: + return f"callback_settings.{definition.name}.callback_type {definition.backend!r} is not a known callback" + unresolved: Final = next( + ( + key + for key, value in definition.callback_vars.items() + if value == "" or value.startswith(_SECRET_REFERENCE_PREFIX) + ), + None, + ) + if unresolved is not None: + return ( + f"callback_settings.{definition.name}.{unresolved} is empty or could not be resolved " + "from the secret manager" + ) + return _definition_backend_error(definition) + + +def _definition_backend_error(definition: CallbackDefinition) -> str | None: + from pydantic import ValidationError + + from litellm.proxy._types import AddTeamCallback + + try: + validated: Final = AddTeamCallback.model_validate( + {"callback_name": definition.backend, "callback_vars": dict(definition.callback_vars)} + ) + except ValidationError as e: + return f"callback_settings.{definition.name}: " + "; ".join(error["msg"] for error in e.errors()) + error: Final = callback_config_error(definition.backend, validated.callback_vars) + return None if error is None else f"callback_settings.{definition.name}: {error}" + + +def callback_reference_error(logging_entries: object) -> str | None: + """A ``callback_settings`` reference is a name only, and the only entry for its backend. + + Every entry's vars merge into one map per backend at request time, so a value on the + reference or on a sibling naming the same backend could point the definition's + credentials somewhere else. + """ + entries: Final = stored_callback_entries(logging_entries) + backends: Final = tuple(None if name is None else callback_backend_name(name) for name, _ in entries) + return next((error for error in _reference_errors(entries, backends) if error is not None), None) + + +def _reference_errors( + entries: Sequence[tuple[str | None, Mapping[str, str]]], backends: Sequence[str | None] +) -> Iterator[str | None]: + for (name, callback_vars), backend in zip(entries, backends, strict=True): + if name is None or callback_definition(name) is None: + continue + if callback_vars: + yield f"callback_settings.{name} is referenced by name only, so the entry may not carry callback_vars" + elif backends.count(backend) > 1: + yield ( + f"callback_settings.{name} already provides the {backend} callback for this key, " + f"so no other {backend} entry may sit next to it" + ) + + +def logging_entry_keys(logging_entries: object) -> tuple[LoggingEntryKey, ...]: + """Every ``logging`` entry as a comparable value: name, event type and sorted vars, so no field changes unnoticed.""" + entries: Final = _object_sequence(logging_entries) or () + return tuple( + (_entry_callback_name(entry), _entry_callback_type(entry), tuple(sorted(_entry_callback_vars(entry).items()))) + for entry in entries + ) def callback_config_error( @@ -279,6 +425,9 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str entries: Final = _object_sequence(metadata.get("logging")) if entries is None: return None + reference_error: Final = callback_reference_error(entries) + if reference_error is not None: + return reference_error entry_error: Final = next( (error for entry in entries if (error := _logging_entry_error(entry)) is not None), None, @@ -318,6 +467,14 @@ def _entry_callback_name(entry: object) -> str | None: return None +def _entry_callback_type(entry: object) -> str | None: + match entry: + case {"callback_type": str() as callback_type}: + return callback_type + case _: + return None + + def _conflicting_entry_setting_error( entry: tuple[str | None, Mapping[str, str]], earlier: Sequence[tuple[str | None, Mapping[str, str]]] ) -> str | None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..2fb049cba43 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -77,6 +77,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_utils import get_request_route from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.callback_config_validation import callback_backend_name, callback_definition from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, get_metadata_variable_name_from_kwargs, @@ -1243,6 +1244,69 @@ def get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallba _get_validated_callback_metadata: Final = get_validated_callback_metadata +_STRING_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def resolved_key_callbacks(logging_entries: object, *, source: str) -> tuple[AddTeamCallback, ...]: + """A key's ``logging`` entries as the backends see them, each ``callback_settings`` reference swapped for its definition. + + A reference is honoured only while no other entry on the key names its backend: every + entry's vars merge into one map per backend, so a sibling could point the definition's + credentials at its own host. The reference's own ``callback_vars`` are never read. + """ + if not isinstance(logging_entries, Sequence) or isinstance(logging_entries, (str, bytes)): + return () + entries: Final = tuple( + _STRING_KEYED_MAPPING.validate_python(item) + for item in _OBJECT_LIST_ADAPTER.validate_python(logging_entries) + if isinstance(item, Mapping) + ) + backends: Final = tuple(_entry_backend(entry) for entry in entries) + return tuple( + callback + for entry, backend in zip(entries, backends, strict=True) + if (callback := _resolved_key_callback(entry, source=source, shared_backend=backends.count(backend) > 1)) + is not None + ) + + +def _entry_backend(entry: Mapping[str, object]) -> str | None: + callback_name: Final = entry.get("callback_name") + return callback_backend_name(callback_name) if isinstance(callback_name, str) else None + + +def _resolved_key_callback(entry: Mapping[str, object], *, source: str, shared_backend: bool) -> AddTeamCallback | None: + definition: Final = callback_definition(entry.get("callback_name")) + if definition is None: + return get_validated_callback_metadata(item=dict(entry), source=source) + if shared_backend: + verbose_proxy_logger.warning( + "Ignoring %s callback metadata naming callback_settings.%s: another entry on this key names the %s callback", + source, + _sanitize_for_log(definition.name), + _sanitize_for_log(definition.backend), + ) + return None + reference: Final = get_validated_callback_metadata(item={**entry, "callback_vars": {}}, source=source) + if reference is None: + return None + return reference.model_copy( + update={"callback_name": definition.backend, "callback_vars": dict(definition.callback_vars)} + ) + + +def _validated_callbacks(logging_entries: Iterable[object], *, source: str) -> tuple[AddTeamCallback, ...]: + return tuple( + callback + for item in logging_entries + if ( + callback := get_validated_callback_metadata( + item=dict(_STRING_KEYED_MAPPING.validate_python(item)), source=source + ) + ) + is not None + ) + class KeyAndTeamLoggingSettings: """ @@ -1276,10 +1340,7 @@ def get_dynamic_logging_metadata( # Key-based callbacks ######################################################################################### if key_dynamic_logging_settings is not None: - for item in key_dynamic_logging_settings: - callback = get_validated_callback_metadata(item=item, source="key-level") - if callback is None: - continue + for callback in resolved_key_callbacks(key_dynamic_logging_settings, source="key-level"): callback_settings_obj = convert_key_logging_metadata_to_callback( data=callback, team_callback_settings_obj=callback_settings_obj, @@ -1443,8 +1504,11 @@ def resolve_tenant_otel_destinations( disabled: Final = _dynamically_disabled_backends(user_api_key_dict, request_headers) callbacks: Final = tuple( callback - for item in entries - if (callback := get_validated_callback_metadata(item=item, source="otel-destination")) is not None + for callback in ( + resolved_key_callbacks(entries, source="otel-destination") + if key_entries is not None + else _validated_callbacks(entries, source="otel-destination") + ) if callback.callback_name.lower() not in disabled ) return tuple( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index a76dad8d02a..15db352da3e 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -76,6 +76,8 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( # noq check_permissions_caller_permission, generate_key_helper_fn, prepare_metadata_fields, + raise_on_invalid_key_logging_config, + raise_on_unauthorized_callback_reference, ) 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 @@ -594,6 +596,11 @@ 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") + raise_on_invalid_key_logging_config(data.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + raise_on_unauthorized_callback_reference( + data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + 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 d56250a5504..c71cef6c662 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -67,7 +67,11 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( evict_and_broadcast, publish_auth_cache_invalidation, ) -from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error +from litellm.proxy.common_utils.callback_config_validation import ( + callback_definition, + logging_entry_keys, + logging_metadata_config_error, +) from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, encrypt_callback_vars, @@ -433,6 +437,7 @@ _KEY_METADATA_REQUEST_FIELDS: Final = frozenset( _DECODED_JSON: Final = TypeAdapter(object) +_STRING_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object]) def _decode_json_string_column(column: str, value: object) -> object: @@ -933,12 +938,49 @@ def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) - raise HTTPException(status_code=400, detail={"error": error}) +def raise_on_unauthorized_callback_reference( + metadata: Mapping[str, object] | None, + user_api_key_dict: UserAPIKeyAuth, + *, + existing_metadata: Mapping[str, object] | None = None, +) -> None: + """Only a proxy admin may point a key at a ``callback_settings`` definition or change the logging of a key holding one. + + The definition carries the operator's credentials, so the team admins and members who + can otherwise edit key metadata may only send such a key's ``logging`` list back unchanged. + """ + if metadata is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + requested: Final = logging_entry_keys(metadata.get("logging")) + stored: Final = logging_entry_keys(_stored_logging_entries(existing_metadata)) + if requested == stored or not any(callback_definition(name) is not None for name, _, _ in (*requested, *stored)): + return + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Only a proxy admin can point a key at a callback_settings definition " + "or change the logging entries of a key that holds one" + ) + }, + ) + + +def _stored_logging_entries(existing_metadata: Mapping[str, object] | None) -> object: + """The stored ``logging`` list with its callback vars decrypted, since ciphertext never equals the plaintext sent back.""" + if existing_metadata is None: + return None + decrypted: Final = _DECODED_JSON.validate_python(decrypt_callback_vars(existing_metadata)) + return _STRING_KEYED_MAPPING.validate_python(decrypted).get("logging") if isinstance(decrypted, Mapping) else None + + def common_key_access_checks( user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest | UpdateKeyRequest, llm_router: Router | None, premium_user: bool, user_id: str | None = None, + existing_metadata: Mapping[str, object] | None = None, ) -> Literal[True]: """ Check if user is allowed to make a key request, for this key @@ -965,6 +1007,11 @@ def common_key_access_checks( llm_router=llm_router, premium_user=premium_user, ) + raise_on_unauthorized_callback_reference( + data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + user_api_key_dict, + existing_metadata=existing_metadata, + ) return True @@ -3773,6 +3820,7 @@ async def _validate_update_key_data( user_id=user_api_key_dict.user_id if acting_as_team_admin else existing_key_row.user_id, llm_router=llm_router, premium_user=premium_user, + existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict ) await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( @@ -6652,6 +6700,11 @@ async def regenerate_key_fn( _key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict user_api_key_dict, ) + raise_on_unauthorized_callback_reference( + None if data is None else data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + user_api_key_dict, + existing_metadata=_key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict + ) # check if user has permission to regenerate key await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f5369bc34b0..b83d5f3f4a5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -168,6 +168,7 @@ from litellm.proxy._types import ( TransformRequestBody, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.callback_config_validation import callback_definitions_error from litellm.proxy.common_utils.callback_utils import ( is_sensitive_callback_key, normalize_callback_names, @@ -5862,6 +5863,16 @@ def _configured_db_callbacks(litellm_settings: Mapping[str, object], setting_key return tuple((setting_key, callback) for callback in callbacks if isinstance(callback, str)) +_CALLBACK_SETTINGS: Final = TypeAdapter(Mapping[str, object]) + + +def _check_callback_definitions(callback_settings: Mapping[str, object]) -> None: + """Refuse to boot on a ``callback_settings`` definition a key could never use.""" + definition_error: Final = callback_definitions_error(callback_settings) + if definition_error is not None: + raise ValueError(definition_error) + + class ProxyConfig: """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. @@ -6785,6 +6796,7 @@ class ProxyConfig: ## Callback settings callback_settings: Final = config.get("callback_settings", {}) if callback_settings: + _check_callback_definitions(_CALLBACK_SETTINGS.validate_python(callback_settings)) litellm.callback_settings = callback_settings ## LITELLM MODULE SETTINGS (e.g. litellm.drop_params=True,..) 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..3e1659db66f 100644 --- a/tests/unit/proxy/common_utils/test_callback_config_validation.py +++ b/tests/unit/proxy/common_utils/test_callback_config_validation.py @@ -3,11 +3,16 @@ from typing import Final import pytest +import litellm from litellm.proxy.common_utils.callback_config_validation import ( + callback_backend_name, callback_config_error, + callback_definition, + callback_definitions_error, conflicting_callback_var_error, conflicting_span_scope_error, cross_entry_family_error, + logging_entry_keys, logging_metadata_config_error, ) @@ -499,3 +504,158 @@ 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 + + +_ARIZE_DEFINITION: Final = { + "callback_type": "arize", + "arize_api_key": "resolved-space-a-key", + "arize_space_id": "space-a", +} + + +@pytest.fixture +def arize_space_a(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callback_settings", {"arize_space_a": dict(_ARIZE_DEFINITION)}) + + +@pytest.mark.parametrize( + ("callback_settings", "named_fields"), + [ + ({"arize": {"callback_type": "arize", "arize_api_key": "k"}}, ("callback_settings.arize", "named after")), + ( + {"space_a": {"callback_type": "not_a_callback"}}, + ("callback_settings.space_a.callback_type", "not_a_callback"), + ), + ({"space_a": {"callback_type": "arize", "arize_api_key": None}}, ("callback_settings.space_a.arize_api_key",)), + ({"space_a": {"callback_type": "arize", "arize_api_key": ""}}, ("callback_settings.space_a.arize_api_key",)), + ( + {"space_a": {"callback_type": "arize", "arize_space_id": "os.environ/UNSET"}}, + ("callback_settings.space_a.arize_space_id", "secret manager"), + ), + ({"space_a": {"callback_type": "arize", "arize_api_kye": "k"}}, ("callback_settings.space_a", "arize_api_kye")), + ( + {"space_a": {"callback_type": "arize", "arize_otlp_protocol": "carrier-pigeon"}}, + ("callback_settings.space_a", "arize_otlp_protocol"), + ), + ], +) +def test_a_definition_a_key_could_never_use_fails_config_load( + callback_settings: dict[str, object], named_fields: tuple[str, ...] +) -> None: + error: Final = callback_definitions_error(callback_settings) + assert error is not None + for fragment in named_fields: + assert fragment in error + + +def test_a_definition_error_names_the_field_and_never_its_value() -> None: + error: Final = callback_definitions_error( + { + "space_a": { + "callback_type": "arize", + "arize_otlp_protocol": "carrier-pigeon", + "arize_api_key": "sk-super-secret", + } + } + ) + assert error is not None and "arize_otlp_protocol" in error + assert "sk-super-secret" not in error + + +@pytest.mark.usefixtures("otel_v2_on") +def test_callback_settings_without_a_referenceable_definition_pass_config_load() -> None: + assert ( + callback_definitions_error( + { + "otel": {"message_logging": False, "excluded_services": ["postgres"]}, + "my_api": {"callback_type": "generic_api", "url": "os.environ/UNSET"}, + "space_a": { + "callback_type": "arize", + "arize_api_key": "k", + "arize_space_id": 12345, + "excluded_services": "postgres,redis", + }, + } + ) + is None + ) + + +@pytest.mark.usefixtures("arize_space_a") +def test_a_reference_resolves_to_its_definition_and_a_plain_name_does_not() -> None: + definition: Final = callback_definition("arize_space_a") + assert definition is not None + assert (definition.backend, dict(definition.callback_vars)) == ( + "arize", + {"arize_api_key": "resolved-space-a-key", "arize_space_id": "space-a"}, + ) + assert callback_definition("arize") is None + assert callback_definition(None) is None + assert (callback_backend_name("arize_space_a"), callback_backend_name("langfuse")) == ("arize", "langfuse") + + +@pytest.mark.usefixtures("arize_space_a") +def test_a_reference_entry_may_not_carry_callback_vars() -> None: + error: Final = logging_metadata_config_error( + { + "logging": [ + { + "callback_name": "arize_space_a", + "callback_type": "success", + "callback_vars": {"arize_space_id": "attacker-space"}, + } + ] + } + ) + assert error is not None and "arize_space_a" in error and "callback_vars" in error + + +@pytest.mark.usefixtures("arize_space_a") +def test_a_reference_is_the_only_entry_for_its_backend() -> None: + sibling: Final = {"callback_name": "arize", "callback_type": "success", "callback_vars": {"arize_space_id": "x"}} + reference: Final = {"callback_name": "arize_space_a", "callback_type": "success"} + error: Final = logging_metadata_config_error({"logging": [reference, sibling]}) + assert error is not None and "arize_space_a" in error and "arize" in error + assert logging_metadata_config_error({"logging": [sibling]}) is None + + +@pytest.mark.usefixtures("arize_space_a") +def test_a_reference_sits_next_to_entries_for_other_backends() -> None: + assert ( + logging_metadata_config_error( + { + "logging": [ + {"callback_name": "arize_space_a", "callback_type": "success"}, + {"callback_name": "arize_space_a", "callback_type": "failure", "callback_vars": {}}, + {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}, + ] + } + ) + is not None + ), "two references to one definition still share its backend" + assert ( + logging_metadata_config_error( + { + "logging": [ + {"callback_name": "arize_space_a", "callback_type": "success"}, + {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}, + ] + } + ) + is None + ) + + +def test_logging_entry_keys_change_with_every_field_a_caller_could_edit() -> None: + entry: Final = {"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {"b": "2", "a": "1"}} + reordered: Final = { + "callback_type": "success", + "callback_vars": {"a": "1", "b": "2"}, + "callback_name": "arize_space_a", + } + assert logging_entry_keys([entry]) == logging_entry_keys([reordered]) + assert logging_entry_keys([entry]) != logging_entry_keys([{**entry, "callback_type": "failure"}]) + assert logging_entry_keys([entry]) != logging_entry_keys([{**entry, "callback_vars": {"a": "1"}}]) + assert logging_entry_keys([entry]) != logging_entry_keys([{**entry, "callback_name": "arize"}]) + assert logging_entry_keys([entry]) != logging_entry_keys([entry, entry]) + assert logging_entry_keys(None) == logging_entry_keys("not-a-list") == () diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 120cf8a2b84..d803493d7cd 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -16,6 +16,8 @@ from fastapi.testclient import TestClient from pydantic import ValidationError from pytest_mock import MockerFixture +import litellm + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( LiteLLM_UserTable, @@ -5496,3 +5498,42 @@ async def test_fips_user_update_hashes_off_the_event_loop(pbkdf2_on_event_loop: assert pbkdf2_on_event_loop == [False] assert data["password"].startswith("pbkdf2:") + + +@pytest.mark.asyncio +async def test_new_user_cannot_hand_its_auto_created_key_a_callback_reference_for_a_non_admin( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """``/user/new`` writes ``metadata`` onto the key it auto-creates, so it runs the same admin gate as ``/key/generate``.""" + from litellm.proxy import proxy_server + from litellm.proxy.utils import PrismaClient + + monkeypatch.setattr( + litellm, + "callback_settings", + {"arize_space_a": {"callback_type": "arize", "arize_api_key": "k", "arize_space_id": "s"}}, + ) + table: Final = SimpleNamespace( + find_first=AsyncMock(return_value=None), + count=AsyncMock(return_value=0), + create=AsyncMock(), + upsert=AsyncMock(), + ) + mocker.patch("prisma.Prisma", return_value=SimpleNamespace(litellm_usertable=table)) + database: Final = PrismaClient( + database_url="postgresql://unused:unused@localhost/unused", + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", database) + reference: Final = {"logging": [{"callback_name": "arize_space_a", "callback_type": "success"}]} + + with pytest.raises(ProxyException) as exc: + await new_user( + data=NewUserRequest(user_email="bob@example.test", metadata=reference, send_invite_email=False), + user_api_key_dict=UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert str(exc.value.code) == "403" + assert "callback_settings" in exc.value.message + table.create.assert_not_awaited() + table.upsert.assert_not_awaited() 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 77dc6eddf51..460de971563 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import ( jwt_key_mapping_cache_key, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, @@ -87,6 +88,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _validate_max_budget, _validate_reset_spend_value, _validate_update_key_data, + raise_on_unauthorized_callback_reference, can_modify_verification_token, check_org_key_model_specific_limits, check_team_key_model_specific_limits, @@ -24243,3 +24245,187 @@ 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") + + +_REFERENCE_ENTRY: Final = {"callback_name": "arize_space_a", "callback_type": "success"} +_LANGFUSE_ENTRY: Final = { + "callback_name": "langfuse", + "callback_type": "success", + "callback_vars": {"langfuse_public_key": "pk", "langfuse_secret_key": "sk-plain"}, +} + + +@pytest.fixture +def arize_space_a(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm, + "callback_settings", + {"arize_space_a": {"callback_type": "arize", "arize_api_key": "k", "arize_space_id": "s"}}, + ) + + +def _non_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-alice", user_id="alice") + + +@pytest.mark.usefixtures("arize_space_a") +class TestCallbackReferenceWriteGate: + def test_a_proxy_admin_may_point_a_key_at_a_definition(self) -> None: + admin: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + assert raise_on_unauthorized_callback_reference({"logging": [_REFERENCE_ENTRY]}, admin) is None + assert ( + raise_on_unauthorized_callback_reference( + {"logging": []}, admin, existing_metadata={"logging": [_REFERENCE_ENTRY]} + ) + is None + ) + + @pytest.mark.parametrize( + "metadata", + [ + {"logging": [_REFERENCE_ENTRY]}, + {"logging": [_LANGFUSE_ENTRY, _REFERENCE_ENTRY]}, + {"logging": [{**_REFERENCE_ENTRY, "callback_vars": {}}]}, + ], + ) + def test_a_non_admin_cannot_point_a_key_at_a_definition(self, metadata: Mapping[str, object]) -> None: + with pytest.raises(HTTPException) as exc: + raise_on_unauthorized_callback_reference(metadata, _non_admin()) + assert exc.value.status_code == 403 + assert "callback_settings" in str(exc.value.detail) + + @pytest.mark.parametrize( + "metadata", + [ + {"note": "edited"}, + {"logging": []}, + {"logging": [{**_REFERENCE_ENTRY, "callback_type": "failure"}]}, + {"logging": [{**_REFERENCE_ENTRY, "callback_name": "arize"}]}, + {"logging": [_REFERENCE_ENTRY, _LANGFUSE_ENTRY]}, + {"logging": [{**_LANGFUSE_ENTRY, "callback_vars": {"langfuse_host": "https://attacker.example"}}]}, + ], + ) + def test_a_non_admin_cannot_change_the_logging_of_a_key_holding_a_reference( + self, metadata: Mapping[str, object] + ) -> None: + with pytest.raises(HTTPException) as exc: + raise_on_unauthorized_callback_reference( + metadata, _non_admin(), existing_metadata={"logging": [_REFERENCE_ENTRY]} + ) + assert exc.value.status_code == 403 + + def test_a_non_admin_may_edit_the_rest_of_the_metadata_when_the_logging_is_sent_back_unchanged(self) -> None: + stored: Final = encrypt_callback_vars({"note": "original", "logging": [_REFERENCE_ENTRY, _LANGFUSE_ENTRY]}) + assert ( + raise_on_unauthorized_callback_reference( + {"note": "edited", "logging": [_REFERENCE_ENTRY, _LANGFUSE_ENTRY]}, + _non_admin(), + existing_metadata=stored, + ) + is None + ) + + @pytest.mark.parametrize("existing_metadata", [None, {}, {"logging": [_LANGFUSE_ENTRY]}]) + def test_a_key_without_a_reference_keeps_the_logging_rules_it_had( + self, existing_metadata: Mapping[str, object] | None + ) -> None: + changed: Final = [ + {**_LANGFUSE_ENTRY, "callback_type": "failure"}, + {"callback_name": "datadog", "callback_vars": {}}, + ] + assert ( + raise_on_unauthorized_callback_reference( + {"logging": changed}, _non_admin(), existing_metadata=existing_metadata + ) + is None + ) + assert ( + raise_on_unauthorized_callback_reference( + None, _non_admin(), existing_metadata={"logging": [_REFERENCE_ENTRY]} + ) + is None + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("arize_space_a") +async def test_generate_key_with_a_callback_reference_is_admin_only() -> None: + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(metadata={"logging": [_REFERENCE_ENTRY]}), + user_api_key_dict=_non_admin(), + litellm_changed_by=None, + team_table=None, + ) + assert exc.value.status_code == 403 + assert "callback_settings" in str(exc.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("arize_space_a") +async def test_generate_key_rejects_callback_vars_on_a_reference_for_everyone() -> None: + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest( + metadata={"logging": [{**_REFERENCE_ENTRY, "callback_vars": {"arize_space_id": "mine"}}]} + ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + litellm_changed_by=None, + team_table=None, + ) + assert exc.value.status_code == 400 + assert "callback_vars" in str(exc.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("arize_space_a") +async def test_update_key_compares_the_sent_logging_against_the_stored_row() -> None: + existing_key_row: Final = MagicMock() + existing_key_row.user_id = "alice" + existing_key_row.token = "hashed_token" + existing_key_row.team_id = None + existing_key_row.metadata = {"logging": [_REFERENCE_ENTRY]} + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=UpdateKeyRequest(key="sk-test-key", metadata={"note": "edited"}), + existing_key_row=existing_key_row, + user_api_key_dict=_non_admin(), + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + assert exc.value.status_code == 403 + assert "callback_settings" in str(exc.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("arize_space_a") +async def test_regenerate_key_cannot_drop_a_reference_for_a_non_admin(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import regenerate_key_fn + + key_row: Final = MagicMock() + key_row.token = "hashed_token" + key_row.user_id = "alice" + key_row.team_id = None + key_row.metadata = {"logging": [_REFERENCE_ENTRY]} + prisma: Final = MagicMock() + prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_row) + prisma.db.litellm_verificationtoken.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-master") + + with pytest.raises(ProxyException) as exc: + await regenerate_key_fn( + key=None, + data=RegenerateKeyRequest(key="sk-test", metadata={"note": "edited"}), + user_api_key_dict=_non_admin(), + litellm_changed_by=None, + ) + assert str(exc.value.code) == "403" + assert "callback_settings" in str(exc.value.message) + prisma.db.litellm_verificationtoken.update.assert_not_awaited() diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 712b4134c05..9e567a67aea 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -6593,3 +6593,68 @@ async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path, return with pytest.raises(ValidationError): await ProxyConfig().load_config(router=None, config_file_path=str(config)) + + +_CALLBACK_DEFINITION_CONFIG: Final = """ +model_list: [] +callback_settings: + arize_space_a: + callback_type: arize + arize_api_key: os.environ/QA_ARIZE_SPACE_A_KEY + arize_space_id: os.environ/QA_ARIZE_SPACE_A_SPACE +""" + + +@pytest.fixture +def config_without_db(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "callback_settings", {}) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("config_without_db") +async def test_a_callback_definition_resolves_its_secrets_at_boot( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("QA_ARIZE_SPACE_A_KEY", "space-a-key-v1") + monkeypatch.setenv("QA_ARIZE_SPACE_A_SPACE", "space-a") + config_file: Final = tmp_path / "callbacks.yaml" + config_file.write_text(_CALLBACK_DEFINITION_CONFIG) + + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + assert litellm.callback_settings["arize_space_a"] == { + "callback_type": "arize", + "arize_api_key": "space-a-key-v1", + "arize_space_id": "space-a", + } + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("config_without_db") +async def test_a_callback_definition_whose_secret_is_missing_fails_boot_naming_the_field( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("QA_ARIZE_SPACE_A_KEY", "space-a-key-v1") + monkeypatch.delenv("QA_ARIZE_SPACE_A_SPACE", raising=False) + config_file: Final = tmp_path / "callbacks.yaml" + config_file.write_text(_CALLBACK_DEFINITION_CONFIG) + + with pytest.raises(ValueError, match=r"callback_settings\.arize_space_a\.arize_space_id"): + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + assert litellm.callback_settings == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("config_without_db") +async def test_a_callback_definition_with_a_misspelled_field_fails_boot(tmp_path: Path) -> None: + config_file: Final = tmp_path / "callbacks.yaml" + config_file.write_text( + "model_list: []\ncallback_settings:\n arize_space_a:\n callback_type: arize\n arize_api_kye: k\n" + ) + + with pytest.raises(ValueError, match="arize_api_kye"): + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 97c21fcc777..0d7fe4089c8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -11443,3 +11443,111 @@ def test_caller_tags_reach_litellm_metadata_whether_metadata_is_a_dict_or_a_json LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data) assert set(data["litellm_metadata"]["tags"]) == {"key-tag", "client-tag"} assert data["metadata"] == caller_metadata + + +_REFERENCE_ENTRY: Final = {"callback_name": "arize_space_a", "callback_type": "success"} +_ARIZE_SPACE_A: Final = { + "callback_type": "arize", + "arize_api_key": "resolved-space-a-key", + "arize_space_id": "space-a", + "arize_otlp_protocol": "http/protobuf", + "excluded_services": "postgres,redis", +} + + +@pytest.fixture +def arize_space_a(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callback_settings", {"arize_space_a": dict(_ARIZE_SPACE_A)}) + + +@pytest.fixture +def otel_v2_on(monkeypatch: pytest.MonkeyPatch): + from litellm.integrations.otel.model.config import is_otel_v2_enabled + + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("ARIZE_ENDPOINT", "https://arize.internal.example/v1") + monkeypatch.delenv("ARIZE_HTTP_ENDPOINT", raising=False) + is_otel_v2_enabled.cache_clear() + yield + is_otel_v2_enabled.cache_clear() + + +@pytest.mark.usefixtures("arize_space_a") +def test_a_key_logging_entry_naming_a_definition_runs_its_backend_with_the_definition_vars() -> None: + auth: Final = UserAPIKeyAuth(api_key="hashed-key", metadata={"logging": [_REFERENCE_ENTRY]}, team_metadata={}) + + result: Final = get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock()) + + assert result is not None + assert result.success_callback == ["arize"] + assert not result.failure_callback + assert result.callback_vars == { + key: value for key, value in _ARIZE_SPACE_A.items() if key not in ("callback_type", "excluded_services") + }, "excluded_services is a destination-only var, which the destination test below covers" + + +@pytest.mark.usefixtures("arize_space_a") +def test_the_vars_stored_on_a_reference_entry_never_reach_the_backend() -> None: + auth: Final = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"logging": [{**_REFERENCE_ENTRY, "callback_vars": {"arize_space_id": "attacker-space"}}]}, + team_metadata={}, + ) + + result: Final = get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock()) + + assert result is not None and result.callback_vars is not None + assert result.callback_vars["arize_space_id"] == "space-a" + + +@pytest.mark.usefixtures("arize_space_a") +def test_a_reference_is_skipped_when_another_entry_on_the_key_names_its_backend() -> None: + """A stored inline entry for the same backend would merge its host or space with the definition's + credentials, so the definition's secrets stay out of the request instead.""" + sibling: Final = { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {"arize_space_id": "attacker-space", "arize_api_key": "attacker-key"}, + } + auth: Final = UserAPIKeyAuth( + api_key="hashed-key", metadata={"logging": [_REFERENCE_ENTRY, sibling]}, team_metadata={} + ) + + result: Final = get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock()) + + assert result is not None + assert result.callback_vars == sibling["callback_vars"] + + +@pytest.mark.usefixtures("arize_space_a") +def test_a_team_logging_entry_naming_a_definition_is_not_resolved() -> None: + auth: Final = UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata={"logging": [_REFERENCE_ENTRY]}) + + assert get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock()) is None + + +@pytest.mark.usefixtures("arize_space_a", "otel_v2_on") +def test_a_reference_resolves_to_the_same_otel_destination_as_the_inline_entry() -> None: + from litellm.proxy.litellm_pre_call_utils import resolve_tenant_otel_destinations + + inline: Final = { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {key: value for key, value in _ARIZE_SPACE_A.items() if key != "callback_type"}, + } + referencing_key: Final = UserAPIKeyAuth( + api_key="hashed-key", team_id="team-a", metadata={"logging": [_REFERENCE_ENTRY]}, team_metadata={} + ) + inline_key: Final = UserAPIKeyAuth( + api_key="other-key", team_id="team-a", metadata={"logging": [inline]}, team_metadata={} + ) + plain_key_on_the_same_team: Final = UserAPIKeyAuth( + api_key="plain-key", team_id="team-a", metadata={}, team_metadata={} + ) + + resolved: Final = resolve_tenant_otel_destinations(referencing_key) + + assert resolved == resolve_tenant_otel_destinations(inline_key) + assert len(resolved) == 1 + assert resolved[0].headers["api_key"] == "resolved-space-a-key" + assert resolve_tenant_otel_destinations(plain_key_on_the_same_team) == ()