mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): scope the body model_id ban to Bedrock, cover the queue route, keep stored creds on test_connection
model_id only picks the upstream model on Bedrock. SageMaker reads it as an inference component and Vertex as a private endpoint id, so the ban now checks the model group's providers and fails closed when the router cannot resolve the model. /queue/chat/completions hands its body to the router but sits in no route group, so it is named as an LLM route for the ban. /health/test_connection no longer drops a deployment's stored credentials when its litellm_params pin a model_id.
This commit is contained in:
parent
0351c0e10d
commit
be22f1f52b
4 changed files with 129 additions and 21 deletions
|
|
@ -41,6 +41,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
from litellm.types.router import (
|
||||
CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS,
|
||||
Deployment,
|
||||
ModelGroupInfo,
|
||||
reject_server_owned_wif_fields,
|
||||
server_owned_wif_fields_present,
|
||||
)
|
||||
|
|
@ -211,6 +212,18 @@ def _is_param_allowed(
|
|||
return False
|
||||
|
||||
|
||||
def _model_group_info(model: str, llm_router: Router | None) -> ModelGroupInfo | None:
|
||||
"""The router's group for ``model``, or for its provider when ``model`` only matches a wildcard deployment."""
|
||||
if llm_router is None:
|
||||
return None
|
||||
registry: Final = llm_router.registry()
|
||||
model_info: Final = registry.group(model)
|
||||
if model_info is not None:
|
||||
return model_info
|
||||
provider: Final = model.split("/", 1)[0]
|
||||
return registry.group(provider) if provider in provider_list else None
|
||||
|
||||
|
||||
def _allow_model_level_clientside_configurable_parameters(
|
||||
model: str, param: str, request_body_value: object, llm_router: Router | None
|
||||
) -> bool:
|
||||
|
|
@ -220,19 +233,7 @@ def _allow_model_level_clientside_configurable_parameters(
|
|||
- check if 'clientside_configurable_parameters' is set for model
|
||||
-
|
||||
"""
|
||||
if llm_router is None:
|
||||
return False
|
||||
# check if model is set
|
||||
registry: Final = llm_router.registry()
|
||||
model_info = registry.group(model)
|
||||
if model_info is None:
|
||||
# check if wildcard model is set
|
||||
if model.split("/", 1)[0] in provider_list:
|
||||
model_info = registry.group(model.split("/", 1)[0])
|
||||
|
||||
if model_info is None:
|
||||
return False
|
||||
|
||||
model_info: Final = _model_group_info(model, llm_router)
|
||||
if model_info is None or model_info.configurable_clientside_auth_params is None:
|
||||
return False
|
||||
|
||||
|
|
@ -319,6 +320,9 @@ def route_manages_mcp_servers(route: str | None) -> bool:
|
|||
|
||||
|
||||
_LLM_CALL_ONLY_BANNED_PARAMS: Final[frozenset[str]] = frozenset(("model_id",))
|
||||
# Hands its body to the router like /chat/completions but belongs to no route group.
|
||||
_UNLISTED_LLM_ROUTES: Final[frozenset[str]] = frozenset(("/queue/chat/completions",))
|
||||
_MODEL_ID_STEERING_PROVIDERS: Final[frozenset[str]] = frozenset(("bedrock",))
|
||||
|
||||
|
||||
def route_calls_an_llm(route: str | None) -> bool:
|
||||
|
|
@ -327,7 +331,16 @@ def route_calls_an_llm(route: str | None) -> bool:
|
|||
Admin routes such as ``/health?model_id=`` and ``/spend/logs/ui?model_id=`` take it as a deployment
|
||||
id instead. ``None`` is a body checked with no route, which callers only do before a provider call.
|
||||
"""
|
||||
return route is None or RouteChecks.is_llm_api_route(route)
|
||||
return route is None or route in _UNLISTED_LLM_ROUTES or RouteChecks.is_llm_api_route(route)
|
||||
|
||||
|
||||
def model_id_steers_the_upstream_model(model: str, llm_router: Router | None) -> bool:
|
||||
"""Bedrock runs the body's ``model_id`` in place of the deployment's model. SageMaker reads it as an
|
||||
inference component on the deployment's own endpoint and Vertex as a private endpoint in its own
|
||||
project, so neither leaves what the deployment already reaches. A model the router cannot resolve
|
||||
fails closed."""
|
||||
group: Final = _model_group_info(model, llm_router)
|
||||
return group is None or any(provider in _MODEL_ID_STEERING_PROVIDERS for provider in group.providers)
|
||||
|
||||
|
||||
_NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body")
|
||||
|
|
@ -522,7 +535,9 @@ def _check_banned_params(
|
|||
for param in _BANNED_REQUEST_BODY_PARAMS:
|
||||
if param not in body:
|
||||
continue
|
||||
if not calls_an_llm and param in _LLM_CALL_ONLY_BANNED_PARAMS:
|
||||
if param in _LLM_CALL_ONLY_BANNED_PARAMS and not (
|
||||
calls_an_llm and model_id_steers_the_upstream_model(model, llm_router)
|
||||
):
|
||||
continue
|
||||
if general_settings.get("allow_client_side_credentials") is True:
|
||||
# Proxy-wide opt-in: every banned param is permitted, exit
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ from litellm.proxy.auth.auth_exception_handler import (
|
|||
)
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
_BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check
|
||||
_LLM_CALL_ONLY_BANNED_PARAMS, # pyright: ignore[reportPrivateUsage] # model_id pins a deployment, it is not a credential
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils import http_parsing_utils
|
||||
|
|
@ -172,7 +173,7 @@ def _request_inherits_config_credentials(
|
|||
bringing its own ``api_key`` describes its own connection when matched only by
|
||||
model name. A blank name or key is no name or key: ``load_credentials_from_list``
|
||||
resolves nothing from it, so it must not cost the request the credentials it
|
||||
would otherwise be probed with.
|
||||
would otherwise be probed with. A deployment-pinned ``model_id`` names no credential either.
|
||||
"""
|
||||
requested_credential: Final = request_params.get("litellm_credential_name")
|
||||
if requested_credential and requested_credential != config_params.get("litellm_credential_name"):
|
||||
|
|
@ -181,7 +182,9 @@ def _request_inherits_config_credentials(
|
|||
return False
|
||||
if allow_client_side_credentials:
|
||||
return True
|
||||
return not any(param in request_params for param in _BANNED_REQUEST_BODY_PARAMS)
|
||||
return not any(
|
||||
param in request_params for param in _BANNED_REQUEST_BODY_PARAMS if param not in _LLM_CALL_ONLY_BANNED_PARAMS
|
||||
)
|
||||
|
||||
|
||||
def _admin_api_key_from_request(
|
||||
|
|
|
|||
|
|
@ -368,7 +368,11 @@ _SERVICE_ACCOUNT_TOKEN_REFERENCE: Final = "oidc/file//var/run/secrets/kubernetes
|
|||
|
||||
|
||||
def _user_config_with(litellm_params: dict[str, str]) -> dict[str, object]:
|
||||
return {"model_list": [{"model_name": "devstral", "litellm_params": {"model": "mistral/devstral-latest", **litellm_params}}]}
|
||||
return {
|
||||
"model_list": [
|
||||
{"model_name": "devstral", "litellm_params": {"model": "mistral/devstral-latest", **litellm_params}}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _user_config_defaulting_to(litellm_params: dict[str, str]) -> dict[str, object]:
|
||||
|
|
@ -388,7 +392,11 @@ def _user_config_search_tool_with(litellm_params: dict[str, str]) -> dict[str, o
|
|||
"search_tools": [
|
||||
{
|
||||
"search_tool_name": "web",
|
||||
"litellm_params": {"search_provider": "tavily", "api_base": "https://attacker.example", **litellm_params},
|
||||
"litellm_params": {
|
||||
"search_provider": "tavily",
|
||||
"api_base": "https://attacker.example",
|
||||
**litellm_params,
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -509,7 +517,9 @@ def test_a_call_cannot_carry_an_oidc_api_key_in_litellm_params(route: str | None
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ["/model/new", "/model/update", "/model/deployment-1/update", "/health/test_connection"])
|
||||
@pytest.mark.parametrize(
|
||||
"route", ["/model/new", "/model/update", "/model/deployment-1/update", "/health/test_connection"]
|
||||
)
|
||||
def test_a_deployment_write_may_carry_an_oidc_api_key_in_litellm_params(route: str):
|
||||
"""A proxy admin configures an oidc/ api_key on a deployment through these routes, and the model management role
|
||||
gate refuses anyone else with a 403."""
|
||||
|
|
@ -3672,6 +3682,7 @@ def test_request_cannot_select_chatgpt_auth_file(allow_client_side_credentials:
|
|||
model="chatgpt-account-a",
|
||||
)
|
||||
|
||||
|
||||
class TestIsRequestBodySafeBlocksBedrockProjectOverride:
|
||||
"""``aws_bedrock_project_id`` pins a deployment to a Bedrock project so
|
||||
that project's data-retention policy applies to its requests. A
|
||||
|
|
@ -5037,6 +5048,58 @@ def test_client_model_id_is_allowed_when_the_deployment_opts_in() -> None:
|
|||
)
|
||||
|
||||
|
||||
def test_the_queue_route_rejects_a_client_model_id_like_every_other_llm_route() -> None:
|
||||
with pytest.raises(ValueError, match="model_id"):
|
||||
is_request_body_safe(
|
||||
request_body={"model": "claude-sonnet-4-6", "model_id": "us.amazon.nova-micro-v1:0"},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="claude-sonnet-4-6",
|
||||
route="/queue/chat/completions",
|
||||
)
|
||||
|
||||
|
||||
def _router_with(*deployments: tuple[str, str]) -> Router:
|
||||
return Router(
|
||||
model_list=[{"model_name": model_name, "litellm_params": {"model": model}} for model_name, model in deployments]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"router, model, model_id",
|
||||
[
|
||||
(_router_with(("sm", "sagemaker_chat/my-endpoint")), "sm", "my-inference-component"),
|
||||
(_router_with(("vx", "vertex_ai/private")), "vx", "1234567890"),
|
||||
],
|
||||
ids=["sagemaker-inference-component", "vertex-private-endpoint"],
|
||||
)
|
||||
def test_model_id_stays_allowed_where_it_does_not_pick_the_bedrock_model(
|
||||
router: Router, model: str, model_id: str
|
||||
) -> None:
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={"model": model, "model_id": model_id, "messages": [{"role": "user", "content": "hi"}]},
|
||||
general_settings={},
|
||||
llm_router=router,
|
||||
model=model,
|
||||
route="/v1/chat/completions",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_model_id_is_rejected_when_the_model_group_mixes_in_a_bedrock_deployment() -> None:
|
||||
router: Final = _router_with(("sm", "sagemaker_chat/my-endpoint"), ("sm", "bedrock/us.anthropic.claude-sonnet-4-6"))
|
||||
with pytest.raises(ValueError, match="model_id"):
|
||||
is_request_body_safe(
|
||||
request_body={"model": "sm", "model_id": "us.amazon.nova-micro-v1:0"},
|
||||
general_settings={},
|
||||
llm_router=router,
|
||||
model="sm",
|
||||
route="/v1/chat/completions",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ["/health", "/spend/logs/ui", "/model/block", "/model/unblock"])
|
||||
def test_admin_routes_keep_model_id_as_a_deployment_id(route: str) -> None:
|
||||
assert (
|
||||
|
|
@ -5430,7 +5493,9 @@ async def test_a_non_numeric_max_request_size_mb_is_ignored_without_premium():
|
|||
def test_chatgpt_model_configuration_reaches_role_authorization(
|
||||
allow_client_side_credentials: bool, route: str
|
||||
) -> None:
|
||||
body: Final = {"litellm_params": {"model": "chatgpt/gpt-6-luna", "chatgpt_auth_file": "/tokens/account-a/auth.json"}}
|
||||
body: Final = {
|
||||
"litellm_params": {"model": "chatgpt/gpt-6-luna", "chatgpt_auth_file": "/tokens/account-a/auth.json"}
|
||||
}
|
||||
assert is_request_body_safe(
|
||||
request_body=body,
|
||||
general_settings={"allow_client_side_credentials": allow_client_side_credentials},
|
||||
|
|
|
|||
|
|
@ -6252,3 +6252,28 @@ async def test_unsaved_chatgpt_account_probe_requires_proxy_admin(
|
|||
await call
|
||||
assert error.value.code == "403"
|
||||
assert len(upstream.calls) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_params, inherits",
|
||||
[
|
||||
(
|
||||
{"model": "bedrock/us.anthropic.claude-sonnet-4-6", "model_id": "arn:aws:bedrock:us-east-1:1:profile/p"},
|
||||
True,
|
||||
),
|
||||
({"model": "bedrock/us.anthropic.claude-sonnet-4-6", "api_base": "https://example.invalid"}, False),
|
||||
],
|
||||
ids=["deployment-pinned model_id keeps the stored credentials", "a client api_base brings its own"],
|
||||
)
|
||||
def test_connection_probe_keeps_stored_credentials_behind_a_pinned_model_id(
|
||||
request_params: dict[str, object], inherits: bool
|
||||
) -> None:
|
||||
assert (
|
||||
_health_endpoints_module._request_inherits_config_credentials( # pyright: ignore[reportPrivateUsage] # the decision under test
|
||||
config_params={"model": "bedrock/us.anthropic.claude-sonnet-4-6", "aws_access_key_id": "AKIA"},
|
||||
request_params=request_params,
|
||||
allow_client_side_credentials=False,
|
||||
selected_by_id=False,
|
||||
)
|
||||
is inherits
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue