Merge remote-tracking branch 'origin' into litellm_access_group_rename

This commit is contained in:
yuneng-jiang
2026-02-13 17:04:59 -08:00
4 changed files with 236 additions and 5 deletions
+13
View File
@@ -585,7 +585,20 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if is_proxy_admin:
return UserAPIKeyAuth(
api_key=None,
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id=user_id,
team_id=team_id,
team_alias=(
team_object.team_alias
if team_object is not None
else None
),
team_metadata=team_object.metadata
if team_object is not None
else None,
org_id=org_id,
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
)
@@ -1099,6 +1099,7 @@ def create_pass_through_route(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
subpath: str = "", # captures sub-paths when include_subpath=True
custom_body: Optional[dict] = None, # accepted for signature compatibility with URL-based path; not forwarded because chat_completion_pass_through_endpoint does not support it
):
return await chat_completion_pass_through_endpoint(
fastapi_response=fastapi_response,
@@ -1115,6 +1116,7 @@ def create_pass_through_route(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
subpath: str = "", # captures sub-paths when include_subpath=True
custom_body: Optional[dict] = None,
):
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
InitPassThroughEndpointHelpers,
@@ -1189,11 +1191,16 @@ def create_pass_through_route(
)
if query_params:
final_query_params.update(query_params)
final_custom_body = (
custom_body_data
if isinstance(custom_body_data, dict) or custom_body_data is None
else None
)
# When a caller (e.g. bedrock_proxy_route) supplies a pre-built
# body, use it instead of the body parsed from the raw request.
if custom_body is not None:
final_custom_body = custom_body
else:
final_custom_body = (
custom_body_data
if isinstance(custom_body_data, dict) or custom_body_data is None
else None
)
return await pass_through_request( # type: ignore
request=request,
@@ -422,3 +422,77 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget():
assert result.user_tpm_limit == 1000
assert result.user_rpm_limit == 100
assert result.user_email == "test@example.com"
def test_proxy_admin_jwt_auth_includes_identity_fields():
"""
Test that the proxy admin early-return path in JWT auth populates
user_id, team_id, team_alias, team_metadata, org_id, and end_user_id.
Regression test: previously the is_proxy_admin branch only set user_role
and parent_otel_span, discarding all identity fields resolved from the JWT.
This caused blank Team Name and Internal User in Request Logs UI.
"""
from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth
team_object = LiteLLM_TeamTable(
team_id="team-123",
team_alias="my-team",
metadata={"tags": ["prod"], "env": "production"},
)
# Simulate the proxy admin early-return path (user_api_key_auth.py ~line 586)
result = UserAPIKeyAuth(
api_key=None,
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="user-abc",
team_id="team-123",
team_alias=(
team_object.team_alias if team_object is not None else None
),
team_metadata=team_object.metadata if team_object is not None else None,
org_id="org-456",
end_user_id="end-user-789",
parent_otel_span=None,
)
assert result.user_role == LitellmUserRoles.PROXY_ADMIN
assert result.user_id == "user-abc"
assert result.team_id == "team-123"
assert result.team_alias == "my-team"
assert result.team_metadata == {"tags": ["prod"], "env": "production"}
assert result.org_id == "org-456"
assert result.end_user_id == "end-user-789"
assert result.api_key is None
def test_proxy_admin_jwt_auth_handles_no_team_object():
"""
Test that the proxy admin early-return path works correctly when
team_object is None (user has admin role but no team association).
"""
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
team_object = None
result = UserAPIKeyAuth(
api_key=None,
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
team_id=None,
team_alias=(
team_object.team_alias if team_object is not None else None
),
team_metadata=team_object.metadata if team_object is not None else None,
org_id=None,
end_user_id=None,
parent_otel_span=None,
)
assert result.user_role == LitellmUserRoles.PROXY_ADMIN
assert result.user_id == "admin-user"
assert result.team_id is None
assert result.team_alias is None
assert result.team_metadata is None
assert result.org_id is None
assert result.end_user_id is None
@@ -2087,6 +2087,143 @@ async def test_add_litellm_data_to_request_adds_headers_to_metadata():
assert "headers" in result["proxy_server_request"]
@pytest.mark.asyncio
async def test_create_pass_through_route_custom_body_url_target():
"""
Test that the URL-based endpoint_func created by create_pass_through_route
accepts a custom_body parameter and forwards it to pass_through_request,
taking precedence over the request-parsed body.
This verifies the fix for issue #16999 where bedrock_proxy_route passes
custom_body=data to the endpoint function, which previously crashed with:
TypeError: endpoint_func() got an unexpected keyword argument 'custom_body'
"""
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
create_pass_through_route,
)
unique_path = "/test/path/unique/custom_body_url"
endpoint_func = create_pass_through_route(
endpoint=unique_path,
target="https://bedrock-agent-runtime.us-east-1.amazonaws.com",
custom_headers={"Content-Type": "application/json"},
_forward_headers=True,
)
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request"
) as mock_pass_through, patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route"
) as mock_is_registered, patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.get_registered_pass_through_route"
) as mock_get_registered, patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._parse_request_data_by_content_type"
) as mock_parse_request:
mock_pass_through.return_value = MagicMock()
mock_is_registered.return_value = True
mock_get_registered.return_value = None
# Simulate the request parser returning a different body
mock_parse_request.return_value = (
{}, # query_params_data
{"parsed_from_request": True}, # custom_body_data (from request)
None, # file_data
False, # stream
)
mock_request = MagicMock(spec=Request)
mock_request.url = MagicMock()
mock_request.url.path = unique_path
mock_request.path_params = {}
mock_request.query_params = QueryParams({})
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.api_key = "test-key"
# The caller-supplied body (e.g. from bedrock_proxy_route)
bedrock_body = {
"retrievalQuery": {"text": "What is in the knowledge base?"},
}
# Call endpoint_func with custom_body — this is the call that
# used to crash with TypeError before the fix
await endpoint_func(
request=mock_request,
fastapi_response=MagicMock(),
user_api_key_dict=mock_user_api_key_dict,
custom_body=bedrock_body,
)
mock_pass_through.assert_called_once()
call_kwargs = mock_pass_through.call_args[1]
# The critical assertion: custom_body takes precedence over
# the body parsed from the raw request
assert call_kwargs["custom_body"] == bedrock_body
@pytest.mark.asyncio
async def test_create_pass_through_route_no_custom_body_falls_back():
"""
Test that the URL-based endpoint_func falls back to the request-parsed body
when custom_body is not provided.
This ensures the default pass-through behavior is preserved only the
Bedrock proxy route (and similar callers) supply a pre-built body.
"""
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
create_pass_through_route,
)
unique_path = "/test/path/unique/no_custom_body"
endpoint_func = create_pass_through_route(
endpoint=unique_path,
target="http://example.com/api",
custom_headers={},
)
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request"
) as mock_pass_through, patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route"
) as mock_is_registered, patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.get_registered_pass_through_route"
) as mock_get_registered, patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._parse_request_data_by_content_type"
) as mock_parse_request:
mock_pass_through.return_value = MagicMock()
mock_is_registered.return_value = True
mock_get_registered.return_value = None
request_parsed_body = {"key": "from_request"}
mock_parse_request.return_value = (
{}, # query_params_data
request_parsed_body, # custom_body_data
None, # file_data
False, # stream
)
mock_request = MagicMock(spec=Request)
mock_request.url = MagicMock()
mock_request.url.path = unique_path
mock_request.path_params = {}
mock_request.query_params = QueryParams({})
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.api_key = "test-key"
# Call without custom_body — should use the request-parsed body
await endpoint_func(
request=mock_request,
fastapi_response=MagicMock(),
user_api_key_dict=mock_user_api_key_dict,
)
mock_pass_through.assert_called_once()
call_kwargs = mock_pass_through.call_args[1]
# Should fall back to the body parsed from the request
assert call_kwargs["custom_body"] == request_parsed_body
def test_build_full_path_with_root_default():
"""
Test _build_full_path_with_root with default root path (/)