From 2461139593a4749439af9dae7d4e36e3958908fd Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Wed, 29 Apr 2026 16:25:13 -0700 Subject: [PATCH] fix(proxy): inherit caller identity in passthrough batch managed-object 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. --- .../anthropic_passthrough_logging_handler.py | 10 +++- .../vertex_passthrough_logging_handler.py | 10 +++- ...t_anthropic_passthrough_logging_handler.py | 51 +++++++++++++++++++ .../test_vertex_ai_batch_passthrough.py | 51 +++++++++++++++++++ 4 files changed, 118 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 216eb61a9d1..c42faa59cf0 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -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, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 86dd23e12ca..6a138532617 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -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, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 3def76b825e..c16c42decc0 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -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 ): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index 756e5fa5bcf..efa26a61bf5 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -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