diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 1cae87aed31..cea30ffad52 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -168,14 +168,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: "AgentResponse", request: Request) -> None: @@ -559,7 +558,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") @@ -673,7 +672,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..8a795214750 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,11 +46,11 @@ 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) - 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 ( @@ -70,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 new file mode 100644 index 00000000000..460b348e188 --- /dev/null +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -0,0 +1,224 @@ +"""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, 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 + +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_GuardrailsTableWhereInput, + 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", + "_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: + return False + 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 + 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(object_type: str) -> bool: + from litellm.proxy import proxy_server + + 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("models"): + return False + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + 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 + 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 + 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 + + +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, + ) + 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 + 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 + 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 + + +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("agents"): + 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 + 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} + 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) +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 b68d4a68b79..e50a3a5a1e7 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2305,8 +2305,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/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/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 49c0135ff10..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,37 @@ 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( + 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=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=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]: @@ -101,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 @@ -390,12 +415,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)}"}, @@ -439,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( @@ -654,12 +681,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)}"}, @@ -699,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( @@ -801,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/proxy_server.py b/litellm/proxy/proxy_server.py index fdce8c868d1..be1914ea50c 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: """ @@ -7094,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) @@ -7278,13 +7276,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 c85325b6fa9..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." } @@ -320,112 +321,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 as e: + requested_model: Final = data.get("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 ( + 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, # 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, + 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/litellm/router.py b/litellm/router.py index efd3b5a527e..00c6c4c8b6f 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: # noqa: BLE001 # best-effort restore: a second failure must not abort the reload + 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 new file mode 100644 index 00000000000..f0fbdea4e85 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -0,0 +1,476 @@ +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"] + + +@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 + self.agent_name = agent_name + self.object_permission = None + self.spend = 0.0 + + 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 +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_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) + + 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 + 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 +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_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: + 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": {}, + "status": "active", + }.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_first = AsyncMock( + return_value=FakeGuardrailRow(guardrail_id, guardrail_name) + ) + prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( + 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) + + 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 + 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) + + +@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_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 + + 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/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index 3240ad20edb..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( @@ -430,3 +477,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=ReconcileOutcome(still_desired=None, live_after=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"] == { + "model_name": 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_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7fdfbea843f..75aa716bb85 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11100,6 +11100,60 @@ async def test_moderations_reraises_proxy_exception_unwrapped(): 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() + + +@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): diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 02e4bddcee0..0523e796543 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -50,6 +50,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 @@ -106,3 +107,79 @@ 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 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(): + 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) + mock_router.is_recognized_model = Mock(return_value=False) + mock_router.team_public_model_names = [] + 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_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) + + 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_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 fc3b14592cd..1e716f7c148 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1119,6 +1119,126 @@ async def test_route_request_rejects_chat_completion_without_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] == {"model_name": 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 table.find_many_wheres == [{"model_name": model_name}, {"model_id": model_name}] + + +@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 == [] + @pytest.mark.asyncio async def test_route_request_routing_group_name_passes_model_gate(): from unittest.mock import AsyncMock, patch @@ -1141,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 == [] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d2d5d9268be..65debae9a16 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7643,6 +7643,95 @@ 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 == [] + + def test_restore_re_adds_popped_deployment(self): + 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 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