diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 43139b18162..e0ed076bcfc 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -89,6 +89,7 @@ def agents_table(prisma_client: PrismaClient) -> AgentTableClient: class AgentRegistry: def __init__(self): self.agent_list: list[AgentResponse] = [] + self.config_agents: tuple[AgentConfig, ...] = () def reset_agent_list(self): self.agent_list = [] @@ -117,9 +118,21 @@ class AgentRegistry: return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest() def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None): + """ + Register the agents declared in config.yaml and remember them for later rebuilds. + + A config entry is skipped when its ``agent_name`` is already registered, so a + database record always wins over a config entry that reuses its name and the + registry never holds two agents under one name. Enforcing that here rather than + in the caller keeps the guarantee independent of the order the two sources load + in. Passing ``None`` leaves the remembered agents untouched; passing an empty + sequence clears them. + """ if agent_config is None: return None + self.config_agents = tuple(agent_config) + for agent_config_item in agent_config: if not isinstance(agent_config_item, dict): raise ValueError("agent_config must be a list of dictionaries") @@ -129,6 +142,9 @@ class AgentRegistry: if not all([agent_name, agent_card_params]): continue + if any(agent.agent_name == agent_name for agent in self.agent_list): + continue + # create a stable hash id for config item config_hash = self._create_agent_id(agent_config_item) @@ -139,26 +155,29 @@ class AgentRegistry: agent_config: Sequence[AgentConfig] | None = None, db_agents: list[dict[str, Any]] | None = None, ): + """ + Rebuild the registry from the DB rows plus the agents declared in config.yaml. + + ``agent_config`` defaults to the agents remembered by the last + ``load_agents_from_config`` call, so a periodic DB reload does not drop + config-defined agents. + + The DB rows are registered first so that a config entry reusing one of their + names is dropped by ``load_agents_from_config``, mirroring how config-declared + MCP servers are unioned under the database registry. Name lookups and + deregistration both address a single agent, so the registry must never hold two + under one name. + """ self.reset_agent_list() - if agent_config: - for agent_config_item in agent_config: - if not isinstance(agent_config_item, dict): - raise ValueError("agent_config must be a list of dictionaries") - - self.register_agent( - agent_config=AgentResponse( - agent_id=self._create_agent_id(agent_config_item), - **agent_config_item, - ) - ) # type: ignore - if db_agents: for db_agent in db_agents: if not isinstance(db_agent, dict): raise ValueError("db_agents must be a list of dictionaries") self.register_agent(agent_config=AgentResponse(**db_agent)) # type: ignore + + self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents) return self.agent_list ########################################################### diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a60ea2da019..314393c2fde 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -654,8 +654,6 @@ from fastapi.security import OAuth2PasswordBearer from fastapi.security.api_key import APIKeyHeader from fastapi.staticfiles import StaticFiles -from litellm.types.agents import AgentConfig - # import enterprise folder enterprise_router = APIRouter() try: @@ -2000,7 +1998,6 @@ config_passthrough_endpoints: Optional[List[Dict[str, Any]]] = None log_file = "api_log.json" worker_config = None master_key: Optional[str] = None -config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None shared_aiohttp_session: Optional["ClientSession"] = None # Global shared session for connection reuse @@ -5093,8 +5090,8 @@ class ProxyConfig: global_mcp_tool_registry.load_tools_from_config(mcp_tools_config, config_file_path=config_file_path) ## AGENTS - agent_config = config.get("agent_list", None) - if agent_config: + agent_config = config.get("agents", config.get("agent_list", None)) + if agent_config is not None: from litellm.proxy.agent_endpoints.agent_registry import ( global_agent_registry, ) @@ -6824,7 +6821,7 @@ class ProxyConfig: try: db_agents = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client) - AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents, agent_config=config_agents) + 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 - {}".format(str(e)) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index ddd9cc09c8e..f506535b5e1 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -112,3 +112,170 @@ async def test_update_agent_in_db_preserves_explicit_static_headers_and_extra_he assert update_data["static_headers"] == '{"Authorization": "Bearer xyz"}' assert update_data["extra_headers"] == ["X-Custom-Header"] + + +def test_load_agents_from_db_and_config_retains_previously_loaded_config_agents(): + """ + A DB reload calls this without an explicit agent_config. It resets the + registry, so it must fall back to the agents remembered from config.yaml + instead of dropping them. + """ + registry = AgentRegistry() + registry.load_agents_from_config( + [ + { + "agent_name": "config-agent", + "agent_card_params": _sample_agent_card_params(), + } + ] + ) + + registry.load_agents_from_db_and_config( + db_agents=[ + { + "agent_id": "db-id", + "agent_name": "db-agent", + "agent_card_params": _sample_agent_card_params(), + } + ] + ) + + assert sorted(agent.agent_name for agent in registry.get_agent_list()) == [ + "config-agent", + "db-agent", + ] + + +def test_load_agents_from_db_and_config_skips_incomplete_config_entries(): + """Config entries missing agent_card_params are skipped, not registered half-built.""" + registry = AgentRegistry() + registry.load_agents_from_config([{"agent_name": "no-card"}]) + + registry.load_agents_from_db_and_config(db_agents=None) + + assert registry.get_agent_list() == [] + + +@pytest.mark.parametrize( + "db_agent_names", + [ + ("shared-name", "other-db-agent"), + ("other-db-agent", "shared-name"), + ], + ids=["colliding-db-row-first", "colliding-db-row-last"], +) +def test_load_agents_from_db_and_config_lets_db_rows_win_a_name_collision(db_agent_names): + """ + A config entry that reuses a DB agent's name must not be registered next to it. + + Name lookups and deregistration both address a single agent, so two entries + under one name would shadow the DB row for routing and let a single delete + drop both. Parametrized over both DB row orders because the registry is a + list and a first-match lookup would otherwise pass on ordering luck. + """ + registry = AgentRegistry() + registry.load_agents_from_config( + [ + { + "agent_name": "shared-name", + "agent_card_params": _sample_agent_card_params(), + }, + { + "agent_name": "config-only", + "agent_card_params": _sample_agent_card_params(), + }, + ] + ) + + registry.load_agents_from_db_and_config( + db_agents=[ + { + "agent_id": f"db-id-{name}", + "agent_name": name, + "agent_card_params": _sample_agent_card_params(), + } + for name in db_agent_names + ] + ) + + registered = registry.get_agent_list() + assert sorted(agent.agent_name for agent in registered) == [ + "config-only", + "other-db-agent", + "shared-name", + ] + + shared = registry.get_agent_by_name("shared-name") + assert shared is not None + assert shared.agent_id == "db-id-shared-name", "the DB row must win the shared name, not the config entry" + + registry.deregister_agent("shared-name") + assert [agent.agent_name for agent in registry.get_agent_list() if agent.agent_name == "shared-name"] == [], ( + "deleting the DB agent must not leave a shadowed config duplicate behind" + ) + + +def test_load_agents_from_config_skips_a_name_already_held_by_a_db_agent(): + """ + The one-agent-per-name rule is enforced by the loader, not by the call order. + + A config load that lands on a registry already holding DB rows must not append a + colliding entry, otherwise the duplicate is reachable any time the two sources are + loaded in the other order. + """ + registry = AgentRegistry() + registry.load_agents_from_db_and_config( + db_agents=[ + { + "agent_id": "db-id", + "agent_name": "shared-name", + "agent_card_params": _sample_agent_card_params(), + } + ] + ) + + registry.load_agents_from_config( + [ + { + "agent_name": "shared-name", + "agent_card_params": _sample_agent_card_params(), + } + ] + ) + + assert [agent.agent_id for agent in registry.get_agent_list()] == ["db-id"] + + +def test_load_agents_from_config_registers_one_agent_per_name_within_the_config(): + """Two config entries sharing a name collapse to the first, keeping name lookups unambiguous.""" + registry = AgentRegistry() + registry.load_agents_from_config( + [ + {"agent_name": "dupe", "agent_card_params": _sample_agent_card_params()}, + {"agent_name": "dupe", "agent_card_params": {**_sample_agent_card_params(), "url": "http://second"}}, + ] + ) + + registered = registry.get_agent_list() + assert len(registered) == 1 + assert registered[0].agent_card_params["url"] == "http://localhost" + + +def test_load_agents_from_config_with_an_empty_list_clears_the_remembered_agents(): + """ + An explicitly empty config must forget the previously loaded agents. + + Otherwise the next DB rebuild replays agents the operator removed from config.yaml. + ``None`` still means "no opinion" and leaves them alone. + """ + registry = AgentRegistry() + registry.load_agents_from_config( + [{"agent_name": "removed-agent", "agent_card_params": _sample_agent_card_params()}] + ) + assert registry.config_agents != () + + registry.load_agents_from_config([]) + assert registry.config_agents == () + + registry.load_agents_from_db_and_config(db_agents=None) + assert registry.get_agent_list() == [], "a removed config agent must not come back on the next rebuild" diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 8d1d8185e4d..bcaab92ebb3 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2354,3 +2354,133 @@ async def test_ProxyConfig_load_config_redacts_secret_litellm_setting_keeps_plai assert "num_retries=7" in rendered, ( f"non-secret num_retries value was over-redacted; expected it visible in {rendered!r}" ) + + +# --------------------------------------------------------------------------- +# ProxyConfig agents from config.yaml +# --------------------------------------------------------------------------- + + +@pytest.fixture +def clean_agent_registry(): + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + 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 = () + try: + yield global_agent_registry + finally: + global_agent_registry.agent_list = original_agents + global_agent_registry.config_agents = original_config_agents + + +def _config_agent(agent_name: str) -> Dict[str, Any]: + return { + "agent_name": agent_name, + "agent_card_params": { + "name": "Config Agent", + "url": "http://localhost:10001", + "protocolVersion": "1.0", + }, + } + + +class _FakeAgentRow: + """Stand-in for a prisma agent record: supports dict() and .object_permission.""" + + def __init__(self, agent_id: str, agent_name: str) -> None: + self.agent_id = agent_id + self.agent_name = agent_name + self.object_permission = None + self.spend = 0.0 + + def __iter__(self): + return iter( + { + "agent_id": self.agent_id, + "agent_name": self.agent_name, + "agent_card_params": {"name": self.agent_name, "url": "http://db-agent"}, + "litellm_params": {}, + }.items() + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("config_key", ["agents", "agent_list"]) +async def test_ProxyConfig__init_non_llm_configs_registers_agents_from_config(clean_agent_registry, config_key): + """The documented ``agents:`` key must register agents, as must the legacy ``agent_list:``.""" + await ProxyConfig()._init_non_llm_configs( + config={config_key: [_config_agent("config-agent")]}, + config_file_path=None, + ) + + assert [agent.agent_name for agent in clean_agent_registry.get_agent_list()] == ["config-agent"] + + +@pytest.mark.asyncio +async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_agent_registry): + """A DB reload rebuilds the registry; config-defined agents must survive it alongside DB rows.""" + await ProxyConfig()._init_non_llm_configs( + config={"agents": [_config_agent("config-agent")]}, + config_file_path=None, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_agentstable.find_many = AsyncMock(return_value=[_FakeAgentRow("db-id", "db-agent")]) + + await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client) + + assert sorted(agent.agent_name for agent in clean_agent_registry.get_agent_list()) == [ + "config-agent", + "db-agent", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "config, expected_agent_names", + [ + ({"agents": [], "agent_list": [_config_agent("legacy-agent")]}, []), + ( + { + "agents": [_config_agent("documented-agent")], + "agent_list": [_config_agent("legacy-agent")], + }, + ["documented-agent"], + ), + ({"agent_list": [_config_agent("legacy-agent")]}, ["legacy-agent"]), + ], + ids=["empty-agents-wins", "populated-agents-wins", "agent_list-alone-still-works"], +) +async def test_ProxyConfig__init_non_llm_configs_prefers_agents_key_by_presence( + clean_agent_registry, config, expected_agent_names +): + """ + ``agents`` outranks the legacy ``agent_list`` whenever the key is present. + + Selecting on truthiness instead would silently register the legacy entries + for a config that spells out ``agents: []``. + """ + await ProxyConfig()._init_non_llm_configs(config=config, config_file_path=None) + + assert [agent.agent_name for agent in clean_agent_registry.get_agent_list()] == expected_agent_names + + +@pytest.mark.asyncio +async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembered_agents(clean_agent_registry): + """ + An explicitly empty ``agents:`` must reach the registry, not be skipped as falsy. + + Skipping it leaves the previously remembered agents in place, so the next DB + rebuild replays agents the operator deleted from config.yaml. + """ + clean_agent_registry.load_agents_from_config([_config_agent("stale-agent")]) + assert clean_agent_registry.config_agents != () + + await ProxyConfig()._init_non_llm_configs(config={"agents": []}, config_file_path=None) + + assert clean_agent_registry.config_agents == () + clean_agent_registry.load_agents_from_db_and_config(db_agents=None) + assert clean_agent_registry.get_agent_list() == []