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 == []