diff --git a/litellm/constants.py b/litellm/constants.py index 1014b472c61..3f99676bbde 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2,7 +2,11 @@ import os import sys from typing import List, Literal, Optional -from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none +from litellm.litellm_core_utils.env_utils import ( + get_env_float, + get_env_int, + get_env_int_or_none, +) DEFAULT_HEALTH_CHECK_PROMPT = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) @@ -1478,6 +1482,11 @@ PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_ PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) +MODEL_CHANGE_PUBSUB_ENABLED = os.getenv("MODEL_CHANGE_PUBSUB_ENABLED", "true").lower() in ["true", "1"] +MODEL_CHANGE_PUBSUB_CHANNEL = os.getenv("MODEL_CHANGE_PUBSUB_CHANNEL", "litellm:model_changes") +MODEL_CHANGE_PUBSUB_POLL_TIMEOUT_SECONDS = get_env_float("MODEL_CHANGE_PUBSUB_POLL_TIMEOUT_SECONDS", 0.5) +MODEL_CHANGE_PUBSUB_RECONNECT_SECONDS = get_env_float("MODEL_CHANGE_PUBSUB_RECONNECT_SECONDS", 5.0) + # APScheduler Configuration - MEMORY LEAK FIX # These settings prevent memory leaks in APScheduler's normalize() and _apply_jitter() functions APSCHEDULER_COALESCE = os.getenv("APSCHEDULER_COALESCE", "True").lower() in [ diff --git a/litellm/litellm_core_utils/env_utils.py b/litellm/litellm_core_utils/env_utils.py index 3a64f44fb25..bc124eb1ae5 100644 --- a/litellm/litellm_core_utils/env_utils.py +++ b/litellm/litellm_core_utils/env_utils.py @@ -21,6 +21,21 @@ def get_env_int(env_var: str, default: int) -> int: return default +def get_env_float(env_var: str, default: float) -> float: + """Parse an environment variable as a float, falling back to default on invalid values. + + Same forgiving behaviour as `get_env_int`, so a typo in a tuning knob cannot crash + the process at import time. + """ + raw = os.getenv(env_var) + if raw is None: + return default + try: + return float(raw.strip()) + except (ValueError, TypeError): + return default + + def get_env_int_or_none(env_var: str) -> int | None: """Parse an environment variable as an integer, returning None when it is unset or unusable. diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 5e3ff8eb7f8..809f3310322 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -22,6 +22,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( reload_serving_verdict, clear_cache, ) +from litellm.proxy.model_change_broadcast import broadcast_model_change from litellm.proxy.utils import PrismaClient from litellm.repositories.model_repository import ModelRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -419,6 +420,7 @@ async def create_model_group( live_before_reload = live_model_ids_snapshot() await clear_cache() + await broadcast_model_change(operation="updated") _raise_http_if_reload_degraded_serving( before=live_before_reload, written_models=updated_pairs, @@ -679,6 +681,7 @@ async def update_access_group( # Clear cache and reload models to pick up the access group changes live_before_reload = live_model_ids_snapshot() await clear_cache() + await broadcast_model_change(operation="updated") _raise_http_if_reload_degraded_serving( before=live_before_reload, written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()), @@ -781,6 +784,7 @@ async def delete_access_group( # Clear cache and reload models to pick up the access group changes live_before_reload = live_model_ids_snapshot() await clear_cache() + await broadcast_model_change(operation="updated") _raise_http_if_reload_degraded_serving( before=live_before_reload, written_models=removed_pairs, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b6422d7f5ae..a2b4bb38667 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team as _legacy_update_team, ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log +from litellm.proxy.model_change_broadcast import broadcast_model_change from litellm.proxy.utils import PrismaClient from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ModelTableRepository @@ -325,6 +326,7 @@ async def patch_model( # Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates) live_before_reload = live_model_ids_snapshot() await clear_cache() + await broadcast_model_change(operation="updated", model_id=model_id) ## CREATE AUDIT LOG ## asyncio.create_task( @@ -430,6 +432,7 @@ async def _set_model_blocked_status( live_before_reload = live_model_ids_snapshot() await clear_cache() + await broadcast_model_change(operation="updated", model_id=data.model_id) asyncio.create_task( create_object_audit_log( @@ -1186,6 +1189,8 @@ async def delete_model( if llm_router is not None: llm_router.delete_deployment(id=model_info.id) + await broadcast_model_change(operation="deleted", model_id=model_info.id) + # Runs after the row delete so the sibling check sees post-delete state. if model_params.model_info.team_id is not None: await _remove_unbacked_team_models( @@ -1370,6 +1375,10 @@ async def add_new_model( prisma_client=prisma_client, ) await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + await broadcast_model_change( + operation="created", + model_id=model_response.model_id if model_response is not None else None, + ) # don't let failed slack alert block the /model/new response _alerting = general_settings.get("alerting", []) or [] if "slack" in _alerting: @@ -1543,6 +1552,7 @@ async def update_model( # Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates) live_before_reload = live_model_ids_snapshot() await clear_cache() + await broadcast_model_change(operation="updated", model_id=_model_id) ## CREATE AUDIT LOG ## asyncio.create_task( create_object_audit_log( diff --git a/litellm/proxy/model_change_broadcast.py b/litellm/proxy/model_change_broadcast.py new file mode 100644 index 00000000000..b5f25ef4ca2 --- /dev/null +++ b/litellm/proxy/model_change_broadcast.py @@ -0,0 +1,230 @@ +""" +Cross-pod invalidation for model CRUD. + +Each pod serves reads (`/model/info`, request routing) from its own in-memory +`llm_router`, reconciled with the DB every `PROXY_CONFIG_RELOAD_INTERVAL_SECONDS`. +A write handled by one pod is therefore invisible to its siblings for up to a full +interval, so a model deleted in the Admin UI keeps showing up on refreshes the load +balancer sends elsewhere. + +The pod that handled the write publishes a notification on the coordination Redis; +subscribers react by re-running the same DB reconcile the timer already performs. +The notification is a trigger, never a payload: siblings always derive state from +the DB, so a dropped, duplicated or reordered message can only cost latency, and the +periodic reconcile remains the backstop. +""" + +import asyncio +import contextlib +from collections.abc import Mapping +from typing import Awaitable, Callable, Final, Literal, Protocol + +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.constants import ( + MODEL_CHANGE_PUBSUB_CHANNEL, + MODEL_CHANGE_PUBSUB_ENABLED, + MODEL_CHANGE_PUBSUB_POLL_TIMEOUT_SECONDS, + MODEL_CHANGE_PUBSUB_RECONNECT_SECONDS, +) + +ModelChangeOperation = Literal["created", "updated", "deleted"] + +PROCESS_ID: Final[str] = str(uuid.uuid4()) + + +class ModelChangeNotification(BaseModel): + model_config = ConfigDict(frozen=True) + + operation: ModelChangeOperation + model_id: str | None = None + origin: str + + +class RedisPubSubConnection(Protocol): + async def subscribe(self, *channels: str) -> None: ... + + async def get_message( + self, + ignore_subscribe_messages: bool = ..., + timeout: float | None = ..., + ) -> Mapping[str, object] | None: ... + + async def aclose(self) -> None: ... + + +class RedisPubSubClient(Protocol): + async def publish(self, channel: str, message: str) -> int: ... + + def pubsub(self) -> RedisPubSubConnection: ... + + +class RedisPubSubBackend(Protocol): + """The slice of `RedisCache` this module needs.""" + + def check_and_fix_namespace(self, key: str) -> str: ... + + def init_async_client(self) -> RedisPubSubClient: ... + + +def _coordination_redis() -> RedisPubSubBackend | None: + from litellm.proxy.proxy_server import redis_usage_cache + + return redis_usage_cache + + +async def broadcast_model_change( + operation: ModelChangeOperation, + model_id: str | None = None, + redis_cache: RedisPubSubBackend | None = None, +) -> None: + """ + Tell sibling pods that the model table changed. Never raises: the write it + follows has already succeeded, and the periodic reconcile still converges. + """ + if not MODEL_CHANGE_PUBSUB_ENABLED: + return + + backend = redis_cache if redis_cache is not None else _coordination_redis() + if backend is None: + return + + notification = ModelChangeNotification(operation=operation, model_id=model_id, origin=PROCESS_ID) + try: + client = backend.init_async_client() + await client.publish( + backend.check_and_fix_namespace(MODEL_CHANGE_PUBSUB_CHANNEL), + notification.model_dump_json(), + ) + except Exception as e: # noqa: BLE001 # no redis error may fail a write that already succeeded + verbose_proxy_logger.warning( + "Could not broadcast model change (%s, model_id=%s) to other pods: %s. " + "They will pick it up on their next config reload.", + operation, + model_id, + str(e), + ) + + +def _parse_notification(payload: object) -> ModelChangeNotification | None: + raw = payload.decode() if isinstance(payload, bytes) else payload + if not isinstance(raw, str): + return None + try: + return ModelChangeNotification.model_validate_json(raw) + except ValidationError: + return None + + +class ModelChangeSubscriber: + """ + Listens for model-change notifications and re-runs `reconcile` for each burst. + + `reconcile` is injected (the proxy passes `ProxyConfig.add_deployment`), so the + subscriber owns no knowledge of how the router is rebuilt. + """ + + def __init__( + self, + redis_cache: RedisPubSubBackend, + reconcile: Callable[[], Awaitable[None]], + origin: str = PROCESS_ID, + poll_timeout_seconds: float = MODEL_CHANGE_PUBSUB_POLL_TIMEOUT_SECONDS, + reconnect_seconds: float = MODEL_CHANGE_PUBSUB_RECONNECT_SECONDS, + ) -> None: + self._redis_cache = redis_cache + self._reconcile = reconcile + self._origin = origin + self._poll_timeout_seconds = poll_timeout_seconds + self._reconnect_seconds = reconnect_seconds + + @property + def channel(self) -> str: + return self._redis_cache.check_and_fix_namespace(MODEL_CHANGE_PUBSUB_CHANNEL) + + async def listen_forever(self) -> None: + while True: + try: + await self.listen_once() + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 # a subscriber that dies leaves this pod on the reload interval + verbose_proxy_logger.warning( + "Model-change subscriber dropped off %s: %s. Reconnecting in %ss; the " + "periodic config reload keeps this pod converging meanwhile.", + self.channel, + str(e), + self._reconnect_seconds, + ) + await asyncio.sleep(self._reconnect_seconds) + + async def listen_once(self) -> None: + """ + One connection's lifetime. Notifications are coalesced: a burst of writes + (deleting several models in a row) triggers a single reconcile once the channel + goes quiet, instead of one full DB reload per message. + """ + connection = self._redis_cache.init_async_client().pubsub() + await connection.subscribe(self.channel) + verbose_proxy_logger.info("Subscribed to model changes from other pods on %s", self.channel) + try: + pending = False + while True: + message = await connection.get_message( + ignore_subscribe_messages=True, + timeout=self._poll_timeout_seconds, + ) + if message is None: + if pending: + pending = False + await self._reconcile() + continue + pending = pending or self._is_foreign_change(message) + finally: + with contextlib.suppress(Exception): + await connection.aclose() + + def _is_foreign_change(self, message: Mapping[str, object]) -> bool: + if message.get("type") != "message": + return False + notification = _parse_notification(message.get("data")) + if notification is None: + verbose_proxy_logger.debug("Ignoring unparseable model-change notification: %s", message.get("data")) + return False + return notification.origin != self._origin + + +class ModelChangeSubscriberHandle: + """ + Owns the subscriber task for the lifetime of the process. Holding the reference + matters: an unreferenced asyncio task can be garbage collected mid-flight. + """ + + def __init__(self) -> None: + self._task: asyncio.Task[None] | None = None + + def start( + self, + redis_cache: RedisPubSubBackend | None, + reconcile: Callable[[], Awaitable[None]], + ) -> None: + if not MODEL_CHANGE_PUBSUB_ENABLED or redis_cache is None: + return + self.stop() + subscriber = ModelChangeSubscriber(redis_cache=redis_cache, reconcile=reconcile) + self._task = asyncio.create_task(subscriber.listen_forever()) + + def stop(self) -> None: + if self._task is None: + return + self._task.cancel() + self._task = None + + @property + def is_running(self) -> bool: + return self._task is not None and not self._task.done() + + +model_change_subscriber = ModelChangeSubscriberHandle() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c72e3d4ee5b..c535112e8b2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -213,7 +213,7 @@ def generate_feedback_box(): import contextlib from collections import defaultdict from contextlib import asynccontextmanager -from functools import lru_cache +from functools import lru_cache, partial import litellm import litellm._redis @@ -468,6 +468,7 @@ from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableRequestMetricsMiddleware, BillingRecorder, ) +from litellm.proxy.model_change_broadcast import model_change_subscriber from litellm.proxy.plugin_routes import ( register_plugins_from_config, ) @@ -797,6 +798,9 @@ def cleanup_router_config_variables(): async def proxy_shutdown_event(): global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") + + model_change_subscriber.stop() + if prisma_client: verbose_proxy_logger.debug("Disconnecting from Prisma") await prisma_client.disconnect() @@ -8122,6 +8126,16 @@ class ProxyStartupEvent: # this will load all existing models on proxy startup await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + ### REACT TO MODEL CHANGES MADE ON OTHER PODS ### + model_change_subscriber.start( + redis_cache=redis_usage_cache, + reconcile=partial( + proxy_config.add_deployment, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ), + ) + ### GET STORED CREDENTIALS ### scheduler.add_job( proxy_config.get_credentials, diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 1e8add52f74..f952c0ab3b7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -1938,6 +1939,100 @@ class TestAddAndDeleteModelLifecycle: assert str(exc_info.value.code) == "400" +class TestModelCrudNotifiesOtherPods: + """Regression: a model deleted (or added) on one pod stayed in every sibling pod's + in-memory router until its next config reload, so the Admin UI kept showing the + deleted model on refreshes the load balancer sent to another pod.""" + + @pytest.mark.asyncio + async def test_add_and_delete_publish_model_change_notifications(self): + import fakeredis.aioredis + + from litellm.constants import MODEL_CHANGE_PUBSUB_CHANNEL + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + add_new_model, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( + delete_model as delete_model_endpoint, + ) + from litellm.proxy.model_change_broadcast import ModelChangeNotification + + class FakeRedisCache: + def __init__(self, client: fakeredis.aioredis.FakeRedis) -> None: + self._client = client + + def check_and_fix_namespace(self, key: str) -> str: + return key + + def init_async_client(self) -> fakeredis.aioredis.FakeRedis: + return self._client + + model_id = "notify-siblings-model-1" + admin_user = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="notify-model", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={"id": model_id}, + created_by="test-admin", + updated_by="test-admin", + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) + + mock_proxy_config = MagicMock() + mock_proxy_config.add_deployment = AsyncMock() + + mock_router = MagicMock() + mock_router.get_model_ids.return_value = [model_id] + + redis_cache = FakeRedisCache(fakeredis.aioredis.FakeRedis(server=fakeredis.FakeServer())) + sibling_pod = redis_cache.init_async_client().pubsub() + await sibling_pod.subscribe(MODEL_CHANGE_PUBSUB_CHANNEL) + + _PS = "litellm.proxy.proxy_server" + _ENCRYPT = "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.proxy_config", mock_proxy_config), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.general_settings", {}), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", mock_router), + patch(f"{_PS}.redis_usage_cache", redis_cache), + patch(_ENCRYPT, side_effect=lambda value, **kwargs: value), + ): + await add_new_model( + model_params=Deployment( + model_name="notify-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"), + model_info={"id": model_id}, + ), + user_api_key_dict=admin_user, + ) + await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=admin_user, + ) + + notifications = [] + deadline = asyncio.get_event_loop().time() + 3.0 + while len(notifications) < 2 and asyncio.get_event_loop().time() < deadline: + message = await sibling_pod.get_message(ignore_subscribe_messages=True, timeout=0.05) + if message is not None: + notifications.append(ModelChangeNotification.model_validate_json(message["data"])) + + assert [notification.operation for notification in notifications] == ["created", "deleted"] + assert {notification.model_id for notification in notifications} == {model_id} + + class TestDeleteTeamBYOKModelGhost: """Regression for issue #22594. diff --git a/tests/test_litellm/proxy/test_model_change_broadcast.py b/tests/test_litellm/proxy/test_model_change_broadcast.py new file mode 100644 index 00000000000..387874b91c5 --- /dev/null +++ b/tests/test_litellm/proxy/test_model_change_broadcast.py @@ -0,0 +1,266 @@ +import asyncio +import json + +import fakeredis.aioredis +import pytest + +from litellm.constants import MODEL_CHANGE_PUBSUB_CHANNEL +from litellm.proxy.model_change_broadcast import ( + ModelChangeNotification, + ModelChangeSubscriber, + ModelChangeSubscriberHandle, + broadcast_model_change, +) + + +class FakeRedisBackend: + """Stands in for `RedisCache`, exposing only what the broadcast module uses.""" + + def __init__(self, client: fakeredis.aioredis.FakeRedis, namespace: str | None = None) -> None: + self._client = client + self._namespace = namespace + + def check_and_fix_namespace(self, key: str) -> str: + if self._namespace is None: + return key + return f"{self._namespace}:{key}" + + def init_async_client(self) -> fakeredis.aioredis.FakeRedis: + return self._client + + +class ExplodingRedisBackend: + def check_and_fix_namespace(self, key: str) -> str: + return key + + def init_async_client(self) -> fakeredis.aioredis.FakeRedis: + raise ConnectionError("redis is down") + + +def _backend(namespace: str | None = None) -> FakeRedisBackend: + server = fakeredis.FakeServer() + return FakeRedisBackend(fakeredis.aioredis.FakeRedis(server=server), namespace=namespace) + + +async def _drain(pubsub, expected: int, timeout: float = 2.0) -> list[dict]: + deadline = asyncio.get_event_loop().time() + timeout + messages: list[dict] = [] + while len(messages) < expected and asyncio.get_event_loop().time() < deadline: + message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=0.05) + if message is not None: + messages.append(message) + return messages + + +@pytest.mark.asyncio +async def test_broadcast_publishes_notification_on_channel(): + backend = _backend() + pubsub = backend.init_async_client().pubsub() + await pubsub.subscribe(MODEL_CHANGE_PUBSUB_CHANNEL) + + await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=backend) + + messages = await _drain(pubsub, expected=1) + assert len(messages) == 1 + notification = ModelChangeNotification.model_validate_json(messages[0]["data"]) + assert notification.operation == "deleted" + assert notification.model_id == "model-1" + assert notification.origin != "" + + +@pytest.mark.asyncio +async def test_broadcast_respects_redis_namespace(): + backend = _backend(namespace="tenant-a") + pubsub = backend.init_async_client().pubsub() + await pubsub.subscribe(f"tenant-a:{MODEL_CHANGE_PUBSUB_CHANNEL}") + + await broadcast_model_change(operation="created", model_id="model-1", redis_cache=backend) + + assert len(await _drain(pubsub, expected=1)) == 1 + + +@pytest.mark.asyncio +async def test_broadcast_never_raises_when_redis_is_unavailable(): + await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=ExplodingRedisBackend()) + + +@pytest.mark.asyncio +async def test_broadcast_is_a_noop_without_a_coordination_redis(): + await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=None) + + +class ReconcileSpy: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self) -> None: + self.calls += 1 + + +async def _run_subscriber(subscriber: ModelChangeSubscriber) -> "asyncio.Task[None]": + task = asyncio.create_task(subscriber.listen_once()) + await asyncio.sleep(0.2) + return task + + +async def _stop(task: "asyncio.Task[None]") -> None: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +async def _wait_for_reconciles(spy: ReconcileSpy, expected: int, timeout: float = 3.0) -> None: + deadline = asyncio.get_event_loop().time() + timeout + while spy.calls < expected and asyncio.get_event_loop().time() < deadline: + await asyncio.sleep(0.05) + + +@pytest.mark.asyncio +async def test_subscriber_reconciles_on_a_change_from_another_pod(): + """The regression: a model deleted on pod A must reach pod B's router without + waiting for the periodic config reload.""" + backend = _backend() + spy = ReconcileSpy() + subscriber = ModelChangeSubscriber( + redis_cache=backend, + reconcile=spy, + origin="pod-b", + poll_timeout_seconds=0.05, + ) + task = await _run_subscriber(subscriber) + + await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=backend) + await _wait_for_reconciles(spy, expected=1) + await _stop(task) + + assert spy.calls == 1 + + +@pytest.mark.asyncio +async def test_subscriber_ignores_its_own_changes(): + backend = _backend() + spy = ReconcileSpy() + origin = "pod-a" + subscriber = ModelChangeSubscriber( + redis_cache=backend, + reconcile=spy, + origin=origin, + poll_timeout_seconds=0.05, + ) + task = await _run_subscriber(subscriber) + + own_notification = ModelChangeNotification(operation="deleted", model_id="model-1", origin=origin) + await backend.init_async_client().publish(subscriber.channel, own_notification.model_dump_json()) + await asyncio.sleep(0.5) + await _stop(task) + + assert spy.calls == 0 + + +@pytest.mark.asyncio +async def test_subscriber_coalesces_a_burst_into_one_reconcile(): + backend = _backend() + spy = ReconcileSpy() + subscriber = ModelChangeSubscriber( + redis_cache=backend, + reconcile=spy, + origin="pod-b", + poll_timeout_seconds=0.05, + ) + task = await _run_subscriber(subscriber) + + for model_id in ("model-1", "model-2", "model-3"): + await broadcast_model_change(operation="deleted", model_id=model_id, redis_cache=backend) + await _wait_for_reconciles(spy, expected=1) + await asyncio.sleep(0.3) + await _stop(task) + + assert spy.calls == 1 + + +@pytest.mark.asyncio +async def test_handle_reconciles_while_running_and_stops_cleanly(): + backend = _backend() + spy = ReconcileSpy() + handle = ModelChangeSubscriberHandle() + + async def publish_from_another_pod(model_id: str) -> None: + notification = ModelChangeNotification(operation="deleted", model_id=model_id, origin="pod-a") + await backend.init_async_client().publish( + backend.check_and_fix_namespace(MODEL_CHANGE_PUBSUB_CHANNEL), + notification.model_dump_json(), + ) + + handle.start(redis_cache=backend, reconcile=spy) + assert handle.is_running + await asyncio.sleep(0.2) + await publish_from_another_pod("model-1") + await _wait_for_reconciles(spy, expected=1) + + handle.stop() + await asyncio.sleep(0.1) + assert not handle.is_running + assert spy.calls == 1 + + await publish_from_another_pod("model-2") + await asyncio.sleep(0.5) + assert spy.calls == 1 + + +@pytest.mark.asyncio +async def test_handle_is_a_noop_without_a_coordination_redis(): + handle = ModelChangeSubscriberHandle() + handle.start(redis_cache=None, reconcile=ReconcileSpy()) + assert not handle.is_running + + +@pytest.mark.asyncio +async def test_subscriber_ignores_unparseable_payloads(): + backend = _backend() + spy = ReconcileSpy() + subscriber = ModelChangeSubscriber( + redis_cache=backend, + reconcile=spy, + origin="pod-b", + poll_timeout_seconds=0.05, + ) + task = await _run_subscriber(subscriber) + client = backend.init_async_client() + + await client.publish(subscriber.channel, "not json") + await client.publish(subscriber.channel, json.dumps({"operation": "exploded"})) + await asyncio.sleep(0.5) + await _stop(task) + + assert spy.calls == 0 + + +@pytest.mark.asyncio +async def test_subscriber_resubscribes_after_a_failed_reconcile(): + """A reconcile that blows up (DB hiccup) must not take the subscriber down for + the rest of the pod's life.""" + backend = _backend() + calls: list[int] = [] + + async def flaky_reconcile() -> None: + calls.append(len(calls)) + if len(calls) == 1: + raise RuntimeError("db unavailable") + + subscriber = ModelChangeSubscriber( + redis_cache=backend, + reconcile=flaky_reconcile, + origin="pod-b", + poll_timeout_seconds=0.05, + reconnect_seconds=0.05, + ) + task = asyncio.create_task(subscriber.listen_forever()) + await asyncio.sleep(0.2) + + await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=backend) + await asyncio.sleep(0.5) + await broadcast_model_change(operation="deleted", model_id="model-2", redis_cache=backend) + await asyncio.sleep(0.5) + await _stop(task) + + assert len(calls) == 2