scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(proxy): push config sync to pods via redis pub/sub

Browse source 
After any management write to a DB-backed config table, publish an
invalidation event on the coordination Redis; every pod runs a
subscriber that debounces, jitters, and triggers an immediate
add_deployment plus get_credentials resync. The interval polls stay
as slow reconciliation fallback and behavior without Redis is
unchanged since publish and subscribe both no-op.
This commit is contained in:
mateo-berri 2026-07-31 20:08:02 -07:00
parent b5cfc2ca00
commit 629d58443e
12 changed files with 978 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

258
litellm/proxy/common_utils/config_sync_pubsub.py Normal file
View file

@ -0,0 +1,258 @@
import asyncio
import json
import random
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
_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",
}
)
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)
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",
"_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,
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,
) -> None:
self._redis_cache = redis_cache
self._resync_callbacks = resync_callbacks
self._debounce_seconds = debounce_seconds
self._jitter_max_seconds = jitter_max_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._task: asyncio.Task[None] | 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._drain_pending(pubsub)
await self._run_resync_callbacks()
@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)

5
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -60,6 +60,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
@ -4240,6 +4244,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 ConfigRepository(prisma_client).table.find_many()

7
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -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.management_endpoints.common_utils import _is_user_team_admin
from litellm.proxy.management_endpoints.team_endpoints import (
@ -809,6 +813,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)

27
litellm/proxy/proxy_server.py
View file

@ -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,12 @@ async def proxy_startup_event(app: FastAPI):
except Exception as e:
verbose_proxy_logger.error(f"Error stopping DB health watchdog task: {e}")
if proxy_config.config_sync_subscriber is not None:
try:
await proxy_config.config_sync_subscriber.stop()
except Exception as e:
verbose_proxy_logger.error(f"Error stopping config sync subscriber: {e}")
await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues]
@ -3837,6 +3845,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):
@ -6495,7 +6504,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}"
@ -6590,7 +6599,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")
@ -8165,6 +8174,20 @@ class ProxyStartupEvent:
)
await proxy_config.get_credentials(prisma_client=prisma_client)
if redis_usage_cache is not None and proxy_config.config_sync_subscriber is None:
async def _resync_config_from_db() -> None:
await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
async def _resync_credentials_from_db() -> None:
await proxy_config.get_credentials(prisma_client=prisma_client)
proxy_config.config_sync_subscriber = ConfigSyncSubscriber(
redis_cache=redis_usage_cache,
resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db),
)
proxy_config.config_sync_subscriber.start()
if store_model_in_db is not True:
await proxy_config.init_mcp_servers_from_db()
if prisma_client is not None:

2
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -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,

11
litellm/proxy/utils.py
View file

@ -119,6 +119,10 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.config_sync_pubsub import (
coordination_redis_cache,
publish_config_change,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.create_views import (
create_missing_views,
@ -2971,9 +2975,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_change(redis_cache=coordination_redis_cache(), object_type=param_name)
async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None:

6
litellm/repositories/credentials_repository.py
View file

@ -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]:

6
litellm/repositories/model_repository.py
View file

@ -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]:

7
litellm/repositories/table_repositories.py
View file

@ -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):

2
ruff.toml
View file

@ -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",
]

600
tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py Normal file
View file

@ -0,0 +1,600 @@
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,
ConfigSyncSubscriber,
_CONFIG_SYNCED_TABLE_NAMES,
_PublishOnWriteActions,
_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",
}
)
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 _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
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,
)
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 test_invalidate_config_param_publishes_param_name() -> None:
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("environment_variables")
finally:
_set_redis_usage_cache(previous_cache)
assert len(client.published) == 1
channel, message = client.published[0]
assert channel == CONFIG_SYNC_CHANNEL
assert json.loads(message) == {"object_type": "environment_variables"}
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 == []

58
tests/test_litellm/repositories/test_repositories.py
View file

@ -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(
@ -2188,6 +2206,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,
@ -2197,9 +2218,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
@ -2208,8 +2231,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 = [
@ -2226,7 +2269,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(