diff --git a/litellm/constants.py b/litellm/constants.py index 32b47bc6b46..6f8bdc74278 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -253,6 +253,7 @@ MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_OAUTH2_TOKEN_EX MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = get_env_int("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", 200) MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = get_env_int("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", 3600) MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = get_env_int("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", 60) +CALLBACK_SETTINGS_SECRET_CACHE_TTL_SECONDS: Final = get_env_int("CALLBACK_SETTINGS_SECRET_CACHE_TTL_SECONDS", 60) # mcp_tool_permissions entry that grants every current and future tool on a server MCP_ALL_TOOLS_WILDCARD: Final = "*" diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 302978d17a4..2760e9f0546 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -10,6 +10,9 @@ from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType from typing import Final, 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 +24,33 @@ _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" +_STRING_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def is_named_callback_definition(definition: object) -> bool: + """Whether a ``callback_settings`` entry names a backend a key can reference. + + Such an entry keeps its ``os.environ/`` values past config load, so the secret + manager is read on each request that uses it rather than once at boot. + """ + if not isinstance(definition, Mapping): + return False + callback_type: Final = _STRING_KEYED_MAPPING.validate_python(definition).get("callback_type") + return isinstance(callback_type, str) and callback_type != _GENERIC_API_CALLBACK_TYPE + + +def named_callback_definition(callback_name: str) -> Mapping[str, object] | None: + raw: Final[object] = litellm.callback_settings.get(callback_name) + if not is_named_callback_definition(raw): + return None + return _STRING_KEYED_MAPPING.validate_python(raw) + + +def callback_backend_name(callback_name: str) -> str: + """The backend a key's ``logging`` entry runs on: the definition's ``callback_type`` for a reference.""" + definition: Final = named_callback_definition(callback_name) + return callback_name if definition is None else str(definition["callback_type"]) def callback_config_error( @@ -286,7 +316,7 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str if entry_error is not None: return entry_error entry_vars: Final = tuple(_entry_callback_vars(entry) for entry in entries) - named_vars: Final = tuple(zip((_entry_callback_name(entry) for entry in entries), entry_vars)) + named_vars: Final = tuple(_backend_entry(entry) for entry in entries) return next( ( error @@ -300,6 +330,21 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str ) +def _backend_entry(entry: object) -> tuple[str | None, Mapping[str, str]]: + """An entry as the backend sees it: a reference carries its definition's content settings.""" + callback_name: Final = _entry_callback_name(entry) + callback_vars: Final = _entry_callback_vars(entry) + definition: Final = None if callback_name is None else named_callback_definition(callback_name) + if callback_name is None or definition is None: + return (callback_name, callback_vars) + inherited: Final = { + key: value + for key, value in definition.items() + if key in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) and isinstance(value, str) + } + return (str(definition["callback_type"]), MappingProxyType({**inherited, **callback_vars})) + + 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 () @@ -369,7 +414,7 @@ def _logging_entry_error(entry: object) -> str | None: return destination_var_error callback_type_raw: Final = typed_entry.get("callback_type") return callback_config_error( - callback_name, + callback_backend_name(callback_name), _entry_callback_vars(typed_entry), callback_type_raw if isinstance(callback_type_raw, str) else None, ) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index d1787f3ff8b..d6a46ec002a 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,19 +1,21 @@ import copy import json import os -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from itertools import accumulate -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, Protocol, TypeAlias -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( BAGGAGE_MAPPED_TAGS_METADATA_KEY, + CALLBACK_SETTINGS_SECRET_CACHE_TTL_SECONDS, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH, @@ -27,7 +29,8 @@ from litellm.litellm_core_utils.core_helpers import ( get_or_create_metadata_bucket, ) from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams +from litellm.proxy._types import AddTeamCallback, CommonProxyErrors, LiteLLMPromptInjectionParams +from litellm.proxy.common_utils.callback_config_validation import named_callback_definition from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -796,3 +799,122 @@ def _decrypt_or_passthrough(key: str, value: object) -> object: inner: Final = value[len(_CALLBACK_VAR_ENCRYPTED_PREFIX) :] decrypted: Final = decrypt_value_helper(value=inner, key=key, exception_type="debug", return_original_value=False) return decrypted if decrypted is not None else value + + +_STRING_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_SEQUENCE: Final = TypeAdapter(tuple[object, ...]) +_SECRET_VALUE: Final[TypeAdapter[str | bool | None]] = TypeAdapter(str | bool | None) +_SECRET_REFERENCE_PREFIX: Final = "os.environ/" + + +class _SecretCache(Protocol): + def get_cache(self, key: str) -> object: ... + def set_cache(self, key: str, value: object) -> None: ... + def flush_cache(self) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _CachedSecret: + """Wraps the value so ``InMemoryCache.get_cache`` cannot json-decode a numeric or boolean-looking string.""" + + value: str | bool + + +_SECRET_CACHE: Final = InMemoryCache(max_size_in_memory=256, default_ttl=CALLBACK_SETTINGS_SECRET_CACHE_TTL_SECONDS) + + +def _secret_cache() -> _SecretCache: + return _SECRET_CACHE + + +def resolve_named_callback(callback: AddTeamCallback) -> AddTeamCallback: + """Swap a reference to a ``callback_settings`` definition for the backend it names. + + The definition's values go through the secret manager on every call, so a rotated + secret is picked up by the next request. The key keeps only the reference name + and may add values the definition does not set. + """ + definition: Final = named_callback_definition(callback.callback_name) + if definition is None: + return callback + backend: Final = str(definition["callback_type"]) + callback_vars: Final = { + **{ + key: _resolved_definition_value(callback.callback_name, key, value) + for key, value in definition.items() + if key != "callback_type" + }, + **callback.callback_vars, + } + try: + return AddTeamCallback(callback_name=backend, callback_type=callback.callback_type, callback_vars=callback_vars) + except (ValidationError, ValueError): + rejected: Final = ", ".join(_rejected_callback_vars(backend, callback_vars)) + raise ValueError( + f"callback_settings.{callback.callback_name} sets {rejected} to a value the {backend} callback rejects" + ) from None + + +def _rejected_callback_vars(backend: str, callback_vars: Mapping[str, str]) -> tuple[str, ...]: + return tuple(key for key, value in callback_vars.items() if not _accepts_callback_var(backend, key, value)) + + +def _accepts_callback_var(backend: str, key: str, value: str) -> bool: + try: + AddTeamCallback(callback_name=backend, callback_vars={key: value}) + except (ValidationError, ValueError): + return False + return True + + +def named_callback_resolution_error(entries: object) -> str | None: + """Why a key's ``callback_settings`` references cannot be resolved right now, or None.""" + items: Final = ( + _OBJECT_SEQUENCE.validate_python(entries) + if isinstance(entries, Sequence) and not isinstance(entries, (str, bytes)) + else () + ) + return next((error for item in items if (error := _reference_resolution_error(item)) is not None), None) + + +def _reference_resolution_error(item: object) -> str | None: + if not isinstance(item, Mapping): + return None + entry: Final = _STRING_KEYED_MAPPING.validate_python(item) + callback_name: Final = entry.get("callback_name") + if not isinstance(callback_name, str) or named_callback_definition(callback_name) is None: + return None + try: + resolve_named_callback(AddTeamCallback.model_validate(entry)) + except (ValidationError, ValueError) as e: + return str(e) + return None + + +def _resolved_definition_value(definition_name: str, key: str, value: object) -> str: + if value is None: + raise ValueError(f"callback_settings.{definition_name}.{key} is not set") + raw: Final = value if isinstance(value, str) else str(value) + if not raw.startswith(_SECRET_REFERENCE_PREFIX): + return raw + secret: Final = _secret_manager_value(raw) + if secret is None: + raise ValueError(f"callback_settings.{definition_name}.{key} could not be resolved from the secret manager") + return secret if isinstance(secret, str) else str(secret).lower() + + +def _secret_manager_value(reference: str) -> str | bool | None: + """One secret-manager read per reference per TTL per worker, so a slow manager does not stall every request. + + Misses are not cached, so a secret created after the first read is picked up on the next request. + """ + cache: Final = _secret_cache() if CALLBACK_SETTINGS_SECRET_CACHE_TTL_SECONDS > 0 else None + cached: Final = cache.get_cache(reference) if cache is not None else None + if isinstance(cached, _CachedSecret): + return cached.value + secret: Final = _SECRET_VALUE.validate_python(get_secret(reference, default_value=reference)) + if secret is None or secret in ("", reference): + return None + if cache is not None: + cache.set_cache(reference, _CachedSecret(secret)) + return secret diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..ed88c4d0151 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -80,6 +80,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, + resolve_named_callback, strip_callback_config, ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports @@ -1241,6 +1242,23 @@ def get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallba return None +def get_resolved_key_callback_metadata(item: Mapping[str, object], *, source: str) -> AddTeamCallback | None: + """A key-level entry, with a ``callback_settings`` reference swapped for the backend it names.""" + callback: Final = get_validated_callback_metadata(item=dict(item), source=source) + if callback is None: + return None + try: + return resolve_named_callback(callback) + except (PydanticValidationError, ValueError) as e: + verbose_proxy_logger.warning( + "Ignoring %s callback metadata naming callback_settings.%s: %s", + source, + _sanitize_for_log(callback.callback_name), + _sanitize_for_log(str(e)), + ) + return None + + _get_validated_callback_metadata: Final = get_validated_callback_metadata @@ -1277,7 +1295,7 @@ def get_dynamic_logging_metadata( ######################################################################################### 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") + callback = get_resolved_key_callback_metadata(item=item, source="key-level") if callback is None: continue callback_settings_obj = convert_key_logging_metadata_to_callback( @@ -1444,7 +1462,7 @@ def resolve_tenant_otel_destinations( callbacks: Final = tuple( callback for item in entries - if (callback := get_validated_callback_metadata(item=item, source="otel-destination")) is not None + if (callback := _destination_callback(item, key_level=key_entries is not None)) is not None if callback.callback_name.lower() not in disabled ) return tuple( @@ -1456,6 +1474,12 @@ def resolve_tenant_otel_destinations( ) +def _destination_callback(item: Mapping[str, object], *, key_level: bool) -> AddTeamCallback | None: + if key_level: + return get_resolved_key_callback_metadata(item=item, source="otel-destination") + return get_validated_callback_metadata(item=dict(item), source="otel-destination") + + def _tenant_destination( name: str, callbacks: Sequence[AddTeamCallback], service_name: str | None ) -> "OtelDestination | None": diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c8895c01f7b..adabc6e7329 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -67,10 +67,15 @@ 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 ( + logging_metadata_config_error, + named_callback_definition, + stored_callback_entries, +) from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, encrypt_callback_vars, + named_callback_resolution_error, ) from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -929,6 +934,47 @@ 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 +) -> None: + """Only a proxy admin may point a key at a ``callback_settings`` definition. + + The definition carries the operator's secret-manager references, so a key that + names it must not be writable by the team admins and members who can otherwise + edit key metadata, and must not override the values the definition sets. + """ + references: Final = tuple( + (callback_name, callback_vars, definition) + for callback_name, callback_vars in stored_callback_entries((metadata or {}).get("logging")) + if callback_name is not None and (definition := named_callback_definition(callback_name)) is not None + ) + if not references: + return + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + raise HTTPException( + status_code=403, + detail={"error": "Only a proxy admin can point a key at a callback_settings definition"}, + ) + overridden: Final = next( + (error for reference in references if (error := _overridden_definition_var_error(reference)) is not None), + None, + ) + if overridden is not None: + raise HTTPException(status_code=400, detail={"error": overridden}) + + +def _overridden_definition_var_error( + reference: tuple[str, Mapping[str, str], Mapping[str, object]], +) -> str | None: + callback_name, callback_vars, definition = reference + variable_name: Final = next((name for name in callback_vars if name in definition), None) + if variable_name is None: + return None + return ( + f"callback_vars.{variable_name} is set by callback_settings.{callback_name} and cannot be overridden on the key" + ) + + def common_key_access_checks( user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest | UpdateKeyRequest, @@ -961,6 +1007,7 @@ def common_key_access_checks( llm_router=llm_router, premium_user=premium_user, ) + raise_on_unauthorized_callback_reference(data.metadata, user_api_key_dict) return True @@ -6462,6 +6509,7 @@ def _check_regenerate_guardrail_opt_out( user_api_key_dict, existing_metadata=existing_metadata, ) + raise_on_unauthorized_callback_reference(data.metadata, user_api_key_dict) @router.post( @@ -8347,10 +8395,16 @@ async def key_health( # Check if logging is configured in metadata if key_metadata and "logging" in key_metadata: - logging_statuses: Final = await test_key_logging( - user_api_key_dict=user_api_key_dict, - request=request, - key_logging=decrypt_callback_vars(key_metadata)["logging"], + key_logging: Final = decrypt_callback_vars(key_metadata)["logging"] + resolution_error: Final = named_callback_resolution_error(key_logging) + logging_statuses: Final = ( + await test_key_logging(user_api_key_dict=user_api_key_dict, request=request, key_logging=key_logging) + if resolution_error is None + else LoggingCallbackStatus( + callbacks=[name for name, _ in stored_callback_entries(key_logging) if name is not None], + status="unhealthy", + details=resolution_error, + ) ) health_status["logging_callbacks"] = logging_statuses diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 698c1e6e752..463e148d318 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -167,6 +167,7 @@ from litellm.proxy._types import ( TransformRequestBody, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.callback_config_validation import is_named_callback_definition from litellm.proxy.common_utils.callback_utils import ( is_sensitive_callback_key, normalize_callback_names, @@ -5904,6 +5905,9 @@ def _configured_db_callbacks(litellm_settings: Mapping[str, object], setting_key return tuple((setting_key, callback) for callback in callbacks if isinstance(callback, str)) +_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object]) + + class ProxyConfig: """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. @@ -6300,7 +6304,21 @@ class ProxyConfig: for key, value in config.items() } + def _resolved_callback_setting(self, definition: object, depth: int, max_depth: int) -> object: + if is_named_callback_definition(definition) or not isinstance(definition, dict): + return definition + return _SETTINGS_ADAPTER.validate_python( + self._check_for_os_environ_vars( + config=_SETTINGS_ADAPTER.validate_python(definition), depth=depth, max_depth=max_depth + ) + ) + def _resolved_config_value(self, key: str, value: object, depth: int, max_depth: int) -> object: + if key == "callback_settings" and depth == 0 and isinstance(value, dict): + return { + name: self._resolved_callback_setting(definition, depth=depth + 2, max_depth=max_depth) + for name, definition in _SETTINGS_ADAPTER.validate_python(value).items() + } if isinstance(value, dict): return self._check_for_os_environ_vars(config=value, depth=depth + 1, max_depth=max_depth) if isinstance(value, list): @@ -6336,10 +6354,11 @@ class ProxyConfig: if litellm.secret_manager_client is not None: return - general_settings: Final = config.get("general_settings") - if not isinstance(general_settings, dict): + raw_general_settings: Final = config.get("general_settings") + if not isinstance(raw_general_settings, dict): return + general_settings: Final = _SETTINGS_ADAPTER.validate_python(raw_general_settings) raw_system: Final = general_settings.get("key_management_system") key_management_system: Final = ( get_secret(raw_system) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 914eb1f2508..eb4be956d50 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -492,3 +492,12 @@ def pytest_sessionfinish() -> None: _close_handler_if_needed(getattr(litellm, name, None)) _run_coroutine_if_needed(close_litellm_async_clients()) enable_socket() + + +@pytest.fixture(autouse=True) +def _flush_callback_settings_secret_cache(): + from litellm.proxy.common_utils.callback_utils import _SECRET_CACHE + + _SECRET_CACHE.flush_cache() + yield + _SECRET_CACHE.flush_cache() diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index 6dd9ff35026..b8c8f504edc 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -3950,6 +3950,67 @@ class TestCallbackTypeFilter: assert resolve_tenant_otel_destinations(self._auth("failure")) == () +class TestCallbackSettingsReference: + """A key that names a ``callback_settings`` definition exports to the backend that + definition configures, with its secrets read from the secret manager, while a + sibling key on the same team without the reference exports nowhere.""" + + _SETTINGS: Final = { + "arize_space_a": { + "callback_type": "arize", + "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY", + "arize_space_id": "os.environ/LIT10061_ARIZE_SPACE_ID", + "excluded_services": "postgres,redis", + } + } + + @pytest.fixture(autouse=True) + def _v2_on_with_definition(self, monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setattr(litellm, "callback_settings", self._SETTINGS) + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "space-a-key") + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + is_otel_v2_enabled.cache_clear() + yield + is_otel_v2_enabled.cache_clear() + + @staticmethod + def _team_key(metadata: dict) -> UserAPIKeyAuth: + return UserAPIKeyAuth(team_id="team-a", metadata=metadata, team_metadata={}) + + def test_the_referencing_key_exports_to_the_definitions_backend(self): + auth = self._team_key( + {"logging": [{"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {}}]} + ) + + (destination,) = resolve_tenant_otel_destinations(auth) + + assert destination.callback_name == "arize" + assert destination.headers["api_key"] == "space-a-key" + assert destination.headers["arize-space-id"] == "space-a" + assert destination.excluded_db_systems == frozenset({"postgresql", "redis"}) + + def test_a_sibling_key_on_the_same_team_exports_nowhere(self): + assert resolve_tenant_otel_destinations(self._team_key({})) == () + assert resolve_tenant_otel_destinations(self._team_key({"logging": []})) == () + + def test_disabling_the_backend_on_the_key_drops_the_referenced_destination(self, monkeypatch): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(litellm, "allow_dynamic_callback_disabling", True) + auth = UserAPIKeyAuth( + team_id="team-a", + metadata={ + "logging": [{"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {}}], + "litellm_disabled_callbacks": ["arize"], + }, + team_metadata={}, + ) + + assert resolve_tenant_otel_destinations(auth) == () + + class TestTenantConfigAgreement: """The destination resolver and ``convert_key_logging_metadata_to_callback`` read the same stored config, so they must not read it two different ways.""" 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..952d375fc6c 100644 --- a/tests/unit/proxy/common_utils/test_callback_config_validation.py +++ b/tests/unit/proxy/common_utils/test_callback_config_validation.py @@ -3,6 +3,7 @@ from typing import Final import pytest +import litellm from litellm.proxy.common_utils.callback_config_validation import ( callback_config_error, conflicting_callback_var_error, @@ -499,3 +500,70 @@ 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 + + +class TestCallbackReferenceEntries: + @staticmethod + def _entry(callback_name: str, **callback_vars: str) -> dict[str, object]: + return { + "logging": [{"callback_name": callback_name, "callback_type": "success", "callback_vars": callback_vars}] + } + + @pytest.fixture + def otel_v2_on(self, monkeypatch) -> Iterator[None]: + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + from litellm.integrations.otel.model.config import is_otel_v2_enabled + + is_otel_v2_enabled.cache_clear() + monkeypatch.setattr( + litellm, + "callback_settings", + { + "arize_space_a": { + "callback_type": "arize", + "arize_api_key": "os.environ/A", + "arize_space_id": "s", + "capture_message_content": "no_content", + }, + "langfuse_team_b": { + "callback_type": "langfuse", + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + }, + }, + ) + try: + yield + finally: + is_otel_v2_enabled.cache_clear() + + def test_a_reference_accepts_the_options_its_backend_supports(self, otel_v2_on): + assert logging_metadata_config_error(self._entry("arize_space_a", excluded_services="postgres,redis")) is None + assert logging_metadata_config_error(self._entry("arize_space_a", arize_otlp_protocol="grpc")) is None + + def test_a_reference_rejects_the_options_its_backend_does_not_support(self, otel_v2_on): + error: Final = logging_metadata_config_error(self._entry("langfuse_team_b", excluded_services="postgres")) + + assert error is not None and "excluded_services" in error and "'langfuse'" in error + + def test_an_inline_entry_cannot_contradict_the_content_setting_its_reference_inherits(self, otel_v2_on): + reference: Final = {"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {}} + inline: Final = { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {"capture_message_content": "span_only"}, + } + + error: Final = logging_metadata_config_error({"logging": [reference, inline]}) + + assert error is not None and "capture_message_content" in error + + def test_an_inline_entry_may_repeat_the_content_setting_its_reference_inherits(self, otel_v2_on): + reference: Final = {"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {}} + inline: Final = { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {"capture_message_content": "no_content"}, + } + + assert logging_metadata_config_error({"logging": [reference, inline]}) 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..200c675263d 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -11,7 +11,7 @@ import litellm from litellm.caching.caching import DualCache from litellm.constants import MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import AddTeamCallback, UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( _serialize_scan_metadata_header, add_guardrail_scan_id, @@ -21,8 +21,10 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, initialize_callbacks_on_proxy, + named_callback_resolution_error, normalize_callback_names, process_callback, + resolve_named_callback, sanitize_openai_provider_metadata, strip_callback_config, ) @@ -167,9 +169,7 @@ def test_add_policy_to_applied_policies_header_uses_litellm_metadata_bucket(): "litellm_metadata": {}, } - add_policy_to_applied_policies_header( - request_data=request_data, policy_name="global-baseline" - ) + add_policy_to_applied_policies_header(request_data=request_data, policy_name="global-baseline") assert request_data["litellm_metadata"]["applied_policies"] == ["global-baseline"] assert "applied_policies" not in request_data.get("metadata", {}) @@ -324,9 +324,7 @@ def test_initialize_callbacks_on_proxy_instantiates_compression_interception( lambda litellm_settings, callback_specific_params: dummy_callback, ) - original_callbacks = ( - list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] - ) + original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] litellm.callbacks = [] try: initialize_callbacks_on_proxy( @@ -379,20 +377,11 @@ def test_encrypt_callback_vars_round_trip(monkeypatch): enc_vars = encrypted["logging"][0]["callback_vars"] assert enc_vars["langfuse_secret_key"] != "sk-lf-secret" assert enc_vars["langfuse_public_key"] != "pk-lf-public" - assert ( - encrypted["callback_settings"]["callback_vars"]["langsmith_api_key"] - != "ls-api-key" - ) + assert encrypted["callback_settings"]["callback_vars"]["langsmith_api_key"] != "ls-api-key" decrypted = decrypt_callback_vars(encrypted) - assert ( - decrypted["logging"][0]["callback_vars"] - == original["logging"][0]["callback_vars"] - ) - assert ( - decrypted["callback_settings"]["callback_vars"] - == original["callback_settings"]["callback_vars"] - ) + assert decrypted["logging"][0]["callback_vars"] == original["logging"][0]["callback_vars"] + assert decrypted["callback_settings"]["callback_vars"] == original["callback_settings"]["callback_vars"] def test_encrypt_callback_vars_is_idempotent(monkeypatch): @@ -415,10 +404,7 @@ def test_decrypt_callback_vars_passes_through_legacy_plaintext(monkeypatch): plaintext = _sample_metadata() decrypted = decrypt_callback_vars(plaintext) # legacy rows decrypt-fail and fall through unchanged - assert ( - decrypted["logging"][0]["callback_vars"]["langfuse_secret_key"] - == "sk-lf-secret" - ) + assert decrypted["logging"][0]["callback_vars"]["langfuse_secret_key"] == "sk-lf-secret" def test_callback_vars_helpers_handle_edge_shapes(monkeypatch): @@ -519,9 +505,7 @@ def test_initialize_callbacks_on_proxy_lakera_ignores_non_dict_callback_settings SimpleNamespace(prisma_client=None), ) - original_callbacks = ( - list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] - ) + original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] litellm.callbacks = [] try: # A non-dict value must be ignored (init_params stays {}), not **-unpacked. @@ -539,9 +523,7 @@ def test_initialize_callbacks_on_proxy_lakera_ignores_non_dict_callback_settings @pytest.mark.parametrize("bad_root", [None, True]) -def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( - monkeypatch, bad_root -): +def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root(monkeypatch, bad_root): """Regression: a blank `callback_settings:` key in YAML loads as None (and `callback_settings: true` as a bool); load_config forwards that value verbatim as callback_specific_params. Membership tests like @@ -558,9 +540,7 @@ def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( CompressionInterceptionLogger, ) - original_callbacks = ( - list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] - ) + original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] litellm.callbacks = [] try: initialize_callbacks_on_proxy( @@ -570,9 +550,7 @@ def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( litellm_settings={}, callback_specific_params=bad_root, ) - assert any( - isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks - ) + assert any(isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks) finally: litellm.callbacks = original_callbacks @@ -624,7 +602,7 @@ def test_strip_callback_config_passes_through_non_dicts(value): _PROBE_MODULE_NAME = "custom_callback_probe" -_PROBE_MODULE_SOURCE = ''' +_PROBE_MODULE_SOURCE = """ from litellm.integrations.custom_logger import CustomLogger @@ -645,7 +623,7 @@ def log_event_fn(kwargs, response_obj, start_time, end_time): NOT_A_CALLBACK = "some-plain-string" proxy_handler_instance = FloorMaxTokens() -''' +""" @pytest.fixture @@ -659,9 +637,7 @@ def probe_config_path(tmp_path): """ (tmp_path / f"{_PROBE_MODULE_NAME}.py").write_text(_PROBE_MODULE_SOURCE) - original_callbacks = ( - list(litellm.callbacks) if isinstance(litellm.callbacks, list) else litellm.callbacks - ) + original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else litellm.callbacks litellm.callbacks = [] ProxyLogging._callback_capabilities_cache.clear() try: @@ -687,7 +663,7 @@ def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_p silently never run the hook. Config load must fail instead.""" entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" - with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: + with pytest.raises(ValueError, match="litellm_settings\\.callbacks entry") as exc_info: _load_callbacks([entry], probe_config_path) message = str(exc_info.value) @@ -705,12 +681,10 @@ def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_p ("NOT_A_CALLBACK", "str 'some-plain-string'"), ], ) -def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( - probe_config_path, attribute, expected_fragment -): +def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values(probe_config_path, attribute, expected_fragment): entry = f"{_PROBE_MODULE_NAME}.{attribute}" - with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: + with pytest.raises(ValueError, match="litellm_settings\\.callbacks entry") as exc_info: _load_callbacks([entry], probe_config_path) message = str(exc_info.value) @@ -722,7 +696,7 @@ def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( def test_initialize_callbacks_on_proxy_rejects_class_valued_non_list_value(probe_config_path): entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" - with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: + with pytest.raises(ValueError, match="litellm_settings\\.callbacks entry") as exc_info: _load_callbacks(entry, probe_config_path) assert entry in str(exc_info.value) @@ -775,3 +749,158 @@ def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_con assert len(litellm.callbacks) == 1 assert isinstance(litellm.callbacks[0], CustomLogger) + + +_ARIZE_REFERENCE_SETTINGS: Final = { + "arize_space_a": { + "callback_type": "arize", + "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY", + "arize_space_id": "os.environ/LIT10061_ARIZE_SPACE_ID", + }, + "otel": {"message_logging": False}, + "audit_sink": {"callback_type": "generic_api", "endpoint": "https://sink.local", "headers": {}}, +} +_REFERENCE_ENTRY: Final = {"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {}} + + +def test_named_callback_resolution_error_names_the_secret_it_could_not_read(monkeypatch): + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + monkeypatch.delenv("LIT10061_ARIZE_API_KEY", raising=False) + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + + error: Final = named_callback_resolution_error([_REFERENCE_ENTRY]) + + assert error is not None + assert "callback_settings.arize_space_a.arize_api_key" in error + assert "space-a" not in error + + +def test_named_callback_resolution_error_is_none_once_every_secret_reads(monkeypatch): + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "key") + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + + assert named_callback_resolution_error([_REFERENCE_ENTRY]) is None + assert named_callback_resolution_error("not-a-list") is None + + +@pytest.mark.parametrize("callback_name", ["otel", "audit_sink", "arize"]) +def test_callback_settings_without_a_backend_are_not_references(monkeypatch, callback_name): + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + callback: Final = AddTeamCallback(callback_name=callback_name, callback_vars={}) + + assert resolve_named_callback(callback) is callback + + +def test_literal_definition_values_never_reach_the_secret_manager(monkeypatch): + monkeypatch.setattr( + litellm, + "callback_settings", + { + "arize_space_a": { + "callback_type": "arize", + "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY", + "arize_space_id": "space-literal", + } + }, + ) + looked_up: Final[list[str]] = [] + + def fake_get_secret(secret_name, default_value=None): + looked_up.append(secret_name) + return "key-from-manager" + + monkeypatch.setattr("litellm.proxy.common_utils.callback_utils.get_secret", fake_get_secret) + + resolved: Final = resolve_named_callback(AddTeamCallback(callback_name="arize_space_a", callback_vars={})) + + assert looked_up == ["os.environ/LIT10061_ARIZE_API_KEY"] + assert resolved.callback_vars == {"arize_api_key": "key-from-manager", "arize_space_id": "space-literal"} + + +def test_a_definition_value_the_backend_rejects_is_reported_by_field_without_its_values(monkeypatch): + monkeypatch.setattr( + litellm, + "callback_settings", + { + "arize_space_a": { + "callback_type": "arize", + "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY", + "arize_space_id": "os.environ/LIT10061_ARIZE_SPACE_ID", + "arize_otlp_protocol": "carrier-pigeon", + } + }, + ) + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "super-secret-key") + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + + error: Final = named_callback_resolution_error([_REFERENCE_ENTRY]) + + assert error is not None + assert "callback_settings.arize_space_a" in error and "arize_otlp_protocol" in error + assert "super-secret-key" not in error and "space-a" not in error and "carrier-pigeon" not in error + + +@pytest.mark.parametrize(("env_value", "expected"), [("false", "false"), ("true", "true"), ("0", "0")]) +def test_a_definition_secret_that_reads_as_a_boolean_or_falsy_value_is_kept(monkeypatch, env_value, expected): + monkeypatch.setattr( + litellm, + "callback_settings", + { + "arize_space_a": { + "callback_type": "arize", + "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY", + "arize_space_id": "os.environ/LIT10061_ARIZE_SPACE_ID", + "turn_off_message_logging": "os.environ/LIT10061_TURN_OFF_MESSAGE_LOGGING", + } + }, + ) + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "key") + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + monkeypatch.setenv("LIT10061_TURN_OFF_MESSAGE_LOGGING", env_value) + + resolved: Final = resolve_named_callback(AddTeamCallback(callback_name="arize_space_a", callback_vars={})) + + assert resolved.callback_vars["turn_off_message_logging"] == expected + assert named_callback_resolution_error([_REFERENCE_ENTRY]) is None + + +def test_a_secret_manager_read_is_reused_within_the_ttl_and_a_miss_is_not_cached(monkeypatch): + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + reads: Final[list[str]] = [] + + def fake_get_secret(secret_name, default_value=None): + reads.append(secret_name) + return "key" if secret_name.endswith("API_KEY") else default_value + + monkeypatch.setattr("litellm.proxy.common_utils.callback_utils.get_secret", fake_get_secret) + + first_error: Final = named_callback_resolution_error([_REFERENCE_ENTRY]) + second_error: Final = named_callback_resolution_error([_REFERENCE_ENTRY]) + + assert first_error is not None and second_error is not None + assert "callback_settings.arize_space_a.arize_space_id" in first_error + assert reads == [ + "os.environ/LIT10061_ARIZE_API_KEY", + "os.environ/LIT10061_ARIZE_SPACE_ID", + "os.environ/LIT10061_ARIZE_SPACE_ID", + ], "the api key read is served from the cache, the missing space id is read again" + + +@pytest.mark.parametrize("space_id", ["12345", "0", "1.5", "true", "[1]"]) +def test_a_cached_secret_keeps_its_string_value_on_later_reads(monkeypatch, space_id): + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + reads: Final[list[str]] = [] + + def fake_get_secret(secret_name, default_value=None): + reads.append(secret_name) + return space_id if secret_name.endswith("SPACE_ID") else "key" + + monkeypatch.setattr("litellm.proxy.common_utils.callback_utils.get_secret", fake_get_secret) + reference: Final = AddTeamCallback(callback_name="arize_space_a", callback_vars={}) + + first: Final = resolve_named_callback(reference).callback_vars["arize_space_id"] + second: Final = resolve_named_callback(reference).callback_vars["arize_space_id"] + + assert (first, second) == (space_id, space_id) + assert reads.count("os.environ/LIT10061_ARIZE_SPACE_ID") == 1, "the second read must come from the cache" diff --git a/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py index 34fa164fd6d..c8183e860f2 100644 --- a/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py @@ -7,7 +7,7 @@ import pytest from fastapi.testclient import TestClient -from typing import cast +from typing import Final, cast from fastapi import FastAPI @@ -51,9 +51,7 @@ class TestCallbackManagementEndpoints: client = TestClient(app) # Make request to list callbacks endpoint - response = client.get( - "/callbacks/list", headers={"Authorization": "Bearer sk-9876"} - ) + response = client.get("/callbacks/list", headers={"Authorization": "Bearer sk-9876"}) # Verify response assert response.status_code == 200 @@ -86,9 +84,7 @@ class TestCallbackManagementEndpoints: litellm._async_success_callback.append("langfuse") # Make request to list callbacks endpoint - response = client.get( - "/callbacks/list", headers={"Authorization": "Bearer sk-9876"} - ) + response = client.get("/callbacks/list", headers={"Authorization": "Bearer sk-9876"}) # Verify response assert response.status_code == 200 @@ -117,9 +113,7 @@ class TestCallbackManagementEndpoints: litellm.callbacks.append("datadog") # Make request to list callbacks endpoint - response = client.get( - "/callbacks/list", headers={"Authorization": "Bearer sk-9876"} - ) + response = client.get("/callbacks/list", headers={"Authorization": "Bearer sk-9876"}) # Verify response assert response.status_code == 200 @@ -133,10 +127,7 @@ class TestCallbackManagementEndpoints: # The categorization logic should deduplicate properly assert len([cb for cb in response_data["success"] if cb == "datadog"]) <= 1 assert len([cb for cb in response_data["failure"] if cb == "datadog"]) <= 1 - assert ( - len([cb for cb in response_data["success_and_failure"] if cb == "datadog"]) - <= 1 - ) + assert len([cb for cb in response_data["success_and_failure"] if cb == "datadog"]) <= 1 # Verify the response structure is correct assert isinstance(response_data["success"], list) @@ -154,9 +145,7 @@ class TestCallbackManagementEndpoints: litellm.callbacks.append("prometheus") # Make request to list callbacks endpoint - response = client.get( - "/callbacks/list", headers={"Authorization": "Bearer sk-9876"} - ) + response = client.get("/callbacks/list", headers={"Authorization": "Bearer sk-9876"}) # Verify response assert response.status_code == 200 @@ -168,24 +157,16 @@ class TestCallbackManagementEndpoints: proxy_internal_callbacks = ["_PROXY_VirtualKeyModelMaxBudgetLimiter"] response_data["success_and_failure"] = [ - cb - for cb in response_data["success_and_failure"] - if cb not in proxy_internal_callbacks + cb for cb in response_data["success_and_failure"] if cb not in proxy_internal_callbacks ] # Verify callbacks are properly categorized - assert ( - "prometheus" in response_data["success_and_failure"] - ) # callbacks list items go to success_and_failure + assert "prometheus" in response_data["success_and_failure"] # callbacks list items go to success_and_failure assert "langfuse" in response_data["success"] assert "datadog" in response_data["failure"] # Verify no duplicates - all_callbacks = ( - response_data["success"] - + response_data["failure"] - + response_data["success_and_failure"] - ) + all_callbacks = response_data["success"] + response_data["failure"] + response_data["success_and_failure"] assert len(set(all_callbacks)) == len(all_callbacks) def test_alist_callbacks_empty_response_structure(self): @@ -194,9 +175,7 @@ class TestCallbackManagementEndpoints: client = TestClient(app) # Make request to list callbacks endpoint - response = client.get( - "/callbacks/list", headers={"Authorization": "Bearer sk-9876"} - ) + response = client.get("/callbacks/list", headers={"Authorization": "Bearer sk-9876"}) # Verify response structure assert response.status_code == 200 @@ -214,9 +193,7 @@ class TestCallbackManagementEndpoints: client = TestClient(app) # Make request to get callback configs endpoint - response = client.get( - "/callbacks/configs", headers={"Authorization": "Bearer sk-9876"} - ) + response = client.get("/callbacks/configs", headers={"Authorization": "Bearer sk-9876"}) # Verify response assert response.status_code == 200 @@ -243,12 +220,9 @@ class TestCallbackManagementEndpoints: # Check if at least one callback has detailed parameter configuration has_detailed_params = any( - config.get("dynamic_params") and len(config.get("dynamic_params", {})) > 0 - for config in response_data + config.get("dynamic_params") and len(config.get("dynamic_params", {})) > 0 for config in response_data ) - assert ( - has_detailed_params - ), "Expected at least one callback to have detailed parameter configuration" + assert has_detailed_params, "Expected at least one callback to have detailed parameter configuration" galileo_config = next( (config for config in response_data if config.get("id") == "galileo"), @@ -481,3 +455,100 @@ class TestNewRelicKeyLoggingValidation: ) finally: is_otel_v2_enabled.cache_clear() + + +_ARIZE_REFERENCE_SETTINGS: Final = { + "arize_space_a": { + "callback_type": "arize", + "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY", + "arize_space_id": "os.environ/LIT10061_ARIZE_SPACE_ID", + } +} + + +class TestCallbackReferenceKeyWrites: + @staticmethod + def _metadata(callback_name, **callback_vars): + return { + "logging": [{"callback_name": callback_name, "callback_type": "success", "callback_vars": callback_vars}] + } + + @pytest.mark.parametrize("role", ["internal_user", "team", None]) + def test_only_a_proxy_admin_can_point_a_key_at_a_definition(self, monkeypatch, role): + from fastapi import HTTPException + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + raise_on_unauthorized_callback_reference, + ) + + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + with pytest.raises(HTTPException) as exc: + raise_on_unauthorized_callback_reference(self._metadata("arize_space_a"), UserAPIKeyAuth(user_role=role)) + assert exc.value.status_code == 403 + assert "callback_settings" in str(exc.value.detail) + + def test_a_proxy_admin_can_point_a_key_at_a_definition_and_add_its_own_values(self, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + raise_on_unauthorized_callback_reference, + ) + + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + assert ( + raise_on_unauthorized_callback_reference( + self._metadata("arize_space_a", turn_off_message_logging="true"), + UserAPIKeyAuth(user_role="proxy_admin"), + ) + is None + ) + + def test_a_key_cannot_override_a_value_its_definition_sets(self, monkeypatch): + from fastapi import HTTPException + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + raise_on_unauthorized_callback_reference, + ) + + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + with pytest.raises(HTTPException) as exc: + raise_on_unauthorized_callback_reference( + self._metadata("arize_space_a", arize_api_key="inline-key"), UserAPIKeyAuth(user_role="proxy_admin") + ) + assert exc.value.status_code == 400 + assert "arize_api_key" in str(exc.value.detail) + + def test_an_inline_callback_entry_is_not_gated(self, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + raise_on_unauthorized_callback_reference, + ) + + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + assert ( + raise_on_unauthorized_callback_reference( + self._metadata("arize", arize_api_key="inline-key", arize_space_id="space-b"), + UserAPIKeyAuth(user_role="internal_user"), + ) + is None + ) + assert raise_on_unauthorized_callback_reference(None, UserAPIKeyAuth(user_role="internal_user")) is None + + def test_a_proxy_admin_can_add_a_backend_option_to_a_reference(self, monkeypatch): + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + raise_on_invalid_key_logging_config, + raise_on_unauthorized_callback_reference, + ) + + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + metadata: Final = self._metadata("arize_space_a", excluded_services="postgres") + try: + assert raise_on_invalid_key_logging_config(metadata) is None + assert raise_on_unauthorized_callback_reference(metadata, UserAPIKeyAuth(user_role="proxy_admin")) is None + finally: + is_otel_v2_enabled.cache_clear() 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 b453f985512..1d119d254fe 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -24240,3 +24240,28 @@ 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") + + +@pytest.mark.asyncio +async def test_key_health_reports_a_callback_settings_secret_it_cannot_read(monkeypatch): + from unittest.mock import MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import key_health + + monkeypatch.setattr( + litellm, + "callback_settings", + {"arize_space_b": {"callback_type": "arize", "arize_api_key": "os.environ/LIT10061_UNSET_ARIZE_API_KEY"}}, + ) + monkeypatch.delenv("LIT10061_UNSET_ARIZE_API_KEY", raising=False) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", + metadata={"logging": [{"callback_name": "arize_space_b", "callback_type": "success", "callback_vars": {}}]}, + ) + + health = await key_health(request=MagicMock(), user_api_key_dict=user_api_key_dict) + + assert health["key"] == "unhealthy" + assert health["logging_callbacks"]["callbacks"] == ["arize_space_b"] + assert "callback_settings.arize_space_b.arize_api_key" in str(health["logging_callbacks"]["details"]) diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 712b4134c05..d8e671131f5 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -1597,6 +1597,34 @@ def test_ProxyConfig__check_for_os_environ_vars_substitutes(monkeypatch): assert out == {"a": "secret-value", "b": 2, "nested": {"c": "secret-value"}} +def test_ProxyConfig__check_for_os_environ_vars_keeps_named_callback_definitions_as_references(monkeypatch): + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "boot-time-value") + monkeypatch.setenv("LIT10061_SINK_TOKEN", "sink-token") + pc = ProxyConfig() + cfg = { + "callback_settings": { + "arize_space_a": {"callback_type": "arize", "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY"}, + "audit_sink": { + "callback_type": "generic_api", + "endpoint": "https://sink.local", + "headers": {"Authorization": "os.environ/LIT10061_SINK_TOKEN"}, + }, + "otel": {"message_logging": False}, + }, + "general_settings": {"master_key": "os.environ/LIT10061_ARIZE_API_KEY"}, + } + + out = pc._check_for_os_environ_vars(cfg) + + assert out["callback_settings"]["arize_space_a"] == { + "callback_type": "arize", + "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY", + } + assert out["callback_settings"]["audit_sink"]["headers"] == {"Authorization": "sink-token"} + assert out["callback_settings"]["otel"] == {"message_logging": False} + assert out["general_settings"]["master_key"] == "boot-time-value" + + def test_ProxyConfig__check_for_os_environ_vars_missing_env_returns_none(monkeypatch): monkeypatch.delenv("NONEXISTENT_TEST_VAR_X", raising=False) pc = ProxyConfig() diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..4e2b488d3ec 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -111,7 +111,8 @@ class TestGetMetadataVariableName: assert get_metadata_variable_name(request) == "litellm_metadata" @pytest.mark.parametrize( - "path", ["/v1beta/interactions", "/interactions", "/v1beta/interactions/v1_abc123", "/interactions/v1_abc123/cancel"] + "path", + ["/v1beta/interactions", "/interactions", "/v1beta/interactions/v1_abc123", "/interactions/v1_abc123/cancel"], ) def test_returns_litellm_metadata_for_interactions_routes(self, path): request = self._make_request(path) @@ -807,7 +808,9 @@ async def test_interactions_spend_logs_metadata_survives_the_router_metadata_buc # The router then adds its deployment fields to the same bucket, and the spend log # resolves that bucket. - data["litellm_metadata"].update({"model_group": "gemini-3.1-flash-lite", "deployment": "gemini/gemini-3.1-flash-lite"}) + data["litellm_metadata"].update( + {"model_group": "gemini-3.1-flash-lite", "deployment": "gemini/gemini-3.1-flash-lite"} + ) resolved = get_litellm_metadata_from_kwargs( {"litellm_params": {"metadata": data.get("metadata") or {}, "litellm_metadata": data["litellm_metadata"]}} ) @@ -2949,6 +2952,103 @@ def test_get_dynamic_logging_metadata_ignores_env_reference_from_key_metadata( assert result is None +_ARIZE_REFERENCE_SETTINGS: Final = { + "arize_space_a": { + "callback_type": "arize", + "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY", + "arize_space_id": "os.environ/LIT10061_ARIZE_SPACE_ID", + "excluded_services": "postgres,redis", + } +} + + +def _key_auth_referencing(callback_name: str, **callback_vars: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="test-key", + team_id="team-a", + metadata={ + "logging": [{"callback_name": callback_name, "callback_type": "success", "callback_vars": callback_vars}] + }, + team_metadata={}, + ) + + +def test_key_callback_reference_picks_up_a_rotated_secret_once_the_cache_ttl_passes(monkeypatch): + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.proxy.common_utils import callback_utils + + now: Final = [1_000.0] + monkeypatch.setattr(callback_utils, "_SECRET_CACHE", InMemoryCache(default_ttl=60, clock=lambda: now[0])) + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "first-key") + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + + first: Final = get_dynamic_logging_metadata( + user_api_key_dict=_key_auth_referencing("arize_space_a"), proxy_config=MagicMock() + ) + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "rotated-key") + within_ttl: Final = get_dynamic_logging_metadata( + user_api_key_dict=_key_auth_referencing("arize_space_a"), proxy_config=MagicMock() + ) + now[0] += 61 + after_ttl: Final = get_dynamic_logging_metadata( + user_api_key_dict=_key_auth_referencing("arize_space_a"), proxy_config=MagicMock() + ) + + assert first is not None and within_ttl is not None and after_ttl is not None + assert first.success_callback == ["arize"] + assert first.callback_vars == {"arize_api_key": "first-key", "arize_space_id": "space-a"} + assert within_ttl.callback_vars["arize_api_key"] == "first-key" + assert after_ttl.callback_vars["arize_api_key"] == "rotated-key" + + +def test_key_callback_reference_keeps_the_values_the_key_adds(monkeypatch): + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "first-key") + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + + result: Final = get_dynamic_logging_metadata( + user_api_key_dict=_key_auth_referencing("arize_space_a", turn_off_message_logging="true"), + proxy_config=MagicMock(), + ) + + assert result is not None + assert result.callback_vars["turn_off_message_logging"] == "true" + assert result.callback_vars["arize_api_key"] == "first-key" + + +def test_key_callback_reference_with_an_unreadable_secret_logs_nowhere(monkeypatch): + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + monkeypatch.delenv("LIT10061_ARIZE_API_KEY", raising=False) + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + + result: Final = get_dynamic_logging_metadata( + user_api_key_dict=_key_auth_referencing("arize_space_a"), proxy_config=MagicMock() + ) + + assert result is None + + +def test_team_callback_reference_is_not_resolved(monkeypatch): + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + monkeypatch.setenv("LIT10061_ARIZE_API_KEY", "first-key") + monkeypatch.setenv("LIT10061_ARIZE_SPACE_ID", "space-a") + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="test-key", + team_id="team-a", + metadata={}, + team_metadata={ + "logging": [{"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {}}] + }, + ) + + result: Final = get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=MagicMock()) + + assert result is not None + assert result.success_callback == ["arize_space_a"] + assert "arize_api_key" not in result.callback_vars + + def test_get_num_retries_from_request(): """ Test LiteLLMProxyRequestSetup._get_num_retries_from_request method @@ -3758,53 +3858,87 @@ async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypa "surface,headers,body,child_id,trace_id", [ ( - "chat/completions", {"X-Session-Id": "child-01"}, {"metadata": {"session_id": "body-001"}}, - "child-01", "parent-01", + "chat/completions", + {"X-Session-Id": "child-01"}, + {"metadata": {"session_id": "body-001"}}, + "child-01", + "parent-01", ), ("responses", {"x-vendor-session-id": "child-01"}, {}, "child-01", "parent-01"), ("messages", {"session-id": "child-01", "user-agent": "codex_cli_rs/1.0"}, {}, "child-01", "parent-01"), ("chat/completions", {"x-litellm-session-id": "short", "x-session-id": "child-01"}, {}, "short", "short"), ("responses", {"x-litellm-trace-id": "trace-01", "x-session-id": "child-01"}, {}, "child-01", "trace-01"), ( - "messages", {"x-litellm-trace-id": "trace-01", "x-litellm-session-id": "child-01"}, - {}, "child-01", "trace-01", + "messages", + {"x-litellm-trace-id": "trace-01", "x-litellm-session-id": "child-01"}, + {}, + "child-01", + "trace-01", ), ("responses", {}, {"metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("messages", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("chat/completions", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("messages", {"baggage": "session.id=child-01"}, {}, "child-01", "parent-01"), ( - "responses", {"baggage": "session.id=child-01"}, {"metadata": {"session_id": "body-001"}}, - "body-001", "parent-01", + "responses", + {"baggage": "session.id=child-01"}, + {"metadata": {"session_id": "body-001"}}, + "body-001", + "parent-01", ), ( - "chat/completions", {}, {"litellm_session_id": "root-001", "metadata": {"session_id": "body-001"}}, - "root-001", "parent-01", + "chat/completions", + {}, + {"litellm_session_id": "root-001", "metadata": {"session_id": "body-001"}}, + "root-001", + "parent-01", ), ( - "responses", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": "active-1"}}, - "active-1", "parent-01", + "responses", + {}, + {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": "active-1"}}, + "active-1", + "parent-01", ), ( - "chat/completions", {}, + "chat/completions", + {}, {"metadata": {"session_id": "active-1"}, "litellm_metadata": {"session_id": "body-001"}}, - "active-1", "parent-01", + "active-1", + "parent-01", ), ( - "chat/completions", {}, {"metadata": {"session_id": ""}, "litellm_metadata": {"session_id": "body-001"}}, - None, "parent-01", + "chat/completions", + {}, + {"metadata": {"session_id": ""}, "litellm_metadata": {"session_id": "body-001"}}, + None, + "parent-01", ), ( - "messages", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": ""}}, - None, "parent-01", + "messages", + {}, + {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": ""}}, + None, + "parent-01", ), ( - "messages", {}, {"litellm_session_id": "root-001", "litellm_metadata": {"session_id": "active-1"}, - "metadata": {"session_id": "body-001", "user_id": "user_abc_account__session_anthro-1"}}, - "anthro-1", "parent-01", + "messages", + {}, + { + "litellm_session_id": "root-001", + "litellm_metadata": {"session_id": "active-1"}, + "metadata": {"session_id": "body-001", "user_id": "user_abc_account__session_anthro-1"}, + }, + "anthro-1", + "parent-01", + ), + ( + "messages", + {"x-session-id": "child-01"}, + {"metadata": {"user_id": {"session_id": "body-001"}}}, + "child-01", + "parent-01", ), - ("messages", {"x-session-id": "child-01"}, {"metadata": {"user_id": {"session_id": "body-001"}}}, - "child-01", "parent-01"), ("chat/completions", {}, {}, None, "parent-01"), ("responses", {"x-session-id": "bad"}, {}, None, "parent-01"), ("messages", {}, {"metadata": {"session_id": 123}}, None, "parent-01"), @@ -4346,9 +4480,7 @@ def _bedrock_passthrough_data( ), ], ) -def test_bedrock_request_metadata_thread_id_becomes_session_id( - headers: dict[str, str], client_body: dict[str, object] -): +def test_bedrock_request_metadata_thread_id_becomes_session_id(headers: dict[str, str], client_body: dict[str, object]): sent_body: Final = copy.deepcopy(client_body) data: Final = _bedrock_passthrough_data(client_body) @@ -4428,7 +4560,9 @@ def test_bedrock_request_metadata_without_usable_session_id_sets_no_session( @pytest.mark.parametrize( ("body_tags", "header_tags", "expected"), [ - pytest.param({"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key"), + pytest.param( + {"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key" + ), pytest.param({"sessionId": "body-session"}, {"sessionId": "header-session"}, "body-session", id="body_wins"), ], ) @@ -11226,9 +11360,7 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request _per_user_env(monkeypatch) - load = AsyncMock( - return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}} - ) + load = AsyncMock(return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}}) monkeypatch.setattr(upc, "aget_user_provider_credentials", load) router = Router( model_list=[ @@ -11282,9 +11414,7 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request _per_user_env(monkeypatch) - load = AsyncMock( - return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}} - ) + load = AsyncMock(return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}}) monkeypatch.setattr(upc, "aget_user_provider_credentials", load) router = Router( model_list=[ diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 3336cce9aaf..13cfa8f5054 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -50,6 +50,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/litellm_core_utils/prompt_templates/factory.py", "litellm/litellm_core_utils/prompt_templates/image_handling.py", "litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py", + "litellm/proxy/common_utils/callback_utils.py", "litellm/litellm_core_utils/specialty_caches/service_trace_id_cache.py", "litellm/llms/azure/common_utils.py", "litellm/llms/bedrock/base_aws_llm.py",