feat(agents): authenticate Entra identities and delegated requests

litellm/proxy/_experimental/mcp_server/idp_token_exchange.py

@@ -170,6 +170,11 @@ async def identity_from_subject_token(
         return _refusal_for(denied, denied.message)
     except Exception as denied:  # noqa: BLE001  # auth_jwt raises a plain Exception on signature and claim failures
         return _refusal_for(denied, denied)
+    if result.get("agent_id") is not None:
+        return SubjectTokenRefusal(
+            error="invalid_request",
+            description="Agent tokens require direct JWT authentication; this exchange supports users only",
+        )
     user_id: Final = result["user_id"]
     if user_id is None:
         return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows")

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -5,7 +5,7 @@
 from datetime import datetime
 from traceback import walk_tb
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Literal
+from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict
 from uuid import uuid4
 
 import anyio
@@ -14,6 +14,7 @@
 from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
 from pydantic import ValidationError
 from starlette.datastructures import Headers
+from typing_extensions import ReadOnly
 
 from litellm._logging import verbose_logger
 from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT
@@ -63,7 +64,27 @@
     from litellm.proxy.utils import ProxyLogging
 from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
 from litellm.types.mcp import MCPAuth
-from litellm.types.utils import CallTypes
+from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall
+
+
+class _MCPModelMetadata(TypedDict):
+    model_group: ReadOnly[str]
+
+
+def _stamp_mcp_tool_metadata(logging_obj: "LiteLLMLoggingObj | None", server_id: str, tool_name: str) -> None:
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
+
+    if logging_obj is None:
+        return
+    server: Final = global_mcp_server_manager.get_mcp_server_by_id(
+        server_id
+    ) or global_mcp_server_manager.get_mcp_server_by_name(server_id)
+    metadata: Final[StandardLoggingMCPToolCall] = {
+        "name": tool_name,
+        "mcp_server_name": server.name if server is not None else server_id,
+    }
+    logging_obj.model_call_details["mcp_tool_call_metadata"] = metadata
+
 
 MCP_AVAILABLE: bool = True
 try:
@@ -1193,6 +1214,12 @@ async def call_tool_rest_api(
                     },
                 )
 
+            data["model"] = f"MCP: {tool_name}"
+            model_metadata: Final[_MCPModelMetadata] = {
+                **(data.get("metadata") or MappingProxyType({})),
+                "model_group": f"MCP: {tool_name}",
+            }
+            data["metadata"] = model_metadata
             proxy_base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
             _request_start_time: Final = datetime.now()  # noqa: DTZ005  # naive to match the tool start time below
             try:
@@ -1226,6 +1253,8 @@ async def call_tool_rest_api(
                 if "metadata" in data and "user_api_key_auth" in data["metadata"]:
                     data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"]
 
+                _stamp_mcp_tool_metadata(logging_obj, server_id, tool_name)
+
                 # Resolve allowed MCP servers with IP filtering
                 (
                     allowed_mcp_servers,

litellm/proxy/_types.py

@@ -3320,6 +3320,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView):  # the expected response ob
     # single-owner so its meaning stays trustworthy.
     mcp_session_resource_server_id: str | None = Field(default=None, exclude=True)
     mcp_toolset_id: str | None = Field(default=None, exclude=True)
+    authenticated_by_custom_auth: bool = Field(default=False, exclude=True)
     via_virtual_key: bool = Field(
         default=False,
         exclude=True,
@@ -3381,6 +3382,7 @@ def check_api_key(cls, values):
         values.pop("mcp_session_resource_server_id", None)
         values.pop("mcp_toolset_id", None)
         values.pop("via_virtual_key", None)
+        values.pop("authenticated_by_custom_auth", None)
         values.pop("agent_caller", None)
         values.pop("managed_agent_context", None)
         values.pop("managed_agent_policy", None)

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -722,6 +722,8 @@ async def invoke_agent_a2a(
                 detail=f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.",
             )
 
+        user_api_key_dict.invoked_agent_id = agent.agent_id
+
         _enforce_inbound_trace_id(agent, request)
 
         # Get backend URL and agent name
@@ -759,6 +761,10 @@ async def invoke_agent_a2a(
         if "metadata" not in body:
             body["metadata"] = {}
         body["metadata"]["agent_id"] = agent.agent_id
+        body["metadata"]["model_group"] = f"a2a_agent/{agent_name}"
+        body["metadata"]["model_info"] = {  # mutable-ok: request hooks mutate metadata before JSON logging
+            "id": agent.agent_id
+        }
         body["agent_id"] = agent.agent_id
 
         body.update(
@@ -862,6 +868,7 @@ async def invoke_agent_a2a(
             # results written by the unified_guardrail hook are captured.
             logging_obj._defer_async_logging = True
             response = await asend_message(
+                model=f"a2a_agent/{agent_name}",
                 request=a2a_request,
                 api_base=agent_url,
                 litellm_params=litellm_params,

litellm/proxy/agent_endpoints/a2a_routing.py

@@ -57,7 +57,7 @@ async def route_a2a_agent_request(
         user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
         or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
     )
-    if not is_admin:
+    if not is_admin or agent.identity_managed:
         is_allowed: Final = await AgentRequestHandler.is_agent_allowed(
             agent_id=agent.agent_id,
             user_api_key_auth=user_api_key_dict,

litellm/proxy/agent_endpoints/auth/agent_permission_handler.py

@@ -177,11 +177,10 @@ async def is_agent_allowed(
 
         registered: Final = global_agent_registry.get_agent_by_id(agent_id)
         registry_managed: Final = isinstance(registered, AgentResponse) and registered.identity_managed
-        if registry_managed or (registered is None and prisma_client is not None):
+        if registry_managed or prisma_client is not None:
             target: Final = await AgentIdentityStore.from_client(prisma_client).agent(agent_id)
             if isinstance(target, AgentIdentityFailure):
-                if registry_managed:
-                    raise_identity_failure(target)
+                raise_identity_failure(target)
             elif target is None and registry_managed:
                 return False
             elif isinstance(target, AgentResponse) and target.identity_managed:
@@ -200,6 +199,7 @@ async def is_agent_allowed(
                     if key_hash
                     and managed_agent_policy(user_api_key_auth) is None
                     and not user_api_key_auth.is_session_token
+                    and not user_api_key_auth.authenticated_by_custom_auth
                     else user_api_key_auth
                 )
                 fresh_auth: Final = authority.model_copy(
@@ -678,14 +678,44 @@ async def _managed_actor_agent_access(auth: UserAPIKeyAuth) -> AgentAccess:
     return RestrictedAgentAccess(capped.intersection(human_ids))
 
 
-async def verified_human_agent_grants(user_id: str | None, team_id: str | None = None) -> frozenset[str]:
+async def _verified_human_agent_sources(
+    user_id: str | None, *, allowed_team_ids: frozenset[str] | None = None
+) -> tuple[tuple[str | None, frozenset[str]], ...]:
     from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
 
     if user_id is None:
-        return frozenset()
+        return ()
     human: Final = await MCPRequestHandler.reload_admitted_user(user_id, requires_fresh_policy=True)
-    sources: Final = await MCPRequestHandler.admitted_subject_sources(
-        human, allowed_team_ids=frozenset((team_id,)) if team_id else frozenset()
+    sources: Final = await MCPRequestHandler.admitted_subject_sources(human, allowed_team_ids=allowed_team_ids)
+    access: Final = await asyncio.gather(*(_strict_agent_access(source) for source in sources))
+    return tuple((source.team_id, _granted_ids(grant)) for source, grant in zip(sources, access, strict=True))
+
+
+async def verified_human_agent_grants(user_id: str | None, team_id: str | None = None) -> frozenset[str]:
+    sources: Final = await _verified_human_agent_sources(
+        user_id, allowed_team_ids=frozenset((team_id,)) if team_id else frozenset()
+    )
+    return frozenset().union(*(grants for source, grants in sources if source is None or source == team_id))
+
+
+async def resolve_delegated_agent_team(
+    user_id: str | None,
+    agent_id: str,
+    team_id: str | None,
+    *,
+    explicit_team: bool,
+    allowed_team_ids: frozenset[str] | None = None,
+) -> str | None:
+    sources: Final = await _verified_human_agent_sources(user_id)
+    if any(source is None and agent_id in grants for source, grants in sources):
+        return team_id
+    granting_teams: Final = frozenset(
+        source
+        for source, grants in sources
+        if source is not None and agent_id in grants and (allowed_team_ids is None or source in allowed_team_ids)
     )
-    human_access: Final = await asyncio.gather(*(_strict_agent_access(source) for source in sources))
-    return frozenset().union(*(_granted_ids(access) for access in human_access))
+    if team_id in granting_teams:
+        return team_id
+    if not explicit_team and granting_teams:
+        return min(granting_teams)
+    raise HTTPException(403, "Select a team that grants access to this agent using x-litellm-team-id")

litellm/proxy/agent_endpoints/auth/managed_authorization.py

@@ -1,11 +1,129 @@
-from typing import Final
+from collections.abc import Mapping
+from itertools import product
+from types import MappingProxyType
+from typing import Annotated, Final, Literal
 
-from litellm.proxy._types import UserAPIKeyAuth
+from pydantic import Field, TypeAdapter, ValidationError
+
+from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth
 from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore
 from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure
 from litellm.types.agents import AgentResponse
 from litellm.types.proxy.agent_identity import AgentIdentityFailure, ManagedAgentContext
 
+_MANAGED_REALTIME_ROUTES: Final = frozenset(("/realtime", "/v1/realtime", "/openai/v1/realtime"))
+_MANAGED_MODEL_ROUTES: Final = frozenset(
+    f"{prefix}/{operation}"
+    for prefix, operation in product(
+        ("", "/v1"),
+        (
+            "chat/completions",
+            "completions",
+            "embeddings",
+            "responses",
+            "messages",
+            "messages/count_tokens",
+            "images/generations",
+            "images/edits",
+            "audio/transcriptions",
+            "audio/speech",
+            "moderations",
+            "rerank",
+            "ocr",
+        ),
+    )
+) | frozenset(
+    (
+        "/openai/v1/responses",
+        "/v2/rerank",
+        "/claude_code_gateway/v1/messages",
+        "/claude_code_gateway/v1/messages/count_tokens",
+        "/cursor/chat/completions",
+    )
+)
+_MANAGED_MODEL_PATHS: Final = (
+    "/engines/{model:path}/chat/completions",
+    "/engines/{model:path}/completions",
+    "/engines/{model:path}/embeddings",
+    "/openai/deployments/{model:path}/chat/completions",
+    "/openai/deployments/{model:path}/completions",
+    "/openai/deployments/{model:path}/embeddings",
+    "/openai/deployments/{model:path}/images/generations",
+    "/openai/deployments/{model:path}/images/edits",
+    "/v1beta/models/{model_name:path}:countTokens",
+    "/v1beta/models/{model_name:path}:generateContent",
+    "/v1beta/models/{model_name:path}:streamGenerateContent",
+    "/models/{model_name:path}:countTokens",
+    "/models/{model_name:path}:generateContent",
+    "/models/{model_name:path}:streamGenerateContent",
+)
+_MANAGED_MCP_ROUTES: Final = tuple(
+    route for route in LiteLLMRoutes.mcp_inference_routes.value if route not in ("/token", "/introspect")
+)
+
+
+_MODEL_ROUTE_KINDS: Final[
+    Mapping[str, Literal["image_generation", "image_edit", "moderation", "speech", "body", "path"]]
+] = MappingProxyType(
+    {
+        "/images/generations": "image_generation",
+        "/images/edits": "image_edit",
+        "/moderations": "moderation",
+        "/audio/transcriptions": "moderation",
+        "/audio/speech": "speech",
+        "/rerank": "body",
+        "/messages/count_tokens": "body",
+        ":countTokens": "path",
+    }
+)
+
+
+def managed_agent_route_allowed(route: str, method: str | None) -> bool:
+    from litellm.proxy.auth.route_checks import RouteChecks
+
+    if route in ("/agents", "/v1/agents"):
+        return method in (None, "GET", "HEAD")
+    if route in _MANAGED_REALTIME_ROUTES:
+        return method in (None, "GET")
+    if route in _MANAGED_MODEL_ROUTES or RouteChecks.check_route_access(route, _MANAGED_MODEL_PATHS):
+        return method in (None, "POST")
+    return RouteChecks.check_route_access(route, _MANAGED_MCP_ROUTES) or RouteChecks.check_route_access(
+        route, LiteLLMRoutes.agent_inference_routes.value
+    )
+
+
+def managed_inference_request(
+    route: str,
+    body: Mapping[str, object],
+    settings: Mapping[str, object],
+    cli_model: str | None,
+    path_model: object = None,
+    query_model: object = None,
+) -> dict[str, object]:
+    from litellm.proxy.auth.route_checks import RouteChecks
+
+    if route in _MANAGED_REALTIME_ROUTES:
+        model: Final = query_model or body.get("model")
+        if not isinstance(model, str) or not model:
+            raise_identity_failure(
+                AgentIdentityFailure(message="Managed inference requires an explicit or configured model")
+            )
+        return {**body, "model": model}  # mutable-ok: centralized auth hooks add request tags and budget metadata
+    if route not in _MANAGED_MODEL_ROUTES and not RouteChecks.check_route_access(route, _MANAGED_MODEL_PATHS):
+        return dict(body)  # mutable-ok: centralized auth hooks add request tags and budget metadata
+    from litellm.proxy.common_utils.http_parsing_utils import resolve_inference_model
+
+    kind: Final = next((kind for suffix, kind in _MODEL_ROUTE_KINDS.items() if route.endswith(suffix)), "completion")
+    endpoint_model: Final = path_model or (
+        query_model if route.endswith(("/completions", "/embeddings", "/images/generations", "/images/edits")) else None
+    )
+    effective: Final = resolve_inference_model(body.get("model"), settings, cli_model, endpoint_model, kind=kind)
+    if not isinstance(effective, str) or not effective:
+        raise_identity_failure(
+            AgentIdentityFailure(message="Managed inference requires an explicit or configured model")
+        )
+    return {**body, "model": effective}  # mutable-ok: centralized auth hooks add request tags and budget metadata
+
 
 def managed_agent_policy(auth: "UserAPIKeyAuth | None") -> AgentResponse | None:
     """The admitted managed policy, or ``None`` when the subject was never admitted as a managed agent.
@@ -82,3 +200,53 @@ def actor_admission_failure(
     if context.mode == "delegated" and not context.user_id:
         return AgentIdentityFailure(message="A verified human subject is required")
     return None
+
+
+_INVOCATION_COST: Final = TypeAdapter(Annotated[float, Field(ge=0, allow_inf_nan=False)])
+
+
+def invocation_target(route: str, body: Mapping[str, object]) -> str | None:
+    components: Final = tuple(route.strip("/").split("/"))
+    path: Final = components[1:] if components and components[0] == "v1" else components
+    if len(path) >= 2 and path[0] == "a2a":
+        return path[1] or None
+    model: Final = body.get("model")
+    return model.removeprefix("a2a/") or None if isinstance(model, str) and model.startswith("a2a/") else None
+
+
+async def prepare_agent_invocation(
+    auth: UserAPIKeyAuth, target_name: str, store: AgentIdentityStore | None, *, billable: bool = True
+) -> None:
+    from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentRequestHandler
+    from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
+
+    registered: Final = await get_agent_with_read_through(target_name)
+    if registered is None:
+        return
+    registered_managed: Final = registered.identity_managed or registered.identity is not None
+    if store is None and registered_managed:
+        raise_identity_failure(
+            AgentIdentityFailure(code="policy_unavailable", message="Managed agent policy requires a database")
+        )
+    target: Final = await store.agent(registered.agent_id) if store is not None else None
+    if isinstance(target, AgentIdentityFailure):
+        raise_identity_failure(target)
+    if target is None and registered_managed:
+        raise_identity_failure(AgentIdentityFailure(message="Invoked agent no longer exists"))
+    effective: Final = target if target is not None else registered
+    if not effective.identity_managed and auth.managed_agent_policy is None:
+        return
+    if not await AgentRequestHandler.is_agent_allowed(effective.agent_id, auth):
+        raise_identity_failure(AgentIdentityFailure(message="The caller is not permitted to invoke this agent"))
+    auth.invoked_agent_id = effective.agent_id
+    auth.invoked_agent_policy = effective
+    if auth.agent_id is None and effective.identity_managed:
+        auth.billing_agent_policy = effective
+    raw_fee: Final = (effective.litellm_params or MappingProxyType({})).get("cost_per_query", 0.0) if billable else 0.0
+    try:
+        fee: Final = _INVOCATION_COST.validate_python(raw_fee)
+    except ValidationError:
+        raise_identity_failure(
+            AgentIdentityFailure(code="policy_unavailable", message="Agent invocation price is invalid")
+        )
+    auth.agent_invocation_cost = fee