fix: align agent endpoint and routing permission checks with existing pattern

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")

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

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
 

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