diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 4be4c2d5e78..1e92754857b 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -159,6 +159,6 @@ def transform_anthropic_messages_request( "model", None ) # do not pass model in request body to vertex ai - sanitize_vertex_anthropic_output_params(anthropic_messages_request) + sanitize_vertex_anthropic_output_params(anthropic_messages_request, model) return anthropic_messages_request diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py index a33ad677789..280cc1c888a 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py @@ -10,23 +10,38 @@ keeps the parent module's import surface narrow. """ -# Keys inside ``output_config`` that Vertex AI Claude does not accept. -# Add an entry only when a 400 "Extra inputs are not permitted" is -# reproducible against the live Vertex endpoint. +# Keys inside ``output_config`` that Vertex AI Claude rejects regardless of +# the target model. Add an entry only when a 400 "Extra inputs are not +# permitted" is reproducible against the live Vertex endpoint for every model. VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS: frozenset = frozenset() -def sanitize_vertex_anthropic_output_params(data: dict) -> None: +def _model_accepts_output_config_effort(model: str) -> bool: + """Whether ``model`` accepts ``output_config.effort`` on Vertex. + + Opus/Sonnet 4.6+ advertise ``supports_output_config`` (or a reasoning + effort level) and accept it; Haiku 4.5 advertises neither and 400s on + ``output_config.effort: Extra inputs are not permitted``. Imported lazily + so this stays a leaf module (see module docstring). + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + return AnthropicConfig._model_supports_effort_param(model) + + +def sanitize_vertex_anthropic_output_params(data: dict, model: str) -> None: """ Strip Vertex-unsupported keys from ``output_config`` / ``output_format`` in-place; forward whatever remains. Behavior: - * ``output_config`` containing only unsupported keys (e.g. ``effort`` - alone) is removed entirely so the request body has no empty dict. - * ``output_config`` containing a mix of supported + unsupported keys - has the unsupported subset filtered out and the rest forwarded. - * ``output_config`` that is supported in full passes through unchanged. + * ``output_config.effort`` is dropped for models that don't accept it + (e.g. Haiku 4.5) and forwarded for those that do (Opus/Sonnet 4.6+). + Clients like Claude Code inject it into every Messages payload, so the + gate has to live here rather than rely on the caller. + * Keys in ``VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS`` are always filtered. + * ``output_config`` left empty after filtering is removed so the request + body has no empty dict. * ``output_format`` is forwarded as-is (Vertex AI Claude accepts it). * Non-dict values for ``output_config`` are dropped to avoid sending malformed payloads downstream. @@ -37,11 +52,19 @@ def sanitize_vertex_anthropic_output_params(data: dict) -> None: if not isinstance(output_config, dict): data.pop("output_config", None) return - sanitized = { - k: v - for k, v in output_config.items() - if k not in VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS - } + + drop_keys = set(VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS) + if "effort" in output_config and not _model_accepts_output_config_effort(model): + from litellm._logging import verbose_logger + + verbose_logger.debug( + "Dropping unsupported output_config.effort for vertex_ai model=%s " + "(no supports_output_config in the model map)", + model, + ) + drop_keys.add("effort") + + sanitized = {k: v for k, v in output_config.items() if k not in drop_keys} if sanitized: data["output_config"] = sanitized else: diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 4627d9f6df3..c852909d475 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -106,7 +106,7 @@ def transform_request( data.pop("model", None) # vertex anthropic doesn't accept 'model' parameter - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, model) tools = optional_params.get("tools") tool_search_used = self.is_tool_search_used(tools) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 38976f79aa3..5301179faba 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -532,6 +532,7 @@ async def common_checks( # noqa: PLR0915 route=route, request_headers=_safe_get_request_headers(request=request), request_query_params=_safe_get_request_query_params(request=request), + llm_router=llm_router, ) # 1. If team is blocked diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 86265270357..956262944ea 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1244,7 +1244,9 @@ def _route_uses_model_routing_sources(route: str) -> bool: def _extract_models_from_managed_resource_id( - resource_id: Any, resource_id_field: Optional[str] = None + resource_id: Any, + resource_id_field: Optional[str] = None, + llm_router: Optional[Router] = None, ) -> List[str]: if not isinstance(resource_id, str) or not resource_id: return [] @@ -1301,16 +1303,18 @@ def _extract_models_from_managed_resource_id( ) if resource_id_field == "video_id": + model_id = decode_video_id_with_provider(resource_id).get("model_id") _append_model_candidates( candidates=candidates, - value=decode_video_id_with_provider(resource_id).get("model_id"), + value=_resolve_model_id_with_router(model_id, llm_router), ) else: + model_id = decode_character_id_with_provider(resource_id).get( + "model_id" + ) _append_model_candidates( candidates=candidates, - value=decode_character_id_with_provider(resource_id).get( - "model_id" - ), + value=_resolve_model_id_with_router(model_id, llm_router), ) except Exception as e: verbose_proxy_logger.debug( @@ -1320,11 +1324,26 @@ def _extract_models_from_managed_resource_id( return _dedupe_model_candidates(candidates) +def _resolve_model_id_with_router( + model_id: Optional[str], llm_router: Optional[Router] +) -> Optional[str]: + if model_id is None or llm_router is None: + return model_id + try: + return llm_router.resolve_model_name_from_model_id(model_id) or model_id + except Exception as e: + verbose_proxy_logger.debug( + "Unable to resolve model_id from managed resource ID: %s", str(e) + ) + return model_id + + def _extract_model_candidates_from_request( request_data: dict, route: str, request_headers: Optional[Mapping[str, Any]] = None, request_query_params: Optional[Mapping[str, Any]] = None, + llm_router: Optional[Router] = None, ) -> List[str]: candidates: List[str] = [] uses_model_routing_sources = _route_uses_model_routing_sources(route=route) @@ -1374,7 +1393,9 @@ def _extract_model_candidates_from_request( _append_model_candidates( candidates, _extract_models_from_managed_resource_id( - request_data.get(field), resource_id_field=field + request_data.get(field), + resource_id_field=field, + llm_router=llm_router, ), ) @@ -1396,12 +1417,14 @@ def get_model_from_request( route: str, request_headers: Optional[Mapping[str, Any]] = None, request_query_params: Optional[Mapping[str, Any]] = None, + llm_router: Optional[Router] = None, ) -> Optional[Union[str, List[str]]]: candidates = _extract_model_candidates_from_request( request_data=request_data, route=route, request_headers=request_headers, request_query_params=request_query_params, + llm_router=llm_router, ) model = _format_model_candidates(candidates) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2f9a1411400..2e6cd1f8e70 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -146,12 +146,14 @@ def _get_model_from_request_context( request_data: dict, route: str, request: Optional[Request], + llm_router: Optional[Any] = None, ) -> Optional[Union[str, List[str]]]: return get_model_from_request( request_data=request_data, route=route, request_headers=_safe_get_request_headers(request=request), request_query_params=_safe_get_request_query_params(request=request), + llm_router=llm_router, ) @@ -1034,6 +1036,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data=request_data, route=route, request=request, + llm_router=llm_router, ) skip_budget_checks = False if model is not None and llm_router is not None: @@ -1451,6 +1454,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data=request_data, route=route, request=request, + llm_router=llm_router, ) skip_budget_checks = False if model is not None and llm_router is not None: @@ -1579,6 +1583,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data=request_data, route=route, request=request, + llm_router=llm_router, ) current_models = _get_model_names_for_budget_checks( model=current_model @@ -2159,6 +2164,7 @@ def _should_skip_budget_checks( request_data=request_data, route=route, request=request, + llm_router=llm_router, ) if model is not None and llm_router is not None: return _is_model_cost_zero(model=model, llm_router=llm_router) @@ -2475,6 +2481,7 @@ async def _enforce_key_and_fallback_model_access( request_data=request_data, route=route, request=request, + llm_router=llm_router, ) if model is not None: @@ -2616,6 +2623,7 @@ async def _run_post_custom_auth_checks( request_data=request_data, route=route, request=request, + llm_router=llm_router, ) current_models = _get_model_names_for_budget_checks(model=current_model) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0e645013b92..97f65912297 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -894,7 +894,12 @@ async def _common_key_generation_helper( # noqa: PLR0915 user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) - if not _is_proxy_admin: + _org_inherited_from_team = ( + team_table is not None + and team_table.organization_id is not None + and data.organization_id == team_table.organization_id + ) + if not _is_proxy_admin and not _org_inherited_from_team: await _validate_caller_can_assign_key_org( user_api_key_dict=user_api_key_dict, organization_id=data.organization_id, 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 3be26eb572d..a94672f9487 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 @@ -114,6 +114,13 @@ def _create_anthropic_response_logging_payload( handles streaming and non-streaming responses """ + # Only record complete_streaming_response for actual streaming responses. + # perform_redaction scrubs this field only when stream is True, so setting + # it on a non-streaming response would bypass message redaction. + if logging_obj.model_call_details.get("stream") is True: + logging_obj.model_call_details["complete_streaming_response"] = ( + litellm_model_response + ) try: # Get custom_llm_provider from logging object if available (e.g., azure_ai for Azure Anthropic) custom_llm_provider = logging_obj.model_call_details.get( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 9d68132b37d..2cc4a34dd57 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1030,6 +1030,9 @@ async def pass_through_request( # noqa: PLR0915 ) if stream: + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + if is_multipart: response = ( await HttpPassThroughEndpointHelpers.make_multipart_http_request( @@ -1108,6 +1111,9 @@ async def pass_through_request( # noqa: PLR0915 verbose_proxy_logger.debug("response.headers= %s", response.headers) if _is_streaming_response(response) is True: + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + try: response.raise_for_status() except httpx.HTTPStatusError as e: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 200a17e2368..eb8af3b073e 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -72,7 +72,7 @@ async def reserve_budget_for_request( return None if route in {"/models", "/v1/models", "/utils/token_counter"}: return None - if get_model_from_request(request_body, route) is None: + if get_model_from_request(request_body, route, llm_router=llm_router) is None: return None counters = await _get_budget_counters( @@ -797,7 +797,7 @@ def estimate_request_max_cost( route: str, llm_router: Optional[Router], ) -> Optional[float]: - model = get_model_from_request(request_body, route) + model = get_model_from_request(request_body, route, llm_router=llm_router) if model is None: return None diff --git a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py index 455c72ff636..5ab0319da47 100644 --- a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py +++ b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py @@ -318,6 +318,7 @@ def test_handle_logging_anthropic_collected_chunks(all_chunks): from litellm.types.utils import ModelResponse litellm_logging_obj = Mock() + litellm_logging_obj.model_call_details = {} pass_through_logging_obj = Mock() sent_args = { diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index b8cd65d3c99..6f4bb4e59c2 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -313,6 +313,40 @@ def test_transform_anthropic_messages_request_removes_scope_from_cache_control() assert result["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral" +def test_messages_request_strips_effort_for_haiku_45(): + """Regression: Claude Code (``claude --model claude-haiku-4.5``) sends + ``output_config.effort`` in its default Messages payload. Haiku 4.5 on + Vertex rejects it with 400 ``output_config.effort: Extra inputs are not + permitted``, so the pass-through must strip it for Haiku while keeping it + for Opus/Sonnet 4.6+.""" + config = VertexAIPartnerModelsAnthropicMessagesConfig() + messages = [{"role": "user", "content": "Hello"}] + + haiku_result = config.transform_anthropic_messages_request( + model="claude-haiku-4-5@20251001", + messages=messages, + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "output_config": {"effort": "high"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert "output_config" not in haiku_result + + opus_result = config.transform_anthropic_messages_request( + model="claude-opus-4-6", + messages=messages, + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "output_config": {"effort": "high"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert opus_result["output_config"] == {"effort": "high"} + + def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instance(): """ Regression test: repeated provider config lookups for the same Vertex Claude model diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index d89d09a4e63..ac2368130d8 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -675,28 +675,60 @@ def test_sanitize_vertex_anthropic_output_params_unit(): sanitize_vertex_anthropic_output_params, ) + supported = "claude-opus-4-6" + # No-op when output_config absent. data: dict = {"max_tokens": 8} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert data == {"max_tokens": 8} - # Effort-only → preserved (Vertex 4.6/4.7 accept it on rawPredict). + # Effort-only on a supporting model → preserved (Vertex 4.6/4.7 accept it). data = {"output_config": {"effort": "high"}} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert data["output_config"] == {"effort": "high"} # Format-only → preserved unchanged. fmt = {"format": {"type": "json_schema", "schema": {"type": "object"}}} data = {"output_config": dict(fmt)} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert data["output_config"] == fmt - # Mixed → both effort and format kept (no current Vertex-unsupported keys). + # Mixed on a supporting model → both effort and format kept. data = {"output_config": {"format": fmt["format"], "effort": "high"}} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert data["output_config"] == {"format": fmt["format"], "effort": "high"} # Non-dict → dropped defensively. data = {"output_config": "garbage"} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert "output_config" not in data + + +def test_sanitize_strips_effort_for_haiku_45(): + """Regression: Haiku 4.5 on Vertex does not support ``output_config.effort`` + and 400s with ``Extra inputs are not permitted``. Claude Code injects + ``effort`` into every Messages payload, so the helper must strip it for + models that don't advertise output_config support while leaving it intact + for Opus/Sonnet 4.6+.""" + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.output_params_utils import ( + sanitize_vertex_anthropic_output_params, + ) + + haiku = "claude-haiku-4-5@20251001" + + # Effort-only → output_config removed entirely (no empty dict on the wire). + data: dict = {"output_config": {"effort": "high"}, "max_tokens": 8} + sanitize_vertex_anthropic_output_params(data, haiku) + assert "output_config" not in data + assert data["max_tokens"] == 8 + + # Mixed → effort stripped, format preserved. + fmt = {"type": "json_schema", "schema": {"type": "object"}} + data = {"output_config": {"effort": "high", "format": fmt}} + sanitize_vertex_anthropic_output_params(data, haiku) + assert data["output_config"] == {"format": fmt} + + # Same payload on a supporting model keeps effort untouched. + data = {"output_config": {"effort": "high"}} + sanitize_vertex_anthropic_output_params(data, "vertex_ai/claude-opus-4-6") + assert data["output_config"] == {"effort": "high"} diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 2d40db9017e..f37eda86f9b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -382,6 +382,62 @@ def test_get_model_from_request_extracts_video_id_model(): ) +def test_get_model_from_request_resolves_video_id_model_with_router(): + from litellm.types.videos.utils import encode_video_id_with_provider + + provider_video_id = ( + "projects/test-project/locations/us-central1/publishers/google/models/" + "veo-3.1-generate-001/operations/operation-id" + ) + video_id = encode_video_id_with_provider( + video_id=provider_video_id, + provider="vertex_ai", + model_id="veo-3.1-generate-001", + ) + llm_router = MagicMock() + llm_router.resolve_model_name_from_model_id.return_value = ( + "gcp/google/veo-3.1-generate-001" + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=llm_router, + ) + == "gcp/google/veo-3.1-generate-001" + ) + llm_router.resolve_model_name_from_model_id.assert_called_once_with( + "veo-3.1-generate-001" + ) + + +def test_get_model_from_request_resolves_character_id_model_with_router(): + from litellm.types.videos.utils import encode_character_id_with_provider + + character_id = encode_character_id_with_provider( + character_id="character-provider-id", + provider="vertex_ai", + model_id="veo-3.1-generate-001", + ) + llm_router = MagicMock() + llm_router.resolve_model_name_from_model_id.return_value = ( + "gcp/google/veo-3.1-generate-001" + ) + + assert ( + get_model_from_request( + request_data={"character_id": character_id}, + route="/v1/videos/characters/{character_id}", + llm_router=llm_router, + ) + == "gcp/google/veo-3.1-generate-001" + ) + llm_router.resolve_model_name_from_model_id.assert_called_once_with( + "veo-3.1-generate-001" + ) + + def test_get_model_from_request_only_runs_media_decoders_for_matching_fields(): with ( patch( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 8fb242372e3..cda22da6ebd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -3094,6 +3094,249 @@ async def test_generate_key_with_object_permission(): assert "object_permission" not in key_data +@pytest.mark.asyncio +async def test_generate_key_team_member_inherits_org_skips_membership_check(): + """Regression: a team member creating a key for an org-scoped team must not + be blocked by the org-membership check. + + When ``organization_id`` is inherited from the key's team (via + ``apply_enterprise_key_management_params`` -> ``add_team_organization_id``), + the caller already passed team-level authorization. Requiring an explicit + ``LiteLLM_OrganizationMembership`` row on top of that broke the normal admin + workflow (admins only add users to teams). This asserts the org-membership + check is skipped when the org id came from the caller's team. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + ) + + org_id = "org-from-team" + + # Team belongs to an org; caller is a team member but NOT an explicit member + # of that organization (the regression scenario). + mock_team_table = MagicMock() + mock_team_table.organization_id = org_id + mock_team_table.metadata = None + + mock_validate_org = AsyncMock() + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "alice", + "team_id": "team-1", + } + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_search_tools_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org", + mock_validate_org, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object", + new_callable=AsyncMock, + return_value=MagicMock(litellm_budget_table=None), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ), + ): + result = await _common_key_generation_helper( + data=GenerateKeyRequest( + user_id="alice", + team_id="team-1", + organization_id=org_id, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ), + litellm_changed_by=None, + team_table=mock_team_table, + ) + + # Key creation proceeded for the team member ... + mock_generate_key.assert_awaited_once() + assert result is not None + # ... and the org-membership check was bypassed because organization_id was + # inherited from the caller's team. + mock_validate_org.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_generate_key_foreign_org_without_team_still_enforces_membership(): + """VERIA-55: a caller assigning a key to an organization that was NOT + inherited from a team must still pass the org-membership check. + + This guards the IDOR fix: ``team_table is None`` (or an org id that does not + match the team) means the org id did not come from team context, so the + explicit membership validation must run. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + ) + + foreign_org_id = "someone-elses-org" + + mock_validate_org = AsyncMock() + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "alice", + "team_id": None, + } + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org", + mock_validate_org, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object", + new_callable=AsyncMock, + return_value=MagicMock(litellm_budget_table=None), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ), + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + user_id="alice", + organization_id=foreign_org_id, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ), + litellm_changed_by=None, + team_table=None, + ) + + # No team context -> the org-membership check must still run. + mock_validate_org.assert_awaited_once() + assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id + + +@pytest.mark.asyncio +async def test_generate_key_foreign_org_with_mismatched_team_still_enforces_membership(): + """VERIA-55: when a team is present but its organization_id differs from the + organization_id on the key request, the org-membership check must still run.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + ) + + team_org_id = "other-org" + foreign_org_id = "someone-elses-org" + + mock_team_table = MagicMock() + mock_team_table.organization_id = team_org_id + mock_team_table.metadata = None + + mock_validate_org = AsyncMock() + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "alice", + "team_id": "team-1", + } + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_search_tools_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm_enterprise.proxy.management_endpoints.key_management_endpoints.apply_enterprise_key_management_params", + side_effect=lambda data, team_table: data, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org", + mock_validate_org, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object", + new_callable=AsyncMock, + return_value=MagicMock(litellm_budget_table=None), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ), + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + user_id="alice", + team_id="team-1", + organization_id=foreign_org_id, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ), + litellm_changed_by=None, + team_table=mock_team_table, + ) + + mock_validate_org.assert_awaited_once() + assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id + + # ============================================ # Organization Key Limit Tests # ============================================ 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 0a9e3031030..f8b6fbde3dc 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 @@ -1043,3 +1043,335 @@ def test_collapse_returns_none_for_interleaved_block_indexes(self): AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) is None ) + + +class TestStreamFalseDeduplication: + """ + Regression tests for the duplicate-callback bug where a streaming pass-through + request had stream=False hardcoded on its Logging object. + + Before the fix: + - logging_obj.stream was always False for pass-through requests + - _is_assembled_stream_success() checked `self.stream is not True` and returned + False immediately, so has_dispatched_final_stream_success was never set + - Any second dispatch_success_handlers call went through unchecked + + After the fix: + - pass_through_endpoints.py sets logging_obj.stream = True after detecting stream + - _create_anthropic_response_logging_payload sets complete_streaming_response on + model_call_details so callbacks see the correct assembled response state + - _is_assembled_stream_success returns True, dedup guard fires on first dispatch + """ + + @staticmethod + def _sse(event, data): + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + @staticmethod + def _make_logging_obj(stream: bool = False) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "hello"}], + stream=stream, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1245", + ) + return logging_obj + + @staticmethod + def _build_chunks(): + frames = [ + TestStreamFalseDeduplication._sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet-20241022", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello"}, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_stop", {"type": "content_block_stop", "index": 0} + ), + TestStreamFalseDeduplication._sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ), + TestStreamFalseDeduplication._sse("message_stop", {"type": "message_stop"}), + ] + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + return PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames) + + def test_complete_streaming_response_set_on_model_call_details(self): + """ + After the fix, _create_anthropic_response_logging_payload must set + complete_streaming_response on logging_obj.model_call_details so that + callbacks like _PROXY_track_cost_callback see the assembled response + instead of None. + + Before the fix: model_call_details had no complete_streaming_response key. + The log showed: "kwargs stream: True + complete streaming response: None" + """ + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + ) + + # pass_through_request sets the stream flag before the streaming handler + # reconstructs the response; mirror that here. + logging_obj = self._make_logging_obj(stream=True) + logging_obj.model_call_details["stream"] = True + all_chunks = list(self._build_chunks()) + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-sonnet-20241022", "stream": True}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + all_chunks=all_chunks, + end_time=datetime.now(), + ) + + # The assembled response must be stored on model_call_details so callbacks + # can identify this as a completed streaming call, not an in-progress one. + assert ( + logging_obj.model_call_details.get("complete_streaming_response") + is not None + ), "complete_streaming_response must be set on model_call_details after assembly" + + # The returned result must match what was stored + assert result["result"] is logging_obj.model_call_details.get( + "complete_streaming_response" + ) + + def test_dedup_guard_fires_when_stream_true_on_logging_obj(self): + """ + When logging_obj.stream is True (set by pass_through_endpoints.py after + detecting a streaming request), dispatch_success_handlers must set + has_dispatched_final_stream_success=True on the first call so that any + second call is a no-op. + + This is the _is_assembled_stream_success gate: with stream=False it + always returned False and the guard was permanently disabled. + """ + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + ) + from litellm.types.utils import ModelResponse + + # Simulate what pass_through_endpoints.py now does after stream detection + logging_obj = self._make_logging_obj(stream=False) + logging_obj.stream = True # fix applied + logging_obj.model_call_details["stream"] = True + + # Simulate what _create_anthropic_response_logging_payload now does + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + assert logging_obj._is_assembled_stream_success(result=mock_response) is True + + # First dispatch sets the flag + assert not logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + logging_obj.model_call_details["has_dispatched_final_stream_success"] = True + + # Second dispatch would be blocked — simulate the guard check + would_skip = bool( + logging_obj._is_assembled_stream_success(result=mock_response) + and logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + ) + assert would_skip is True, ( + "Dedup guard must block a second dispatch_success_handlers call for the " + "same assembled streaming response" + ) + + def test_sse_fallback_path_sets_stream_true_for_dedup(self): + """ + When a nominally non-streaming request receives an SSE response + (_is_streaming_response returns True), the fallback branch in + pass_through_endpoints.py must set logging_obj.stream = True so the + dedup guard activates. + + Before the fix the fallback path never set stream=True, so + _is_assembled_stream_success always returned False and duplicate + callback dispatches were never blocked. + """ + from litellm.types.utils import ModelResponse + + # logging_obj starts with stream=False, as created before the request + logging_obj = self._make_logging_obj(stream=False) + assert logging_obj._is_assembled_stream_success(result=MagicMock()) is False + + # Simulate what the SSE fallback branch in pass_through_endpoints.py now does + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + # With stream=True the dedup guard must be active + assert logging_obj._is_assembled_stream_success(result=mock_response) is True + + logging_obj.model_call_details["has_dispatched_final_stream_success"] = True + + would_skip = bool( + logging_obj._is_assembled_stream_success(result=mock_response) + and logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + ) + assert would_skip is True + + def test_stream_false_logging_obj_bypasses_dedup_guard(self): + """ + Demonstrates the pre-fix state: with stream=False on the logging object, + _is_assembled_stream_success always returns False regardless of whether + complete_streaming_response is set. This means the dedup guard can never + fire, so duplicate dispatches go through unchecked. + + This test documents the old broken behavior so the fix is clearly justified. + """ + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + # With stream=False, _is_assembled_stream_success returns False even though + # complete_streaming_response is present — the guard is permanently disabled. + assert logging_obj._is_assembled_stream_success(result=mock_response) is False + + +class TestNonStreamingResponseRedaction: + """ + Regression tests ensuring _create_anthropic_response_logging_payload only sets + complete_streaming_response for streaming responses. perform_redaction scrubs + that field exclusively when model_call_details["stream"] is True, so storing it + on a non-streaming response would deliver the unredacted response to logging + callbacks when message logging is disabled. + """ + + @staticmethod + def _make_logging_obj(stream: bool) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "hello"}], + stream=stream, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1245", + ) + # pass_through_request mirrors the stream flag onto model_call_details, + # which is the key perform_redaction inspects. + logging_obj.model_call_details["stream"] = stream + return logging_obj + + def test_non_streaming_does_not_set_complete_streaming_response(self): + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + response = ModelResponse(model="claude-3-5-sonnet-20241022") + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert ( + "complete_streaming_response" not in logging_obj.model_call_details + ), "non-streaming responses must not populate complete_streaming_response" + + def test_streaming_sets_complete_streaming_response(self): + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=True) + response = ModelResponse(model="claude-3-5-sonnet-20241022") + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert ( + logging_obj.model_call_details.get("complete_streaming_response") + is response + ) + + def test_non_streaming_response_is_redacted_when_message_logging_off(self): + from litellm.litellm_core_utils.redact_messages import ( + redact_message_input_output_from_logging, + ) + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + response = ModelResponse( + model="claude-3-5-sonnet-20241022", + choices=[Choices(message=Message(role="assistant", content="secret"))], + ) + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + logging_obj.model_call_details["litellm_params"] = { + "metadata": {"headers": {"x-litellm-enable-message-redaction": True}} + } + + redacted = redact_message_input_output_from_logging( + model_call_details=logging_obj.model_call_details, + result=response, + ) + + leaked = logging_obj.model_call_details.get("complete_streaming_response") + assert leaked is None + assert redacted.choices[0].message.content == "redacted-by-litellm" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 61299e2662a..89c57bcced6 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1050,6 +1050,131 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): assert metadata["user_api_key_user_id"] == "test-user-id" +@pytest.mark.asyncio +async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): + """ + Regression: a streaming pass-through request must flag its logging object as + streaming (logging_obj.stream and model_call_details["stream"]) before the + response is dispatched, so cost/success callbacks treat it as a stream and the + streaming dedup guard fires instead of double-logging. + """ + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock( + return_value={"model": "claude-3", "stream": True} + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {} + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def _empty_chunks(*args, **kwargs): + return + yield # pragma: no cover + + mock_chunk_processor.return_value = _empty_chunks() + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/v1/messages" + mock_request.body = AsyncMock( + return_value=b'{"model": "claude-3", "stream": true}' + ) + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="http://target-api.com/v1/messages", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + async_client.send.assert_awaited_once() + assert async_client.send.call_args.kwargs["stream"] is True + + mock_chunk_processor.assert_called_once() + logging_obj = mock_chunk_processor.call_args.kwargs[ + "litellm_logging_obj" + ] + assert logging_obj.stream is True + assert logging_obj.model_call_details["stream"] is True + + +@pytest.mark.asyncio +async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): + """ + Regression: a request that is not flagged as streaming up front but whose + upstream response comes back as an SSE stream (content-type text/event-stream) + must still flag its logging object as streaming before dispatch. Otherwise the + cost/success callbacks treat the assembled stream as a non-stream and the dedup + guard never fires, double-logging the request. + """ + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock( + return_value={"model": "claude-3"} + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {"content-type": "text/event-stream"} + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.request = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def _empty_chunks(*args, **kwargs): + return + yield # pragma: no cover + + mock_chunk_processor.return_value = _empty_chunks() + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/v1/messages" + mock_request.body = AsyncMock(return_value=b'{"model": "claude-3"}') + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="http://target-api.com/v1/messages", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=False, + ) + + async_client.request.assert_awaited_once() + + mock_chunk_processor.assert_called_once() + logging_obj = mock_chunk_processor.call_args.kwargs[ + "litellm_logging_obj" + ] + assert logging_obj.stream is True + assert logging_obj.model_call_details["stream"] is True + + @pytest.mark.asyncio async def test_create_pass_through_endpoint(): """