From 9aa3afc8d74a6afdc5ad225ba81cd73f7f8c2137 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 4 Mar 2026 19:02:17 -0800 Subject: [PATCH] fix: wire enable_tracing UI toggle to backend and optimize agent lookup - Map enable_tracing form field to require_trace_id_on_calls_by_agent in buildAgentDataFromForm and parseAgentForForm so the tracing toggle actually takes effect end-to-end - Replace O(n) linear scan in AgentRegistry.get_agent_by_id with O(1) dict lookup via _agent_by_id index, reducing latency in auth hot path - Update test cleanup to preserve _agent_by_id consistency Co-Authored-By: Claude Opus 4.6 --- .../proxy/agent_endpoints/agent_registry.py | 18 +++++++++--------- tests/test_litellm/test_a2a_registry_lookup.py | 4 +++- .../src/components/agents/agent_config.ts | 2 ++ 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 159c9fb93d9..e4f246c8a0e 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -15,14 +15,21 @@ from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest class AgentRegistry: def __init__(self): self.agent_list: List[AgentResponse] = [] + self._agent_by_id: Dict[str, AgentResponse] = {} def reset_agent_list(self): self.agent_list = [] + self._agent_by_id = {} def register_agent(self, agent_config: AgentResponse): self.agent_list.append(agent_config) + if agent_config.agent_id: + self._agent_by_id[agent_config.agent_id] = agent_config def deregister_agent(self, agent_name: str): + for agent in self.agent_list: + if agent.agent_name == agent_name and agent.agent_id: + self._agent_by_id.pop(agent.agent_id, None) self.agent_list = [ agent for agent in self.agent_list if agent.agent_name != agent_name ] @@ -358,16 +365,9 @@ class AgentRegistry: agent_id: str, ) -> Optional[AgentResponse]: """ - Get an agent by its ID from the database + Get an agent by its ID using O(1) dict lookup """ - try: - for agent in self.agent_list: - if agent.agent_id == agent_id: - return agent - - return None - except Exception as e: - raise Exception(f"Error getting agent from DB: {str(e)}") + return self._agent_by_id.get(agent_id) def get_agent_by_name(self, agent_name: str) -> Optional[AgentResponse]: """ diff --git a/tests/test_litellm/test_a2a_registry_lookup.py b/tests/test_litellm/test_a2a_registry_lookup.py index 9938f10a43f..c84ada2566d 100644 --- a/tests/test_litellm/test_a2a_registry_lookup.py +++ b/tests/test_litellm/test_a2a_registry_lookup.py @@ -56,8 +56,9 @@ def test_a2a_registry_integration(): # Register and test original_agents = global_agent_registry.agent_list.copy() + original_agent_by_id = global_agent_registry._agent_by_id.copy() global_agent_registry.register_agent(test_agent) - + try: litellm.completion( model="a2a/test-agent", @@ -68,6 +69,7 @@ def test_a2a_registry_integration(): assert "registry-url.example.com" in str(e) or "APIConnectionError" in str(type(e).__name__) finally: global_agent_registry.agent_list = original_agents + global_agent_registry._agent_by_id = original_agent_by_id except ImportError: pytest.skip("Registry not available (not in proxy context)") diff --git a/ui/litellm-dashboard/src/components/agents/agent_config.ts b/ui/litellm-dashboard/src/components/agents/agent_config.ts index 3fc80373025..ac71f5ae99c 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_config.ts +++ b/ui/litellm-dashboard/src/components/agents/agent_config.ts @@ -278,6 +278,7 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { if (values.cost_per_query) params.cost_per_query = parseFloat(values.cost_per_query); if (values.input_cost_per_token) params.input_cost_per_token = parseFloat(values.input_cost_per_token); if (values.output_cost_per_token) params.output_cost_per_token = parseFloat(values.output_cost_per_token); + if (values.enable_tracing !== undefined) params.require_trace_id_on_calls_by_agent = values.enable_tracing; if (Object.keys(params).length > 0) { agentData.litellm_params = params; @@ -316,5 +317,6 @@ export const parseAgentForForm = (agent: any) => { cost_per_query: agent.litellm_params?.cost_per_query, input_cost_per_token: agent.litellm_params?.input_cost_per_token, output_cost_per_token: agent.litellm_params?.output_cost_per_token, + enable_tracing: agent.litellm_params?.require_trace_id_on_calls_by_agent, }; };