diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 60367178c7f..7d06568792b 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -1,5 +1,4 @@ import hashlib -import json from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone from typing import Any, Protocol, TypedDict @@ -114,8 +113,15 @@ class AgentRegistry: public_agent_list.append(agent) return public_agent_list - def _create_agent_id(self, agent_config: AgentConfig) -> str: - return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest() + def _create_agent_id(self, agent_name: str) -> str: + """ + Derive a config agent's id from its name, which is unique across the registry. + + Hashing the whole config entry instead would fold every mutable value into the + identity, so rotating a secret in ``static_headers`` would mint a new agent id + and silently invalidate the permissions granted on the old one. + """ + return hashlib.sha256(agent_name.encode()).hexdigest() def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None): """ @@ -139,14 +145,14 @@ class AgentRegistry: agent_name = agent_config_item.get("agent_name") agent_card_params = agent_config_item.get("agent_card_params") - if not all([agent_name, agent_card_params]): + if not agent_name or not 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) + config_hash = self._create_agent_id(agent_name) self.register_agent(agent_config=AgentResponse(agent_id=config_hash, **agent_config_item)) # type: ignore 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 f506535b5e1..fc8b6b78d68 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -279,3 +279,54 @@ def test_load_agents_from_config_with_an_empty_list_clears_the_remembered_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" + + +def test_config_agent_id_survives_a_static_header_secret_rotation(): + """ + A config agent's id must not move when a value inside the entry changes. + + The id used to be a sha256 of the whole config entry with `os.environ/...` + placeholders already resolved, so rotating an upstream token minted a brand new + agent id: every permission granted on the old id stopped matching and the keys + that had been secured lost access while unrestricted keys kept theirs. + """ + registry_before = AgentRegistry() + registry_before.load_agents_from_config( + [ + { + "agent_name": "rotating-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + ] + ) + + registry_after = AgentRegistry() + registry_after.load_agents_from_config( + [ + { + "agent_name": "rotating-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v2"}, + } + ] + ) + + agent_before = registry_before.get_agent_by_name("rotating-agent") + agent_after = registry_after.get_agent_by_name("rotating-agent") + assert agent_before is not None and agent_after is not None + assert agent_before.agent_id == agent_after.agent_id + + +def test_config_agents_with_different_names_get_different_ids(): + """Distinct config agents must stay distinguishable for permission grants.""" + registry = AgentRegistry() + registry.load_agents_from_config( + [ + {"agent_name": "agent-alpha", "agent_card_params": _sample_agent_card_params()}, + {"agent_name": "agent-beta", "agent_card_params": _sample_agent_card_params()}, + ] + ) + + agent_ids = {agent.agent_id for agent in registry.get_agent_list()} + assert len(agent_ids) == 2