diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 82135cb1b74..b9f2906e683 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -4303,6 +4303,14 @@ async def _prefetch_referenced_auth_objects( ) +async def _prefetch_callback_settings_secrets(user_api_key_dict: UserAPIKeyAuth) -> None: + """Read the secrets this key's ``callback_settings`` references need in a thread, before the sync path resolves them.""" + from litellm.proxy.common_utils.callback_utils import prefetch_named_callback_secrets + from litellm.proxy.litellm_pre_call_utils import KeyAndTeamLoggingSettings + + await prefetch_named_callback_secrets(KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)) + + def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> None: """Anchor the OTLP destinations this key or team overrides its traces to. @@ -4396,6 +4404,7 @@ async def user_api_key_auth( user_api_key_auth_obj.agent_caller = agent_caller_from_headers( safe_get_request_headers(request), user_api_key_auth_obj ) + await _prefetch_callback_settings_secrets(user_api_key_auth_obj) _seed_request_destinations(user_api_key_auth_obj, request) # A body that never parsed is authenticated (so the trace carries identity @@ -4884,6 +4893,7 @@ async def authorize_internal_virtual_key( ) if recovered is not None: return recovered + await _prefetch_callback_settings_secrets(auth) _seed_request_destinations(auth, request) auth.request_route = route request.state.principal = _resolve_request_principal(request, auth) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index d6a46ec002a..09a01361b4c 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,9 +1,12 @@ +import asyncio import copy import json import os from collections.abc import Callable, Iterable, Mapping, Sequence +from contextvars import ContextVar from dataclasses import dataclass from itertools import accumulate +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, Protocol, TypeAlias from pydantic import TypeAdapter, ValidationError @@ -30,7 +33,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker 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.callback_config_validation import named_callback_definition, stored_callback_entries from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -827,16 +830,68 @@ def _secret_cache() -> _SecretCache: return _SECRET_CACHE +_REQUEST_SECRETS: Final[ContextVar[Mapping[str, str | bool]]] = ContextVar( + "callback_settings_request_secrets", default=MappingProxyType({}) +) + + +async def prefetch_named_callback_secrets(logging_entries: object) -> None: + """Read the secrets a key's ``callback_settings`` references need before the request's synchronous path does. + + Each read runs in a worker thread, so a slow secret manager delays only this + request and never the event loop. The values land in a ``ContextVar`` that the + resolution on this request task consults before it falls back to a blocking read. + """ + references: Final = tuple( + reference for reference in _definition_secret_references(logging_entries) if _known_secret(reference) is None + ) + if not references: + return + secrets: Final = await asyncio.gather( + *(asyncio.to_thread(_read_secret_manager_value, reference) for reference in references) + ) + _REQUEST_SECRETS.set( + MappingProxyType( + { + **_REQUEST_SECRETS.get(), + **{ + reference: secret + for reference, secret in zip(references, secrets, strict=True) + if secret is not None + }, + } + ) + ) + + +def _definition_secret_references(logging_entries: object) -> tuple[str, ...]: + return tuple( + dict.fromkeys( + value + for callback_name, _ in stored_callback_entries(logging_entries) + if callback_name is not None and (definition := named_callback_definition(callback_name)) is not None + for key, value in definition.items() + if key != "callback_type" and isinstance(value, str) and value.startswith(_SECRET_REFERENCE_PREFIX) + ) + ) + + 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. + and may add values the definition does not set, never replace one it does. """ definition: Final = named_callback_definition(callback.callback_name) if definition is None: return callback + overridden: Final = next((key for key in callback.callback_vars if key in definition), None) + if overridden is not None: + raise ValueError( + f"callback_vars.{overridden} is set by callback_settings.{callback.callback_name} " + "and cannot be overridden on the key" + ) backend: Final = str(definition["callback_type"]) callback_vars: Final = { **{ @@ -904,17 +959,28 @@ def _resolved_definition_value(definition_name: str, key: str, value: object) -> 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. + """The value a reference resolves to, read from the secret manager only when nothing on hand already has it.""" + known: Final = _known_secret(reference) + return known if known is not None else _read_secret_manager_value(reference) + + +def _known_secret(reference: str) -> str | bool | None: + """What this request prefetched, else what this worker read within the TTL. Misses are not cached, so a secret created after the first read is picked up on the next request. """ + prefetched: Final = _REQUEST_SECRETS.get().get(reference) + if prefetched is not None: + return prefetched 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 + return cached.value if isinstance(cached, _CachedSecret) else None + + +def _read_secret_manager_value(reference: str) -> str | bool | None: 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)) + if CALLBACK_SETTINGS_SECRET_CACHE_TTL_SECONDS > 0: + _secret_cache().set_cache(reference, _CachedSecret(secret)) return secret diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d2aac9cae38..4ee09adeac2 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 @@ -590,6 +592,8 @@ async def new_user( data=data, user_api_key_dict=user_api_key_dict, ) + raise_on_invalid_key_logging_config(data.metadata) + raise_on_unauthorized_callback_reference(data.metadata, user_api_key_dict) data_json = data.json() data_json = _update_internal_new_user_params(data_json, data) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 44e25c29cb2..abf9ee751c0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -76,6 +76,7 @@ from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, encrypt_callback_vars, named_callback_resolution_error, + prefetch_named_callback_secrets, ) from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -971,18 +972,39 @@ def raise_on_unauthorized_callback_reference( if overridden is not None: raise HTTPException(status_code=400, detail={"error": overridden}) return - stored: Final = _callback_references(existing_metadata or {}) - if _reference_keys(references) == _reference_keys(stored): + stored: Final = _admin_only_entry_keys(existing_metadata or {}) + if _admin_only_entry_keys(metadata) == stored: return - raise HTTPException( - status_code=403, - detail={ - "error": ( - "Only a proxy admin can point a key at a callback_settings definition" - if not stored - else "Only a proxy admin can change or remove a key's callback_settings reference" - ) - }, + unknown: Final = next(iter(_unknown_callback_names(metadata)), None) + raise HTTPException(status_code=403, detail={"error": _callback_reference_denied(unknown, stored=bool(stored))}) + + +def _callback_reference_denied(unknown_callback_name: str | None, *, stored: bool) -> str: + if stored: + return "Only a proxy admin can change or remove a key's callback_settings reference" + if unknown_callback_name is None: + return "Only a proxy admin can point a key at a callback_settings definition" + return ( + f"Only a proxy admin can store a logging entry for {unknown_callback_name!r}, a callback the proxy does not " + f"know, since a later callback_settings.{unknown_callback_name} definition would hand the key its secrets" + ) + + +def _admin_only_entry_keys(metadata: Mapping[str, object]) -> tuple[tuple[str, tuple[tuple[str, str], ...]], ...]: + """References to a definition, plus names no callback owns yet, which a later definition would turn into one.""" + return ( + *_reference_keys(_callback_references(metadata)), + *((callback_name, ()) for callback_name in _unknown_callback_names(metadata)), + ) + + +def _unknown_callback_names(metadata: Mapping[str, object]) -> tuple[str, ...]: + return tuple( + callback_name + for callback_name, _ in stored_callback_entries(metadata.get("logging")) + if callback_name is not None + and callback_name not in litellm._known_custom_logger_compatible_callbacks # pyright: ignore[reportPrivateUsage] # the registry of callback names has no public accessor + and named_callback_definition(callback_name) is None ) @@ -8437,6 +8459,7 @@ async def key_health( # Check if logging is configured in metadata if key_metadata and "logging" in key_metadata: key_logging: Final = decrypt_callback_vars(key_metadata)["logging"] + await prefetch_named_callback_secrets(key_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) diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index 200c675263d..db2df56b03e 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -1,6 +1,7 @@ import copy import json import sys +import threading from types import ModuleType, SimpleNamespace from typing import Final from unittest.mock import patch @@ -23,6 +24,7 @@ from litellm.proxy.common_utils.callback_utils import ( initialize_callbacks_on_proxy, named_callback_resolution_error, normalize_callback_names, + prefetch_named_callback_secrets, process_callback, resolve_named_callback, sanitize_openai_provider_metadata, @@ -792,6 +794,48 @@ def test_callback_settings_without_a_backend_are_not_references(monkeypatch, cal assert resolve_named_callback(callback) is callback +async def test_prefetched_definition_secrets_are_read_in_a_thread_once_and_reused_by_the_sync_path(monkeypatch): + """The auth phase prefetches off the event loop, so request-time resolution never blocks on the secret manager.""" + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + monkeypatch.setattr("litellm.proxy.common_utils.callback_utils.CALLBACK_SETTINGS_SECRET_CACHE_TTL_SECONDS", 0) + reads: Final[list[tuple[str, int]]] = [] + + def fake_get_secret(secret_name, default_value=None): + reads.append((secret_name, threading.get_ident())) + return {"os.environ/LIT10061_ARIZE_API_KEY": "key-a", "os.environ/LIT10061_ARIZE_SPACE_ID": "space-a"}[ + secret_name + ] + + monkeypatch.setattr("litellm.proxy.common_utils.callback_utils.get_secret", fake_get_secret) + + await prefetch_named_callback_secrets([_REFERENCE_ENTRY]) + resolved: Final = resolve_named_callback(AddTeamCallback.model_validate(_REFERENCE_ENTRY)) + + assert resolved.callback_vars == {"arize_api_key": "key-a", "arize_space_id": "space-a"} + assert sorted(name for name, _ in reads) == [ + "os.environ/LIT10061_ARIZE_API_KEY", + "os.environ/LIT10061_ARIZE_SPACE_ID", + ], "the sync path must reuse the prefetched values instead of reading again" + assert all(thread != threading.get_ident() for _, thread in reads), ( + "prefetch reads must leave the event loop thread" + ) + + +def test_a_key_cannot_override_a_definition_value_at_request_time(monkeypatch): + """A stored entry that pins its own value for a definition field is refused, so a key written before the definition existed cannot redirect it.""" + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + monkeypatch.setattr("litellm.proxy.common_utils.callback_utils.get_secret", lambda name, default_value=None: "x") + + with pytest.raises( + ValueError, match=r"callback_vars\.arize_space_id is set by callback_settings\.arize_space_a" + ) as exc: + resolve_named_callback( + AddTeamCallback(callback_name="arize_space_a", callback_vars={"arize_space_id": "attacker-space"}) + ) + + assert "attacker-space" not in str(exc.value) + + def test_literal_definition_values_never_reach_the_secret_manager(monkeypatch): monkeypatch.setattr( litellm, 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 4701ec501f6..7451ed13367 100644 --- a/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_callback_management_endpoints.py @@ -554,6 +554,46 @@ class TestCallbackReferenceKeyWrites: is_otel_v2_enabled.cache_clear() +class TestUnknownCallbackNamesOnKeys: + """A name no callback owns is admin-only, since a later ``callback_settings`` definition under it would hand the key its secrets.""" + + _UNKNOWN: Final = { + "logging": [{"callback_name": "arize_space_b", "callback_type": "success", "callback_vars": {}}], + } + + @pytest.fixture(autouse=True) + def definitions(self, monkeypatch) -> None: + monkeypatch.setattr(litellm, "callback_settings", _ARIZE_REFERENCE_SETTINGS) + + @staticmethod + def _check(metadata, role: str, existing_metadata=None): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + raise_on_unauthorized_callback_reference, + ) + + return raise_on_unauthorized_callback_reference( + metadata, UserAPIKeyAuth(user_role=role), existing_metadata=existing_metadata + ) + + @pytest.mark.parametrize("role", ["internal_user", "team", None]) + def test_a_non_admin_cannot_store_a_name_no_callback_owns(self, role): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + self._check(self._UNKNOWN, role) + assert exc.value.status_code == 403 + assert "arize_space_b" in str(exc.value.detail) + + def test_a_proxy_admin_can_store_a_name_no_callback_owns(self): + assert self._check(self._UNKNOWN, "proxy_admin") is None + + def test_a_non_admin_can_resend_a_stored_unknown_name_unchanged(self): + assert ( + self._check({**self._UNKNOWN, "note": "edited"}, "internal_user", existing_metadata=self._UNKNOWN) is None + ) + + class TestCallbackReferenceKeyUpdates: """A non-admin may edit a referencing key only by sending its stored reference back unchanged.""" 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 d8aa0239e2d..d57f17c5945 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1199,6 +1199,73 @@ async def test_new_user_non_admin_omits_permissions_succeeds(mocker): assert result is not None +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller_role", "expected_status"), + [(LitellmUserRoles.ORG_ADMIN, 403), (LitellmUserRoles.INTERNAL_USER, 403), (LitellmUserRoles.PROXY_ADMIN, None)], +) +async def test_new_user_gates_the_auto_created_key_callback_settings_reference_like_key_generate( + mocker, monkeypatch, caller_role, expected_status +): + """``/user/new`` copies ``metadata`` onto the key it creates, so pointing it at a definition is admin-only there too.""" + import litellm + from litellm.proxy.management_endpoints.internal_user_endpoints import new_user + + monkeypatch.setattr( + litellm, + "callback_settings", + {"arize_space_a": {"callback_type": "arize", "arize_api_key": "os.environ/LIT10061_ARIZE_API_KEY"}}, + ) + mock_prisma_client = mocker.MagicMock() + + async def mock_count(*args, **kwargs): + return 5 + + mock_prisma_client.db.litellm_usertable.count = mock_count + + async def mock_check(*_args, **_kwargs): + return None + + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_email", + mock_check, + ) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id", + mock_check, + ) + mock_license_check = mocker.MagicMock() + mock_license_check.is_over_limit.return_value = False + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check) + helper_calls: list[dict] = [] + + async def stub_helper(**kwargs): + helper_calls.append(kwargs) + return {"user_id": "alice", "key": "sk-alice", "expires": None} + + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", + stub_helper, + ) + data = NewUserRequest( + user_email="alice@example.com", + user_role=LitellmUserRoles.INTERNAL_USER, + metadata={"logging": [{"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {}}]}, + ) + caller = UserAPIKeyAuth(user_id="caller", user_role=caller_role) + + if expected_status is None: + assert await new_user(data=data, user_api_key_dict=caller) is not None + assert len(helper_calls) == 1 + return + with pytest.raises(ProxyException) as exc: + await new_user(data=data, user_api_key_dict=caller) + assert exc.value.code == str(expected_status) + assert "callback_settings" in str(exc.value.message) + assert helper_calls == [], "nothing may be persisted once the reference is refused" + + @pytest.mark.asyncio async def test_new_user_admin_can_set_permissions(mocker): """`new_user` accepts a PROXY_ADMIN caller for any shape of