diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py new file mode 100644 index 00000000000..092f9f29606 --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -0,0 +1,145 @@ +from collections.abc import Mapping +from itertools import product +from typing import Final + +from litellm.proxy._types import LiteLLMRoutes +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") +) + + +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} + if route not in _MANAGED_MODEL_ROUTES and not RouteChecks.check_route_access(route, _MANAGED_MODEL_PATHS): + return dict(body) + from litellm.proxy.common_utils.http_parsing_utils import resolve_inference_model + + kind: Final = ( + "image_generation" + if route.endswith("/images/generations") + else "image_edit" + if route.endswith("/images/edits") + else "moderation" + if route.endswith(("/moderations", "/audio/transcriptions")) + else "speech" + if route.endswith("/audio/speech") + else "body" + if route.endswith(("/rerank", "/messages/count_tokens")) + else "path" + if route.endswith(":countTokens") + else "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} + + +def actor_admission_failure( + agent: AgentResponse, + context: ManagedAgentContext | None, +) -> AgentIdentityFailure | None: + if not agent.enabled or agent.identity is None or not agent.identity.active: + return AgentIdentityFailure(message="Agent execution is disabled") + if context is None: + return AgentIdentityFailure(message="This agent requires its bound identity provider token") + if context.agent_id != agent.agent_id or context.binding_revision != agent.identity.revision: + return AgentIdentityFailure(message="Agent identity changed during authentication; retry") + if agent.execution_mode not in (context.mode, "both"): + return AgentIdentityFailure(message="Agent is not enabled for this execution mode") + if context.mode == "delegated" and not context.user_id: + return AgentIdentityFailure(message="A verified human subject is required") + return None + + +def invocation_target(route: str, body: Mapping[str, object]) -> str | None: + model: Final = body.get("model") + if isinstance(model, str) and model.startswith("a2a/"): + return model.removeprefix("a2a/") or None + components: Final = tuple(route.strip("/").split("/")) + path: Final = components[1:] if components and components[0] == "v1" else components + return path[1] if len(path) >= 2 and path[0] == "a2a" else None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 64b0c6c1967..b3be5fb78c9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -93,6 +93,7 @@ from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_bod from litellm.proxy.common_utils.http_parsing_utils import ( get_client_requested_model, get_tags_from_request_body, + resolve_inference_model, ) from litellm.proxy.common_utils.openai_error_payload import ( LITELLM_CALL_ID_HEADER, @@ -2068,11 +2069,12 @@ class ProxyBaseLLMRequestProcessing: if isinstance(model, str): reject_url_valued_destination("model", model) - self.data["model"] = ( - general_settings.get("completion_model", None) # server default - or user_model # model name passed via cli args - or model # for azure deployments - or self.data.get("model", None) # default passed in http request + self.data["model"] = resolve_inference_model( + self.data.get("model"), + general_settings, + user_model, + model, + kind="image_edit" if route_type == "aimage_edit" else "completion", ) # override with user settings, these are params passed via cli diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 1c2bd7ea217..be62448e3ed 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -2,7 +2,7 @@ import json import re from collections.abc import Collection, Mapping from types import MappingProxyType, UnionType -from typing import Annotated, Any, Final, Union, get_args, get_origin +from typing import Annotated, Any, Final, Literal, Union, get_args, get_origin import orjson from fastapi import Request, UploadFile, status @@ -25,6 +25,39 @@ _FORM_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-www-form- _ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required}) +def resolve_inference_model( + body_model: object, + settings: Mapping[str, object], + cli_model: str | None, + endpoint_model: object = None, + *, + kind: Literal[ + "completion", "image_generation", "image_edit", "moderation", "speech", "body", "path" + ] = "completion", +) -> object: + match kind: + case "image_generation": + return cli_model or endpoint_model or settings.get("image_generation_model") or body_model + case "image_edit": + return ( + settings.get("completion_model") + or cli_model + or endpoint_model + or settings.get("image_generation_model") + or body_model + ) + case "moderation": + return cli_model or settings.get("moderation_model") or body_model + case "speech": + return cli_model or body_model + case "body": + return body_model + case "path": + return endpoint_model + case "completion": + return settings.get("completion_model") or cli_model or endpoint_model or body_model + + def _normalize_media_type(content_type: str) -> str: """Return the bare media type per RFC 7231: strip params, trim, lowercase.""" if not content_type: diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index b9580ba3948..16dc38575da 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -21,6 +21,7 @@ from litellm.proxy.common_request_processing import ( from litellm.proxy.common_utils.http_parsing_utils import ( coerce_numeric_form_fields, numeric_form_fields, + resolve_inference_model, ) from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, @@ -118,14 +119,9 @@ async def image_generation( if isinstance(model, str): reject_url_valued_destination("model", model) - data["model"] = ( - model - or general_settings.get("image_generation_model", None) # server default - or user_model # model name passed via cli args - or data.get("model", None) # default passed in http request + data["model"] = resolve_inference_model( + data.get("model"), general_settings, user_model, model, kind="image_generation" ) - if user_model: - data["model"] = user_model ### MODEL ALIAS MAPPING ### # check if model name in model alias map @@ -324,12 +320,6 @@ async def image_edit_api( if "prompt" not in data: data["prompt"] = None - data["model"] = ( - model - or general_settings.get("image_generation_model", None) # server default - or user_model # model name passed via cli args - or data.get("model", None) # default passed in http request - ) ######################################################### # Process request ######################################################### @@ -346,7 +336,7 @@ async def image_edit_api( general_settings=general_settings, proxy_config=proxy_config, select_data_generator=select_data_generator, - model=None, + model=model, user_model=user_model, user_temperature=user_temperature, user_request_timeout=user_request_timeout, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d2f9a4d7d93..5de9d3d73aa 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -419,6 +419,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, check_file_size_under_limit, get_form_data, + resolve_inference_model, ) from litellm.proxy.common_utils.load_config_utils import get_config_from_bucket from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations @@ -12139,13 +12140,7 @@ async def moderations( proxy_config=proxy_config, ) - data["model"] = ( - general_settings.get("moderation_model", None) # server default - or user_model # model name passed via cli args - or data.get("model") # default passed in http request - ) - if user_model: - data["model"] = user_model + data["model"] = resolve_inference_model(data.get("model"), general_settings, user_model, kind="moderation") ### CALL HOOKS ### - modify incoming data / reject request before calling the model data = await proxy_logging_obj.pre_call_hook( @@ -12399,13 +12394,7 @@ async def audio_transcriptions( if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id - data["model"] = ( - general_settings.get("moderation_model", None) # server default - or user_model # model name passed via cli args - or data.get("model", None) # default passed in http request - ) - if user_model: - data["model"] = user_model + data["model"] = resolve_inference_model(data.get("model"), general_settings, user_model, kind="moderation") router_model_names: Final = llm_router.model_names if llm_router is not None else [] diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py new file mode 100644 index 00000000000..7e1e4e6c51f --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py @@ -0,0 +1,184 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.proxy.agent_endpoints.auth.managed_authorization import ( + actor_admission_failure, + invocation_target, +) +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import AgentIdentityBinding, AgentIdentityFailure, ManagedAgentContext + +BINDING: Final = AgentIdentityBinding( + agent_id="agent", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + service_principal_id="principal", + issuer="issuer", + revision="current", +) + + +def agent(**overrides: object) -> AgentResponse: + return AgentResponse.model_validate( + { + "agent_id": "agent", + "agent_name": "Agent", + "agent_card_params": {}, + "identity": BINDING, + "identity_managed": True, + "execution_mode": "both", + **overrides, + } + ) + + +@pytest.mark.parametrize( + "state", + [ + {"enabled": False}, + {"identity": None}, + {"identity": BINDING.model_copy(update={"active": False})}, + {"execution_mode": "delegated"}, + ], +) +def test_keys_cannot_bypass_lifecycle_or_delegated_only_mode(state: dict[str, object]) -> None: + assert isinstance(actor_admission_failure(agent(**state), None), AgentIdentityFailure) + + +@pytest.mark.parametrize("mode", ["autonomous", "both", "delegated"]) +def test_keys_cannot_impersonate_an_entra_bound_agent(mode: str) -> None: + assert isinstance(actor_admission_failure(agent(execution_mode=mode), None), AgentIdentityFailure) + + +@pytest.mark.parametrize( + "context", + [ + ManagedAgentContext(agent_id="agent", binding_revision="previous", mode="autonomous"), + ManagedAgentContext(agent_id="another", binding_revision="current", mode="autonomous"), + ManagedAgentContext(agent_id="agent", binding_revision="current", mode="delegated"), + ], +) +def test_stale_binding_and_unverified_delegation_cannot_pass_admission(context: ManagedAgentContext) -> None: + assert isinstance(actor_admission_failure(agent(), context), AgentIdentityFailure) + + +@pytest.mark.parametrize( + "route,body,expected", + [ + ("/a2a/agent", {}, "agent"), + ("/v1/a2a/agent/", {}, "agent"), + ("/v1/chat/completions", {"model": "a2a/Readable name"}, "Readable name"), + ("/v1/chat/completions", {"model": "a2a/"}, None), + ("/v1/chat/completions", {"model": "ordinary-model"}, None), + ("/a2a", {}, None), + ], +) +def test_invocation_routes_resolve_the_same_target(route: str, body: dict[str, object], expected: str | None) -> None: + assert invocation_target(route, body) == expected + + +def test_execution_mode_must_match_verified_token_mode() -> None: + context: Final = ManagedAgentContext(agent_id="agent", binding_revision="current", mode="autonomous") + failure: Final = actor_admission_failure(agent(execution_mode="delegated"), context) + assert isinstance(failure, AgentIdentityFailure) + assert "execution mode" in failure.message + + +@pytest.mark.parametrize("mode,user", [("autonomous", None), ("delegated", "verified-human")]) +def test_matching_identity_revision_and_execution_mode_pass_admission(mode: str, user: str | None) -> None: + context: Final = ManagedAgentContext.model_validate( + {"agent_id": "agent", "binding_revision": "current", "mode": mode, "user_id": user} + ) + assert actor_admission_failure(agent(), context) is None + + +@pytest.mark.parametrize( + "route,method,allowed", + [ + ("/v1/agents", "GET", True), + ("/v1/agents", "POST", False), + ("/v1/chat/completions", "POST", True), + ("/v1/chat/completions", "DELETE", False), + ("/openai/deployments/model/chat/completions", "POST", True), + ("/engines/openai/model/chat/completions", "POST", True), + ("/openai/deployments/openai/model/images/generations", "POST", True), + ("/openai/deployments/openai/model/images/edits", "POST", True), + ("/v1beta/models/gemini-model:generateContent", "POST", True), + ("/v1/realtime", "GET", True), + ("/v1/realtime", "POST", False), + ("/v1/realtime/client_secrets", "POST", False), + ("/mcp/tools/call", "POST", True), + ("/a2a/target/message/send", "POST", True), + ("/v1/agents/target", "PATCH", False), + ("/v1/responses/other-response", "GET", False), + ("/v1/files", "GET", False), + ("/v1/files", "POST", False), + ("/openai/v1/files", "GET", False), + ("/anthropic/v1/files", "GET", False), + ], +) +def test_managed_route_scope_excludes_provider_resources(route: str, method: str, allowed: bool) -> None: + from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_route_allowed + + assert managed_agent_route_allowed(route, method) is allowed + + +@pytest.mark.parametrize( + "route,body,settings,cli_model,path_model,expected", + [ + ("/v1/chat/completions", {"model": "body"}, {"completion_model": "default"}, "cli", "path", "default"), + ("/v1/moderations", {"model": "body"}, {"moderation_model": "default"}, "cli", None, "cli"), + ("/v1/audio/speech", {"model": "body"}, {"completion_model": "ignored"}, None, None, "body"), + ("/openai/deployments/path/embeddings", {"model": "body"}, {}, None, "path", "path"), + ("/v1/messages/count_tokens", {"model": "body"}, {"completion_model": "ignored"}, "cli", None, "body"), + ("/mcp/tools/call", {}, {"completion_model": "ignored"}, "cli", None, None), + ("/v1/images/generations", {"model": "image"}, {"completion_model": "text"}, None, None, "image"), + ("/v1/images/generations", {}, {"image_generation_model": "image"}, None, None, "image"), + ("/v1/images/edits", {}, {"image_generation_model": "image"}, None, None, "image"), + ("/v1/rerank", {"model": "reranker"}, {"completion_model": "text"}, "cli", None, "reranker"), + ("/v1beta/models/path:countTokens", {"model": "body"}, {"completion_model": "text"}, "cli", "path", "path"), + ], +) +def test_managed_inference_resolves_dispatch_precedence(route, body, settings, cli_model, path_model, expected): + from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_inference_request + + assert managed_inference_request(route, body, settings, cli_model, path_model).get("model") == expected + + +def test_managed_inference_without_any_model_cannot_skip_model_grants(): + from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_inference_request + + with pytest.raises(HTTPException, match="explicit or configured model"): + managed_inference_request("/v1/moderations", {}, {}, None) + + +@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/images/generations", "/v1/images/edits"]) +def test_managed_inference_query_model_takes_precedence_over_body(route: str): + from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_inference_request + + assert managed_inference_request(route, {"model": "body"}, {}, None, query_model="query")["model"] == "query" + + +def test_managed_inference_ignores_unsupported_query_model(): + from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_inference_request + + assert ( + managed_inference_request("/v1/messages", {"model": "body"}, {}, None, query_model="query")["model"] == "body" + ) + + +@pytest.mark.parametrize("route", ["/realtime", "/v1/realtime", "/openai/v1/realtime"]) +def test_managed_realtime_requires_a_model_and_ignores_completion_defaults(route: str) -> None: + from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_inference_request + + with pytest.raises(HTTPException, match="explicit or configured model"): + managed_inference_request(route, {}, {"completion_model": "allowed-default"}, "cli") + assert ( + managed_inference_request(route, {"model": "requested"}, {"completion_model": "allowed-default"}, "cli")[ + "model" + ] + == "requested" + ) diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index 7929a0b21af..d11d0ee3a58 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -1210,3 +1210,28 @@ class TestCoerceNumericFormFields: numeric_fields=self.numeric_fields, ) assert result == {"n": 3, "temperature": None, "image": buffer} + + +@pytest.mark.parametrize( + "kind,settings,cli,path,body,expected", + [ + ("completion", {"completion_model": "default"}, "cli", "path", "body", "default"), + ("completion", {}, "cli", "path", "body", "cli"), + ("completion", {}, None, "path", "body", "path"), + ("completion", {}, None, None, "body", "body"), + ("image_generation", {"completion_model": "text", "image_generation_model": "image"}, None, None, "body", "image"), + ("image_generation", {"image_generation_model": "image"}, "cli", "path", "body", "cli"), + ("image_generation", {"image_generation_model": "image"}, None, "path", "body", "path"), + ("image_edit", {"completion_model": "text", "image_generation_model": "image"}, None, None, "body", "text"), + ("image_edit", {"image_generation_model": "image"}, None, "path", "body", "path"), + ("image_edit", {"image_generation_model": "image"}, None, None, "body", "image"), + ("moderation", {"moderation_model": "mod"}, "cli", None, "body", "cli"), + ("speech", {"completion_model": "text"}, None, None, "body", "body"), + ("body", {"completion_model": "text"}, "cli", None, "body", "body"), + ("path", {"completion_model": "text"}, "cli", "path", "body", "path"), + ], +) +def test_shared_inference_model_selection_preserves_handler_precedence(kind, settings, cli, path, body, expected): + from litellm.proxy.common_utils.http_parsing_utils import resolve_inference_model + + assert resolve_inference_model(body, settings, cli, path, kind=kind) == expected