From 292161f766ca7ac88cf3cdbef0c4b599a5576a88 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:25:10 -0700 Subject: [PATCH 1/6] fix(proxy): read through to the DB on registry misses so just-created models, guardrails, and agents resolve on sibling replicas --- .../proxy/agent_endpoints/a2a_endpoints.py | 15 +- litellm/proxy/agent_endpoints/a2a_routing.py | 6 +- .../common_utils/registry_read_through.py | 143 +++++++++++ .../proxy/guardrails/guardrail_endpoints.py | 6 +- ...model_access_group_management_endpoints.py | 30 ++- litellm/proxy/route_llm_request.py | 232 ++++++++++-------- ruff.toml | 2 +- .../test_registry_read_through.py | 231 +++++++++++++++++ .../test_access_group_management.py | 96 ++++++++ .../proxy/test_route_a2a_models.py | 75 ++++++ .../proxy/test_route_llm_request.py | 121 +++++++++ 11 files changed, 843 insertions(+), 114 deletions(-) create mode 100644 litellm/proxy/common_utils/registry_read_through.py create mode 100644 tests/test_litellm/proxy/common_utils/test_registry_read_through.py diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 27780aeb994..a4e1ac126d9 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -152,14 +152,13 @@ def _jsonrpc_error( ) -def _get_agent(agent_id: str): +async def _get_agent(agent_id: str) -> "AgentResponse | None": """Look up an agent by ID or name. Returns None if not found.""" - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.common_utils.registry_read_through import ( + get_agent_with_read_through, + ) - agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) - if agent is None: - agent = global_agent_registry.get_agent_by_name(agent_name=agent_id) - return agent + return await get_agent_with_read_through(agent_id) def _enforce_inbound_trace_id(agent: Any, request: Request) -> None: @@ -531,7 +530,7 @@ async def get_agent_card( ) try: - agent: Final = _get_agent(agent_id) + agent: Final = await _get_agent(agent_id) if agent is None: raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found") @@ -645,7 +644,7 @@ async def invoke_agent_a2a( params.pop(key) # Find the agent - agent: Final = _get_agent(agent_id) + agent: Final = await _get_agent(agent_id) if agent is None: return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 038b6b4a840..2228735d805 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -25,10 +25,12 @@ async def route_a2a_agent_request( Returns None if not an A2A request (allows normal routing to continue). """ # Import here to avoid circular imports - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, ) + from litellm.proxy.common_utils.registry_read_through import ( + get_agent_with_read_through, + ) from litellm.proxy.route_llm_request import ( ROUTE_ENDPOINT_MAPPING, ProxyModelNotFoundError, @@ -44,7 +46,7 @@ async def route_a2a_agent_request( agent_name: Final = model_name[4:] # Look up agent in registry - agent: Final = global_agent_registry.get_agent_by_name(agent_name) + agent: Final = await get_agent_with_read_through(agent_name) if agent is None: verbose_proxy_logger.error("[A2A] Agent '%s' not found in registry", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py new file mode 100644 index 00000000000..b78106205d4 --- /dev/null +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -0,0 +1,143 @@ +"""Read-through recovery for in-memory registries in multi-replica deployments. + +A management write (POST /model/new, /guardrails, /v1/agents) lands on one +replica and reaches Postgres, but sibling replicas only refresh their in-memory +registries on the periodic config reload or the Redis config-sync resync, both +of which lag by seconds. A request that uses the new object immediately can +land on a sibling that has never heard of it and fail with a 400/404. + +On a registry miss, callers here fetch the missing object from the DB and load +it into the local registry before giving up. A short negative-result TTL keeps +repeated lookups of genuinely unknown names from hammering the DB. +""" + +import asyncio +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.agents import AgentResponse + +READ_THROUGH_MISS_TTL_SECONDS: Final = 2.0 + + +class RegistryReadThrough: + __slots__ = ("_lock", "_miss_ttl_seconds", "_recent_misses", "_resync") + + def __init__( + self, + resync: Callable[[str], Awaitable[bool]], + miss_ttl_seconds: float = READ_THROUGH_MISS_TTL_SECONDS, + ) -> None: + self._resync = resync + self._miss_ttl_seconds = miss_ttl_seconds + self._lock = asyncio.Lock() + self._recent_misses = InMemoryCache(max_size_in_memory=1000) + + async def attempt(self, key: str) -> bool: + if self._recent_misses.get_cache(key) is not None: + return False + async with self._lock: + if self._recent_misses.get_cache(key) is not None: + return False + try: + found: Final = await self._resync(key) + except Exception as e: # noqa: BLE001 # a failed read-through must surface the original miss error, not a 500 + verbose_proxy_logger.warning("registry read-through for %r failed: %s", key, e) + return False + if not found: + self._recent_misses.set_cache(key, True, ttl=self._miss_ttl_seconds) + return found + + +def _db_backed_registries_enabled() -> bool: + from litellm.proxy import proxy_server + + return proxy_server.prisma_client is not None and proxy_server.store_model_in_db is True + + +async def _resync_model_deployments(model_name: str) -> bool: + from litellm.proxy import proxy_server + from litellm.repositories.model_repository import ModelRepository + + if not _db_backed_registries_enabled(): + return False + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + rows: Final = await ModelRepository(prisma_client).table.find_many( + where={"OR": [{"model_name": model_name}, {"model_id": model_name}]} + ) + if not rows: + return False + if proxy_server.llm_router is None: + await proxy_server.proxy_config.add_deployment( + prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj + ) + return proxy_server.llm_router is not None + proxy_server.proxy_config._add_deployment(db_models=rows) + proxy_server.llm_model_list = proxy_server.llm_router.get_model_list() + return True + + +async def _resync_guardrails(guardrail_name: str) -> bool: + from litellm.proxy import proxy_server + + if not _db_backed_registries_enabled(): + return False + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + await proxy_server.proxy_config._init_guardrails_in_db(prisma_client=prisma_client) + return _initialized_guardrail(guardrail_name) is not None + + +async def _resync_agents(agent_id_or_name: str) -> bool: + from litellm.proxy import proxy_server + + if not _db_backed_registries_enabled(): + return False + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + await proxy_server.proxy_config._init_agents_in_db(prisma_client=prisma_client) + return _agent_from_registry(agent_id_or_name) is not None + + +model_registry_read_through: Final = RegistryReadThrough(resync=_resync_model_deployments) +guardrail_registry_read_through: Final = RegistryReadThrough(resync=_resync_guardrails) +agent_registry_read_through: Final = RegistryReadThrough(resync=_resync_agents) + + +def _agent_from_registry(agent_id_or_name: str) -> "AgentResponse | None": + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + by_id: Final = global_agent_registry.get_agent_by_id(agent_id=agent_id_or_name) + if by_id is not None: + return by_id + return global_agent_registry.get_agent_by_name(agent_name=agent_id_or_name) + + +async def get_agent_with_read_through(agent_id_or_name: str) -> "AgentResponse | None": + agent: Final = _agent_from_registry(agent_id_or_name) + if agent is not None: + return agent + if not await agent_registry_read_through.attempt(agent_id_or_name): + return None + return _agent_from_registry(agent_id_or_name) + + +def _initialized_guardrail(guardrail_name: str) -> "CustomGuardrail | None": + from litellm.proxy.guardrails import guardrail_endpoints + + return guardrail_endpoints.GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(guardrail_name=guardrail_name) + + +async def get_initialized_guardrail_with_read_through(guardrail_name: str) -> "CustomGuardrail | None": + active: Final = _initialized_guardrail(guardrail_name) + if active is not None: + return active + if not await guardrail_registry_read_through.attempt(guardrail_name): + return None + return _initialized_guardrail(guardrail_name) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 761d8aabc8a..dff70ccf68d 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2244,8 +2244,12 @@ async def apply_guardrail( litellm_logging_obj = None start_time: Final = datetime.now(timezone.utc) + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + try: - active_guardrail: Final[CustomGuardrail | None] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( + active_guardrail: Final[CustomGuardrail | None] = await get_initialized_guardrail_with_read_through( guardrail_name=request.guardrail_name ) if active_guardrail is None: 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 7051f705a03..75d33c6c40a 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -7,7 +7,10 @@ Endpoints here: import json from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final + +if TYPE_CHECKING: + from litellm.router import Router from fastapi import APIRouter, Depends, HTTPException @@ -52,6 +55,23 @@ def validate_models_exist(model_names: list[str], llm_router) -> tuple[bool, lis return (len(missing) == 0, missing) +async def _missing_models_after_read_through( + model_names: Sequence[str], llm_router: "Router | None" +) -> tuple[str, ...]: + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + model_registry_read_through, + ) + + _, missing = validate_models_exist(model_names=list(model_names), llm_router=llm_router) + if not missing: + return () + for name in missing: + await model_registry_read_through.attempt(name) + _, still_missing = validate_models_exist(model_names=list(model_names), llm_router=proxy_server.llm_router) + return tuple(still_missing) + + def add_access_group_to_deployment(model_info: dict[str, Any], access_group: str) -> tuple[dict[str, Any], bool]: """ Add an access group to a deployment's model_info. @@ -369,12 +389,12 @@ async def create_model_group( # Validate model_names exist in router (only if using model_names path) if not use_model_ids and has_model_names: assert data.model_names is not None - all_valid, missing_models = validate_models_exist( + missing_models: Final = await _missing_models_after_read_through( model_names=data.model_names, llm_router=llm_router, ) - if not all_valid: + if missing_models: raise HTTPException( status_code=400, detail={"error": f"Model(s) not found: {', '.join(missing_models)}"}, @@ -633,12 +653,12 @@ async def update_access_group( # Validation: Check if all new models exist (only if using model_names path) if not use_model_ids and has_model_names: assert data.model_names is not None - all_valid, missing_models = validate_models_exist( + missing_models: Final = await _missing_models_after_read_through( model_names=data.model_names, llm_router=llm_router, ) - if not all_valid: + if missing_models: raise HTTPException( status_code=400, detail={"error": f"Model(s) not found: {', '.join(missing_models)}"}, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index dd8deed57f1..657e0cbcafa 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -313,112 +313,150 @@ async def add_shared_session_to_data(data: dict) -> None: pass +RouteType = Literal[ + "acompletion", + "atext_completion", + "aembedding", + "aimage_generation", + "aspeech", + "atranscription", + "amoderation", + "arerank", + "aresponses", + "aget_responses", + "adelete_responses", + "acancel_responses", + "acompact_responses", + "acreate_response_reply", + "alist_input_items", + "_arealtime", # private function for realtime API + "acreate_realtime_client_secret", + "arealtime_calls", + "acreate_realtime_transcription_session", + "_aresponses_websocket", # private function for responses WebSocket mode + "aimage_edit", + "agenerate_content", + "agenerate_content_stream", + "allm_passthrough_route", + "acreate_batch", + "aretrieve_batch", + "alist_batches", + "afile_content", + "afile_retrieve", + "acreate_fine_tuning_job", + "acancel_fine_tuning_job", + "alist_fine_tuning_jobs", + "aretrieve_fine_tuning_job", + "avector_store_search", + "avector_store_create", + "avector_store_retrieve", + "avector_store_list", + "avector_store_update", + "avector_store_delete", + "avector_store_file_create", + "avector_store_file_list", + "avector_store_file_retrieve", + "avector_store_file_content", + "avector_store_file_update", + "avector_store_file_delete", + "aocr", + "asearch", + "avideo_generation", + "avideo_list", + "avideo_status", + "avideo_content", + "avideo_remix", + "avideo_create_character", + "avideo_get_character", + "avideo_edit", + "avideo_extension", + "acreate_container", + "alist_containers", + "aretrieve_container", + "adelete_container", + "aupload_container_file", + "alist_container_files", + "aretrieve_container_file", + "adelete_container_file", + "aretrieve_container_file_content", + "acreate_skill", + "alist_skills", + "aget_skill", + "adelete_skill", + "aingest", + "anthropic_messages", + "acreate_interaction", + "aget_interaction", + "adelete_interaction", + "acancel_interaction", + "acreate_agent", + "alist_agents", + "aget_agent", + "adelete_agent", + "alist_agent_versions", + "asend_message", + "call_mcp_tool", + "acancel_batch", + "afile_delete", + "acreate_eval", + "alist_evals", + "aget_eval", + "aupdate_eval", + "adelete_eval", + "acancel_eval", + "acreate_run", + "alist_runs", + "aget_run", + "acancel_run", + "adelete_run", +] + + async def route_request( data: dict, llm_router: LitellmRouter | None, user_model: str | None, - route_type: Literal[ - "acompletion", - "atext_completion", - "aembedding", - "aimage_generation", - "aspeech", - "atranscription", - "amoderation", - "arerank", - "aresponses", - "aget_responses", - "adelete_responses", - "acancel_responses", - "acompact_responses", - "acreate_response_reply", - "alist_input_items", - "_arealtime", # private function for realtime API - "acreate_realtime_client_secret", - "arealtime_calls", - "acreate_realtime_transcription_session", - "_aresponses_websocket", # private function for responses WebSocket mode - "aimage_edit", - "agenerate_content", - "agenerate_content_stream", - "allm_passthrough_route", - "acreate_batch", - "aretrieve_batch", - "alist_batches", - "afile_content", - "afile_retrieve", - "acreate_fine_tuning_job", - "acancel_fine_tuning_job", - "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job", - "avector_store_search", - "avector_store_create", - "avector_store_retrieve", - "avector_store_list", - "avector_store_update", - "avector_store_delete", - "avector_store_file_create", - "avector_store_file_list", - "avector_store_file_retrieve", - "avector_store_file_content", - "avector_store_file_update", - "avector_store_file_delete", - "aocr", - "asearch", - "avideo_generation", - "avideo_list", - "avideo_status", - "avideo_content", - "avideo_remix", - "avideo_create_character", - "avideo_get_character", - "avideo_edit", - "avideo_extension", - "acreate_container", - "alist_containers", - "aretrieve_container", - "adelete_container", - "aupload_container_file", - "alist_container_files", - "aretrieve_container_file", - "adelete_container_file", - "aretrieve_container_file_content", - "acreate_skill", - "alist_skills", - "aget_skill", - "adelete_skill", - "aingest", - "anthropic_messages", - "acreate_interaction", - "aget_interaction", - "adelete_interaction", - "acancel_interaction", - "acreate_agent", - "alist_agents", - "aget_agent", - "adelete_agent", - "alist_agent_versions", - "asend_message", - "call_mcp_tool", - "acancel_batch", - "afile_delete", - "acreate_eval", - "alist_evals", - "aget_eval", - "aupdate_eval", - "adelete_eval", - "acancel_eval", - "acreate_run", - "alist_runs", - "aget_run", - "acancel_run", - "adelete_run", - ], + route_type: RouteType, user_api_key_dict: UserAPIKeyAuth | None = None, ): """ Common helper to route the request """ + try: + return await _route_request_single_attempt( + data=data, + llm_router=llm_router, + user_model=user_model, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + ) + except ProxyModelNotFoundError: + requested_model: Final = data.get("model", "") + if not isinstance(requested_model, str) or not requested_model: + raise + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + model_registry_read_through, + ) + + if not await model_registry_read_through.attempt(requested_model): + raise + return await _route_request_single_attempt( + data=data, + llm_router=proxy_server.llm_router, + user_model=user_model, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + ) + + +async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited provider coroutines; the inferred union keeps route_request's callers typed + data: dict, # noqa: LIT001 # request body is the proxy-wide mutable dict contract shared with route_request + llm_router: LitellmRouter | None, + user_model: str | None, + route_type: RouteType, + user_api_key_dict: UserAPIKeyAuth | None = None, +): raise_if_required_body_param_missing(route_type=route_type, data=data) await add_shared_session_to_data(data) diff --git a/ruff.toml b/ruff.toml index 095e3e24c52..bd3f8334e94 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", "TID251", + "ANN202", "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_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py new file mode 100644 index 00000000000..e1c8f031579 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -0,0 +1,231 @@ +import asyncio +from typing import Final + +import pytest + +from litellm.proxy.common_utils.registry_read_through import RegistryReadThrough + + +class ResyncSpy: + def __init__(self, found: bool = True, error: Exception | None = None) -> None: + self.found = found + self.error = error + self.calls: list[str] = [] + + async def __call__(self, key: str) -> bool: + self.calls.append(key) + if self.error is not None: + raise self.error + return self.found + + +@pytest.mark.asyncio +async def test_attempt_returns_true_when_resync_finds_object(): + spy: Final = ResyncSpy(found=True) + read_through: Final = RegistryReadThrough(resync=spy) + + assert await read_through.attempt("new-model") is True + assert spy.calls == ["new-model"] + + +@pytest.mark.asyncio +async def test_attempt_found_key_is_not_negative_cached(): + spy: Final = ResyncSpy(found=True) + read_through: Final = RegistryReadThrough(resync=spy) + + assert await read_through.attempt("new-model") is True + assert await read_through.attempt("new-model") is True + assert spy.calls == ["new-model", "new-model"] + + +@pytest.mark.asyncio +async def test_missing_key_is_negative_cached_within_ttl(): + spy: Final = ResyncSpy(found=False) + read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + + assert await read_through.attempt("ghost-model") is False + assert await read_through.attempt("ghost-model") is False + assert spy.calls == ["ghost-model"] + + +@pytest.mark.asyncio +async def test_negative_cache_expires_and_resync_runs_again(): + spy: Final = ResyncSpy(found=False) + read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=0.05) + + assert await read_through.attempt("ghost-model") is False + await asyncio.sleep(0.1) + assert await read_through.attempt("ghost-model") is False + assert spy.calls == ["ghost-model", "ghost-model"] + + +@pytest.mark.asyncio +async def test_resync_exception_returns_false_without_negative_caching(): + spy: Final = ResyncSpy(error=RuntimeError("db down")) + read_through: Final = RegistryReadThrough(resync=spy) + + assert await read_through.attempt("new-model") is False + assert await read_through.attempt("new-model") is False + assert spy.calls == ["new-model", "new-model"] + + +@pytest.mark.asyncio +async def test_concurrent_attempts_for_missing_key_resync_once(): + class SlowResyncSpy(ResyncSpy): + async def __call__(self, key: str) -> bool: + await asyncio.sleep(0.05) + return await super().__call__(key) + + spy: Final = SlowResyncSpy(found=False) + read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + + results: Final = await asyncio.gather(*(read_through.attempt("ghost-model") for _ in range(5))) + assert results == [False] * 5 + assert spy.calls == ["ghost-model"] + + +@pytest.mark.asyncio +async def test_distinct_keys_do_not_share_negative_cache(): + spy: Final = ResyncSpy(found=False) + read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + + assert await read_through.attempt("ghost-a") is False + assert await read_through.attempt("ghost-b") is False + assert spy.calls == ["ghost-a", "ghost-b"] + + +class FakeAgentRow: + def __init__(self, agent_id: str, agent_name: str) -> None: + self.agent_id = agent_id + self.agent_name = agent_name + self.object_permission = None + self.spend = 0.0 + + def __iter__(self): + return iter( + { + "agent_id": self.agent_id, + "agent_name": self.agent_name, + "agent_card_params": {"name": self.agent_name, "url": "http://db-agent"}, + "litellm_params": {}, + }.items() + ) + + +@pytest.fixture +def clean_agent_registry(): + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + original_agents: Final = list(global_agent_registry.agent_list) + original_config_agents: Final = getattr(global_agent_registry, "config_agents", ()) + global_agent_registry.agent_list = [] + global_agent_registry.config_agents = () + try: + yield global_agent_registry + finally: + global_agent_registry.agent_list = original_agents + global_agent_registry.config_agents = original_config_agents + + +@pytest.mark.asyncio +async def test_get_agent_with_read_through_recovers_agent_created_on_sibling_replica( + clean_agent_registry, monkeypatch +): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent_id: Final = "read-through-db-agent-id" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_agentstable.find_many = AsyncMock( + return_value=[FakeAgentRow(agent_id, "read-through-db-agent")] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + assert clean_agent_registry.get_agent_by_id(agent_id=agent_id) is None + agent: Final = await get_agent_with_read_through(agent_id) + + assert agent is not None + assert agent.agent_id == agent_id + + +@pytest.mark.asyncio +async def test_get_agent_with_read_through_returns_none_for_unknown_agent(clean_agent_registry, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_agentstable.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + assert await get_agent_with_read_through("agent-nobody-created") is None + + +class FakeGuardrailRow: + def __init__(self, guardrail_id: str, guardrail_name: str) -> None: + self.guardrail_id = guardrail_id + self.guardrail_name = guardrail_name + + def __iter__(self): + return iter( + { + "guardrail_id": self.guardrail_id, + "guardrail_name": self.guardrail_name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "default_on": True, + "blocked_words": [{"keyword": "secret", "action": "BLOCK"}], + }, + "guardrail_info": {}, + }.items() + ) + + +@pytest.mark.asyncio +async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sibling_replica(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + guardrail_id: Final = "read-through-db-guardrail-id" + guardrail_name: Final = "read-through-db-guardrail" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[FakeGuardrailRow(guardrail_id, guardrail_name)] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + try: + guardrail: Final = await get_initialized_guardrail_with_read_through(guardrail_name=guardrail_name) + assert guardrail is not None + assert guardrail.guardrail_name == guardrail_name + finally: + IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id) + + +@pytest.mark.asyncio +async def test_get_guardrail_with_read_through_returns_none_for_unknown_guardrail(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + assert await get_initialized_guardrail_with_read_through(guardrail_name="guardrail-nobody-created") is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index 3240ad20edb..1722d8c377b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -430,3 +430,99 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): assert response.models_updated == 1 mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_create_access_group_read_through_recovers_model_created_on_sibling_replica(): + """Regression: an access group referencing a model that another replica just wrote + to the DB must be created instead of 400ing until the periodic config reload.""" + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + create_model_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + NewModelGroupRequest, + ) + + from types import SimpleNamespace + + model_name = "e2e-ag-sibling-replica-model" + db_row = SimpleNamespace( + model_id=f"{model_name}-id", + model_name=model_name, + litellm_params={"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "hi"}, + model_info={}, + blocked=False, + ) + + mock_router = Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=[[db_row], [], [db_row]]) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch( + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock(return_value=None), + ), + ): + response = await create_model_group( + data=NewModelGroupRequest(access_group="replica-lag-group", model_names=[model_name]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response.models_updated == 1 + assert response.model_names == [model_name] + assert mock_prisma.db.litellm_proxymodeltable.find_many.await_args_list[0].kwargs["where"] == { + "OR": [{"model_name": model_name}, {"model_id": model_name}] + } + + +@pytest.mark.asyncio +async def test_create_access_group_model_missing_everywhere_still_400s(): + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + create_model_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + NewModelGroupRequest, + ) + + model_name = "e2e-ag-model-nobody-created" + mock_router = Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + with pytest.raises(HTTPException) as exc_info: + await create_model_group( + data=NewModelGroupRequest(access_group="ghost-group", model_names=[model_name]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert exc_info.value.status_code == 400 + assert model_name in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 616fa62cda5..22f99d03d21 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -49,6 +49,7 @@ async def test_route_a2a_model_bypasses_router(): ) mock_registry = Mock() + mock_registry.get_agent_by_id = Mock(return_value=None) mock_registry.get_agent_by_name = Mock(return_value=mock_agent) # Mock litellm.acompletion to verify it's called @@ -104,3 +105,77 @@ async def test_route_non_a2a_model_raises_error_if_not_in_router(): user_model=None, route_type="acompletion", ) + + +class _DbAgentRow: + def __init__(self, agent_id: str, agent_name: str) -> None: + self.agent_id = agent_id + self.agent_name = agent_name + self.object_permission = None + self.spend = 0.0 + + def __iter__(self): + return iter( + { + "agent_id": self.agent_id, + "agent_name": self.agent_name, + "agent_card_params": {"name": self.agent_name, "url": "http://sibling-db-agent.example.com"}, + "litellm_params": {}, + }.items() + ) + + +def _router_without_models(): + mock_router = Mock() + mock_router.model_names = [] + mock_router.deployment_names = [] + mock_router.has_model_id = Mock(return_value=False) + mock_router.model_group_alias = None + mock_router.router_general_settings = Mock(pass_through_all_models=False) + mock_router.default_deployment = None + mock_router.pattern_router = Mock(patterns=[]) + mock_router.map_team_model = Mock(return_value=None) + return mock_router + + +@pytest.mark.asyncio +async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_replica(monkeypatch): + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + agent_name = "a2a-sibling-replica-agent" + prisma_client = Mock() + prisma_client.db.litellm_agentstable.find_many = AsyncMock( + return_value=[_DbAgentRow("a2a-sibling-replica-agent-id", agent_name)] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + original_agents = list(global_agent_registry.agent_list) + original_config_agents = getattr(global_agent_registry, "config_agents", ()) + global_agent_registry.agent_list = [] + global_agent_registry.config_agents = () + + data = { + "model": f"a2a/{agent_name}", + "messages": [{"role": "user", "content": "Hello"}], + } + mock_acompletion = AsyncMock(return_value={"id": "read-through-response"}) + + try: + with patch("litellm.acompletion", mock_acompletion): + await route_request( + data=data, + llm_router=_router_without_models(), + user_model=None, + route_type="acompletion", + ) + finally: + global_agent_registry.agent_list = original_agents + global_agent_registry.config_agents = original_config_agents + + mock_acompletion.assert_called_once() + call_kwargs = mock_acompletion.call_args.kwargs + assert call_kwargs["model"] == f"a2a/{agent_name}" + assert call_kwargs["api_base"] == "http://sibling-db-agent.example.com" + prisma_client.db.litellm_agentstable.find_many.assert_awaited() diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 3ae0e1e7d18..ebd52c448bc 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1091,3 +1091,124 @@ async def test_route_request_rejects_chat_completion_without_messages(): assert exc_info.value.status_code == 400 assert exc_info.value.param == "messages" llm_router.acompletion.assert_not_called() + + +class FakeProxyModelTable: + def __init__(self, rows): + self.rows = rows + self.find_many_wheres = [] + + async def find_many(self, where=None, **kwargs): + self.find_many_wheres.append(where) + return list(self.rows) + + +def _fake_prisma_client_with_models(rows): + from types import SimpleNamespace + + table = FakeProxyModelTable(rows) + return SimpleNamespace(db=SimpleNamespace(litellm_proxymodeltable=table)), table + + +def _db_model_row(model_name: str, mock_response: str): + from types import SimpleNamespace + + return SimpleNamespace( + model_id=f"{model_name}-id", + model_name=model_name, + litellm_params={"model": "openai/gpt-4o", "api_key": "fake", "mock_response": mock_response}, + model_info={}, + blocked=False, + ) + + +@pytest.mark.asyncio +async def test_route_request_read_through_recovers_model_created_on_sibling_replica(monkeypatch): + """Regression: a model written to the DB by another replica must be served on + first request instead of 400ing until the periodic config reload.""" + import litellm + import litellm.proxy.proxy_server as proxy_server + + model_name = "e2e-sibling-replica-model" + router = litellm.Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + fake_prisma, table = _fake_prisma_client_with_models([_db_model_row(model_name, "hello-from-db")]) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + llm_call = await route_request( + data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]}, + llm_router=router, + user_model=None, + route_type="acompletion", + ) + response = await llm_call + + assert response.choices[0].message.content == "hello-from-db" + assert len(table.find_many_wheres) == 1 + assert table.find_many_wheres[0] == {"OR": [{"model_name": model_name}, {"model_id": model_name}]} + + +@pytest.mark.asyncio +async def test_route_request_unknown_model_raises_and_hits_db_once_within_ttl(monkeypatch): + import litellm + import litellm.proxy.proxy_server as proxy_server + + model_name = "e2e-model-nobody-created" + router = litellm.Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + fake_prisma, table = _fake_prisma_client_with_models([]) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + data = {"model": model_name, "messages": [{"role": "user", "content": "hi"}]} + with pytest.raises(ProxyModelNotFoundError): + await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion") + with pytest.raises(ProxyModelNotFoundError): + await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion") + + assert len(table.find_many_wheres) == 1 + + +@pytest.mark.asyncio +async def test_route_request_read_through_disabled_without_store_model_in_db(monkeypatch): + import litellm + import litellm.proxy.proxy_server as proxy_server + + model_name = "e2e-config-only-proxy-model" + router = litellm.Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + fake_prisma, table = _fake_prisma_client_with_models([_db_model_row(model_name, "should-not-load")]) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", False) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with pytest.raises(ProxyModelNotFoundError): + await route_request( + data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]}, + llm_router=router, + user_model=None, + route_type="acompletion", + ) + + assert table.find_many_wheres == [] From ac2db91b0616c156c99b9e5494a005f4e50e74d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:02:12 -0700 Subject: [PATCH 2/6] fix(proxy): single-row read-through resyncs and reload-race hardening Resync registry misses with single-row DB fetches (guardrail by unique name, agent by unique id or name, model by name then id) instead of full-table loads, and bound them with a global budget of 20 resyncs per 5s window per registry that fails closed without negative-caching the key. Access group create/update now trust the reconcile outcome snapshot captured under the reload lock instead of a post-lock router read, so a concurrent reconcile can no longer surface a false degraded-serving 500. Router.upsert_deployment restores the previously served deployment when the replacement add fails under ignore_invalid_deployments, so a bad update no longer silently drops a healthy deployment from serving. --- .../common_utils/registry_read_through.py | 95 ++++++++++++--- ...model_access_group_management_endpoints.py | 36 ++++-- litellm/proxy/route_llm_request.py | 2 +- litellm/router.py | 27 ++++- ruff-strict-budget.json | 2 +- .../test_registry_read_through.py | 111 +++++++++++++++--- .../test_access_group_management.py | 63 ++++++++-- .../proxy/test_route_a2a_models.py | 24 ++-- .../proxy/test_route_llm_request.py | 4 +- tests/test_litellm/test_router.py | 56 +++++++++ type-discipline-budget.json | 6 +- 11 files changed, 356 insertions(+), 70 deletions(-) diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index b78106205d4..7ace046ae26 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -2,16 +2,16 @@ A management write (POST /model/new, /guardrails, /v1/agents) lands on one replica and reaches Postgres, but sibling replicas only refresh their in-memory -registries on the periodic config reload or the Redis config-sync resync, both -of which lag by seconds. A request that uses the new object immediately can -land on a sibling that has never heard of it and fail with a 400/404. - -On a registry miss, callers here fetch the missing object from the DB and load -it into the local registry before giving up. A short negative-result TTL keeps -repeated lookups of genuinely unknown names from hammering the DB. +registries on the periodic config reload, so a request using the new object +immediately can land on a sibling that has never heard of it and fail 400/404. +On a registry miss, callers here fetch the missing row from the DB and load it +into the local registry before giving up. A short negative-result TTL per key +plus a global resync budget per window bound the DB load from lookups of +genuinely unknown names. """ import asyncio +import time from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Final @@ -19,24 +19,57 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache if TYPE_CHECKING: + from prisma.types import ( + LiteLLM_AgentsTableInclude, + LiteLLM_AgentsTableWhereUniqueInput, + LiteLLM_ProxyModelTableWhereInput, + ) + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.agents import AgentResponse READ_THROUGH_MISS_TTL_SECONDS: Final = 2.0 +READ_THROUGH_RESYNC_WINDOW_SECONDS: Final = 5.0 +READ_THROUGH_MAX_RESYNCS_PER_WINDOW: Final = 20 class RegistryReadThrough: - __slots__ = ("_lock", "_miss_ttl_seconds", "_recent_misses", "_resync") + __slots__ = ( + "_lock", + "_max_resyncs_per_window", + "_miss_ttl_seconds", + "_recent_misses", + "_resync", + "_resync_window_seconds", + "_window_resyncs", + "_window_started_at", + ) def __init__( self, resync: Callable[[str], Awaitable[bool]], miss_ttl_seconds: float = READ_THROUGH_MISS_TTL_SECONDS, + max_resyncs_per_window: int = READ_THROUGH_MAX_RESYNCS_PER_WINDOW, + resync_window_seconds: float = READ_THROUGH_RESYNC_WINDOW_SECONDS, ) -> None: self._resync = resync self._miss_ttl_seconds = miss_ttl_seconds + self._max_resyncs_per_window = max_resyncs_per_window + self._resync_window_seconds = resync_window_seconds self._lock = asyncio.Lock() self._recent_misses = InMemoryCache(max_size_in_memory=1000) + self._window_started_at = float("-inf") + self._window_resyncs = 0 + + def _consume_resync_budget(self) -> bool: + now: Final = time.monotonic() + if now - self._window_started_at >= self._resync_window_seconds: + self._window_started_at = now + self._window_resyncs = 0 + if self._window_resyncs >= self._max_resyncs_per_window: + return False + self._window_resyncs += 1 + return True async def attempt(self, key: str) -> bool: if self._recent_misses.get_cache(key) is not None: @@ -44,6 +77,14 @@ class RegistryReadThrough: async with self._lock: if self._recent_misses.get_cache(key) is not None: return False + if not self._consume_resync_budget(): + verbose_proxy_logger.warning( + "registry read-through for %r skipped: resync budget of %s per %ss exhausted", + key, + self._max_resyncs_per_window, + self._resync_window_seconds, + ) + return False try: found: Final = await self._resync(key) except Exception as e: # noqa: BLE001 # a failed read-through must surface the original miss error, not a 500 @@ -68,9 +109,10 @@ async def _resync_model_deployments(model_name: str) -> bool: return False prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None - rows: Final = await ModelRepository(prisma_client).table.find_many( - where={"OR": [{"model_name": model_name}, {"model_id": model_name}]} - ) + table: Final = ModelRepository(prisma_client).table + name_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_name": model_name} + id_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_id": model_name} + rows: Final = await table.find_many(where=name_filter) or await table.find_many(where=id_filter) if not rows: return False if proxy_server.llm_router is None: @@ -85,24 +127,49 @@ async def _resync_model_deployments(model_name: str) -> bool: async def _resync_guardrails(guardrail_name: str) -> bool: from litellm.proxy import proxy_server + from litellm.proxy.guardrails.guardrail_registry import ( + IN_MEMORY_GUARDRAIL_HANDLER, + GuardrailRegistry, + ) if not _db_backed_registries_enabled(): return False prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None - await proxy_server.proxy_config._init_guardrails_in_db(prisma_client=prisma_client) + row: Final = await GuardrailRegistry().get_guardrail_by_name_from_db( + guardrail_name=guardrail_name, prisma_client=prisma_client + ) + if row is None: + return False + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=row) return _initialized_guardrail(guardrail_name) is not None async def _resync_agents(agent_id_or_name: str) -> bool: from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.agent_registry import ( + agents_table, + global_agent_registry, + ) + from litellm.types.agents import AgentResponse if not _db_backed_registries_enabled(): return False + if _agent_from_registry(agent_id_or_name) is not None: + return True prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None - await proxy_server.proxy_config._init_agents_in_db(prisma_client=prisma_client) - return _agent_from_registry(agent_id_or_name) is not None + table: Final = agents_table(prisma_client) + id_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id_or_name} + name_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_name": agent_id_or_name} + include_permission: Final[LiteLLM_AgentsTableInclude] = {"object_permission": True} + row: Final = await table.find_unique(where=id_filter, include=include_permission) or await table.find_unique( + where=name_filter, include=include_permission + ) + if row is None: + return False + global_agent_registry.register_agent(agent_config=AgentResponse.model_validate(row.model_dump())) + return True model_registry_read_through: Final = RegistryReadThrough(resync=_resync_model_deployments) 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 6b5e1a5d0b5..8e8545a51cc 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -57,20 +57,20 @@ def _model_table(prisma_client: PrismaClient) -> _ModelTableClient: return ModelRepository(prisma_client).table -def validate_models_exist(model_names: list[str], llm_router: "Router | None") -> tuple[bool, list[str]]: +def validate_models_exist(model_names: Sequence[str], llm_router: "Router | None") -> tuple[bool, Sequence[str]]: """ Validate that all requested model names exist in the router. Checks only exact model name matches. Returns: - Tuple[bool, List[str]]: (all_valid, missing_models) + (all_valid, missing_models) """ if llm_router is None: return False, model_names - router_model_names: Final = set(llm_router.get_model_names()) - missing: Final = [m for m in model_names if m not in router_model_names] - return (len(missing) == 0, missing) + router_model_names: Final = frozenset(llm_router.get_model_names()) + missing: Final = tuple(m for m in model_names if m not in router_model_names) + return (not missing, missing) async def _missing_models_after_read_through( @@ -81,12 +81,12 @@ async def _missing_models_after_read_through( model_registry_read_through, ) - _, missing = validate_models_exist(model_names=list(model_names), llm_router=llm_router) + _, missing = validate_models_exist(model_names=model_names, llm_router=llm_router) if not missing: return () for name in missing: await model_registry_read_through.attempt(name) - _, still_missing = validate_models_exist(model_names=list(model_names), llm_router=proxy_server.llm_router) + _, still_missing = validate_models_exist(model_names=model_names, llm_router=proxy_server.llm_router) return tuple(still_missing) @@ -118,13 +118,21 @@ def _raise_http_if_reload_degraded_serving( before: frozenset[str], written_models: Sequence[tuple[str, object]], access_group: str, + still_desired: frozenset[str] | None, + live_after: frozenset[str] | None, ) -> None: """Same verdict as the model-write endpoints, expressed through this file's HTTPException error convention, with the metadata-only obligation: these writes change group membership, not the models themselves, so a row that was already not serving before the reload is never blamed here; only a model this reload stopped serving is reported.""" - missing, collateral = reload_serving_verdict(before=before, written_models=written_models, written_must_serve=False) + missing, collateral = reload_serving_verdict( + before=before, + written_models=written_models, + written_must_serve=False, + still_desired=still_desired, + live_after=live_after, + ) gone: Final = tuple(dict.fromkeys((*missing, *collateral))) if not gone: return @@ -456,11 +464,13 @@ async def create_model_group( live_before_reload: Final = live_model_ids_snapshot() - await clear_cache() + reload_outcome: Final = await clear_cache() _raise_http_if_reload_degraded_serving( before=live_before_reload, written_models=updated_pairs, access_group=data.access_group, + still_desired=reload_outcome.still_desired, + live_after=reload_outcome.live_after, ) verbose_proxy_logger.info( @@ -716,11 +726,13 @@ async def update_access_group( # Clear cache and reload models to pick up the access group changes live_before_reload: Final = live_model_ids_snapshot() - await clear_cache() + reload_outcome: Final = await clear_cache() _raise_http_if_reload_degraded_serving( before=live_before_reload, written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()), access_group=access_group, + still_desired=reload_outcome.still_desired, + live_after=reload_outcome.live_after, ) verbose_proxy_logger.info( @@ -818,11 +830,13 @@ async def delete_access_group( # Clear cache and reload models to pick up the access group changes live_before_reload: Final = live_model_ids_snapshot() - await clear_cache() + reload_outcome: Final = await clear_cache() _raise_http_if_reload_degraded_serving( before=live_before_reload, written_models=removed_pairs, access_group=access_group, + still_desired=reload_outcome.still_desired, + live_after=reload_outcome.live_after, ) verbose_proxy_logger.info( diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 4e7f86b87d5..492c3d750a9 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -458,7 +458,7 @@ async def route_request( async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited provider coroutines; the inferred union keeps route_request's callers typed - data: dict, # noqa: LIT001 # request body is the proxy-wide mutable dict contract shared with route_request + data: dict, # mutable-ok: request body is the proxy-wide mutable dict contract shared with route_request llm_router: LitellmRouter | None, user_model: str | None, route_type: RouteType, diff --git a/litellm/router.py b/litellm/router.py index efd3b5a527e..f3fb7c08d2c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8573,11 +8573,9 @@ class Router: Returns: - The added/updated deployment """ + _deployment_model_id: Final = deployment.model_info.id or "" + _deployment_on_router: Final[Deployment | None] = self.get_deployment(model_id=_deployment_model_id) try: - # check if deployment already exists - _deployment_model_id: Final = deployment.model_info.id or "" - - _deployment_on_router: Final[Deployment | None] = self.get_deployment(model_id=_deployment_model_id) if _deployment_on_router is not None: # deployment with this model_id exists on the router if ( @@ -8628,10 +8626,31 @@ class Router: deployment.model_info.id, e, ) + self._restore_deployment_after_failed_upsert( + previous_deployment=_deployment_on_router, model_id=_deployment_model_id + ) return None else: raise e + def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None: + if previous_deployment is None or self.has_model_id(model_id): + return + try: + self.add_deployment(deployment=previous_deployment) + verbose_router_logger.info( + "Restored deployment %s (id=%s); it keeps serving its previous configuration.", + previous_deployment.model_name, + model_id, + ) + except Exception as restore_error: + verbose_router_logger.warning( + "Could not restore previously served deployment %s (id=%s) after the failed upsert: %s", + previous_deployment.model_name, + model_id, + restore_error, + ) + @staticmethod def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: """The ``litellm.model_cost`` keys a deployment's shared backend info is registered under.""" diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 6882479a344..096e039b8aa 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -12,7 +12,7 @@ "limit": 2017 }, "ANN202": { - "limit": 855 + "limit": 854 }, "ANN204": { "limit": 711 diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index e1c8f031579..31b41d5458e 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -94,6 +94,32 @@ async def test_distinct_keys_do_not_share_negative_cache(): assert spy.calls == ["ghost-a", "ghost-b"] +@pytest.mark.asyncio +async def test_resync_budget_exhausted_blocks_resync_without_negative_caching(): + spy: Final = ResyncSpy(found=False) + read_through: Final = RegistryReadThrough( + resync=spy, miss_ttl_seconds=60.0, max_resyncs_per_window=2, resync_window_seconds=60.0 + ) + + assert await read_through.attempt("ghost-a") is False + assert await read_through.attempt("ghost-b") is False + assert await read_through.attempt("ghost-c") is False + assert spy.calls == ["ghost-a", "ghost-b"] + assert read_through._recent_misses.get_cache("ghost-c") is None + + +@pytest.mark.asyncio +async def test_resync_budget_replenishes_after_window(): + spy: Final = ResyncSpy(found=True) + read_through: Final = RegistryReadThrough(resync=spy, max_resyncs_per_window=1, resync_window_seconds=0.05) + + assert await read_through.attempt("model-a") is True + assert await read_through.attempt("model-b") is False + await asyncio.sleep(0.1) + assert await read_through.attempt("model-b") is True + assert spy.calls == ["model-a", "model-b"] + + class FakeAgentRow: def __init__(self, agent_id: str, agent_name: str) -> None: self.agent_id = agent_id @@ -101,15 +127,15 @@ class FakeAgentRow: self.object_permission = None self.spend = 0.0 - def __iter__(self): - return iter( - { - "agent_id": self.agent_id, - "agent_name": self.agent_name, - "agent_card_params": {"name": self.agent_name, "url": "http://db-agent"}, - "litellm_params": {}, - }.items() - ) + def model_dump(self): + return { + "agent_id": self.agent_id, + "agent_name": self.agent_name, + "agent_card_params": {"name": self.agent_name, "url": "http://db-agent"}, + "litellm_params": {}, + "object_permission": None, + "spend": self.spend, + } @pytest.fixture @@ -138,8 +164,8 @@ async def test_get_agent_with_read_through_recovers_agent_created_on_sibling_rep agent_id: Final = "read-through-db-agent-id" prisma_client: Final = MagicMock() - prisma_client.db.litellm_agentstable.find_many = AsyncMock( - return_value=[FakeAgentRow(agent_id, "read-through-db-agent")] + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + return_value=FakeAgentRow(agent_id, "read-through-db-agent") ) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) @@ -149,6 +175,35 @@ async def test_get_agent_with_read_through_recovers_agent_created_on_sibling_rep assert agent is not None assert agent.agent_id == agent_id + prisma_client.db.litellm_agentstable.find_unique.assert_awaited_once_with( + where={"agent_id": agent_id}, + include={"object_permission": True}, + ) + + +@pytest.mark.asyncio +async def test_get_agent_with_read_through_recovers_agent_by_name(clean_agent_registry, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent_name: Final = "read-through-db-agent-by-name" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + side_effect=[None, FakeAgentRow("read-through-name-lookup-id", agent_name)] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + agent: Final = await get_agent_with_read_through(agent_name) + + assert agent is not None + assert agent.agent_name == agent_name + prisma_client.db.litellm_agentstable.find_unique.assert_awaited_with( + where={"agent_name": agent_name}, + include={"object_permission": True}, + ) @pytest.mark.asyncio @@ -159,11 +214,33 @@ async def test_get_agent_with_read_through_returns_none_for_unknown_agent(clean_ from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through prisma_client: Final = MagicMock() - prisma_client.db.litellm_agentstable.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) assert await get_agent_with_read_through("agent-nobody-created") is None + assert prisma_client.db.litellm_agentstable.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_resync_agents_already_registered_skips_db(clean_agent_registry, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import _resync_agents + + agent_id: Final = "read-through-dedup-agent-id" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + return_value=FakeAgentRow(agent_id, "read-through-dedup-agent") + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + assert await _resync_agents(agent_id) is True + assert await _resync_agents(agent_id) is True + assert prisma_client.db.litellm_agentstable.find_unique.await_count == 1 + assert len(clean_agent_registry.agent_list) == 1 class FakeGuardrailRow: @@ -200,8 +277,11 @@ async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sib guardrail_id: Final = "read-through-db-guardrail-id" guardrail_name: Final = "read-through-db-guardrail" prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock( + return_value=FakeGuardrailRow(guardrail_id, guardrail_name) + ) prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( - return_value=[FakeGuardrailRow(guardrail_id, guardrail_name)] + side_effect=AssertionError("full-table guardrail scan on read-through miss") ) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) @@ -210,6 +290,9 @@ async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sib guardrail: Final = await get_initialized_guardrail_with_read_through(guardrail_name=guardrail_name) assert guardrail is not None assert guardrail.guardrail_name == guardrail_name + prisma_client.db.litellm_guardrailstable.find_unique.assert_awaited_once_with( + where={"guardrail_name": guardrail_name} + ) finally: IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id) @@ -224,7 +307,7 @@ async def test_get_guardrail_with_read_through_returns_none_for_unknown_guardrai ) prisma_client: Final = MagicMock() - prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index 1722d8c377b..c973c6a8346 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -13,6 +13,9 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm import Router +from litellm.proxy.management_endpoints.model_management_endpoints import ( + ReconcileOutcome, +) @pytest.mark.asyncio @@ -121,7 +124,7 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new_callable=AsyncMock, + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): response = await create_model_group( @@ -186,7 +189,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments(): patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new_callable=AsyncMock, + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): response = await create_model_group( @@ -236,7 +239,7 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names(): patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new_callable=AsyncMock, + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): response = await create_model_group( @@ -313,7 +316,7 @@ async def test_create_access_group_invalid_model_id_returns_400(): patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new_callable=AsyncMock, + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): with pytest.raises(HTTPException) as exc_info: @@ -352,7 +355,7 @@ async def test_create_access_group_surfaces_dropped_models(): patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new=AsyncMock(return_value=None), + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): with pytest.raises(HTTPException) as exc_info: @@ -365,6 +368,50 @@ async def test_create_access_group_surfaces_dropped_models(): assert "deploy-A" in str(exc_info.value.detail) + +@pytest.mark.asyncio +async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_read(): + """A concurrent reconcile sampled after the lock is released must not make this + write's reload look like it dropped the tagged model: the verdict has to judge from + the ReconcileOutcome the reload captured under the lock, not a fresh router read.""" + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + create_model_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + NewModelGroupRequest, + ) + + deploy_a = MagicMock(model_id="deploy-A", model_name="gpt-4o", model_info={}) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() + + concurrently_wiped_router = MagicMock() + concurrently_wiped_router.get_model_ids.side_effect = [["deploy-A"], []] + with ( + patch("litellm.proxy.proxy_server.llm_router", concurrently_wiped_router), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch( + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock( + return_value=ReconcileOutcome( + still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"}) + ) + ), + ), + ): + response = await create_model_group( + data=NewModelGroupRequest(access_group="production-models", model_ids=["deploy-A"]), + user_api_key_dict=UserAPIKeyAuth(user_id="test_admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response.models_updated == 1 + assert concurrently_wiped_router.get_model_ids.call_count == 1 + + @pytest.mark.asyncio async def test_tag_deployment_parses_string_model_info_and_refuses_corrupt(): """The model_info column can arrive as its JSON string; tagging must parse it rather @@ -420,7 +467,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new=AsyncMock(return_value=None), + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): response = await delete_access_group( @@ -474,7 +521,7 @@ async def test_create_access_group_read_through_recovers_model_created_on_siblin patch("litellm.proxy.proxy_server.store_model_in_db", True), patch( "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new=AsyncMock(return_value=None), + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): response = await create_model_group( @@ -485,7 +532,7 @@ async def test_create_access_group_read_through_recovers_model_created_on_siblin assert response.models_updated == 1 assert response.model_names == [model_name] assert mock_prisma.db.litellm_proxymodeltable.find_many.await_args_list[0].kwargs["where"] == { - "OR": [{"model_name": model_name}, {"model_id": model_name}] + "model_name": model_name } diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 770f7857265..0523e796543 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -116,15 +116,15 @@ class _DbAgentRow: self.object_permission = None self.spend = 0.0 - def __iter__(self): - return iter( - { - "agent_id": self.agent_id, - "agent_name": self.agent_name, - "agent_card_params": {"name": self.agent_name, "url": "http://sibling-db-agent.example.com"}, - "litellm_params": {}, - }.items() - ) + def model_dump(self): + return { + "agent_id": self.agent_id, + "agent_name": self.agent_name, + "agent_card_params": {"name": self.agent_name, "url": "http://sibling-db-agent.example.com"}, + "litellm_params": {}, + "object_permission": None, + "spend": self.spend, + } def _router_without_models(): @@ -149,8 +149,8 @@ async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_re agent_name = "a2a-sibling-replica-agent" prisma_client = Mock() - prisma_client.db.litellm_agentstable.find_many = AsyncMock( - return_value=[_DbAgentRow("a2a-sibling-replica-agent-id", agent_name)] + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + side_effect=[None, _DbAgentRow("a2a-sibling-replica-agent-id", agent_name)] ) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) @@ -182,4 +182,4 @@ async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_re call_kwargs = mock_acompletion.call_args.kwargs assert call_kwargs["model"] == f"a2a/{agent_name}" assert call_kwargs["api_base"] == "http://sibling-db-agent.example.com" - prisma_client.db.litellm_agentstable.find_many.assert_awaited() + prisma_client.db.litellm_agentstable.find_unique.assert_awaited() diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 8f52fa179fe..15b4f09822a 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1179,7 +1179,7 @@ async def test_route_request_read_through_recovers_model_created_on_sibling_repl assert response.choices[0].message.content == "hello-from-db" assert len(table.find_many_wheres) == 1 - assert table.find_many_wheres[0] == {"OR": [{"model_name": model_name}, {"model_id": model_name}]} + assert table.find_many_wheres[0] == {"model_name": model_name} @pytest.mark.asyncio @@ -1207,7 +1207,7 @@ async def test_route_request_unknown_model_raises_and_hits_db_once_within_ttl(mo with pytest.raises(ProxyModelNotFoundError): await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion") - assert len(table.find_many_wheres) == 1 + assert table.find_many_wheres == [{"model_name": model_name}, {"model_id": model_name}] @pytest.mark.asyncio diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index fb8438ccb01..68ad911e792 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7595,6 +7595,62 @@ def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypa assert len(result) == 1 +class TestUpsertDeploymentRollback: + """ + Regression tests: `upsert_deployment` pops the previous deployment before + re-adding the edited one. When the re-add raises under + `ignore_invalid_deployments=True`, the pop must be rolled back so this pod + keeps serving the previous configuration instead of silently dropping a live + deployment (the "Error upserting deployment" drop behind the access-group + reload 500 in the 2-replica e2e suite). + """ + + def test_failed_upsert_keeps_previous_deployment_serving(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = litellm.Router( + model_list=[ + { + "model_name": "prod-model", + "litellm_params": {"model": "gpt-4o", "api_key": "sk-old"}, + "model_info": {"id": "prod-1", "db_model": True}, + } + ], + ignore_invalid_deployments=True, + ) + + result = router.upsert_deployment( + deployment=Deployment( + model_name="prod-model", + litellm_params=LiteLLM_Params(model="auto_router/broken"), + model_info=ModelInfo(id="prod-1", db_model=True), + ) + ) + + assert result is None + restored = router.get_deployment(model_id="prod-1") + assert restored is not None + assert restored.litellm_params.model == "gpt-4o" + assert [model["model_name"] for model in router.model_list] == ["prod-model"] + + def test_failed_fresh_add_returns_none_without_restore(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = litellm.Router(model_list=[], ignore_invalid_deployments=True) + + result = router.upsert_deployment( + deployment=Deployment( + model_name="fresh-router", + litellm_params=LiteLLM_Params(model="auto_router/broken"), + model_info=ModelInfo(id="fresh-1", db_model=True), + ) + ) + + assert result is None + assert router.get_deployment(model_id="fresh-1") is None + assert router.model_list == [] + + class TestConsumedRequestTagsStamp: """Issue #36621: when a request's tags select a tagged pre-routing strategy, those tags are consumed by the selection; the hook must stamp the rewritten model group so diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f8e481dc142..e7cfff93aa4 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22894 + "limit": 22892 }, "LIT002": { - "limit": 26888 + "limit": 26886 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16700 + "limit": 16699 }, "LIT011": { "limit": 5590 From afeed48a70feb8a903ce50b4fdfe283fd80cfbed Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:43:44 -0700 Subject: [PATCH 3/6] fix(proxy): serialize read-through with reloads, gate db object types The model resync now mutates the router under MODEL_RECONCILE_LOCK, and the agent resync shares the new AGENT_RECONCILE_LOCK with the periodic agent reload, so a reconcile built from a pre-write DB snapshot can no longer evict or duplicate what a read-through just registered. Every resync checks should_load_db_object for its object type, keeping read-through consistent with what the replica is configured to load, and the a2a raise sites tag ProxyModelNotFoundError as non-retryable so an agent miss no longer burns the model resync budget. --- litellm/proxy/agent_endpoints/a2a_routing.py | 4 +- .../proxy/agent_endpoints/agent_registry.py | 1 + .../common_utils/registry_read_through.py | 38 +++--- litellm/proxy/proxy_server.py | 64 +++++----- litellm/proxy/route_llm_request.py | 7 +- .../test_registry_read_through.py | 110 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 26 +++++ .../proxy/test_route_llm_request.py | 36 ++++++ 8 files changed, 234 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 2228735d805..8a795214750 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -50,7 +50,7 @@ async def route_a2a_agent_request( if agent is None: verbose_proxy_logger.error("[A2A] Agent '%s' not found in registry", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) - raise ProxyModelNotFoundError(route=route_name, model_name=model_name) + raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) # Verify the caller is permitted to use this agent (admins bypass the check) is_admin: Final = user_api_key_dict is not None and ( @@ -72,7 +72,7 @@ async def route_a2a_agent_request( if not agent.agent_card_params or "url" not in agent.agent_card_params: verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) - raise ProxyModelNotFoundError(route=route_name, model_name=model_name) + raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) # Inject API base and route to litellm data["api_base"] = agent.agent_card_params["url"] diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 742fdf35b1e..64de6827679 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -600,3 +600,4 @@ class AgentRegistry: global_agent_registry: Final = AgentRegistry() +AGENT_RECONCILE_LOCK: Final = asyncio.Lock() diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 7ace046ae26..e92803d3f51 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -95,17 +95,19 @@ class RegistryReadThrough: return found -def _db_backed_registries_enabled() -> bool: +def _db_backed_registries_enabled(object_type: str) -> bool: from litellm.proxy import proxy_server - return proxy_server.prisma_client is not None and proxy_server.store_model_in_db is True + if proxy_server.prisma_client is None or proxy_server.store_model_in_db is not True: + return False + return proxy_server.should_load_db_object(object_type=object_type) async def _resync_model_deployments(model_name: str) -> bool: from litellm.proxy import proxy_server from litellm.repositories.model_repository import ModelRepository - if not _db_backed_registries_enabled(): + if not _db_backed_registries_enabled("models"): return False prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None @@ -115,13 +117,15 @@ async def _resync_model_deployments(model_name: str) -> bool: rows: Final = await table.find_many(where=name_filter) or await table.find_many(where=id_filter) if not rows: return False - if proxy_server.llm_router is None: + router: Final = proxy_server.llm_router + if router is None: await proxy_server.proxy_config.add_deployment( prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj ) return proxy_server.llm_router is not None - proxy_server.proxy_config._add_deployment(db_models=rows) - proxy_server.llm_model_list = proxy_server.llm_router.get_model_list() + async with proxy_server.MODEL_RECONCILE_LOCK: + proxy_server.proxy_config._add_deployment(db_models=rows) + proxy_server.llm_model_list = router.get_model_list() return True @@ -132,7 +136,7 @@ async def _resync_guardrails(guardrail_name: str) -> bool: GuardrailRegistry, ) - if not _db_backed_registries_enabled(): + if not _db_backed_registries_enabled("guardrails"): return False prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None @@ -148,12 +152,13 @@ async def _resync_guardrails(guardrail_name: str) -> bool: async def _resync_agents(agent_id_or_name: str) -> bool: from litellm.proxy import proxy_server from litellm.proxy.agent_endpoints.agent_registry import ( + AGENT_RECONCILE_LOCK, agents_table, global_agent_registry, ) from litellm.types.agents import AgentResponse - if not _db_backed_registries_enabled(): + if not _db_backed_registries_enabled("agents"): return False if _agent_from_registry(agent_id_or_name) is not None: return True @@ -163,13 +168,16 @@ async def _resync_agents(agent_id_or_name: str) -> bool: id_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id_or_name} name_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_name": agent_id_or_name} include_permission: Final[LiteLLM_AgentsTableInclude] = {"object_permission": True} - row: Final = await table.find_unique(where=id_filter, include=include_permission) or await table.find_unique( - where=name_filter, include=include_permission - ) - if row is None: - return False - global_agent_registry.register_agent(agent_config=AgentResponse.model_validate(row.model_dump())) - return True + async with AGENT_RECONCILE_LOCK: + if _agent_from_registry(agent_id_or_name) is not None: + return True + row: Final = await table.find_unique(where=id_filter, include=include_permission) or await table.find_unique( + where=name_filter, include=include_permission + ) + if row is None: + return False + global_agent_registry.register_agent(agent_config=AgentResponse.model_validate(row.model_dump())) + return True model_registry_read_through: Final = RegistryReadThrough(resync=_resync_model_deployments) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 56036713fa9..dbbd87c7bb3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4120,6 +4120,31 @@ def _swap_in_model_cost_map(new_model_cost_map: dict) -> int: return fetched_model_count +def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool: + """ + Check if an object type should be loaded from the database based on general_settings.supported_db_objects. + + Args: + object_type: Type of object to check (e.g., SupportedDBObjectType.MODELS, "models", etc.) + + Returns: + True if the object should be loaded, False otherwise + """ + supported_db_objects: Final = general_settings.get("supported_db_objects", None) + + if supported_db_objects is None: + return True + + if not isinstance(supported_db_objects, list): + verbose_proxy_logger.warning( + "supported_db_objects is not a list, got %s. Loading all objects.", type(supported_db_objects) + ) + return True + + object_type_str: Final = str(object_type) + return any(str(obj) == object_type_str for obj in supported_db_objects) + + class ProxyConfig: """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. @@ -6522,36 +6547,7 @@ class ProxyConfig: return config def _should_load_db_object(self, object_type: str | SupportedDBObjectType) -> bool: - """ - Check if an object type should be loaded from the database based on general_settings.supported_db_objects. - - Args: - object_type: Type of object to check (e.g., SupportedDBObjectType.MODELS, "models", etc.) - - Returns: - True if the object should be loaded, False otherwise - """ - global general_settings - - # Get the supported_db_objects configuration - supported_db_objects: Final = general_settings.get("supported_db_objects", None) - - # If supported_db_objects is not set, load all objects (default behavior) - if supported_db_objects is None: - return True - - # If supported_db_objects is set, only load specified objects - if not isinstance(supported_db_objects, list): - verbose_proxy_logger.warning( - "supported_db_objects is not a list, got %s. Loading all objects.", type(supported_db_objects) - ) - return True - - # Convert object_type to string for comparison (handles both str and enum) - object_type_str: Final = str(object_type) - - # Check if the object type is in the list (supports both str and enum values) - return any(str(obj) == object_type_str for obj in supported_db_objects) + return should_load_db_object(object_type=object_type) async def _get_models_from_db(self, prisma_client: PrismaClient) -> list | None: """ @@ -7278,13 +7274,17 @@ class ProxyConfig: ) async def _init_agents_in_db(self, prisma_client: PrismaClient): + from litellm.proxy.agent_endpoints.agent_registry import ( + AGENT_RECONCILE_LOCK, + ) from litellm.proxy.agent_endpoints.agent_registry import ( global_agent_registry as AGENT_REGISTRY, ) try: - db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client) - AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents) + async with AGENT_RECONCILE_LOCK: + db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client) + AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:_init_agents_in_db - %s", e) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 492c3d750a9..91a0c68fd58 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -146,7 +146,8 @@ ROUTE_ENDPOINT_MAPPING: Final = { class ProxyModelNotFoundError(HTTPException): - def __init__(self, route: str, model_name: str): + def __init__(self, route: str, model_name: str, retryable_with_model_read_through: bool = True): + self.retryable_with_model_read_through: Final = retryable_with_model_read_through detail: Final = { "error": f"{route}: Invalid model name passed in model={model_name}. Call `/v1/models` to view available models for your key." } @@ -437,9 +438,9 @@ async def route_request( route_type=route_type, user_api_key_dict=user_api_key_dict, ) - except ProxyModelNotFoundError: + except ProxyModelNotFoundError as e: requested_model: Final = data.get("model", "") - if not isinstance(requested_model, str) or not requested_model: + if not e.retryable_with_model_read_through or not isinstance(requested_model, str) or not requested_model: raise from litellm.proxy import proxy_server from litellm.proxy.common_utils.registry_read_through import ( diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index 31b41d5458e..56c71e38ef5 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -312,3 +312,113 @@ async def test_get_guardrail_with_read_through_returns_none_for_unknown_guardrai monkeypatch.setattr(proxy_server, "store_model_in_db", True) assert await get_initialized_guardrail_with_read_through(guardrail_name="guardrail-nobody-created") is None + + +@pytest.mark.asyncio +async def test_resync_model_deployments_mutates_router_under_model_reconcile_lock(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[MagicMock()]) + router: Final = MagicMock() + router.get_model_list.return_value = [] + lock_states: list[bool] = [] + + def record_add_deployment(db_models) -> None: + lock_states.append(proxy_server.MODEL_RECONCILE_LOCK.locked()) + + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", record_add_deployment) + + assert await _resync_model_deployments("lock-scope-model") is True + assert lock_states == [True] + assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + + +@pytest.mark.asyncio +async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock( + side_effect=AssertionError("db hit for an object type this replica does not load") + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["guardrails"]}) + + assert await _resync_model_deployments("gated-out-model") is False + + +@pytest.mark.asyncio +async def test_resync_guardrails_respects_supported_db_objects(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import _resync_guardrails + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock( + side_effect=AssertionError("db hit for an object type this replica does not load") + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]}) + + assert await _resync_guardrails("gated-out-guardrail") is False + + +@pytest.mark.asyncio +async def test_resync_agents_respects_supported_db_objects(clean_agent_registry, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import _resync_agents + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + side_effect=AssertionError("db hit for an object type this replica does not load") + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]}) + + assert await _resync_agents("gated-out-agent") is False + + +@pytest.mark.asyncio +async def test_resync_agents_waits_for_agent_reload_and_skips_duplicate_registration(clean_agent_registry, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.agent_endpoints.agent_registry import AGENT_RECONCILE_LOCK + from litellm.proxy.common_utils.registry_read_through import _resync_agents + from litellm.types.agents import AgentResponse + + agent_id: Final = "reload-race-agent-id" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + side_effect=AssertionError("db hit while the agent reload held the reconcile lock") + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + async with AGENT_RECONCILE_LOCK: + resync_task: Final = asyncio.ensure_future(_resync_agents(agent_id)) + await asyncio.sleep(0.05) + assert not resync_task.done() + clean_agent_registry.register_agent( + agent_config=AgentResponse.model_validate(FakeAgentRow(agent_id, "reload-race-agent").model_dump()) + ) + + assert await resync_task is True + assert len(clean_agent_registry.agent_list) == 1 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5545ee92e84..3fbbcf674da 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11098,3 +11098,29 @@ async def test_moderations_reraises_proxy_exception_unwrapped(): assert exc_info.value.code == "400" assert exc_info.value.param == "metadata" mock_logging.post_call_failure_hook.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch): + from litellm.proxy.agent_endpoints.agent_registry import ( + AGENT_RECONCILE_LOCK, + global_agent_registry, + ) + from litellm.proxy.proxy_server import ProxyConfig + + lock_states: list[bool] = [] + + async def fake_get_all_agents_from_db(prisma_client) -> list: + lock_states.append(AGENT_RECONCILE_LOCK.locked()) + return [] + + def fake_load_agents_from_db_and_config(db_agents) -> None: + lock_states.append(AGENT_RECONCILE_LOCK.locked()) + + monkeypatch.setattr(global_agent_registry, "get_all_agents_from_db", fake_get_all_agents_from_db) + monkeypatch.setattr(global_agent_registry, "load_agents_from_db_and_config", fake_load_agents_from_db_and_config) + + await ProxyConfig()._init_agents_in_db(prisma_client=MagicMock()) + + assert lock_states == [True, True] + assert not AGENT_RECONCILE_LOCK.locked() diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 15b4f09822a..1e716f7c148 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1261,3 +1261,39 @@ async def test_route_request_routing_group_name_passes_model_gate(): assert response == "group_response" spy.assert_called_once_with(**data) + + +@pytest.mark.asyncio +async def test_route_request_a2a_agent_miss_does_not_consume_model_read_through(monkeypatch): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + import litellm + import litellm.proxy.proxy_server as proxy_server + + model_name = "a2a/agent-nobody-created" + router = litellm.Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + fake_prisma, model_table = _fake_prisma_client_with_models([]) + agents_find_unique = AsyncMock(return_value=None) + fake_prisma.db.litellm_agentstable = SimpleNamespace(find_unique=agents_find_unique) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with pytest.raises(ProxyModelNotFoundError): + await route_request( + data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]}, + llm_router=router, + user_model=None, + route_type="acompletion", + ) + + assert agents_find_unique.await_count == 2 + assert model_table.find_many_wheres == [] From 43389e987aa43b09f1b935fa92bc71f41a9953f0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:54:08 -0700 Subject: [PATCH 4/6] test: call _restore_deployment_after_failed_upsert directly for the router coverage gate --- tests/test_litellm/test_router.py | 35 +++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 546a3f33e65..cda495f750a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7698,6 +7698,41 @@ class TestUpsertDeploymentRollback: assert router.get_deployment(model_id="fresh-1") is None assert router.model_list == [] + def test_restore_re_adds_popped_deployment(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = litellm.Router( + model_list=[ + { + "model_name": "prod-model", + "litellm_params": {"model": "gpt-4o", "api_key": "sk-old"}, + "model_info": {"id": "prod-1", "db_model": True}, + } + ], + ignore_invalid_deployments=True, + ) + previous = router.get_deployment(model_id="prod-1") + router.delete_deployment(id="prod-1") + assert router.has_model_id("prod-1") is False + + router._restore_deployment_after_failed_upsert( + previous_deployment=previous, model_id="prod-1" + ) + + restored = router.get_deployment(model_id="prod-1") + assert restored is not None + assert restored.litellm_params.model == "gpt-4o" + + router._restore_deployment_after_failed_upsert( + previous_deployment=previous, model_id="prod-1" + ) + assert len(router.model_list) == 1 + + router._restore_deployment_after_failed_upsert( + previous_deployment=None, model_id="prod-1" + ) + assert len(router.model_list) == 1 + class TestConsumedRequestTagsStamp: """Issue #36621: when a request's tags select a tagged pre-routing strategy, those From e9c01da23389c8394354411cba573e32494c8bba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:54:46 -0700 Subject: [PATCH 5/6] test: drop unused imports in the direct restore test --- tests/test_litellm/test_router.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cda495f750a..65debae9a16 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7699,8 +7699,6 @@ class TestUpsertDeploymentRollback: assert router.model_list == [] def test_restore_re_adds_popped_deployment(self): - from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo - router = litellm.Router( model_list=[ { From f8a23aab09607b400aeaf5862e4afa8997938c3e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 01:46:33 -0700 Subject: [PATCH 6/6] fix: gate guardrail read-through to active rows and serialize it with the reload reconcile --- .../common_utils/registry_read_through.py | 16 +++-- .../proxy/guardrails/guardrail_registry.py | 3 + litellm/proxy/proxy_server.py | 52 ++++++++-------- .../test_registry_read_through.py | 60 +++++++++++++++++-- tests/test_litellm/proxy/test_proxy_server.py | 28 +++++++++ 5 files changed, 125 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index e92803d3f51..460b348e188 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from prisma.types import ( LiteLLM_AgentsTableInclude, LiteLLM_AgentsTableWhereUniqueInput, + LiteLLM_GuardrailsTableWhereInput, LiteLLM_ProxyModelTableWhereInput, ) @@ -132,20 +133,25 @@ async def _resync_model_deployments(model_name: str) -> bool: async def _resync_guardrails(guardrail_name: str) -> bool: from litellm.proxy import proxy_server from litellm.proxy.guardrails.guardrail_registry import ( + GUARDRAIL_RECONCILE_LOCK, IN_MEMORY_GUARDRAIL_HANDLER, - GuardrailRegistry, ) + from litellm.repositories.table_repositories import GuardrailsRepository + from litellm.types.guardrails import Guardrail if not _db_backed_registries_enabled("guardrails"): return False prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None - row: Final = await GuardrailRegistry().get_guardrail_by_name_from_db( - guardrail_name=guardrail_name, prisma_client=prisma_client - ) + active_row_filter: Final[LiteLLM_GuardrailsTableWhereInput] = { + "guardrail_name": guardrail_name, + "status": "active", + } + row: Final = await GuardrailsRepository(prisma_client).table.find_first(where=active_row_filter) if row is None: return False - IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=row) + async with GUARDRAIL_RECONCILE_LOCK: + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=Guardrail(**dict(row))) return _initialized_guardrail(guardrail_name) is not None diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 5f7374581a2..f6d348c1045 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -1,5 +1,6 @@ # litellm/proxy/guardrails/guardrail_registry.py +import asyncio import importlib import os from collections.abc import Callable, Iterator, Mapping @@ -813,4 +814,6 @@ class InMemoryGuardrailHandler: # In Memory Guardrail Handler for LiteLLM Proxy ######################################################## IN_MEMORY_GUARDRAIL_HANDLER: Final = InMemoryGuardrailHandler() + +GUARDRAIL_RECONCILE_LOCK: Final = asyncio.Lock() ######################################################## diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0deb71f917..be1914ea50c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7090,38 +7090,40 @@ class ProxyConfig: async def _init_guardrails_in_db(self, prisma_client: PrismaClient): from litellm.proxy.guardrails.guardrail_registry import ( + GUARDRAIL_RECONCILE_LOCK, IN_MEMORY_GUARDRAIL_HANDLER, Guardrail, GuardrailRegistry, ) try: - guardrails_in_db: Final[list[Guardrail]] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client - ) - verbose_proxy_logger.debug("guardrails from the DB %s", str(guardrails_in_db)) - db_guardrail_ids: Final[set] = set() - for guardrail in guardrails_in_db: - guardrail_id = guardrail.get("guardrail_id") - if guardrail_id: - db_guardrail_ids.add(guardrail_id) - try: - IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( - guardrail=cast(Guardrail, guardrail), - ) - except Exception as e: # noqa: BLE001 # one unloadable row must not stop the remaining guardrails - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - " - "skipping guardrail '%s' (ID: %s): %s: %s", - guardrail.get("guardrail_name"), - guardrail_id, - type(e).__name__, - e, - ) + async with GUARDRAIL_RECONCILE_LOCK: + guardrails_in_db: Final[list[Guardrail]] = await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) + verbose_proxy_logger.debug("guardrails from the DB %s", str(guardrails_in_db)) + db_guardrail_ids: Final[set] = set() + for guardrail in guardrails_in_db: + guardrail_id = guardrail.get("guardrail_id") + if guardrail_id: + db_guardrail_ids.add(guardrail_id) + try: + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( + guardrail=cast(Guardrail, guardrail), + ) + except Exception as e: # noqa: BLE001 # one unloadable row must not stop the remaining guardrails + verbose_proxy_logger.error( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - " + "skipping guardrail '%s' (ID: %s): %s: %s", + guardrail.get("guardrail_name"), + guardrail_id, + type(e).__name__, + e, + ) - # Drop in-memory DB-backed entries whose row was deleted on another - # pod. Config-loaded entries are never touched. - IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids) + # Drop in-memory DB-backed entries whose row was deleted on another + # pod. Config-loaded entries are never touched. + IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - %s", e) diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index 56c71e38ef5..f0fbdea4e85 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -260,6 +260,7 @@ class FakeGuardrailRow: "blocked_words": [{"keyword": "secret", "action": "BLOCK"}], }, "guardrail_info": {}, + "status": "active", }.items() ) @@ -277,7 +278,7 @@ async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sib guardrail_id: Final = "read-through-db-guardrail-id" guardrail_name: Final = "read-through-db-guardrail" prisma_client: Final = MagicMock() - prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock( + prisma_client.db.litellm_guardrailstable.find_first = AsyncMock( return_value=FakeGuardrailRow(guardrail_id, guardrail_name) ) prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( @@ -290,8 +291,8 @@ async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sib guardrail: Final = await get_initialized_guardrail_with_read_through(guardrail_name=guardrail_name) assert guardrail is not None assert guardrail.guardrail_name == guardrail_name - prisma_client.db.litellm_guardrailstable.find_unique.assert_awaited_once_with( - where={"guardrail_name": guardrail_name} + prisma_client.db.litellm_guardrailstable.find_first.assert_awaited_once_with( + where={"guardrail_name": guardrail_name, "status": "active"} ) finally: IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id) @@ -307,13 +308,64 @@ async def test_get_guardrail_with_read_through_returns_none_for_unknown_guardrai ) prisma_client: Final = MagicMock() - prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) assert await get_initialized_guardrail_with_read_through(guardrail_name="guardrail-nobody-created") is None +@pytest.mark.asyncio +async def test_resync_guardrails_never_loads_non_active_rows(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import _resync_guardrails + + pending_name: Final = "pending-review-guardrail" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + assert await _resync_guardrails(pending_name) is False + prisma_client.db.litellm_guardrailstable.find_first.assert_awaited_once_with( + where={"guardrail_name": pending_name, "status": "active"} + ) + + +@pytest.mark.asyncio +async def test_resync_guardrails_syncs_under_guardrail_reconcile_lock(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.common_utils.registry_read_through as read_through_module + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import _resync_guardrails + from litellm.proxy.guardrails.guardrail_registry import ( + GUARDRAIL_RECONCILE_LOCK, + IN_MEMORY_GUARDRAIL_HANDLER, + ) + + guardrail_name: Final = "lock-scope-guardrail" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_first = AsyncMock( + return_value=FakeGuardrailRow("lock-scope-guardrail-id", guardrail_name) + ) + lock_states: list[bool] = [] + + def record_sync(guardrail) -> None: + lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked()) + + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(IN_MEMORY_GUARDRAIL_HANDLER, "sync_guardrail_from_db", record_sync) + monkeypatch.setattr(read_through_module, "_initialized_guardrail", lambda guardrail_name: MagicMock()) + + assert await _resync_guardrails(guardrail_name) is True + assert lock_states == [True] + assert not GUARDRAIL_RECONCILE_LOCK.locked() + + @pytest.mark.asyncio async def test_resync_model_deployments_mutates_router_under_model_reconcile_lock(monkeypatch): from unittest.mock import AsyncMock, MagicMock diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e7d7d4c9323..75aa716bb85 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11126,6 +11126,34 @@ async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(mo assert not AGENT_RECONCILE_LOCK.locked() +@pytest.mark.asyncio +async def test_init_guardrails_in_db_snapshots_and_reconciles_under_guardrail_reconcile_lock(monkeypatch): + from litellm.proxy.guardrails.guardrail_registry import ( + GUARDRAIL_RECONCILE_LOCK, + IN_MEMORY_GUARDRAIL_HANDLER, + GuardrailRegistry, + ) + from litellm.proxy.proxy_server import ProxyConfig + + lock_states: list[bool] = [] + + async def fake_get_all_guardrails_from_db(prisma_client) -> list: + lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked()) + return [] + + def fake_reconcile_db_guardrails(db_guardrail_ids) -> list: + lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked()) + return [] + + monkeypatch.setattr(GuardrailRegistry, "get_all_guardrails_from_db", fake_get_all_guardrails_from_db) + monkeypatch.setattr(IN_MEMORY_GUARDRAIL_HANDLER, "reconcile_db_guardrails", fake_reconcile_db_guardrails) + + await ProxyConfig()._init_guardrails_in_db(prisma_client=MagicMock()) + + assert lock_states == [True, True] + assert not GUARDRAIL_RECONCILE_LOCK.locked() + + class TestEmbeddingsFailureHookRequestData: @pytest.mark.asyncio async def test_failure_hook_gets_post_setup_data_with_logging_obj(self):