diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py new file mode 100644 index 00000000000..76c3066ae83 --- /dev/null +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -0,0 +1,302 @@ +import asyncio +import json +import random +import time +from collections.abc import Awaitable, Callable +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Protocol, cast # noqa: TID251 # untyped prisma/redis boundary needs cast + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + + +class _ConfigSyncPubSub(Protocol): + def subscribe(self, *channels: str) -> Awaitable[object]: ... + + def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Awaitable[object]: ... + + def aclose(self) -> Awaitable[object]: ... + + +class _ConfigSyncPubSubClient(Protocol): + def publish(self, channel: str, message: str) -> Awaitable[int]: ... + + def pubsub(self) -> _ConfigSyncPubSub: ... + + +CONFIG_SYNC_CHANNEL = "litellm_proxy.config_change" +CONFIG_SYNC_DEBOUNCE_SECONDS = 1.0 +CONFIG_SYNC_JITTER_MAX_SECONDS = 5.0 +CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS = 10.0 +_POLL_TIMEOUT_SECONDS = 1.0 +_BACKOFF_INITIAL_SECONDS = 5.0 +_BACKOFF_MAX_SECONDS = 60.0 + +_WRITE_ACTION_NAMES: frozenset[str] = frozenset( + {"create", "create_many", "update", "update_many", "upsert", "delete", "delete_many"} +) + +_CONFIG_SYNCED_TABLE_NAMES: frozenset[str] = frozenset( + { + "litellm_proxymodeltable", + "litellm_credentialstable", + "litellm_guardrailstable", + "litellm_policytable", + "litellm_policyattachmenttable", + "litellm_managedvectorstorestable", + "litellm_managedvectorstoreindextable", + "litellm_mcpservertable", + "litellm_agentstable", + "litellm_prompttable", + "litellm_searchtoolstable", + "litellm_ssoconfig", + "litellm_cacheconfig", + "litellm_configoverrides", + } +) + +_RESYNC_APPLIED_CONFIG_PARAM_NAMES: frozenset[str] = frozenset( + { + "general_settings", + "router_settings", + "litellm_settings", + "model_cost_map_reload_config", + "anthropic_beta_headers_reload_config", + } +) + + +def coordination_redis_cache() -> "RedisCache | None": + from litellm.proxy.proxy_server import redis_usage_cache + + return redis_usage_cache + + +def config_sync_channel(redis_cache: "RedisCache") -> str: + if redis_cache.namespace is None: + return CONFIG_SYNC_CHANNEL + return f"{redis_cache.namespace}:{CONFIG_SYNC_CHANNEL}" + + +def _raw_async_client(redis_cache: "RedisCache") -> object: + return cast( # cast-ok: redis-py generics leave the client type partially unknown + object, + redis_cache.init_async_client(), # pyright: ignore[reportUnknownMemberType] # redis generics + ) + + +def _pubsub_capable_client(redis_cache: "RedisCache") -> _ConfigSyncPubSubClient | None: + from redis.asyncio import Redis + + client = _raw_async_client(redis_cache) + if isinstance(client, Redis): + return cast(_ConfigSyncPubSubClient, client) # cast-ok: protocol view of the standalone redis client + return None + + +@dataclass(frozen=True, slots=True) +class _ConfigChangeMessage: + object_type: str + + +def _config_change_message_json(object_type: str) -> str: + return json.dumps(asdict(_ConfigChangeMessage(object_type=object_type))) + + +async def publish_config_change(redis_cache: "RedisCache | None", object_type: str) -> None: + if redis_cache is None: + return + try: + client = _pubsub_capable_client(redis_cache) + if client is None: + verbose_proxy_logger.debug( + "config sync publish for %s skipped: cluster redis client has no pub/sub support", + object_type, + ) + return + await client.publish(config_sync_channel(redis_cache), _config_change_message_json(object_type)) + except Exception as e: # noqa: BLE001 # best-effort publish; writes must never fail on redis errors + verbose_proxy_logger.warning("config sync publish for %s failed: %s", object_type, e) + + +async def publish_config_change_for_object_type(object_type: str) -> None: + await publish_config_change(redis_cache=coordination_redis_cache(), object_type=object_type) + + +async def publish_config_param_change(param_name: str) -> None: + if param_name not in _RESYNC_APPLIED_CONFIG_PARAM_NAMES: + verbose_proxy_logger.debug( + "config sync publish for %s skipped: no resync callback applies this param outside proxy startup", + param_name, + ) + return + await publish_config_change_for_object_type(param_name) + + +class _PublishOnWriteActions: + __slots__ = ("_actions", "_object_type", "_publish") + + def __init__(self, actions: object, object_type: str, publish: Callable[[str], Awaitable[None]]) -> None: + self._actions = actions + self._object_type = object_type + self._publish = publish + + def __getattr__(self, name: str) -> object: + attribute = cast(object, getattr(self._actions, name)) # cast-ok: getattr on dynamic prisma actions + if name not in _WRITE_ACTION_NAMES: + return attribute + write_action = cast(Callable[..., Awaitable[object]], attribute) # cast-ok: prisma actions are untyped + object_type = self._object_type + publish = self._publish + + async def _write_then_publish( + *args: object, + **kwargs: object, # kwargs-ok: transparent passthrough to untyped prisma action + ) -> object: + result = await write_action(*args, **kwargs) + await publish(object_type) + return result + + return _write_then_publish + + +def wrap_table_actions_for_config_sync( + actions: object, + table_name: str, + publish: Callable[[str], Awaitable[None]] = publish_config_change_for_object_type, +) -> object: + if table_name not in _CONFIG_SYNCED_TABLE_NAMES: + return actions + return _PublishOnWriteActions(actions=actions, object_type=table_name, publish=publish) + + +class ConfigSyncSubscriber: + __slots__ = ( + "_backoff_initial_seconds", + "_backoff_max_seconds", + "_debounce_seconds", + "_jitter_max_seconds", + "_last_resync_at", + "_min_resync_interval_seconds", + "_monotonic", + "_redis_cache", + "_resync_callbacks", + "_rng", + "_sleep", + "_task", + ) + + def __init__( + self, + redis_cache: "RedisCache", + resync_callbacks: tuple[Callable[[], Awaitable[None]], ...], + debounce_seconds: float = CONFIG_SYNC_DEBOUNCE_SECONDS, + jitter_max_seconds: float = CONFIG_SYNC_JITTER_MAX_SECONDS, + min_resync_interval_seconds: float = CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS, + backoff_initial_seconds: float = _BACKOFF_INITIAL_SECONDS, + backoff_max_seconds: float = _BACKOFF_MAX_SECONDS, + rng: random.Random | None = None, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + monotonic: Callable[[], float] = time.monotonic, + ) -> None: + self._redis_cache = redis_cache + self._resync_callbacks = resync_callbacks + self._debounce_seconds = debounce_seconds + self._jitter_max_seconds = jitter_max_seconds + self._min_resync_interval_seconds = min_resync_interval_seconds + self._backoff_initial_seconds = backoff_initial_seconds + self._backoff_max_seconds = backoff_max_seconds + self._rng = rng if rng is not None else random.Random() + self._sleep = sleep + self._monotonic = monotonic + self._task: asyncio.Task[None] | None = None + self._last_resync_at: float | None = None + + def start(self) -> None: + if self._task is not None: + return + self._task = asyncio.create_task(self._run()) + + async def stop(self) -> None: + task = self._task + if task is None: + return + self._task = None + _ = task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def _run(self) -> None: + backoff_seconds = self._backoff_initial_seconds + while True: + try: + client = _pubsub_capable_client(self._redis_cache) + if client is None: + verbose_proxy_logger.warning( + "config sync subscriber disabled: cluster redis client has no pub/sub support; " + "interval polling remains the only sync mechanism" + ) + return + pubsub = client.pubsub() + try: + await pubsub.subscribe(config_sync_channel(self._redis_cache)) + backoff_seconds = self._backoff_initial_seconds + await self._consume(pubsub) + finally: + await self._close_pubsub(pubsub) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect + verbose_proxy_logger.warning( + "config sync subscriber redis error: %s; reconnecting in %.0fs", + e, + backoff_seconds, + ) + await self._sleep(backoff_seconds) + backoff_seconds = min(backoff_seconds * 2, self._backoff_max_seconds) + + async def _consume(self, pubsub: _ConfigSyncPubSub) -> None: + while True: + message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) + if message is None: + continue + await self._sleep(self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds)) + await self._wait_for_min_resync_interval() + await self._drain_pending(pubsub) + await self._run_resync_callbacks() + self._last_resync_at = self._monotonic() + + async def _wait_for_min_resync_interval(self) -> None: + if self._last_resync_at is None: + return + seconds_until_next_resync = self._min_resync_interval_seconds - (self._monotonic() - self._last_resync_at) + if seconds_until_next_resync <= 0: + return + verbose_proxy_logger.debug( + "config sync resync throttled for %.1fs to cap fleet-wide reload rate", + seconds_until_next_resync, + ) + await self._sleep(seconds_until_next_resync) + + @staticmethod + async def _drain_pending(pubsub: _ConfigSyncPubSub) -> None: + while await pubsub.get_message(ignore_subscribe_messages=True, timeout=0) is not None: + pass + + async def _run_resync_callbacks(self) -> None: + for callback in self._resync_callbacks: + try: + await callback() + except Exception as e: # noqa: BLE001 # one failing resync callback must not kill the subscriber + verbose_proxy_logger.warning("config sync resync callback failed: %s", e) + + @staticmethod + async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: + try: + await pubsub.aclose() + except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection + verbose_proxy_logger.debug("config sync pubsub close failed: %s", e) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b910eee7130..c336e9321e0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -61,6 +61,10 @@ from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, encrypt_callback_vars, ) +from litellm.proxy.common_utils.config_sync_pubsub import ( + coordination_redis_cache, + publish_config_change, +) from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -4304,6 +4308,7 @@ async def _rotate_master_key( await tx.litellm_proxymodeltable.create_many( data=new_models, ) + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") # 3. process config table try: config = await _config_table(prisma_client).find_many() diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 4885ec42578..f43c62b68b9 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -38,6 +38,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.config_sync_pubsub import ( + coordination_redis_cache, + publish_config_change, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin @@ -812,6 +816,9 @@ async def delete_team_models( await tx.litellm_proxymodeltable.delete_many(where={"model_id": {"in": model_ids}}) deleted_model_ids.extend(model_ids) + if deleted_model_ids: + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") + if llm_router is not None: for model_id in deleted_model_ids: llm_router.delete_deployment(id=model_id) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 26d5b386cdd..04161d775f8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -301,6 +301,7 @@ from litellm.proxy.common_request_processing import ( create_response, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -546,6 +547,7 @@ from litellm.proxy.utils import ( _get_redoc_url, _is_projected_spend_over_limit, _is_valid_team_configs, + evict_config_param, get_config_param, get_custom_url, get_error_message_str, @@ -1151,6 +1153,8 @@ async def proxy_startup_event(app: FastAPI): except Exception as e: verbose_proxy_logger.error(f"Error stopping DB health watchdog task: {e}") + await proxy_config.stop_config_sync_subscriber() + await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] @@ -3836,6 +3840,7 @@ class ProxyConfig: self._last_semantic_filter_config: Optional[Dict[str, Any]] = None self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None self.worker_registry: List["WorkerRegistryEntry"] = [] + self.config_sync_subscriber: ConfigSyncSubscriber | None = None def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -6216,6 +6221,38 @@ class ProxyConfig: return still_desired_ids + def start_config_sync_subscriber( + self, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + redis_cache: Optional[RedisCache], + ) -> None: + if redis_cache is None or self.config_sync_subscriber is not None: + return + + async def _resync_config_from_db() -> None: + await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + + async def _resync_credentials_from_db() -> None: + await self.get_credentials(prisma_client=prisma_client) + + subscriber = ConfigSyncSubscriber( + redis_cache=redis_cache, + resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db), + ) + self.config_sync_subscriber = subscriber + subscriber.start() + + async def stop_config_sync_subscriber(self) -> None: + subscriber = self.config_sync_subscriber + if subscriber is None: + return + self.config_sync_subscriber = None + try: + await subscriber.stop() + except Exception as e: + verbose_proxy_logger.error(f"Error stopping config sync subscriber: {e}") + async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): """ Use this to read non-llm objects from the db and initialize them @@ -6507,7 +6544,7 @@ class ProxyConfig: }, }, ) - await invalidate_config_param("model_cost_map_reload_config") + await evict_config_param("model_cost_map_reload_config") verbose_proxy_logger.info( f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}" @@ -6602,7 +6639,7 @@ class ProxyConfig: }, }, ) - await invalidate_config_param("anthropic_beta_headers_reload_config") + await evict_config_param("anthropic_beta_headers_reload_config") # Count providers in config provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description") @@ -8208,6 +8245,12 @@ class ProxyStartupEvent: ) await proxy_config.get_credentials(prisma_client=prisma_client) + proxy_config.start_config_sync_subscriber( + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + redis_cache=redis_usage_cache, + ) + if store_model_in_db is not True: await proxy_config.init_mcp_servers_from_db() if prisma_client is not None: diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 8f3f8ad1bfc..e4872a4b6b5 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -21,6 +21,7 @@ from litellm.proxy.config_resolvers.sso import ( SSO_SECRET_FIELDS, resolve_sso_config, ) +from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( SSOConfigRepository, @@ -971,6 +972,7 @@ async def update_sso_settings( "param_value": json.dumps(filtered_env_vars, default=str), }, ) + await invalidate_config_param("environment_variables") except Exception as e: raise HTTPException( status_code=500, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4fd13cc5c6c..e1b8f89149b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -120,6 +120,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( create_missing_views, @@ -2973,9 +2974,14 @@ async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any] return row +async def evict_config_param(param_name: str) -> None: + await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) + + async def invalidate_config_param(param_name: str) -> None: """Evict from both cache layers; call after every LiteLLM_Config write.""" - await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) + await evict_config_param(param_name) + await publish_config_param_change(param_name) async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None: diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index b5a315d233c..8e4b9ac0be7 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -9,6 +9,7 @@ so reads return the stored values verbatim. from typing import Any, Dict, Optional from litellm.models.credentials import CredentialItem +from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync class CredentialsRepository: @@ -25,7 +26,10 @@ class CredentialsRepository: @property def table(self) -> Any: - return self.prisma_client.db.litellm_credentialstable + return wrap_table_actions_for_config_sync( + actions=self.prisma_client.db.litellm_credentialstable, + table_name="litellm_credentialstable", + ) @staticmethod def _to_model(record: Any) -> Optional[CredentialItem]: diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 0da51519964..50a50cc60d6 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -7,6 +7,7 @@ from typing import Any, Dict, List, Optional, Type from litellm.models.model import LiteLLM_ProxyModelTable from litellm.repositories.base_repository import BaseRepository +from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -22,7 +23,10 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): @property def table(self) -> Any: - return self.prisma_client.db.litellm_proxymodeltable + return wrap_table_actions_for_config_sync( + actions=self.prisma_client.db.litellm_proxymodeltable, + table_name="litellm_proxymodeltable", + ) @property def model_class(self) -> Type[LiteLLM_ProxyModelTable]: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 54008c0950c..af8be986831 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -9,6 +9,8 @@ methods; richer repositories live in their own modules. from typing import Any +from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync + class PrismaTableRepository: """Base for repositories that expose a single Prisma table.""" @@ -26,7 +28,10 @@ class PrismaTableRepository: @property def table(self) -> Any: - return getattr(self.prisma_client.db, self.table_name) + return wrap_table_actions_for_config_sync( + actions=getattr(self.prisma_client.db, self.table_name), + table_name=self.table_name, + ) class PolicyRepository(PrismaTableRepository): diff --git a/ruff.toml b/ruff.toml index 2ea9d7260fb..b652e206f41 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,7 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "C901", + "C901", "TID251", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py new file mode 100644 index 00000000000..6872407808c --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -0,0 +1,894 @@ +import asyncio +import json +import random +from typing import Callable, Coroutine, Iterable, List, Optional, Tuple +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from redis.asyncio import Redis + +import litellm +from litellm.proxy.common_utils.config_sync_pubsub import ( + CONFIG_SYNC_CHANNEL, + CONFIG_SYNC_JITTER_MAX_SECONDS, + CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS, + ConfigSyncSubscriber, + _CONFIG_SYNCED_TABLE_NAMES, + _PublishOnWriteActions, + _RESYNC_APPLIED_CONFIG_PARAM_NAMES, + _WRITE_ACTION_NAMES, + publish_config_change, + wrap_table_actions_for_config_sync, +) + +_EXPECTED_WRITE_ACTION_NAMES = ( + "create", + "create_many", + "delete", + "delete_many", + "update", + "update_many", + "upsert", +) + +_EXPECTED_CONFIG_SYNCED_TABLE_NAMES = frozenset( + { + "litellm_agentstable", + "litellm_cacheconfig", + "litellm_configoverrides", + "litellm_credentialstable", + "litellm_guardrailstable", + "litellm_managedvectorstoreindextable", + "litellm_managedvectorstorestable", + "litellm_mcpservertable", + "litellm_policyattachmenttable", + "litellm_policytable", + "litellm_prompttable", + "litellm_proxymodeltable", + "litellm_searchtoolstable", + "litellm_ssoconfig", + } +) + +_EXPECTED_RESYNC_APPLIED_CONFIG_PARAM_NAMES = frozenset( + { + "anthropic_beta_headers_reload_config", + "general_settings", + "litellm_settings", + "model_cost_map_reload_config", + "router_settings", + } +) + +_STARTUP_ONLY_CONFIG_PARAM_NAMES = ("environment_variables",) + + +class _RecordingRedisClient(Redis): + def __init__(self) -> None: + self.published: List[Tuple[str, str]] = [] + + async def publish(self, channel: str, message: str) -> int: + self.published.append((channel, message)) + return 1 + + +class _FailingPublishRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + raise ConnectionError("redis down") + + +class _NotRedisClient: + def __init__(self) -> None: + self.published: List[Tuple[str, str]] = [] + + async def publish(self, channel: str, message: str) -> int: + self.published.append((channel, message)) + return 1 + + +class _QueuePubSub: + def __init__(self, initial_messages: Iterable[str] = ()) -> None: + self.queue: "asyncio.Queue[str]" = asyncio.Queue() + for message in initial_messages: + self.queue.put_nowait(message) + self.subscribed_channels: List[str] = [] + self.closed = False + + async def subscribe(self, *channels: str) -> None: + self.subscribed_channels.extend(channels) + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + if timeout == 0: + try: + return self.queue.get_nowait() + except asyncio.QueueEmpty: + return None + try: + return await asyncio.wait_for(self.queue.get(), timeout) + except asyncio.TimeoutError: + return None + + async def aclose(self) -> None: + self.closed = True + + +class _BrokenPubSub(_QueuePubSub): + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + raise ConnectionError("connection lost") + + +class _CloseFailingBrokenPubSub(_BrokenPubSub): + async def aclose(self) -> None: + raise ConnectionError("close failed") + + +class _EmptyPollsThenMessagePubSub(_QueuePubSub): + def __init__(self, empty_polls: int, initial_messages: Iterable[str] = ()) -> None: + super().__init__(initial_messages=initial_messages) + self.remaining_empty_polls = empty_polls + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + if timeout != 0 and self.remaining_empty_polls > 0: + self.remaining_empty_polls -= 1 + return None + return await super().get_message(ignore_subscribe_messages=ignore_subscribe_messages, timeout=timeout) + + +class _FakeClock: + def __init__(self, now: float = 1000.0) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + +class _ScriptedPubSubRedisClient(Redis): + def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None: + self._scripted_pubsubs = iter(pubsubs) + + def pubsub(self) -> _QueuePubSub: + return next(self._scripted_pubsubs) + + +class _FakeRedisCache: + def __init__(self, client: object, namespace: Optional[str] = None) -> None: + self._client = client + self.namespace = namespace + + def init_async_client(self) -> object: + return self._client + + +class _ExplodingRedisCache: + namespace: Optional[str] = None + + def init_async_client(self) -> object: + raise ConnectionError("cannot connect") + + +def _recording_callback( + events: List[str], name: str, fired: asyncio.Event +) -> Callable[[], Coroutine[None, None, None]]: + async def callback() -> None: + events.append(name) + fired.set() + + return callback + + +async def test_publish_noops_when_redis_cache_is_none() -> None: + await publish_config_change(redis_cache=None, object_type="litellm_proxymodeltable") + + +async def test_publish_sends_object_type_json_on_channel() -> None: + client = _RecordingRedisClient() + cache = _FakeRedisCache(client) + + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") + + assert len(client.published) == 1 + channel, message = client.published[0] + assert channel == "litellm_proxy.config_change" + assert json.loads(message) == {"object_type": "litellm_proxymodeltable"} + + +async def test_publish_uses_namespaced_channel() -> None: + client = _RecordingRedisClient() + cache = _FakeRedisCache(client, namespace="prod-eu") + + await publish_config_change(redis_cache=cache, object_type="litellm_credentialstable") + + assert client.published[0][0] == "prod-eu:litellm_proxy.config_change" + + +async def test_publish_swallows_redis_publish_errors() -> None: + cache = _FakeRedisCache(_FailingPublishRedisClient()) + + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") + + +async def test_publish_swallows_client_init_errors() -> None: + await publish_config_change(redis_cache=_ExplodingRedisCache(), object_type="litellm_proxymodeltable") + + +async def test_publish_skips_clients_without_pubsub_support() -> None: + client = _NotRedisClient() + cache = _FakeRedisCache(client) + + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") + + assert client.published == [] + + +async def test_subscriber_runs_injected_callbacks_in_order_on_message() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + events: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=( + _recording_callback(events, "add_deployment", asyncio.Event()), + _recording_callback(events, "get_credentials", fired), + ), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + ) + + subscriber.start() + pubsub.queue.put_nowait(json.dumps({"object_type": "litellm_proxymodeltable"})) + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert events == ["add_deployment", "get_credentials"] + assert pubsub.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert pubsub.closed is True + + +async def test_burst_within_debounce_window_coalesces_into_one_resync() -> None: + burst = [json.dumps({"object_type": "litellm_proxymodeltable"}) for _ in range(5)] + pubsub = _QueuePubSub(initial_messages=burst) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.05, + jitter_max_seconds=0.0, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + await asyncio.sleep(0.3) + await subscriber.stop() + + assert resyncs == ["resync"] + assert pubsub.queue.empty() + + +async def test_subscriber_subscribes_on_namespaced_channel_and_resyncs() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub]), namespace="prod-eu") + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + ) + + subscriber.start() + pubsub.queue.put_nowait(json.dumps({"object_type": "litellm_proxymodeltable"})) + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert pubsub.subscribed_channels == ["prod-eu:litellm_proxy.config_change"] + assert resyncs == ["resync"] + + +class _MaxJitterRandom(random.Random): + def uniform(self, a: float, b: float) -> float: + return b + + +async def test_debounce_sleep_adds_jitter_from_injected_rng() -> None: + pubsub = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_proxymodeltable"})]) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + sleeps: List[float] = [] + fired = asyncio.Event() + + async def recording_sleep(seconds: float) -> None: + sleeps.append(seconds) + + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback([], "resync", fired),), + debounce_seconds=1.0, + jitter_max_seconds=4.0, + rng=_MaxJitterRandom(), + sleep=recording_sleep, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert sleeps == [5.0] + + +def test_default_jitter_window_is_nonzero() -> None: + assert CONFIG_SYNC_JITTER_MAX_SECONDS > 0 + + +def test_default_min_resync_interval_caps_reload_rate() -> None: + assert CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS > CONFIG_SYNC_JITTER_MAX_SECONDS + + +def _throttled_subscriber( + cache: object, + events: List[str], + fired: asyncio.Event, + clock: _FakeClock, + min_resync_interval_seconds: float = 10.0, +) -> ConfigSyncSubscriber: + async def recording_sleep(seconds: float) -> None: + events.append(f"sleep:{seconds}") + await asyncio.sleep(0) + + async def resync() -> None: + events.append("resync") + fired.set() + + return ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(resync,), + debounce_seconds=0.0, + jitter_max_seconds=0.0, + min_resync_interval_seconds=min_resync_interval_seconds, + sleep=recording_sleep, + monotonic=clock, + ) + + +async def test_resync_arriving_inside_min_interval_waits_out_the_remainder() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + events: List[str] = [] + fired = asyncio.Event() + clock = _FakeClock() + subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + + subscriber.start() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + fired.clear() + clock.now += 4.0 + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert events == ["sleep:0.0", "resync", "sleep:0.0", "sleep:6.0", "resync"] + + +async def test_resync_after_min_interval_elapsed_is_not_throttled() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + events: List[str] = [] + fired = asyncio.Event() + clock = _FakeClock() + subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + + subscriber.start() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + fired.clear() + clock.now += 30.0 + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert events == ["sleep:0.0", "resync", "sleep:0.0", "resync"] + + +async def test_writes_during_the_throttle_wait_collapse_into_the_next_resync() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + events: List[str] = [] + fired = asyncio.Event() + clock = _FakeClock() + subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + + subscriber.start() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + fired.clear() + for _ in range(5): + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + await asyncio.sleep(0.1) + await subscriber.stop() + + assert events.count("resync") == 2 + assert pubsub.queue.empty() + + +async def test_polls_without_messages_do_not_trigger_resyncs() -> None: + pubsub = _EmptyPollsThenMessagePubSub(empty_polls=3, initial_messages=["change"]) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + await asyncio.sleep(0.1) + await subscriber.stop() + + assert pubsub.remaining_empty_polls == 0 + assert resyncs == ["resync"] + + +async def test_failing_pubsub_close_still_reconnects() -> None: + broken = _CloseFailingBrokenPubSub() + healthy = _QueuePubSub(initial_messages=["change"]) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([broken, healthy])) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + backoff_initial_seconds=0.02, + backoff_max_seconds=0.05, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert healthy.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert resyncs == ["resync"] + + +async def test_second_start_does_not_open_a_second_subscription() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + subscriber = ConfigSyncSubscriber(redis_cache=cache, resync_callbacks=(), debounce_seconds=0.01) + + subscriber.start() + task = subscriber._task + subscriber.start() + assert task is not None + assert subscriber._task is task + await asyncio.sleep(0.05) + await subscriber.stop() + + assert pubsub.subscribed_channels == [CONFIG_SYNC_CHANNEL] + + +async def test_redis_error_leads_to_backoff_and_resubscribe() -> None: + broken = _BrokenPubSub() + healthy = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_credentialstable"})]) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([broken, healthy])) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + backoff_initial_seconds=0.02, + backoff_max_seconds=0.05, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + task = subscriber._task + assert task is not None + assert task.done() is False + await subscriber.stop() + + assert broken.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert broken.closed is True + assert healthy.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert resyncs == ["resync"] + + +async def test_failing_resync_callback_does_not_kill_subscriber() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + resyncs: List[str] = [] + fired = asyncio.Event() + + async def failing_callback() -> None: + raise RuntimeError("resync exploded") + + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(failing_callback, _recording_callback(resyncs, "resync", fired)), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + min_resync_interval_seconds=0.0, + ) + + subscriber.start() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + fired.clear() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert resyncs == ["resync", "resync"] + + +async def test_stop_cancels_subscriber_cleanly() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + subscriber = ConfigSyncSubscriber(redis_cache=cache, resync_callbacks=(), debounce_seconds=0.01) + + subscriber.start() + await asyncio.sleep(0.05) + task = subscriber._task + assert task is not None + await subscriber.stop() + + assert task.done() is True + assert subscriber._task is None + assert pubsub.closed is True + await subscriber.stop() + + +async def test_stop_before_start_is_a_noop() -> None: + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([])) + subscriber = ConfigSyncSubscriber(redis_cache=cache, resync_callbacks=()) + + await subscriber.stop() + + +async def test_subscriber_exits_without_callbacks_when_client_lacks_pubsub() -> None: + cache = _FakeRedisCache(_NotRedisClient()) + resyncs: List[str] = [] + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", asyncio.Event()),), + ) + + subscriber.start() + task = subscriber._task + assert task is not None + await asyncio.wait_for(task, timeout=5) + + assert resyncs == [] + + +class _FakeTableActions: + def __init__(self, calls: List[Tuple[str, str]]) -> None: + self._calls = calls + + async def create(self, **kwargs: object) -> object: + self._calls.append(("write", "create")) + return {"id": "m-1"} + + async def find_many(self, **kwargs: object) -> object: + self._calls.append(("read", "find_many")) + return [] + + +class _AllWritesTableActions: + def __init__(self, calls: List[str]) -> None: + self._calls = calls + + def __getattr__(self, name: str) -> Callable[..., Coroutine[None, None, str]]: + async def action(*args: object, **kwargs: object) -> str: + self._calls.append(name) + return name + + return action + + +def _recording_publish(calls: List[Tuple[str, str]]) -> Callable[[str], Coroutine[None, None, None]]: + async def publish(object_type: str) -> None: + calls.append(("publish", object_type)) + + return publish + + +def test_wrapper_passes_through_unsynced_tables() -> None: + actions = object() + + wrapped = wrap_table_actions_for_config_sync(actions=actions, table_name="litellm_spendlogs") + + assert wrapped is actions + + +async def test_wrapper_publishes_table_name_after_write() -> None: + calls: List[Tuple[str, str]] = [] + wrapped = wrap_table_actions_for_config_sync( + actions=_FakeTableActions(calls), + table_name="litellm_proxymodeltable", + publish=_recording_publish(calls), + ) + + result = await wrapped.create(data={"model_name": "gpt-5.2"}) + + assert result == {"id": "m-1"} + assert calls == [("write", "create"), ("publish", "litellm_proxymodeltable")] + + +async def test_wrapper_does_not_publish_on_reads() -> None: + calls: List[Tuple[str, str]] = [] + wrapped = wrap_table_actions_for_config_sync( + actions=_FakeTableActions(calls), + table_name="litellm_proxymodeltable", + publish=_recording_publish(calls), + ) + + result = await wrapped.find_many(where={}) + + assert result == [] + assert calls == [("read", "find_many")] + + +def test_write_action_names_are_pinned() -> None: + assert _WRITE_ACTION_NAMES == frozenset(_EXPECTED_WRITE_ACTION_NAMES) + + +def test_config_synced_table_membership_is_pinned() -> None: + assert _CONFIG_SYNCED_TABLE_NAMES == _EXPECTED_CONFIG_SYNCED_TABLE_NAMES + + +def test_tool_telemetry_table_writes_pass_through_unwrapped() -> None: + actions = object() + + wrapped = wrap_table_actions_for_config_sync(actions=actions, table_name="litellm_tooltable") + + assert wrapped is actions + + +@pytest.mark.parametrize("action_name", _EXPECTED_WRITE_ACTION_NAMES) +async def test_wrapper_publishes_for_every_write_action(action_name: str) -> None: + write_calls: List[str] = [] + publish_calls: List[Tuple[str, str]] = [] + wrapped = wrap_table_actions_for_config_sync( + actions=_AllWritesTableActions(write_calls), + table_name="litellm_guardrailstable", + publish=_recording_publish(publish_calls), + ) + + result = await getattr(wrapped, action_name)(data={}) + + assert result == action_name + assert write_calls == [action_name] + assert publish_calls == [("publish", "litellm_guardrailstable")] + + +async def test_model_repository_write_publishes_via_live_coordination_cache() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import _set_redis_usage_cache + from litellm.repositories.model_repository import ModelRepository + + client = _RecordingRedisClient() + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.update = AsyncMock(return_value={"model_id": "m-1"}) + repository = ModelRepository(prisma_client) + table = repository.table + assert isinstance(table, _PublishOnWriteActions) + + previous_cache = proxy_server.redis_usage_cache + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + await table.update(where={"model_id": "m-1"}, data={"model_name": "gpt-5.2"}) + finally: + _set_redis_usage_cache(previous_cache) + + prisma_client.db.litellm_proxymodeltable.update.assert_awaited_once_with( + where={"model_id": "m-1"}, data={"model_name": "gpt-5.2"} + ) + assert len(client.published) == 1 + channel, message = client.published[0] + assert channel == CONFIG_SYNC_CHANNEL + assert json.loads(message) == {"object_type": "litellm_proxymodeltable"} + + +async def _publish_calls_for_invalidated_param(param_name: str) -> List[Tuple[str, str]]: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import _set_redis_usage_cache + from litellm.proxy.utils import invalidate_config_param + + client = _RecordingRedisClient() + previous_cache = proxy_server.redis_usage_cache + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + await invalidate_config_param(param_name) + finally: + _set_redis_usage_cache(previous_cache) + return client.published + + +@pytest.mark.parametrize("param_name", sorted(_EXPECTED_RESYNC_APPLIED_CONFIG_PARAM_NAMES)) +async def test_invalidate_config_param_publishes_params_a_resync_applies(param_name: str) -> None: + published = await _publish_calls_for_invalidated_param(param_name) + + assert len(published) == 1 + channel, message = published[0] + assert channel == CONFIG_SYNC_CHANNEL + assert json.loads(message) == {"object_type": param_name} + + +@pytest.mark.parametrize("param_name", _STARTUP_ONLY_CONFIG_PARAM_NAMES) +async def test_invalidate_config_param_does_not_publish_startup_only_params(param_name: str) -> None: + published = await _publish_calls_for_invalidated_param(param_name) + + assert published == [] + + +def test_resync_applied_config_param_membership_is_pinned() -> None: + assert _RESYNC_APPLIED_CONFIG_PARAM_NAMES == _EXPECTED_RESYNC_APPLIED_CONFIG_PARAM_NAMES + + +async def test_evict_config_param_does_not_publish() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import _set_redis_usage_cache + from litellm.proxy.utils import evict_config_param + + client = _RecordingRedisClient() + previous_cache = proxy_server.redis_usage_cache + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + await evict_config_param("model_cost_map_reload_config") + finally: + _set_redis_usage_cache(previous_cache) + + assert client.published == [] + + +def _reload_config_prisma_client() -> MagicMock: + config_record = MagicMock() + config_record.param_value = {"interval_hours": 6, "force_reload": True} + prisma_client = MagicMock() + prisma_client.get_generic_data = AsyncMock(return_value=config_record) + prisma_client.db.litellm_config.upsert = AsyncMock(return_value=None) + return prisma_client + + +async def test_model_cost_map_reload_does_not_publish_config_change() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import ProxyConfig, _set_redis_usage_cache + from litellm.proxy.utils import litellm_config_cache + from litellm.utils import _invalidate_model_cost_lowercase_map + + litellm_config_cache.flush_cache() + prisma_client = _reload_config_prisma_client() + client = _RecordingRedisClient() + previous_cache = proxy_server.redis_usage_cache + original_model_cost = litellm.model_cost.copy() + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + with patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map: + mock_get_map.return_value = {"gpt-5.2": {"input_cost_per_token": 0.001}} + await ProxyConfig()._check_and_reload_model_cost_map(prisma_client=prisma_client) + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + _set_redis_usage_cache(previous_cache) + + prisma_client.db.litellm_config.upsert.assert_awaited_once() + assert client.published == [] + + +async def test_anthropic_beta_headers_reload_does_not_publish_config_change() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import ProxyConfig, _set_redis_usage_cache + from litellm.proxy.utils import litellm_config_cache + + litellm_config_cache.flush_cache() + prisma_client = _reload_config_prisma_client() + client = _RecordingRedisClient() + previous_cache = proxy_server.redis_usage_cache + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + with patch("litellm.anthropic_beta_headers_manager.reload_beta_headers_config") as mock_reload: + mock_reload.return_value = {} + await ProxyConfig()._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client) + finally: + _set_redis_usage_cache(previous_cache) + + prisma_client.db.litellm_config.upsert.assert_awaited_once() + assert client.published == [] + + +class _StopFailingSubscriber(ConfigSyncSubscriber): + async def stop(self) -> None: + raise RuntimeError("stop failed") + + +async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> None: + from litellm.proxy.proxy_server import ProxyConfig + + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) + config = ProxyConfig() + prisma_client = MagicMock() + proxy_logging_obj = MagicMock() + calls: List[Tuple[str, object, object]] = [] + + async def fake_add_deployment(prisma_client: object, proxy_logging_obj: object) -> None: + calls.append(("add_deployment", prisma_client, proxy_logging_obj)) + + async def fake_get_credentials(prisma_client: object) -> None: + calls.append(("get_credentials", prisma_client, None)) + + config.add_deployment = fake_add_deployment + config.get_credentials = fake_get_credentials + config.start_config_sync_subscriber( + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + redis_cache=cache, + ) + subscriber = config.config_sync_subscriber + assert subscriber is not None + for callback in subscriber._resync_callbacks: + await callback() + await config.stop_config_sync_subscriber() + + assert calls == [ + ("add_deployment", prisma_client, proxy_logging_obj), + ("get_credentials", prisma_client, None), + ] + assert config.config_sync_subscriber is None + assert subscriber._task is None + + +async def test_proxy_config_does_not_start_subscriber_without_coordination_redis() -> None: + from litellm.proxy.proxy_server import ProxyConfig + + config = ProxyConfig() + + config.start_config_sync_subscriber( + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + redis_cache=None, + ) + + assert config.config_sync_subscriber is None + + +async def test_proxy_config_keeps_the_first_subscriber_on_repeat_start() -> None: + from litellm.proxy.proxy_server import ProxyConfig + + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) + config = ProxyConfig() + + config.start_config_sync_subscriber(prisma_client=MagicMock(), proxy_logging_obj=MagicMock(), redis_cache=cache) + first = config.config_sync_subscriber + config.start_config_sync_subscriber(prisma_client=MagicMock(), proxy_logging_obj=MagicMock(), redis_cache=cache) + second = config.config_sync_subscriber + await config.stop_config_sync_subscriber() + + assert first is not None + assert second is first + + +async def test_proxy_config_shutdown_survives_a_failing_subscriber_stop() -> None: + from litellm.proxy.proxy_server import ProxyConfig + + config = ProxyConfig() + config.config_sync_subscriber = _StopFailingSubscriber( + redis_cache=_FakeRedisCache(_ScriptedPubSubRedisClient([])), + resync_callbacks=(), + ) + + await config.stop_config_sync_subscriber() + + assert config.config_sync_subscriber is None diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 4f1efe484ab..9a3702d2355 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2058,6 +2058,42 @@ async def test_ProxyConfig__add_router_settings_from_db_config_none_router_noop( await pc._add_router_settings_from_db_config() # type: ignore[call-arg] +# --------------------------------------------------------------------------- +# ProxyConfig.add_deployment +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch): + from litellm.proxy import proxy_server + + pc = ProxyConfig() + fake_router = MagicMock() + fake_router.get_model_list = MagicMock(return_value=[]) + fake_prisma = MagicMock() + fake_prisma.db.litellm_config.find_first = AsyncMock( + return_value=SimpleNamespace(param_value={"routing_strategy": "latency-based-routing"}) + ) + + async def fake_get_config(*args, **kwargs): + return {} + + monkeypatch.setattr(pc, "get_config", fake_get_config) + monkeypatch.setattr(pc, "_get_models_from_db", AsyncMock(return_value=[])) + monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock()) + monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock()) + monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "master_key", "sk-master") + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "proxy_config", pc) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + fake_router.update_settings.assert_called_once_with(routing_strategy="latency-based-routing") + + # --------------------------------------------------------------------------- # ProxyConfig._add_general_settings_from_db_config # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index 8253308f393..3f567397e1a 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -307,6 +307,15 @@ class TestModelRepository: client = MockPrismaClient() return ModelRepository(client) + def test_table_is_wrapped_for_config_sync(self, repo): + from litellm.proxy.common_utils.config_sync_pubsub import ( + _PublishOnWriteActions, + ) + + table = repo.table + assert isinstance(table, _PublishOnWriteActions) + assert table._actions is repo.prisma_client.db.litellm_proxymodeltable + @pytest.mark.asyncio @patch( "litellm.repositories.model_repository.encrypt_value_helper", @@ -1313,6 +1322,15 @@ class TestCredentialsRepository: client = MockPrismaClient() return CredentialsRepository(client) + def test_table_is_wrapped_for_config_sync(self, repo): + from litellm.proxy.common_utils.config_sync_pubsub import ( + _PublishOnWriteActions, + ) + + table = repo.table + assert isinstance(table, _PublishOnWriteActions) + assert table._actions is repo.prisma_client.db.litellm_credentialstable + @pytest.mark.asyncio async def test_create(self, repo): record = await repo.create( @@ -2210,6 +2228,9 @@ class TestConfigRepositoryDeepCopy: class TestPrismaTableRepository: def test_table_property_returns_named_delegate(self): + from litellm.proxy.common_utils.config_sync_pubsub import ( + _PublishOnWriteActions, + ) from litellm.repositories.table_repositories import ( AgentsRepository, PolicyRepository, @@ -2219,9 +2240,11 @@ class TestPrismaTableRepository: agents = AgentsRepository(prisma_client) policy = PolicyRepository(prisma_client) - assert agents.table is prisma_client.db.litellm_agentstable - assert policy.table is prisma_client.db.litellm_policytable - assert agents.table is not policy.table + assert isinstance(agents.table, _PublishOnWriteActions) + assert isinstance(policy.table, _PublishOnWriteActions) + assert agents.table._actions is prisma_client.db.litellm_agentstable + assert policy.table._actions is prisma_client.db.litellm_policytable + assert agents.table._actions is not policy.table._actions def test_table_access_raises_without_db(self): from litellm.repositories.table_repositories import SpendLogsRepository @@ -2230,8 +2253,28 @@ class TestPrismaTableRepository: with pytest.raises(RuntimeError, match="No DB Connected"): _ = repo.table + CONFIG_SYNCED_TABLE_NAMES = frozenset( + { + "litellm_agentstable", + "litellm_cacheconfig", + "litellm_configoverrides", + "litellm_guardrailstable", + "litellm_managedvectorstoreindextable", + "litellm_managedvectorstorestable", + "litellm_mcpservertable", + "litellm_policyattachmenttable", + "litellm_policytable", + "litellm_prompttable", + "litellm_searchtoolstable", + "litellm_ssoconfig", + } + ) + def test_each_repository_binds_its_own_table_name(self): import litellm.repositories.table_repositories as tr + from litellm.proxy.common_utils.config_sync_pubsub import ( + _PublishOnWriteActions, + ) prisma_client = MagicMock() repos = [ @@ -2248,7 +2291,14 @@ class TestPrismaTableRepository: assert name.startswith("litellm_") assert name not in seen, f"duplicate table_name {name}" seen.add(name) - assert repo_cls(prisma_client).table is getattr(prisma_client.db, name) + table = repo_cls(prisma_client).table + raw_actions = getattr(prisma_client.db, name) + if name in self.CONFIG_SYNCED_TABLE_NAMES: + assert isinstance(table, _PublishOnWriteActions), name + assert table._actions is raw_actions + else: + assert table is raw_actions, name + assert self.CONFIG_SYNCED_TABLE_NAMES <= seen def _json_path_equals(