scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(proxy): inherit caller identity in passthrough batch managed-object

Browse source 
Read user_id and team_id from the request's litellm_params metadata when
fabricating the UserAPIKeyAuth handed to the managed_files hook, so
batches created via passthrough are attributed to the real requester
instead of a hardcoded fallback. Adds parametrized regression coverage
for both the populated-metadata and empty-kwargs cases.
This commit is contained in:
Ryan Crabbe 2026-04-29 16:25:13 -07:00
parent fc0cc9c581
commit 2461139593
No known key found for this signature in database
4 changed files with 118 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -549,10 +549,16 @@ class AnthropicPassthroughLoggingHandler:
# Create a mock user API key dict for the managed object storage
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
_request_metadata = (kwargs.get("litellm_params", {}) or {}).get(
"metadata", {}
) or {}
user_api_key_dict = UserAPIKeyAuth(
user_id=kwargs.get("user_id", "default-user"),
user_id=_request_metadata.get(
"user_api_key_user_id", "default-user"
),
api_key="",
team_id=None,
team_id=_request_metadata.get("user_api_key_team_id"),
team_alias=None,
user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value
user_email=None,

10
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
View file

@ -849,10 +849,16 @@ class VertexPassthroughLoggingHandler:
# Create a mock user API key dict for the managed object storage
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
_request_metadata = (kwargs.get("litellm_params", {}) or {}).get(
"metadata", {}
) or {}
user_api_key_dict = UserAPIKeyAuth(
user_id=kwargs.get("user_id", "default-user"),
user_id=_request_metadata.get(
"user_api_key_user_id", "default-user"
),
api_key="",
team_id=None,
team_id=_request_metadata.get("user_api_key_team_id"),
team_alias=None,
user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value
user_email=None,

51
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
View file

@ -573,6 +573,57 @@ class TestAnthropicBatchPassthroughCostTracking:
or "claude-sonnet-4-5-20250929" in decoded
)
@pytest.mark.parametrize(
"kwargs,expected_user_id,expected_team_id",
[
(
{
"litellm_params": {
"metadata": {
"user_api_key_user_id": "real-user-123",
"user_api_key_team_id": "team-456",
}
}
},
"real-user-123",
"team-456",
),
({}, "default-user", None),
],
)
def test_store_batch_managed_object_propagates_user_identity_from_metadata(
self,
mock_logging_obj,
kwargs,
expected_user_id,
expected_team_id,
):
"""The fabricated UserAPIKeyAuth must inherit user_id/team_id from the
request's litellm_params.metadata, not the (always-empty) top-level
kwargs lookup. Falls back to "default-user" only when metadata is
absent."""
mock_managed_files_hook = MagicMock()
with (
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl,
patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler.verbose_proxy_logger"
),
):
mock_pl.get_proxy_hook.return_value = mock_managed_files_hook
AnthropicPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id="uoi",
batch_object={"id": "b1", "object": "batch", "status": "validating"},
model_object_id="b1",
logging_obj=mock_logging_obj,
**kwargs,
)
mock_managed_files_hook.store_unified_object_id.assert_called_once()
call_kwargs = mock_managed_files_hook.store_unified_object_id.call_args[1]
assert call_kwargs["user_api_key_dict"].user_id == expected_user_id
assert call_kwargs["user_api_key_dict"].team_id == expected_team_id
def test_batch_creation_handler_failure_status_code(
self, mock_logging_obj, mock_request_body
):

51
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py
View file

@ -264,6 +264,57 @@ class TestVertexAIBatchPassthroughHandler:
# Verify the managed files hook was called
mock_managed_files_hook.store_unified_object_id.assert_called_once()
@pytest.mark.parametrize(
"kwargs,expected_user_id,expected_team_id",
[
(
{
"litellm_params": {
"metadata": {
"user_api_key_user_id": "real-user-123",
"user_api_key_team_id": "team-456",
}
}
},
"real-user-123",
"team-456",
),
({}, "default-user", None),
],
)
def test_store_batch_managed_object_propagates_user_identity_from_metadata(
self,
mock_logging_obj,
mock_managed_files_hook,
kwargs,
expected_user_id,
expected_team_id,
):
"""The fabricated UserAPIKeyAuth must inherit user_id/team_id from the
request's litellm_params.metadata, not the (always-empty) top-level
kwargs lookup. Falls back to "default-user" only when metadata is
absent."""
with (
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl,
patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.verbose_proxy_logger"
),
):
mock_pl.get_proxy_hook.return_value = mock_managed_files_hook
VertexPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id="uoi",
batch_object={"id": "b1", "object": "batch", "status": "validating"},
model_object_id="b1",
logging_obj=mock_logging_obj,
**kwargs,
)
mock_managed_files_hook.store_unified_object_id.assert_called_once()
call_kwargs = mock_managed_files_hook.store_unified_object_id.call_args[1]
assert call_kwargs["user_api_key_dict"].user_id == expected_user_id
assert call_kwargs["user_api_key_dict"].team_id == expected_team_id
def test_batch_cost_calculation_integration(self):
"""Single Vertex AI response → non-zero cost with correct token counts."""
from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage