From cdb29946eb75a792566d8447e1e261387fa38368 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 16 Apr 2026 17:22:28 -0700 Subject: [PATCH] fix: align agent endpoint and routing permission checks with existing pattern --- litellm/proxy/agent_endpoints/a2a_routing.py | 23 +++++++++++++- litellm/proxy/agent_endpoints/endpoints.py | 32 ++++++++++++++------ litellm/proxy/common_request_processing.py | 25 +++++++++------ litellm/proxy/route_llm_request.py | 6 +++- 4 files changed, 64 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 8f951414994..d5f5348731f 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -8,10 +8,17 @@ Looks up agents in the registry and injects their API base URL. from typing import Any, Optional import litellm +from fastapi import HTTPException + from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth -def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]: +async def route_a2a_agent_request( + data: dict, + route_type: str, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, +) -> Optional[Any]: """ Route A2A agent requests directly to litellm with injected API base. @@ -19,6 +26,9 @@ def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]: """ # 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.route_llm_request import ( ROUTE_ENDPOINT_MAPPING, ProxyModelNotFoundError, @@ -40,6 +50,17 @@ def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]: route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError(route=route_name, model_name=model_name) + # Verify the caller is permitted to use this agent + is_allowed = await AgentRequestHandler.is_agent_allowed( + agent_id=agent.agent_id, + user_api_key_auth=user_api_key_dict, + ) + if not is_allowed: + raise HTTPException( + status_code=403, + detail=f"Agent '{agent_name}' is not allowed for your key/team. Contact proxy admin for access.", + ) + # Get API base URL from agent config if not agent.agent_card_params or "url" not in agent.agent_card_params: verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' has no URL configured") diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 64c20d5ed5e..4c5fbaa22c5 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -200,10 +200,9 @@ async def get_agents( for agent in returned_agents: if agent.litellm_params is None: agent.litellm_params = {} - agent.litellm_params[ - "is_public" - ] = litellm.public_agent_groups is not None and ( - agent.agent_id in litellm.public_agent_groups + agent.litellm_params["is_public"] = ( + litellm.public_agent_groups is not None + and (agent.agent_id in litellm.public_agent_groups) ) # Redact sensitive fields for non-admin users @@ -393,6 +392,19 @@ async def get_agent_by_id( """ await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + AgentRequestHandler, + ) + + is_allowed = await AgentRequestHandler.is_agent_allowed( + agent_id=agent_id, user_api_key_auth=user_api_key_dict + ) + if not is_allowed: + raise HTTPException( + status_code=403, + detail=f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.", + ) + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -409,13 +421,13 @@ async def get_agent_by_id( agent_dict = agent_row.model_dump() if agent_row.object_permission is not None: try: - agent_dict[ - "object_permission" - ] = agent_row.object_permission.model_dump() + agent_dict["object_permission"] = ( + agent_row.object_permission.model_dump() + ) except Exception: - agent_dict[ - "object_permission" - ] = agent_row.object_permission.dict() + agent_dict["object_permission"] = ( + agent_row.object_permission.dict() + ) agent = AgentResponse(**agent_dict) # type: ignore else: # Agent found in memory — refresh spend from DB diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c4717ad9cf3..97801baaf0c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -866,9 +866,11 @@ class ProxyBaseLLMRequestProcessing: "Request received by LiteLLM: payload too large to log (%d bytes, limit %d). Keys: %s", len(_payload_str), MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, - list(self.data.keys()) - if isinstance(self.data, dict) - else type(self.data).__name__, + ( + list(self.data.keys()) + if isinstance(self.data, dict) + else type(self.data).__name__ + ), ) else: verbose_proxy_logger.debug( @@ -1054,6 +1056,7 @@ class ProxyBaseLLMRequestProcessing: route_type=route_type, llm_router=llm_router, user_model=user_model, + user_api_key_dict=user_api_key_dict, ) tasks.append(llm_call) @@ -1128,9 +1131,9 @@ class ProxyBaseLLMRequestProcessing: # aliasing/routing, but the OpenAI-compatible response `model` field should reflect # what the client sent. if requested_model_from_client: - self.data[ - "_litellm_client_requested_model" - ] = requested_model_from_client + self.data["_litellm_client_requested_model"] = ( + requested_model_from_client + ) # Streaming: attach a closure that fires after all guardrail # end-of-stream blocks complete. CSW.__anext__ stores the @@ -1731,7 +1734,9 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.debug("inside generator") try: str_so_far = "" - async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook( + async for ( + chunk + ) in proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=response, request_data=request_data, @@ -1959,9 +1964,9 @@ class ProxyBaseLLMRequestProcessing: # Add cache-related fields to **params (handled by Usage.__init__) if cache_creation_input_tokens is not None: - usage_kwargs[ - "cache_creation_input_tokens" - ] = cache_creation_input_tokens + usage_kwargs["cache_creation_input_tokens"] = ( + cache_creation_input_tokens + ) if cache_read_input_tokens is not None: usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index f1590b16c24..17cc4374560 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional from fastapi import HTTPException, status import litellm +from litellm.proxy._types import UserAPIKeyAuth if TYPE_CHECKING: from litellm.router import Router as _Router @@ -314,6 +315,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin "acancel_run", "adelete_run", ], + user_api_key_dict: Optional[UserAPIKeyAuth] = None, ): """ Common helper to route the request @@ -548,7 +550,9 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin route_a2a_agent_request, ) - result = route_a2a_agent_request(data, route_type) + result = await route_a2a_agent_request( + data, route_type, user_api_key_dict=user_api_key_dict + ) if result is not None: return result # Fall through to raise exception below if result is None