From 135ec00b27ac2452611032336c743c57f64ccae7 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:06:09 +0000 Subject: [PATCH 01/43] feat(model_armor): logging_only mode scans completed streams after delivery Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 25 +- .../model_armor/model_armor.py | 68 ++++- .../integrations/test_custom_guardrail.py | 57 +++- .../guardrail_hooks/test_model_armor.py | 289 ++++++++++++++++++ 4 files changed, 428 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 77bf4820a1a..ae6d646b5ae 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -30,6 +30,7 @@ from litellm.types.utils import ( GuardrailStatus, GuardrailTracingDetail, LLMResponseTypes, + ModelResponse, StandardLoggingGuardrailInformation, ) @@ -602,9 +603,7 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) - if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks - else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) if self.uses_apply_guardrail_interface() else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( @@ -883,7 +882,9 @@ class CustomGuardrail(CustomLogger): """logging_only: run apply_guardrail on copies of the logged request/response and record the verdict.""" from litellm.llms import get_guardrail_translation_mapping - if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: + if not self.uses_apply_guardrail_interface(): + return kwargs, result + if not self._event_hook_is_event_type(GuardrailEventHooks.logging_only): return kwargs, result try: translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))() @@ -922,6 +923,8 @@ class CustomGuardrail(CustomLogger): translation: "BaseTranslation", scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: + from litellm.llms import get_guardrail_translation_mapping + optional_params: Final = kwargs.get("optional_params") or {} scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) scratch_request: Final = { @@ -933,8 +936,18 @@ class CustomGuardrail(CustomLogger): "metadata": scratch_metadata, } await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) - await translation.process_output_response( - response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request + response: Final = ( + kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result + ) + if response is None: + return + output_translation: Final = ( + get_guardrail_translation_mapping(CallTypes.acompletion)() + if isinstance(response, ModelResponse) + else translation + ) + await output_translation.process_output_response( + response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request ) def supports_scan_only_tool_results(self) -> bool: diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index fde40111d49..01b05227a91 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,11 +1,13 @@ +import time from collections.abc import AsyncGenerator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional import httpx from fastapi import HTTPException if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel import json @@ -52,6 +54,7 @@ from litellm.types.utils import ( CallTypes, CallTypesLiteral, Choices, + GenericGuardrailAPIInputs, GuardrailStatus, ModelResponse, ModelResponseStream, @@ -118,8 +121,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): Supports: - Pre-call sanitization (sanitizeUserPrompt) - Post-call sanitization (sanitizeModelResponse) + - logging_only: scans the completed response after it reaches the client and + records the verdict in spend logs without blocking """ + use_native_lifecycle_hooks: ClassVar[bool] = True + @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: return [ @@ -128,6 +135,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): GuardrailEventHooks.post_call, GuardrailEventHooks.pre_mcp_call, GuardrailEventHooks.during_mcp_call, + GuardrailEventHooks.logging_only, ] def __init__( @@ -1096,6 +1104,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): add_guardrail_to_applied_guardrails_header, ) + if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: + async for chunk in response: + yield chunk + return + all_chunks: Final[Sequence[object]] = tuple([chunk async for chunk in response]) if not all_chunks or self._is_terminal_error_stream(all_chunks): @@ -1213,6 +1226,59 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): for chunk in all_chunks: yield chunk + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + content: Final = "\n".join(text for text in inputs.get("texts") or () if text) + if not content: + return inputs + + source: Final[Literal["user_prompt", "model_response"]] = ( + "user_prompt" if input_type == "request" else "model_response" + ) + start_time: Final = time.time() + try: + armor_response: Final = await self.make_model_armor_request( + content=content, source=source, request_data=request_data + ) + except ModelArmorAPIError as e: + error_end_time: Final = time.time() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=str(e), + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + guardrail_provider="model_armor", + start_time=start_time, + end_time=error_end_time, + duration=error_end_time - start_time, + ) + raise + + flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) + end_time: Final = time.time() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=self._build_logging_response(armor_response), + request_data=request_data, + guardrail_status="guardrail_flagged" if flagged else "success", + guardrail_provider="model_armor", + start_time=start_time, + end_time=end_time, + duration=end_time - start_time, + ) + if flagged: + raise HTTPException( + status_code=400, + detail=self._build_block_error_detail( + "Response blocked by Model Armor" if input_type == "response" else "Violated content safety policy", + armor_response, + ), + ) + return inputs + @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: """ diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index ddc8439a83a..c0463210445 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2404,7 +2404,7 @@ def test_logging_only_requires_framework_support_or_explicit_declaration( event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, ) -> None: supported: Final = [GuardrailEventHooks.pre_call] - if guardrail_type is _InheritedApplyGuardrail: + if guardrail_type is not CustomGuardrail: guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) assert guardrail.event_hook == event_hook assert supported == [GuardrailEventHooks.pre_call] @@ -2566,7 +2566,7 @@ class TestLoggingOnlyApplyGuardrail: assert [e["guardrail_status"] for e in entries] == ["success"] @pytest.mark.asyncio - async def test_native_lifecycle_hook_guardrail_is_left_alone(self): + async def test_native_lifecycle_hook_guardrail_scans_in_logging_only(self): class _NativeHooks(_ApplyOnlyObserver): use_native_lifecycle_hooks = True @@ -2575,9 +2575,9 @@ class TestLoggingOnlyApplyGuardrail: out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) - assert guardrail.calls == [] - assert out_kwargs is kwargs + assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] assert out_response is response + assert out_kwargs["standard_logging_object"]["guardrail_information"] @pytest.mark.asyncio async def test_aresponses_scans_logged_messages_when_input_is_cleared(self): @@ -2829,3 +2829,52 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: assert response.choices[0].message.content == "filtered response" assert "guardrail_to_apply" not in request_data assert len(_guardrail_entries(request_data)) == 1 + + +class _NativeLifecycleLoggingGuardrail(CustomGuardrail): + """Native lifecycle guardrail that also implements apply_guardrail, like the azure guards.""" + + use_native_lifecycle_hooks: ClassVar[bool] = True + + def __init__(self): + from litellm.types.guardrails import GuardrailEventHooks + + super().__init__( + guardrail_name="native-logging-guardrail", + event_hook=GuardrailEventHooks.logging_only, + ) + self.calls: list = [] + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.calls.append((input_type, list(inputs.get("texts") or []))) + return inputs + + +@pytest.mark.asyncio +async def test_native_lifecycle_guardrail_logging_only_scans_assembled_response(): + """A use_native_lifecycle_hooks guardrail accepts mode logging_only and its + async_logging_hook scans kwargs["async_complete_streaming_response"], not the raw result.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _NativeLifecycleLoggingGuardrail() + assembled = ModelResponse( + choices=[Choices(message=Message(role="assistant", content="assembled stream text"))] + ) + sentinel_result = object() + kwargs = { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + "async_complete_streaming_response": assembled, + } + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=sentinel_result, call_type=CallTypes.acompletion.value + ) + + assert out_result is sentinel_result + assert ("response", ["assembled stream text"]) in guardrail.calls + assert out_kwargs["standard_logging_object"]["guardrail_information"] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 47089b7b1b1..a6ed4e14616 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -4929,3 +4929,292 @@ def test_every_responses_delta_event_is_in_the_scanned_set(): } assert not missing assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES + + +def _clean_armor_response() -> dict: + return { + "sanitizationResult": { + "filterMatchState": "NO_MATCH_FOUND", + "filterResults": {}, + } + } + + +def _flagged_armor_response() -> dict: + return { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}}, + } + } + + +def _logging_only_guardrail() -> ModelArmorGuardrail: + return ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-logging", + event_hook=GuardrailEventHooks.logging_only, + ) + + +def _logged_kwargs() -> dict: + return { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + } + + +def _chat_response(text: str) -> litellm.ModelResponse: + return litellm.ModelResponse( + choices=[ + litellm.types.utils.Choices( + message=litellm.types.utils.Message(role="assistant", content=text) + ) + ] + ) + + +def _stream_chunk(text: str) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content=text) + ) + ] + ) + + +def _metadata_entries(kwargs: dict) -> list: + return kwargs["standard_logging_object"].get("guardrail_information") or [] + + +def test_logging_only_mode_is_accepted_and_keeps_native_hooks(): + guardrail = _logging_only_guardrail() + assert guardrail.event_hook == GuardrailEventHooks.logging_only + assert guardrail.use_native_lifecycle_hooks is True + assert GuardrailEventHooks.logging_only in ModelArmorGuardrail.get_supported_event_hooks() + + post_call_guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-post", + event_hook=GuardrailEventHooks.post_call, + ) + assert post_call_guardrail._deployment_hook_target() is post_call_guardrail + + +@pytest.mark.asyncio +async def test_logging_only_stream_yields_chunks_without_waiting_for_scan(): + """A logging_only guardrail must pass stream chunks straight through; the scan happens + afterwards on the assembled response via async_logging_hook.""" + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock( + side_effect=AssertionError("logging_only must not scan the stream") + ) + + produced = 0 + + async def gen(): + nonlocal produced + for i in range(3): + produced += 1 + yield _stream_chunk(f"chunk-{i} ") + + hook_iter = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=gen(), + request_data={"metadata": {}, "guardrails": ["model-armor-logging"]}, + ) + first = await hook_iter.__anext__() + assert produced == 1 + chunks = [first] + async for chunk in hook_iter: + chunks.append(chunk) + assert len(chunks) == 3 + guardrail.make_model_armor_request.assert_not_awaited() + + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + response = _chat_response("all clear") + kwargs = _logged_kwargs() + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + assert out_result is response + entries = _metadata_entries(out_kwargs) + assert len(entries) >= 1 + entry = entries[-1] + assert entry["guardrail_status"] == "success" + assert entry["guardrail_mode"] == "logging_only" + assert entry["guardrail_provider"] == "model_armor" + + +@pytest.mark.asyncio +async def test_logging_only_records_flagged_verdict_without_altering_response(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + response = _chat_response("flagged output") + kwargs = _logged_kwargs() + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + assert out_result is response + entries = _metadata_entries(out_kwargs) + assert entries[-1]["guardrail_status"] == "guardrail_flagged" + assert entries[-1]["guardrail_mode"] == "logging_only" + + +@pytest.mark.asyncio +async def test_logging_only_records_model_armor_api_error(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock( + side_effect=ModelArmorAPIError("Model Armor API error (upstream 500)") + ) + response = _chat_response("some output") + kwargs = _logged_kwargs() + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + assert out_result is response + entries = _metadata_entries(out_kwargs) + assert entries[-1]["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_logging_only_scans_assembled_responses_api_stream(): + """The terminal ResponseCompletedEvent is an envelope; the scan must run on the + assembled ResponsesAPIResponse kept in kwargs.""" + from openai.types.responses import ResponseOutputMessage, ResponseOutputText + + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + assembled = ResponsesAPIResponse( + id="resp-1", + created_at=1700000000, + output=[ + ResponseOutputMessage( + id="msg-1", + type="message", + role="assistant", + status="completed", + content=[ + ResponseOutputText( + annotations=[], text="assembled output text", type="output_text" + ) + ], + ) + ], + ) + event = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=assembled + ) + + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + kwargs = _logged_kwargs() + del kwargs["messages"] + kwargs["input"] = "hello" + kwargs["async_complete_streaming_response"] = assembled + + out_kwargs, _ = await guardrail.async_logging_hook( + kwargs=kwargs, result=event, call_type="aresponses" + ) + + response_scans = [ + call + for call in guardrail.make_model_armor_request.await_args_list + if call.kwargs.get("source") == "model_response" + ] + assert response_scans, "expected a model_response scan of the assembled response" + assert "assembled output text" in response_scans[0].kwargs["content"] + assert _metadata_entries(out_kwargs) + + +@pytest.mark.asyncio +async def test_logging_only_scans_anthropic_messages_model_response(): + """/v1/messages logs a ModelResponse; the output scan must extract the assistant text.""" + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + kwargs = _logged_kwargs() + kwargs["messages"] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + response = _chat_response("anthropic assembled text") + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="anthropic_messages" + ) + + assert out_result is response + response_scans = [ + call + for call in guardrail.make_model_armor_request.await_args_list + if call.kwargs.get("source") == "model_response" + ] + assert response_scans + assert "anthropic assembled text" in response_scans[0].kwargs["content"] + assert _metadata_entries(out_kwargs) + + +@pytest.mark.asyncio +async def test_logging_only_skips_output_scan_when_no_assembled_response(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + kwargs = _logged_kwargs() + + await guardrail.async_logging_hook(kwargs=kwargs, result=None, call_type="acompletion") + + sources = [call.kwargs.get("source") for call in guardrail.make_model_armor_request.await_args_list] + assert "model_response" not in sources + + +@pytest.mark.asyncio +async def test_native_post_call_mode_ignores_logging_hook(): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-post", + event_hook=GuardrailEventHooks.post_call, + ) + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) + response = _chat_response("some output") + kwargs = _logged_kwargs() + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + assert out_kwargs is kwargs + assert out_result is response + guardrail.make_model_armor_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_apply_guardrail_raises_on_flagged_content(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + request_data = {"metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["forbidden output"]}, + request_data=request_data, + input_type="response", + ) + + assert exc_info.value.status_code == 400 + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert entries[-1]["guardrail_status"] == "guardrail_flagged" From a5cc65f1a9b5f7a674d98f2e45460790952724e5 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:10:35 +0000 Subject: [PATCH 02/43] refactor(custom_guardrail): resolve logging_only output translation in async_logging_hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index ae6d646b5ae..305a5c20764 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -901,8 +901,16 @@ class CustomGuardrail(CustomLogger): for key, value in (litellm_params.get("metadata") or {}).items() if key != "standard_logging_guardrail_information" } + response: Final = ( + kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result + ) + output_translation: Final = ( + get_guardrail_translation_mapping(CallTypes.acompletion)() + if isinstance(response, ModelResponse) + else translation + ) try: - await self._scan_logged_call(kwargs, result, translation, scratch_metadata) + await self._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata) except Exception as e: verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e) recorded: Final = scratch_metadata.get("standard_logging_guardrail_information") @@ -919,12 +927,11 @@ class CustomGuardrail(CustomLogger): async def _scan_logged_call( self, kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract - result: object, + response: object | None, translation: "BaseTranslation", + output_translation: "BaseTranslation", scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: - from litellm.llms import get_guardrail_translation_mapping - optional_params: Final = kwargs.get("optional_params") or {} scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) scratch_request: Final = { @@ -936,16 +943,8 @@ class CustomGuardrail(CustomLogger): "metadata": scratch_metadata, } await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) - response: Final = ( - kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result - ) if response is None: return - output_translation: Final = ( - get_guardrail_translation_mapping(CallTypes.acompletion)() - if isinstance(response, ModelResponse) - else translation - ) await output_translation.process_output_response( response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request ) From f417d7f739fca2b8abed6ad3d4a399e39b2cfee7 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:49:11 +0000 Subject: [PATCH 03/43] fix(model_armor): record logging_only verdicts without raising so both scans run Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_armor/model_armor.py | 12 +---- .../guardrail_hooks/test_model_armor.py | 53 ++++++++++++++++--- 2 files changed, 47 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 01b05227a91..127a7ea6786 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1245,7 +1245,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): armor_response: Final = await self.make_model_armor_request( content=content, source=source, request_data=request_data ) - except ModelArmorAPIError as e: + except (ModelArmorAPIError, httpx.HTTPError) as e: error_end_time: Final = time.time() self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=str(e), @@ -1256,7 +1256,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=error_end_time, duration=error_end_time - start_time, ) - raise + return inputs flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) end_time: Final = time.time() @@ -1269,14 +1269,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=end_time, duration=end_time - start_time, ) - if flagged: - raise HTTPException( - status_code=400, - detail=self._build_block_error_detail( - "Response blocked by Model Armor" if input_type == "response" else "Violated content safety policy", - armor_response, - ), - ) return inputs @staticmethod diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index a6ed4e14616..bd358e84148 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -5203,18 +5203,55 @@ async def test_native_post_call_mode_ignores_logging_hook(): @pytest.mark.asyncio -async def test_apply_guardrail_raises_on_flagged_content(): +async def test_apply_guardrail_records_flagged_without_raising(): guardrail = _logging_only_guardrail() guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) request_data = {"metadata": {}} + inputs = {"texts": ["forbidden output"]} - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs={"texts": ["forbidden output"]}, - request_data=request_data, - input_type="response", - ) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) - assert exc_info.value.status_code == 400 + assert result == inputs entries = request_data["metadata"]["standard_logging_guardrail_information"] assert entries[-1]["guardrail_status"] == "guardrail_flagged" + + +@pytest.mark.asyncio +async def test_logging_only_records_transport_error(): + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(side_effect=httpx.ConnectError("boom")) + response = _chat_response("some output") + kwargs = _logged_kwargs() + + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + assert out_result is response + entries = _metadata_entries(out_kwargs) + failed = [e for e in entries if e["guardrail_status"] == "guardrail_failed_to_respond"] + assert failed + assert all(e["guardrail_provider"] == "model_armor" for e in failed) + + +@pytest.mark.asyncio +async def test_logging_only_flagged_prompt_still_scans_response(): + """A flagged input scan must not abort the output scan; both verdicts are recorded.""" + guardrail = _logging_only_guardrail() + guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + response = _chat_response("flagged output") + kwargs = _logged_kwargs() + + out_kwargs, _ = await guardrail.async_logging_hook( + kwargs=kwargs, result=response, call_type="acompletion" + ) + + sources = [call.kwargs.get("source") for call in guardrail.make_model_armor_request.await_args_list] + assert sources == ["user_prompt", "model_response"] + entries = _metadata_entries(out_kwargs) + flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] + assert len(flagged) == 2 From 39b916f13f3330b5099ff45ede27e59f7278f53a Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 09:07:06 +0000 Subject: [PATCH 04/43] fix(model_armor): decorate apply_guardrail with log_guardrail_information Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/guardrails/guardrail_hooks/model_armor/model_armor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 127a7ea6786..a833de6096d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1226,6 +1226,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): for chunk in all_chunks: yield chunk + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, From 2ca29a9a9157d52149a4c2cd38e98f89c7f57e30 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 09:33:57 +0000 Subject: [PATCH 05/43] fix(model_armor): gate apply_guardrail raise to non-logging_only and require native guardrails to declare logging_only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 4 ++- .../model_armor/model_armor.py | 12 +++++++-- .../integrations/test_custom_guardrail.py | 2 +- .../guardrail_hooks/test_model_armor.py | 27 +++++++++++++++++++ 4 files changed, 41 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 305a5c20764..407445b2828 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -603,7 +603,9 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) if self.uses_apply_guardrail_interface() else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) + if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks + else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index a833de6096d..a0563a7a1c9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,7 +1,7 @@ import time from collections.abc import AsyncGenerator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal import httpx from fastapi import HTTPException @@ -1232,7 +1232,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], - logging_obj: Optional["LiteLLMLoggingObj"] = None, + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: content: Final = "\n".join(text for text in inputs.get("texts") or () if text) if not content: @@ -1270,6 +1270,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=end_time, duration=end_time - start_time, ) + if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only): + raise HTTPException( + status_code=400, + detail=self._build_block_error_detail( + "Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor", + armor_response, + ), + ) return inputs @staticmethod diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index c0463210445..bdb7fad21b3 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2404,7 +2404,7 @@ def test_logging_only_requires_framework_support_or_explicit_declaration( event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, ) -> None: supported: Final = [GuardrailEventHooks.pre_call] - if guardrail_type is not CustomGuardrail: + if guardrail_type is _InheritedApplyGuardrail: guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) assert guardrail.event_hook == event_hook assert supported == [GuardrailEventHooks.pre_call] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index bd358e84148..126e162fec8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -5255,3 +5255,30 @@ async def test_logging_only_flagged_prompt_still_scans_response(): entries = _metadata_entries(out_kwargs) flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] assert len(flagged) == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_raises_on_flagged_when_not_logging_only(): + """The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a + non-logging_only instance must signal the block so flagged text is not returned as clean.""" + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-pre", + event_hook=GuardrailEventHooks.pre_call, + ) + guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + request_data = {"metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["forbidden prompt"]}, + request_data=request_data, + input_type="request", + ) + + assert exc_info.value.status_code == 400 + entries = request_data["metadata"]["standard_logging_guardrail_information"] + flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] + assert len(flagged) == 1 From 88de192dcf55e26f0f2cabb4d769a825dd8dbc7f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:29:04 -0700 Subject: [PATCH 06/43] test: bind management E2E callers and isolate JWT actors --- .github/e2e-stack/assert_tests_ran.py | 15 + .github/e2e-stack/oidc-profile.sh | 5 + .github/e2e-stack/select_tests.py | 2 + .github/workflows/test-e2e-changed.yml | 2 +- .../test_e2e_changed_gate.py | 23 ++ tests/e2e/CONTRIBUTING.md | 6 +- .../e2e/coverage_registry/management_cases.py | 151 +++++++++ tests/e2e/coverage_registry/mgmt.yaml | 8 + tests/e2e/e2e_http.py | 31 +- tests/e2e/idp.py | 254 ++++++++++++++- tests/e2e/junit_properties.py | 3 +- tests/e2e/management/conftest.py | 19 +- tests/e2e/management/jwt_actors.py | 175 ++++++++++ tests/e2e/management/management_client.py | 122 ++++--- .../e2e/management/test_jwt_management_e2e.py | 233 +++++++++++-- tests/e2e/models.py | 28 +- tests/e2e/proxy_client.py | 116 ++++--- tests/e2e/test_e2e_http.py | 10 + tests/e2e/test_idp.py | 122 ++++++- tests/e2e/test_proxy_client.py | 307 +++++++++++++++++- tests/e2e/transport.py | 17 +- tests/e2e/ui/oidcSetup.ts | 30 ++ tests/e2e/ui/playwright.oidc.config.ts | 22 ++ 23 files changed, 1541 insertions(+), 160 deletions(-) create mode 100755 .github/e2e-stack/oidc-profile.sh create mode 100644 tests/e2e/coverage_registry/management_cases.py create mode 100644 tests/e2e/management/jwt_actors.py create mode 100644 tests/e2e/ui/oidcSetup.ts create mode 100644 tests/e2e/ui/playwright.oidc.config.ts diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py index c4348c20873..2303c42f4fb 100644 --- a/.github/e2e-stack/assert_tests_ran.py +++ b/.github/e2e-stack/assert_tests_ran.py @@ -3,6 +3,9 @@ import xml.etree.ElementTree as ET from pathlib import Path from typing import Final +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e")) +from coverage_registry.management_cases import MANAGEMENT_CASES + def main() -> int: selected: Final = tuple(sys.argv[2:]) @@ -16,6 +19,17 @@ def main() -> int: case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) ) missing: Final = tuple(path for path in selected if path not in passed) + required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected) + passed_nodes: Final = frozenset( + prop.get("value") + for case in cases + if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) + for prop in case.findall("./properties/property") + if prop.get("name") == "management_node" + ) + missing_nodes: Final = required_nodes - passed_nodes + for node in sorted(missing_nodes): + _ = sys.stdout.write(f"::error::required management case did not pass: {node}\n") for path in selected: collected: Final = sum(case.get("file") == path for case in cases) skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases) @@ -27,6 +41,7 @@ def main() -> int: if ( selected and not missing + and not missing_nodes and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error")) ): return 0 diff --git a/.github/e2e-stack/oidc-profile.sh b/.github/e2e-stack/oidc-profile.sh new file mode 100755 index 00000000000..84eaaaf8051 --- /dev/null +++ b/.github/e2e-stack/oidc-profile.sh @@ -0,0 +1,5 @@ +#!/usr/bin/env bash +set -euo pipefail +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "${REPO_ROOT}" +exec uv run --no-sync python tests/e2e/idp.py "$@" diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index 238818a0d36..982e93cf642 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -12,6 +12,8 @@ UNSUPPORTED: Final = re.compile( HARNESS: Final = re.compile( r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" r"|^tests/e2e/idp_realm\.json$" + r"|^tests/e2e/management/(management_client|jwt_actors|conftest)\.py$" + r"|^tests/e2e/coverage_registry/management_cases\.py$" r"|^tests/e2e/gateway/" r"|^\.github/e2e-stack/" r"|^\.github/workflows/test-e2e-changed\.yml$" diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index 1db597ff673..c9f08deb36e 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -183,7 +183,7 @@ jobs: log="${RUNNER_TEMP}/e2e-pass-${pass}.log" echo "::group::pass ${pass} of 3" set +e - uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \ + uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v --reruns 0 -p no:cacheprovider \ -o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1 status=$? uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}" diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 588402e3996..101816c7f11 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -81,6 +81,25 @@ def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None assert result.returncode == 1 +@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer")) +def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None: + suite: Final = ET.Element("testsuite") + path: Final = "tests/e2e/management/test_jwt_management_e2e.py" + case: Final = ET.SubElement(suite, "testcase", file=path) + properties: Final = ET.SubElement(case, "properties") + _ = ET.SubElement( + properties, + "property", + name="management_node", + value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]", + ) + report: Final = tmp_path / "report.xml" + ET.ElementTree(suite).write(report) + result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True) + assert result.returncode == 1 + assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout + + def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_path: Path) -> None: env_path: Final = tmp_path / ".env" @@ -141,6 +160,10 @@ def test_changed_suite_files_are_selected_unless_the_stack_cannot_run_them( ( "tests/e2e/proxy_client.py", "tests/e2e/conftest.py", + "tests/e2e/management/management_client.py", + "tests/e2e/management/jwt_actors.py", + "tests/e2e/management/conftest.py", + "tests/e2e/coverage_registry/management_cases.py", "tests/e2e/pytest.ini", "tests/e2e/gateway/stage_mirror_ci_config.yml", ".github/e2e-stack/up.sh", diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 44564a51e26..78c05ea4b30 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -60,7 +60,11 @@ The suites run against a live proxy, so bring one up first by running the litell Keycloak's password grant is a test-only provisioning shortcut, not a production login recommendation. The `litellm-e2e-admin` client adds the proxy's admin scope; the normal client does not. Never reuse this permissive realm outside an isolated test stack. - Management tests can use the shared `idp` and `jwt_identity` fixtures. Each test gets a unique Keycloak group/user and a matching proxy user/team. Setup and fallback cleanup use the master key; the operations and read-backs being tested must explicitly use `caller_key=idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)` (or a member token). See `management/test_jwt_management_e2e.py` for create/read/update/clear/delete and tenant-denial examples. A group claim alone is not database team membership: permission tests explicitly add the member and prove an allowed read before asserting the denied write. + Management tests can bind a credential once with `client.with_caller(Caller(...))`; direct calls, delegated helpers and replica read-backs then retain that caller. Explicit `caller_key` arguments override the binding. Keep the original master-backed client for bootstrap and cleanup. `actor_factory` lazily provisions database roles and tenant memberships, with `database_role` tokens carrying no groups and `group_scoped` actors retaining the existing team route gate. Token minting is explicit through `actor.mint_caller(idp)`. The factory runs requests without backend retries and reports cleanup failures. `coverage_registry/management_cases.py` records exact canary nodes and non-secret actor labels; the CI execution assertion rejects a missing or skipped actor row + + For the opt-in browser profile, start the existing IdP first, then run `.github/e2e-stack/oidc-profile.sh "$PROXY_BASE_URL" `. The wrapper creates a confidential client with an exact `/sso/callback` redirect and S256 PKCE, passes the client secret only through the child process environment, and removes the client on exit. It uses the existing generic OIDC handler with `GENERIC_USER_ID_ATTRIBUTE=sub`. Preserve the IdP's PostgreSQL data across restarts + + `tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping; browser journey specs under `ui/oidc/` are a separate coverage step Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`: diff --git a/tests/e2e/coverage_registry/management_cases.py b/tests/e2e/coverage_registry/management_cases.py new file mode 100644 index 00000000000..15dc7d333c5 --- /dev/null +++ b/tests/e2e/coverage_registry/management_cases.py @@ -0,0 +1,151 @@ +from dataclasses import dataclass +from typing import Final, Literal + +CredentialKind = Literal["master", "idp_admin", "direct_jwt", "virtual_key", "dashboard_session"] +DependencyProfile = Literal["management_only", "real_oidc_browser", "external_provider_required"] + + +@dataclass(frozen=True, slots=True) +class ManagementCase: + node: str + credential_kind: CredentialKind + actor: str + profile: str + method: Literal["GET", "POST"] + path: str + operation_family: str + dependency_profile: DependencyProfile = "management_only" + + +JWT_FILE: Final = "tests/e2e/management/test_jwt_management_e2e.py" +JWT_CLASS: Final = f"{JWT_FILE}::TestJwtManagement" +ACTORS: Final = ( + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", +) +MANAGEMENT_CASES: Final = tuple( + ManagementCase( + node=f"{JWT_CLASS}::test_actor_subject_and_database_role[{role}]", + credential_kind="direct_jwt", + actor=role, + profile="database_role", + method="GET", + path="/user/info", + operation_family="identity", + ) + for role in ACTORS +) + ( + ManagementCase( + node=f"{JWT_CLASS}::test_admin_viewer_reads_but_cannot_update", + credential_kind="direct_jwt", + actor="proxy_admin_viewer", + profile="database_role", + method="POST", + path="/key/update", + operation_family="denial", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[direct_jwt]", + credential_kind="direct_jwt", + actor="proxy_admin", + profile="group_scoped", + method="POST", + path="/key/generate", + operation_family="key_lifecycle", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[virtual_key]", + credential_kind="virtual_key", + actor="proxy_admin", + profile="database_role", + method="POST", + path="/key/generate", + operation_family="key_lifecycle", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_two_actor_sets_keep_tenants_and_keys_isolated", + credential_kind="direct_jwt", + actor="team_member", + profile="group_scoped", + method="GET", + path="/key/info", + operation_family="tenant_isolation", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_member_cannot_write_and_another_team_cannot_read_the_key", + credential_kind="direct_jwt", + actor="team_member", + profile="group_scoped", + method="POST", + path="/key/update", + operation_family="tenant_isolation", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_multi_group_actor_keeps_exact_memberships", + credential_kind="master", + actor="bootstrap", + profile="group_scoped", + method="GET", + path="/team/info", + operation_family="memberships", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_successful_actor_cleanup_removes_owned_state", + credential_kind="master", + actor="bootstrap", + profile="failure_cleanup", + method="GET", + path="/team/info", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[group]", + credential_kind="idp_admin", + actor="idp_admin", + profile="failure_cleanup", + method="POST", + path="/groups", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[user]", + credential_kind="idp_admin", + actor="idp_admin", + profile="failure_cleanup", + method="POST", + path="/users", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_oidc_browser_profile_identity_mapping", + credential_kind="direct_jwt", + actor="internal_user", + profile="oidc_configuration", + method="GET", + path="/protocol/openid-connect/userinfo", + operation_family="oidc_identity", + ), +) + + +def canonical_node(node: str) -> str: + return node if node.startswith("tests/e2e/") else f"tests/e2e/{node}" + + +def case_properties(node: str) -> tuple[tuple[str, str], ...]: + case: Final = next((case for case in MANAGEMENT_CASES if case.node == canonical_node(node)), None) + if case is None: + return () + return ( + ("management_node", case.node), + ("credential_kind", case.credential_kind), + ("actor", case.actor), + ("auth_profile", case.profile), + ("dependency_profile", case.dependency_profile), + ) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d1227fe7c0c..31ad61ba3e2 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -90,3 +90,11 @@ - {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"} - {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"} - {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"} + +- {id: mgmt.user.jwt.database_roles, module: mgmt, tier: P0, surface: api, assertions: [database_roles], source: "auth/handle_jwt.py", rationale: "User-only JWT subjects retain their seeded database roles and memberships"} +- {id: mgmt.key.jwt.viewer_denied, module: mgmt, tier: P0, surface: api, assertions: [viewer_denied], source: "auth/route_checks.py", rationale: "An admin viewer can read a key but cannot update it or change stored state"} +- {id: mgmt.user.oidc.identity_mapping, module: mgmt, tier: P0, surface: api, assertions: [identity_mapping], source: "tests/e2e/idp.py", rationale: "IdP configuration canary only: confidential-client token and userinfo subjects match the seeded user; application SSO is separate"} +- {id: mgmt.team.jwt.tenant_isolation, module: mgmt, tier: P0, surface: api, assertions: [tenant_isolation], source: "auth/handle_jwt.py", rationale: "Isolated team actors read their own key and receive 403 for the other tenant key"} +- {id: mgmt.team.jwt.multiple_memberships, module: mgmt, tier: P0, surface: api, assertions: [multiple_memberships], source: "auth/handle_jwt.py", rationale: "A multi-group actor has exactly the configured memberships without admin scope"} +- {id: mgmt.user.jwt.cleanup, module: mgmt, tier: P0, surface: api, assertions: [cleanup], source: "management_endpoints/internal_user_endpoints.py", rationale: "Owned users teams organizations keys and IdP objects disappear after successful cleanup"} +- {id: mgmt.user.jwt.partial_cleanup, module: mgmt, tier: P0, surface: api, assertions: [partial_cleanup], source: "auth/handle_jwt.py", rationale: "Partial identity setup removes the group and user created before failure"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index ce069720c6e..67370c98274 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -16,9 +16,11 @@ requests itself imports. from __future__ import annotations import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Generator, Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass -from typing import Final, Generator, Generic, Iterator, Literal, NewType, Protocol, TypeVar, cast +from typing import Final, Generic, Literal, NewType, Protocol, TypeVar, cast import pytest import requests @@ -36,8 +38,8 @@ class Headers(BaseModel): class AuthHeaders(Headers): # litellm accepts either; set whichever the call needs, leave the other None. - authorization: str | None = None - x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key") + authorization: str | None = Field(default=None, repr=False) + x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key", repr=False) class AnthropicHeaders(AuthHeaders): @@ -292,6 +294,22 @@ def _params(params: BaseModel | None) -> dict[str, str]: TRANSIENT_STATUSES: frozenset[int] = frozenset({529}) RETRY_ATTEMPTS: int = 3 +_QUALIFICATION: Final[ContextVar[bool]] = ContextVar("e2e_qualification", default=False) + + +def retry_attempts(default: int) -> int: + return 1 if _QUALIFICATION.get() else default + + +@contextmanager +def without_retries() -> Generator[None]: + token: Final = _QUALIFICATION.set(True) + try: + yield + finally: + _QUALIFICATION.reset(token) + + RETRY_BACKOFF_SECONDS: float = 0.5 @@ -319,7 +337,7 @@ def request_with_retry[T: RetryableResponse]( hang should surface as a hang instead of doubling the wall clock. Every retry prints, so flakiness stays visible in the run log instead of vanishing into green.""" - for attempt in range(1, RETRY_ATTEMPTS): + for attempt in range(1, retry_attempts(RETRY_ATTEMPTS)): resp = issue() if resp.status_code not in TRANSIENT_STATUSES: return resp @@ -414,6 +432,7 @@ def get_external[R: BaseModel]( url: str, *, response_type: type[R], + headers: BaseModel | None = None, timeout: float = 30.0, ) -> Result[R]: """GET an absolute URL outside the proxy (e.g. a public /.well-known document). @@ -422,7 +441,7 @@ def get_external[R: BaseModel]( try: resp = requests.get( url, - headers={"Accept": "application/json"}, + headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})}, timeout=timeout, ) except requests.RequestException as exc: diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py index 6d2fc84eb27..12db91bbd88 100644 --- a/tests/e2e/idp.py +++ b/tests/e2e/idp.py @@ -2,11 +2,17 @@ from __future__ import annotations +import base64 import os import secrets +import signal +import subprocess +import sys import warnings from collections.abc import Callable -from dataclasses import dataclass, field +from contextlib import ExitStack +from dataclasses import dataclass, field, replace +from types import FrameType from typing import Final, Literal import pytest @@ -14,11 +20,15 @@ from e2e_http import ( AuthHeaders, ExternalWrite, NetworkError, + NoBody, Result, Success, + UnknownApiError, delete_external, + get_external, post_form_external, post_json_external, + unwrap, ) from pydantic import BaseModel, Field @@ -46,7 +56,9 @@ class TokenGrantForm(BaseModel): grant_type: Literal["password"] = "password" client_id: str username: str - password: str + password: str = Field(repr=False) + client_secret: str | None = Field(default=None, repr=False) + scope: str | None = None class TokenResponse(BaseModel): @@ -63,7 +75,7 @@ class GroupCreateBody(BaseModel): class PasswordCredential(BaseModel): type: Literal["password"] = "password" - value: str + value: str = Field(repr=False) temporary: bool = False @@ -101,8 +113,20 @@ class Identity: user_id: str username: str password: str = field(repr=False) - group: str - group_id: str + groups: tuple[str, ...] + group_ids: tuple[str, ...] + + @property + def group(self) -> str: + if len(self.groups) != 1: + raise ValueError("A single-group identity is required") + return self.groups[0] + + @property + def group_id(self) -> str: + if len(self.group_ids) != 1: + raise ValueError("A single-group identity is required") + return self.group_ids[0] @dataclass(frozen=True, slots=True) @@ -111,6 +135,10 @@ class Keycloak: realm: str admin_username: str admin_password: str = field(repr=False) + strict_cleanup: bool = False + + def with_strict_cleanup(self) -> Keycloak: + return replace(self, strict_cleanup=True) @property def issuer(self) -> str: @@ -150,7 +178,9 @@ class Keycloak: f"group {name}", ) - def create_user(self, *, username: str, email: str, password: str, group: str) -> str: + def create_user( + self, *, username: str, email: str, password: str, group: str | None = None, groups: tuple[str, ...] = () + ) -> str: return created_id( post_json_external( self._admin_url("/users"), @@ -158,7 +188,7 @@ class Keycloak: json=UserCreateBody( username=username, email=email, - groups=(group,), + groups=(group,) if group is not None else groups, credentials=(PasswordCredential(value=password),), ), ), @@ -171,14 +201,28 @@ class Keycloak: def delete_group(self, group_id: str) -> None: self._delete(f"/groups/{group_id}") + def assert_absent(self, kind: Literal["users", "groups", "clients"], resource_id: str) -> None: + result: Final = get_external( + self._admin_url(f"/{kind}/{resource_id}"), + headers=self._admin_headers(), + response_type=NoBody, + ) + assert isinstance(result, UnknownApiError) and result.status_code == 404, ( + f"Owned IdP {kind} still exists: {result}" + ) + def _delete(self, path: str) -> None: try: headers: Final = self._admin_headers() except pytest.fail.Exception as exc: + if self.strict_cleanup: + raise RuntimeError(f"Keycloak cleanup could not authenticate for {path}") from exc warnings.warn(f"Keycloak cleanup could not authenticate for {path}: {exc}", RuntimeWarning, stacklevel=2) return result: Final = delete_external(self._admin_url(path), headers=headers) if result.status_code not in (204, 404): + if self.strict_cleanup: + raise RuntimeError(f"Keycloak cleanup failed for {path}: HTTP {result.status_code}") warnings.warn( f"Keycloak cleanup failed for {path}: HTTP {result.status_code} {result.body[:300]}", RuntimeWarning, @@ -188,15 +232,34 @@ class Keycloak: def provision(self, *, marker: str, group: str, defer: Callable[[Callable[[], object]], None]) -> Identity: """Create `group` and a user in it, credentialed with a password generated for this test alone, and hand back the identity a token can be minted for.""" - group_id: Final = self.create_group(group) - defer(lambda: self.delete_group(group_id)) + return self.provision_groups(marker=marker, groups=(group,), defer=defer) + + def provision_groups( + self, *, marker: str, groups: tuple[str, ...], defer: Callable[[Callable[[], object]], None] + ) -> Identity: + def provision_group(name: str) -> str: + created: Final = self.create_group(name) + defer(lambda: self.delete_group(created)) + return created + + group_ids: Final = tuple(provision_group(group) for group in groups) + return self.provision_user(marker=marker, groups=groups, group_ids=group_ids, defer=defer) + + def provision_user( + self, + *, + marker: str, + groups: tuple[str, ...], + group_ids: tuple[str, ...], + defer: Callable[[Callable[[], object]], None], + ) -> Identity: username: Final = f"e2e-jwt-user-{marker}" password: Final = secrets.token_urlsafe(24) user_id: Final = self.create_user( - username=username, email=f"{username}@example.com", password=password, group=group + username=username, email=f"{username}@example.com", password=password, groups=groups ) defer(lambda: self.delete_user(user_id)) - return Identity(user_id=user_id, username=username, password=password, group=group, group_id=group_id) + return Identity(user_id=user_id, username=username, password=password, groups=groups, group_ids=group_ids) def access_token( self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID, issuer_host: str | None = None @@ -211,6 +274,65 @@ class Keycloak: ) return self._token(result, f"a token for {identity.username}") + def discovery(self) -> Discovery: + return unwrap(get_external(f"{self.issuer}/.well-known/openid-configuration", response_type=Discovery)) + + def browser_client(self, *, callback_url: str, defer: Callable[[Callable[[], object]], None]) -> BrowserClient: + client: Final = BrowserClient( + client_id=f"e2e-browser-{secrets.token_hex(8)}", + secret=secrets.token_urlsafe(32), + callback_url=callback_url, + ) + resource_id: Final = created_id( + post_json_external( + self._admin_url("/clients"), + headers=self._admin_headers(), + json=BrowserClientBody( + clientId=client.client_id, + secret=client.secret, + redirectUris=(callback_url,), + ), + ), + "browser client", + ) + defer(lambda: self._delete(f"/clients/{resource_id}")) + configured: Final = unwrap( + get_external( + self._admin_url(f"/clients/{resource_id}"), + headers=self._admin_headers(), + response_type=BrowserClientBody, + ) + ) + assert configured.redirect_uris == (callback_url,) + assert configured.standard_flow_enabled and not configured.public_client + assert configured.attributes.pkce == "S256" + return client + + def browser_token(self, identity: Identity, client: BrowserClient) -> str: + return self._token( + post_form_external( + self.token_url(self.realm), + form=TokenGrantForm( + client_id=client.client_id, + client_secret=client.secret, + username=identity.username, + password=identity.password, + scope="openid email", + ), + response_type=TokenResponse, + ), + "browser-profile identity mapping", + ) + + def userinfo(self, token: str) -> UserInfo: + return unwrap( + get_external( + f"{self.issuer}/protocol/openid-connect/userinfo", + headers=AuthHeaders(authorization=f"Bearer {token}"), + response_type=UserInfo, + ) + ) + def keycloak_from_env() -> Keycloak: admin_username: Final = os.environ.get(KEYCLOAK_ADMIN_USER_ENV, "").strip() @@ -226,3 +348,113 @@ def keycloak_from_env() -> Keycloak: admin_username=admin_username, admin_password=admin_password, ) + + +class TokenClaims(BaseModel): + sub: str + iss: str + aud: str | tuple[str, ...] + exp: int + scope: str = "" + groups: tuple[str, ...] = () + + +class Discovery(BaseModel): + issuer: str + authorization_endpoint: str + token_endpoint: str + userinfo_endpoint: str + jwks_uri: str + + +class UserInfo(BaseModel): + sub: str + email: str + + +class BrowserAttributes(BaseModel): + pkce: str = Field(default="S256", alias="pkce.code.challenge.method") + + +class AudienceConfig(BaseModel): + audience: str = Field(default="litellm-e2e", alias="included.custom.audience") + access_token: str = Field(default="true", alias="access.token.claim") + id_token: str = Field(default="false", alias="id.token.claim") + + +class AudienceMapper(BaseModel): + name: str = "litellm-audience" + protocol: str = "openid-connect" + mapper: str = Field(default="oidc-audience-mapper", alias="protocolMapper") + config: AudienceConfig = Field(default_factory=AudienceConfig) + + +class BrowserClientBody(BaseModel): + client_id: str = Field(alias="clientId") + secret: str = Field(repr=False) + redirect_uris: tuple[str, ...] = Field(alias="redirectUris") + enabled: bool = True + public_client: bool = Field(default=False, alias="publicClient") + standard_flow_enabled: bool = Field(default=True, alias="standardFlowEnabled") + direct_access_grants_enabled: bool = Field(default=True, alias="directAccessGrantsEnabled") + default_client_scopes: tuple[str, ...] = Field(default=("email", "basic"), alias="defaultClientScopes") + attributes: BrowserAttributes = Field(default_factory=BrowserAttributes) + protocol_mappers: tuple[AudienceMapper, ...] = Field(default=(AudienceMapper(),), alias="protocolMappers") + + +@dataclass(frozen=True, slots=True) +class BrowserClient: + client_id: str + secret: str = field(repr=False) + callback_url: str + + def environment(self, discovery: Discovery) -> dict[str, str]: + return { + "GENERIC_CLIENT_ID": self.client_id, + "GENERIC_CLIENT_SECRET": self.secret, + "GENERIC_USER_ID_ATTRIBUTE": "sub", + "GENERIC_AUTHORIZATION_ENDPOINT": discovery.authorization_endpoint, + "GENERIC_TOKEN_ENDPOINT": discovery.token_endpoint, + "GENERIC_USERINFO_ENDPOINT": discovery.userinfo_endpoint, + "GENERIC_CLIENT_USE_PKCE": "true", + "GENERIC_SCOPE": "openid email", + } + + +def token_claims(token: str) -> TokenClaims: + payload: Final = token.split(".")[1] + return TokenClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))) + + +def run_oidc_profile(proxy_url: str, command: list[str]) -> int: + idp: Final = keycloak_from_env().with_strict_cleanup() + with ExitStack() as cleanup: + + def terminate(signum: int, frame: FrameType | None) -> None: + raise SystemExit(128 + signum) + + previous: Final = signal.signal(signal.SIGTERM, terminate) + cleanup.callback(signal.signal, signal.SIGTERM, previous) + + def defer(callback: Callable[[], object]) -> None: + cleanup.callback(callback) + + client: Final = idp.browser_client(callback_url=f"{proxy_url.rstrip('/')}/sso/callback", defer=defer) + environment: Final = {**os.environ, **client.environment(idp.discovery()), "PROXY_BASE_URL": proxy_url} + with subprocess.Popen(command, env=environment) as child: + try: + return child.wait() + finally: + if child.poll() is None: + child.terminate() + try: + child.wait(timeout=5) + except subprocess.TimeoutExpired: + child.kill() + child.wait() + + +if __name__ == "__main__": + if len(sys.argv) < 3: + raise SystemExit("Usage: idp.py PROXY_URL COMMAND [ARG ...]; requires a running test IdP") + raise SystemExit(run_oidc_profile(sys.argv[1], sys.argv[2:])) diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py index c5971c5362c..b9f5da871ae 100644 --- a/tests/e2e/junit_properties.py +++ b/tests/e2e/junit_properties.py @@ -19,6 +19,7 @@ from __future__ import annotations from collections.abc import Iterable import pytest +from coverage_registry.management_cases import case_properties # Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing # at runtime names this suite's place in the repo. test_junit_properties.py @@ -94,7 +95,7 @@ def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]: ("package", package_from_nodeid(item.nodeid)), ("covers", ",".join(covers_from_item(item))), ("source", source_from_item(item)), - ) + ) + case_properties(item.nodeid) def attach_result_properties(item: pytest.Item) -> None: diff --git a/tests/e2e/management/conftest.py b/tests/e2e/management/conftest.py index bd69c8c0ff3..5a11b634085 100644 --- a/tests/e2e/management/conftest.py +++ b/tests/e2e/management/conftest.py @@ -5,8 +5,14 @@ holds the shared ProxyClient so `resources` / `scoped_key` clean up keys, teams, users, and orgs this suite creates. """ -import pytest +from collections.abc import Generator +from typing import Final +import pytest +from e2e_http import without_retries +from idp import Keycloak +from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory from management_client import ManagementClient, build_client from proxy_client import ProxyClient @@ -21,3 +27,14 @@ def pytest_configure(config: pytest.Config) -> None: @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> ManagementClient: return build_client(proxy) + + +@pytest.fixture +def actor_factory(proxy: ProxyClient, idp: Keycloak) -> Generator[ActorFactory]: + bootstrap: Final = build_client(proxy) + resources: Final = ResourceManager(client=proxy, strict_cleanup=True) + with without_retries(): + try: + yield ActorFactory(bootstrap=bootstrap, idp=idp, resources=resources) + finally: + resources.teardown() diff --git a/tests/e2e/management/jwt_actors.py b/tests/e2e/management/jwt_actors.py new file mode 100644 index 00000000000..909d1652ada --- /dev/null +++ b/tests/e2e/management/jwt_actors.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final, Literal + +from e2e_config import unique_marker +from e2e_http import NoBody, unwrap +from idp import ADMIN_CLIENT_ID, TESTS_CLIENT_ID, Identity, Keycloak +from lifecycle import ResourceManager +from management.management_client import ManagementClient +from models import ( + KeyGenerateBody, + KeyGenerateResponse, + OrgDeleteBody, + OrgDeleteResponse, + OrgMemberAddBody, + OrgMemberEntry, + OrgNewBody, + TeamDeleteBody, + TeamMemberAddBody, + TeamMemberEntry, + TeamNewBody, + UserNewBody, + UserRole, +) +from proxy_client import Caller + +ActorRole = Literal[ + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", +] +ActorProfile = Literal["database_role", "group_scoped"] + + +@dataclass(frozen=True, slots=True) +class Tenant: + organization_id: str + team_id: str + group_id: str + + +@dataclass(frozen=True, slots=True) +class Actor: + identity: Identity + role: ActorRole + global_role: UserRole + profile: ActorProfile + tenants: tuple[Tenant, ...] + + def mint_caller(self, idp: Keycloak) -> Caller: + return Caller( + credential=idp.access_token( + self.identity, client_id=ADMIN_CLIENT_ID if self.role == "proxy_admin" else TESTS_CLIENT_ID + ), + kind="direct_jwt", + role=self.role, + tenant=self.tenants[0].organization_id if self.tenants else None, + ) + + +@dataclass(frozen=True, slots=True) +class ActorFactory: + bootstrap: ManagementClient + idp: Keycloak + resources: ResourceManager + + def __post_init__(self) -> None: + if self.bootstrap.proxy.caller is not None: + raise ValueError("Actor bootstrap requires a separately held master client") + + def key(self, tenant: Tenant | None = None, *, user_id: str | None = None) -> KeyGenerateResponse: + created: Final = unwrap( + self.bootstrap.generate_key( + KeyGenerateBody( + team_id=tenant.team_id if tenant is not None else None, + user_id=user_id, + key_alias=f"e2e-actor-key-{unique_marker()}", + ) + ) + ) + self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key)) + return created + + def tenant(self) -> Tenant: + marker: Final = unique_marker() + organization_id: Final = self.bootstrap.create_org(OrgNewBody(organization_alias=f"e2e-organization-{marker}")) + self.resources.defer( + lambda: unwrap( + self.bootstrap.proxy.transport.delete( + "/organization/delete", + headers=self.bootstrap.proxy.management_headers(), + json=OrgDeleteBody(organization_ids=[organization_id]), + response_type=OrgDeleteResponse, + ) + ) + ) + team_id: Final = self.bootstrap.proxy.create_team( + TeamNewBody(team_alias=f"e2e-team-{marker}", organization_id=organization_id) + ) + self.resources.defer( + lambda: unwrap( + self.bootstrap.proxy.transport.post( + "/team/delete", + headers=self.bootstrap.proxy.management_headers(), + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + ) + ) + self.bootstrap.delete_team_member(team_id, self.bootstrap.user_info().user_id) + group_id: Final = self.idp.create_group(team_id) + self.resources.defer(lambda: self.idp.with_strict_cleanup().delete_group(group_id)) + return Tenant(organization_id=organization_id, team_id=team_id, group_id=group_id) + + def create( + self, role: ActorRole, *, tenants: tuple[Tenant, ...] = (), profile: ActorProfile = "database_role" + ) -> Actor: + if role in ("team_admin", "team_member", "organization_admin") and not tenants: + raise ValueError("A membership actor requires a tenant") + identity: Final = self.idp.with_strict_cleanup().provision_user( + marker=unique_marker(), + groups=tuple(tenant.team_id for tenant in tenants) if profile == "group_scoped" else (), + group_ids=tuple(tenant.group_id for tenant in tenants) if profile == "group_scoped" else (), + defer=self.resources.defer, + ) + global_role: Final[UserRole] = ( + role + if role in ("proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer") + else "internal_user" + ) + self.bootstrap.create_user( + UserNewBody( + user_id=identity.user_id, + user_email=f"{identity.username}@example.com", + user_role=global_role, + auto_create_key=False, + ) + ) + self.resources.defer(lambda: self.bootstrap.delete_user_strict(identity.user_id)) + for tenant in tenants: + unwrap( + self.bootstrap.proxy.transport.post( + "/organization/member_add", + headers=self.bootstrap.proxy.management_headers(), + json=OrgMemberAddBody( + organization_id=tenant.organization_id, + member=OrgMemberEntry( + user_id=identity.user_id, + role="org_admin" if role == "organization_admin" else "internal_user", + ), + ), + response_type=NoBody, + ) + ) + unwrap( + self.bootstrap.proxy.transport.post( + "/team/member_add", + headers=self.bootstrap.proxy.management_headers(), + json=TeamMemberAddBody( + team_id=tenant.team_id, + member=TeamMemberEntry( + user_id=identity.user_id, + role="admin" if role == "team_admin" else "user", + ), + ), + response_type=NoBody, + ) + ) + return Actor(identity=identity, role=role, global_role=global_role, profile=profile, tenants=tenants) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index e17b92a13ed..a0243e868e2 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -7,7 +7,8 @@ llm-only key hitting a management route). from __future__ import annotations import time -from dataclasses import dataclass +import warnings +from dataclasses import dataclass, field, replace import jwt from e2e_config import MASTER_KEY @@ -20,6 +21,7 @@ from e2e_http import ( StreamingResponse, Success, UnknownApiError, + retry_attempts, unwrap, ) from models import ( @@ -81,7 +83,7 @@ from models import ( UserNewResponse, UserUpdateBody, ) -from proxy_client import ProxyClient +from proxy_client import Caller, ProxyClient MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" @@ -98,7 +100,7 @@ class DashboardSession: its bearer on every subsequent call, the claims it renders the signed-in user from, and where it lands the browser.""" - session_key: str + session_key: str = field(repr=False) claims: UiSessionClaims redirect_url: str @@ -106,7 +108,10 @@ class DashboardSession: @dataclass(frozen=True, slots=True) class ManagementClient: proxy: ProxyClient - master_key: str + master_key: str = field(repr=False) + + def with_caller(self, caller: Caller) -> ManagementClient: + return replace(self, proxy=self.proxy.with_caller(caller)) def llm_only_key(self) -> str: return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])) @@ -117,7 +122,7 @@ class ManagementClient: dashboard creates it under the session key their sign-in minted). Returns the outcome rather than unwrapping it, so a caller can poll a route that is only transiently refusing.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) return self.proxy.transport.post( "/key/generate", headers=headers, @@ -131,9 +136,9 @@ class ManagementClient: sign-in minted, never the master key). Returns the outcome rather than unwrapping it, so a caller can poll a route that is only transiently refusing; `update_key_models` is the unwrapping shorthand.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) last: Result[NoBody] = NetworkError(message="/key/update was never attempted") - for attempt in range(_KEY_WRITE_ATTEMPTS): + for attempt in range(retry_attempts(_KEY_WRITE_ATTEMPTS)): last = self.proxy.transport.post( "/key/update", headers=headers, @@ -144,6 +149,7 @@ class ManagementClient: case UnknownApiError(body=error_body) if any( marker in error_body.lower() for marker in _TRANSIENT_BACKEND_MARKERS ): + warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2) time.sleep(0.5 * (attempt + 1)) continue case _: @@ -153,10 +159,10 @@ class ManagementClient: def update_key_models(self, key: str, models: list[str]) -> None: _ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models))) - def key_info_as(self, key: str, *, caller_key: str) -> Result[KeyInfoResponse]: + def key_info_as(self, key: str, *, caller_key: str | None = None) -> Result[KeyInfoResponse]: return self.proxy.transport.get( "/key/info", - headers=self.proxy.transport.bearer(caller_key), + headers=self.proxy.management_headers(caller_key), params=KeyInfoParams(key=key), response_type=KeyInfoResponse, ) @@ -167,7 +173,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/key/delete", - headers=self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key), + headers=self.proxy.management_headers(caller_key), json=KeyDeleteBody(keys=[key]), response_type=NoBody, ) @@ -179,7 +185,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/model/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=ModelDeleteBody(id=model_id), response_type=NoBody, ) @@ -190,7 +196,7 @@ class ManagementClient: Connection button, probing the live provider with the supplied params.""" return self.proxy.transport.post( "/health/test_connection", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=ConnectionTestResponse, timeout=120.0, @@ -200,7 +206,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/key/block", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyBlockBody(key=key), response_type=NoBody, ) @@ -209,7 +215,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/key/regenerate", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyRegenerateBody(key=key, grace_period=grace_period), response_type=KeyGenerateResponse, ) @@ -219,7 +225,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( f"/key/{key}/reset_spend", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyResetSpendBody(reset_to=reset_to), response_type=KeyResetSpendResponse, ) @@ -228,7 +234,7 @@ class ManagementClient: def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]: """GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is who is asking: the master key by default, or a virtual key.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) return self.proxy.transport.get( "/key/list", headers=headers, @@ -266,7 +272,7 @@ class ManagementClient: team_id = unwrap( self.proxy.transport.post( "/team/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=TeamNewResponse, ) @@ -276,10 +282,10 @@ class ManagementClient: def update_team(self, body: TeamUpdateBody) -> None: last: Result[NoBody] | None = None - for attempt in range(5): + for attempt in range(retry_attempts(5)): last = self.proxy.transport.post( "/team/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -289,6 +295,7 @@ class ManagementClient: case UnknownApiError(body=body_text) if ( "connecting to redis" in body_text.lower() or "name resolution" in body_text.lower() ): + warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2) time.sleep(0.5 * (attempt + 1)) continue case _: @@ -299,7 +306,7 @@ class ManagementClient: def delete_team(self, team_id: str) -> None: _ = self.proxy.transport.post( "/team/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) @@ -308,7 +315,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/team/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -320,7 +327,7 @@ class ManagementClient: for entry in unwrap( self.proxy.transport.get( "/team/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=NoBody(), response_type=TeamListResponse, ) @@ -328,14 +335,16 @@ class ManagementClient: ) def team_info_status(self, team_id: str) -> ProbeResult: - return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id)) + return self.proxy.transport.probe( + "/team/info", params=TeamInfoParams(team_id=team_id), headers=self.proxy.management_headers() + ) def _wait_for_team(self, team_id: str) -> None: last: Result[TeamInfoResponse] | None = None - for _ in range(_TEAM_READY_ATTEMPTS): + for _ in range(retry_attempts(_TEAM_READY_ATTEMPTS)): last = self.proxy.transport.get( "/team/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -343,25 +352,29 @@ class ManagementClient: case Success(): return case _: + warnings.warn("Repeating team read while the team becomes available", RuntimeWarning, stacklevel=2) time.sleep(_TEAM_READY_SLEEP_SECONDS) assert last is not None raise AssertionError(last) def add_team_member(self, team_id: str, user_id: str) -> None: last: Result[NoBody] | None = None - for attempt in range(_TEAM_READY_ATTEMPTS): + for attempt in range(retry_attempts(_TEAM_READY_ATTEMPTS)): last = self.proxy.transport.post( "/team/member_add", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)), response_type=NoBody, ) match last: case Success(): return - case UnknownApiError(body=body) if ( - "doesn't exist" in body and attempt + 1 < _TEAM_READY_ATTEMPTS + case UnknownApiError(body=body) if "doesn't exist" in body and attempt + 1 < retry_attempts( + _TEAM_READY_ATTEMPTS ): + warnings.warn( + "Retrying team membership while the team becomes available", RuntimeWarning, stacklevel=2 + ) time.sleep(_TEAM_READY_SLEEP_SECONDS) continue case _: @@ -373,7 +386,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/team/member_delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id), response_type=NoBody, ) @@ -383,7 +396,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/user/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=UserNewResponse, ) @@ -393,7 +406,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/customer/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=CustomerNewBody(user_id=user_id), response_type=CustomerResponse, ) @@ -404,7 +417,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/customer/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=CustomerInfoParams(end_user_id=end_user_id), response_type=CustomerResponse, ) @@ -413,7 +426,7 @@ class ManagementClient: def delete_customer(self, user_id: str) -> None: _ = self.proxy.transport.post( "/customer/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=CustomerDeleteBody(user_ids=[user_id]), response_type=NoBody, ) @@ -422,7 +435,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/user/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -431,7 +444,7 @@ class ManagementClient: def delete_user(self, user_id: str) -> None: _ = self.proxy.transport.post( "/user/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=NoBody, ) @@ -442,17 +455,17 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/user/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=UserDeleteResponse, ) ) - def user_info(self, user_id: str) -> UserInfoResponse: + def user_info(self, user_id: str | None = None) -> UserInfoResponse: return unwrap( self.proxy.transport.get( "/user/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserInfoParams(user_id=user_id), response_type=UserInfoResponse, ) @@ -462,7 +475,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/user/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserListParams(user_ids=user_id), response_type=UserListResponse, ) @@ -472,7 +485,7 @@ class ManagementClient: listing = unwrap( self.proxy.transport.get( "/user/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserListParams(user_ids=user_id), response_type=UserListResponse, ) @@ -483,7 +496,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/organization/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=OrgNewResponse, ) @@ -493,7 +506,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.patch( "/organization/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -502,7 +515,7 @@ class ManagementClient: def delete_org(self, organization_id: str) -> None: _ = self.proxy.transport.delete( "/organization/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=OrgDeleteBody(organization_ids=[organization_id]), response_type=NoBody, ) @@ -511,19 +524,24 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/organization/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=OrgInfoParams(organization_id=organization_id), response_type=OrgInfoResponse, ) ) def org_info_status(self, organization_id: str) -> ProbeResult: - return self.proxy.transport.probe("/organization/info", params=OrgInfoParams(organization_id=organization_id)) + return self.proxy.transport.probe( + "/organization/info", + params=OrgInfoParams(organization_id=organization_id), + headers=self.proxy.management_headers(), + ) + def create_tag(self, body: TagNewBody) -> None: _ = unwrap( self.proxy.transport.post( "/tag/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -532,7 +550,7 @@ class ManagementClient: def delete_tag(self, name: str) -> None: _ = self.proxy.transport.post( "/tag/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TagDeleteBody(name=name), response_type=NoBody, ) @@ -542,7 +560,7 @@ class ManagementClient: unwrap( self.proxy.transport.get( "/tag/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=NoBody(), response_type=TagListResponse, ) @@ -553,7 +571,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/v1/mcp/server", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=McpServerRow, ) @@ -565,7 +583,7 @@ class ManagementClient: return unwrap( self.proxy.transport.put( "/v1/mcp/server", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=McpServerRow, ) @@ -576,7 +594,7 @@ class ManagementClient: unwrap it while a deferred teardown can ignore an already-deleted server.""" return self.proxy.transport.delete( f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=NoBody(), response_type=NoBody, ) diff --git a/tests/e2e/management/test_jwt_management_e2e.py b/tests/e2e/management/test_jwt_management_e2e.py index 22306da8eb8..0d23f954158 100644 --- a/tests/e2e/management/test_jwt_management_e2e.py +++ b/tests/e2e/management/test_jwt_management_e2e.py @@ -2,60 +2,249 @@ from __future__ import annotations -from typing import Final +from typing import Final, Literal import pytest -from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_config import CHEAP_OPENAI_MODEL, PROXY_BASE_URL, unique_marker from e2e_http import UnauthorizedError, UnknownApiError, unwrap -from idp import ADMIN_CLIENT_ID, Identity, Keycloak +from idp import ADMIN_CLIENT_ID, Identity, Keycloak, token_claims from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory, ActorRole from management_client import ManagementClient -from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserNewBody +from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserInfoParams, UserInfoResponse, UserNewBody +from proxy_client import Caller pytestmark = pytest.mark.e2e class TestJwtManagement: + @pytest.mark.parametrize( + "role", + ( + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", + ), + ) + @pytest.mark.covers("mgmt.user.jwt.database_roles") + def test_actor_subject_and_database_role(self, actor_factory: ActorFactory, role: ActorRole) -> None: + tenants: Final = ( + (actor_factory.tenant(),) if role in ("organization_admin", "team_admin", "team_member") else () + ) + actor: Final = actor_factory.create(role, tenants=tenants) + caller: Final = actor.mint_caller(actor_factory.idp) + claims: Final = token_claims(caller.credential) + assert claims.sub == actor.identity.user_id + assert claims.iss == actor_factory.idp.issuer + assert claims.aud == "litellm-e2e" or "litellm-e2e" in claims.aud + assert actor.identity.groups == () + assert ("litellm_proxy_admin" in claims.scope.split()) == (role == "proxy_admin") + stored: Final = actor_factory.bootstrap.user_info(actor.identity.user_id) + assert stored.user_id == actor.identity.user_id + assert stored.user_info.user_role == actor.global_role + bound: Final = actor_factory.bootstrap.with_caller(caller) + own: Final = unwrap( + bound.proxy.transport.get( + "/user/info", + headers=bound.proxy.management_headers(), + params=UserInfoParams(), + response_type=UserInfoResponse, + ) + ) + assert own.user_id == actor.identity.user_id + assert own.user_info.user_role == actor.global_role + for tenant in tenants: + info = actor_factory.bootstrap.team_info(tenant.team_id) + assert info.organization_id == tenant.organization_id + assert {(member.user_id, member.role) for member in info.members_with_roles} == { + (actor.identity.user_id, "admin" if role == "team_admin" else "user") + } + assert { + (member.user_id, member.user_role) + for member in actor_factory.bootstrap.org_info(tenant.organization_id).members + } == {(actor.identity.user_id, "org_admin" if role == "organization_admin" else "internal_user")} + + @pytest.mark.covers("mgmt.key.jwt.viewer_denied") + def test_admin_viewer_reads_but_cannot_update(self, actor_factory: ActorFactory) -> None: + actor: Final = actor_factory.create("proxy_admin_viewer") + viewer: Final = actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp)) + alias: Final = f"e2e-viewer-{unique_marker()}" + key: Final = actor_factory.key().key + unwrap(actor_factory.bootstrap.update_key(KeyUpdateBody(key=key, key_alias=alias))) + assert viewer.proxy.key_info(key).key_alias == alias + denied: Final = viewer.update_key(KeyUpdateBody(key=key, key_alias="forbidden")) + assert isinstance(denied, UnknownApiError) and denied.status_code == 403, f"viewer write was accepted: {denied}" + assert "proxy_admin_viewer" in denied.body and "/key/update" in denied.body + assert actor_factory.bootstrap.proxy.key_info(key).key_alias == alias + + @pytest.mark.covers("mgmt.user.oidc.identity_mapping") + def test_oidc_browser_profile_identity_mapping(self, actor_factory: ActorFactory) -> None: + actor: Final = actor_factory.create("internal_user") + idp: Final = actor_factory.idp.with_strict_cleanup() + discovery: Final = idp.discovery() + assert discovery.issuer == idp.issuer + assert discovery.jwks_uri == idp.jwks_url + callback: Final = f"{PROXY_BASE_URL}/sso/callback" + browser: Final = idp.browser_client(callback_url=callback, defer=actor_factory.resources.defer) + token: Final = idp.browser_token(actor.identity, browser) + assert token_claims(token).sub == actor.identity.user_id + userinfo: Final = idp.userinfo(token) + assert userinfo.sub == actor.identity.user_id + assert userinfo.email == f"{actor.identity.username}@example.com" + assert browser.environment(discovery)["GENERIC_USER_ID_ATTRIBUTE"] == "sub" + @pytest.mark.covers("mgmt.key.jwt.lifecycle") + @pytest.mark.parametrize("credential_kind", ("direct_jwt", "virtual_key")) def test_admin_creates_reads_updates_clears_and_deletes_a_key( - self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager + self, + client: ManagementClient, + idp: Keycloak, + jwt_identity: Identity, + resources: ResourceManager, + actor_factory: ActorFactory, + credential_kind: Literal["direct_jwt", "virtual_key"], ) -> None: - admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + actor: Final = actor_factory.create("proxy_admin") + virtual_key: Final = ( + actor_factory.key(user_id=actor.identity.user_id).key if credential_kind == "virtual_key" else None + ) + admin: Final = ( + virtual_key if virtual_key is not None else idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + ) + bound: Final = client.with_caller(Caller(credential=admin, kind=credential_kind, role="proxy_admin")) alias: Final = f"e2e-jwt-key-{unique_marker()}" created: Final = unwrap( - client.generate_key( + bound.generate_key( KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]), - caller_key=admin, ) ) resources.defer(lambda: client.proxy.delete_key(created.key)) - original: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + original: Final = unwrap(bound.key_info_as(created.key)).info assert original.key_alias == alias and original.team_id == jwt_identity.group assert original.models == [CHEAP_OPENAI_MODEL] updated_alias: Final = f"{alias}-updated" - unwrap( - client.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120), caller_key=admin) - ) - updated: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + unwrap(bound.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120))) + updated: Final = unwrap(bound.key_info_as(created.key)).info assert updated.key_alias == updated_alias and updated.rpm_limit == 120 assert updated.models == [CHEAP_OPENAI_MODEL], "omitted models must preserve the restriction" - unwrap(client.update_key(KeyUpdateBody(key=created.key, models=[]), caller_key=admin)) - cleared: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + unwrap(bound.update_key(KeyUpdateBody(key=created.key, models=[]))) + cleared: Final = unwrap(bound.key_info_as(created.key)).info assert cleared.models == [] and cleared.rpm_limit == 120 - assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 1 - client.delete_key_strict(created.key, caller_key=admin) - assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 0 + assert unwrap(bound.key_list(updated_alias)).total_count == 1 + bound.delete_key_strict(created.key) + assert unwrap(bound.key_list(updated_alias)).total_count == 0 + + @pytest.mark.covers("mgmt.team.jwt.tenant_isolation") + def test_two_actor_sets_keep_tenants_and_keys_isolated(self, actor_factory: ActorFactory) -> None: + first: Final = actor_factory.tenant() + second: Final = actor_factory.tenant() + assert first.organization_id != second.organization_id and first.team_id != second.team_id + actors: Final = tuple( + actor_factory.create("team_member", tenants=(tenant,), profile="group_scoped") for tenant in (first, second) + ) + assert actors[0].identity.user_id != actors[1].identity.user_id + callers: Final = tuple( + actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp)) for actor in actors + ) + keys: Final = tuple(actor_factory.key(tenant) for tenant in (first, second)) + assert keys[0].key != keys[1].key + assert callers[0].proxy.key_info(keys[0].key).team_id == first.team_id + assert callers[1].proxy.key_info(keys[1].key).team_id == second.team_id + for caller, other_key in ((callers[0], keys[1].key), (callers[1], keys[0].key)): + hidden = caller.key_info_as(other_key) + assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403 + assert tuple(actor.identity.groups for actor in actors) == ((first.team_id,), (second.team_id,)) + + @pytest.mark.covers("mgmt.team.jwt.multiple_memberships") + def test_multi_group_actor_keeps_exact_memberships(self, actor_factory: ActorFactory) -> None: + tenants: Final = (actor_factory.tenant(), actor_factory.tenant()) + actor: Final = actor_factory.create("team_member", tenants=tenants, profile="group_scoped") + claims: Final = token_claims(actor.mint_caller(actor_factory.idp).credential) + assert set(claims.groups) == {tenant.team_id for tenant in tenants} + assert "litellm_proxy_admin" not in claims.scope.split() + assert actor.identity.groups == tuple(tenant.team_id for tenant in tenants) + for tenant in tenants: + assert { + (entry.user_id, entry.role) + for entry in actor_factory.bootstrap.team_info(tenant.team_id).members_with_roles + } == {(actor.identity.user_id, "user")} + + @pytest.mark.covers("mgmt.user.jwt.cleanup") + def test_successful_actor_cleanup_removes_owned_state(self, actor_factory: ActorFactory) -> None: + resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True) + factory: Final = ActorFactory(bootstrap=actor_factory.bootstrap, idp=actor_factory.idp, resources=resources) + try: + tenant: Final = factory.tenant() + actor: Final = factory.create("team_member", tenants=(tenant,), profile="group_scoped") + key: Final = factory.key(tenant) + alias: Final = factory.bootstrap.proxy.key_info(key.key).key_alias + assert alias is not None + finally: + resources.teardown() + assert factory.bootstrap.user_count(actor.identity.user_id) == 0 + assert factory.bootstrap.key_alias_count(alias) == 0 + assert factory.bootstrap.team_info_status(tenant.team_id).status_code == 404 + assert factory.bootstrap.org_info_status(tenant.organization_id).status_code == 404 + factory.idp.assert_absent("users", actor.identity.user_id) + factory.idp.assert_absent("groups", tenant.group_id) + + @pytest.mark.parametrize("stage", ("group", "user")) + @pytest.mark.covers("mgmt.user.jwt.partial_cleanup") + def test_partial_setup_removes_previously_created_identities( + self, + actor_factory: ActorFactory, + stage: Literal["group", "user"], + ) -> None: + idp: Final = actor_factory.idp.with_strict_cleanup() + resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True) + marker: Final = unique_marker() + group_id: Final = idp.create_group(f"e2e-partial-{marker}") + resources.defer(lambda: idp.delete_group(group_id)) + try: + identity: Final = ( + idp.provision_user( + marker=marker, + groups=(f"e2e-partial-{marker}",), + group_ids=(group_id,), + defer=resources.defer, + ) + if stage == "user" + else None + ) + if identity is None: + with pytest.raises(pytest.fail.Exception, match="HTTP 409"): + idp.create_group(f"e2e-partial-{marker}") + else: + with pytest.raises(pytest.fail.Exception, match="HTTP 409"): + idp.create_user( + username=identity.username, + email=f"{identity.username}@example.com", + password=identity.password, + groups=identity.groups, + ) + finally: + resources.teardown() + idp.assert_absent("groups", group_id) + if identity is not None: + idp.assert_absent("users", identity.user_id) @pytest.mark.covers("mgmt.key.jwt.member_denied", "mgmt.key.jwt.other_team_denied") def test_member_cannot_write_and_another_team_cannot_read_the_key( self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager ) -> None: admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + bound: Final = client.with_caller(Caller(credential=admin, kind="direct_jwt", role="proxy_admin")) member: Final = idp.access_token(jwt_identity) + member_client: Final = client.with_caller(Caller(credential=member, kind="direct_jwt", role="team_member")) alias: Final = f"e2e-jwt-owned-{unique_marker()}" created: Final = unwrap( client.generate_key(KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group), caller_key=admin) @@ -63,14 +252,14 @@ class TestJwtManagement: resources.defer(lambda: client.proxy.delete_key(created.key)) client.add_team_member(jwt_identity.group, jwt_identity.user_id) - assert unwrap(client.key_info_as(created.key, caller_key=member)).info.key_alias == alias + assert unwrap(member_client.key_info_as(created.key)).info.key_alias == alias - refused: Final = client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"), caller_key=member) + refused: Final = member_client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden")) assert isinstance(refused, UnauthorizedError), f"member write was accepted: {refused}" assert "does not have permissions for endpoint" in refused.body.lower(), ( f"expected a permission denial: {refused}" ) - assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.key_alias == alias + assert unwrap(bound.key_info_as(created.key)).info.key_alias == alias marker: Final = unique_marker() outsider: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer) @@ -88,4 +277,4 @@ class TestJwtManagement: assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403, ( f"another team must not read this key: {hidden}" ) - assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.team_id == jwt_identity.group + assert unwrap(bound.key_info_as(created.key)).info.team_id == jwt_identity.group diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 3cab0334dea..6fca1268ebc 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1091,13 +1091,13 @@ class UiLoginBody(BaseModel): class UiLoginResponse(BaseModel): - token: str + token: str = Field(repr=False) redirect_url: str class UiSessionClaims(BaseModel): user_id: str - key: str + key: str = Field(repr=False) user_role: str login_method: Literal["sso", "username_password"] exp: int @@ -1135,6 +1135,7 @@ class TeamInfoParams(BaseModel): class TeamData(BaseModel): + organization_id: str | None = None team_alias: str | None = None models: list[str] = [] members_with_roles: list[TeamMemberEntry] = [] @@ -1175,6 +1176,7 @@ class UserNewBody(BaseModel): user_email: str user_role: UserRole user_id: str | None = None + auto_create_key: bool | None = None class UserNewResponse(BaseModel): @@ -1187,7 +1189,7 @@ class UserUpdateBody(BaseModel): class UserInfoParams(BaseModel): - user_id: str + user_id: str | None = None class UserData(BaseModel): @@ -1240,16 +1242,36 @@ class OrgInfoParams(BaseModel): organization_id: str +class OrgMembership(BaseModel): + user_id: str + user_role: str + + class OrgInfoResponse(BaseModel): organization_id: str organization_alias: str | None = None models: list[str] = [] + members: tuple[OrgMembership, ...] = () + + +class OrgMemberEntry(BaseModel): + user_id: str + role: Literal["org_admin", "internal_user"] + + +class OrgMemberAddBody(BaseModel): + organization_id: str + member: OrgMemberEntry class OrgDeleteBody(BaseModel): organization_ids: list[str] +class OrgDeleteResponse(RootModel[tuple[OrgInfoResponse, ...]]): + pass + + # ---------- tags (management) ---------- diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 1fe2ec905ef..48a6110dc0b 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -11,14 +11,23 @@ from __future__ import annotations import time import warnings from collections.abc import Callable, Mapping -from dataclasses import dataclass -from functools import reduce +from dataclasses import dataclass, field, replace from datetime import datetime +from functools import reduce from types import MappingProxyType -from typing import Final - -from pydantic import BaseModel +from typing import Final, Literal +from e2e_config import ( + CONTROL_PLANE_BASE_URL, + MASTER_KEY, + POLL_INTERVAL, + POLL_TIMEOUT, + PROXY_BASE_URL, + PROXY_REPLICA_URLS, + REQUEST_TIMEOUT, + SLOW_PROVIDER_TIMEOUT_SECONDS, + settle_propagation, +) from e2e_http import ( AnthropicHeaders, AuthHeaders, @@ -55,6 +64,7 @@ from models import ( KeyInfoParams, KeyInfoResponse, LiteLLMParamsBody, + MemorySummaryResponse, ModelDeleteBody, ModelInfoBody, ModelInfoEntry, @@ -63,7 +73,6 @@ from models import ( ModelNewBody, ModelNewResponse, ModelsListParams, - MemorySummaryResponse, ModelsListResponse, ModelUpdateBody, OcrBody, @@ -76,23 +85,13 @@ from models import ( TeamDeleteBody, TeamNewBody, TeamNewResponse, - UserDeleteBody, - UserDeleteResponse, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, + UserDeleteBody, + UserDeleteResponse, ) -from e2e_config import ( - CONTROL_PLANE_BASE_URL, - MASTER_KEY, - POLL_INTERVAL, - POLL_TIMEOUT, - PROXY_BASE_URL, - PROXY_REPLICA_URLS, - REQUEST_TIMEOUT, - SLOW_PROVIDER_TIMEOUT_SECONDS, - settle_propagation, -) +from pydantic import BaseModel from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -421,11 +420,23 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re ) +CredentialKind = Literal["master", "direct_jwt", "virtual_key", "dashboard_session"] + + +@dataclass(frozen=True, slots=True) +class Caller: + credential: str = field(repr=False) + kind: CredentialKind + role: str + tenant: str | None = None + + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport replicas: Mapping[str, Transport] control_replicas: Mapping[str, Transport] + caller: Caller | None = None poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT @@ -433,13 +444,24 @@ class ProxyClient: model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT + def with_caller(self, caller: Caller) -> ProxyClient: + return replace(self, caller=caller) + + def management_headers(self, caller_key: str | None = None, *, transport: Transport | None = None) -> AuthHeaders: + selected: Final = self.transport if transport is None else transport + if caller_key is not None: + return selected.bearer(caller_key) + if self.caller is not None: + return selected.bearer(self.caller.credential) + return selected.master + # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- def generate_key(self, body: KeyGenerateBody) -> str: return unwrap( self.transport.post( "/key/generate", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=KeyGenerateResponse, ) @@ -448,7 +470,7 @@ class ProxyClient: def delete_key(self, key: str) -> None: _ = self.transport.post( "/key/delete", - headers=self.transport.master, + headers=self.management_headers(), json=KeyDeleteBody(keys=[key]), response_type=NoBody, ) @@ -458,7 +480,7 @@ class ProxyClient: return _ = self.transport.post( "/customer/delete", - headers=self.transport.master, + headers=self.management_headers(), json=CustomerDeleteBody(user_ids=user_ids), response_type=NoBody, ) @@ -467,7 +489,7 @@ class ProxyClient: return unwrap( self.transport.get( "/key/info", - headers=self.transport.master, + headers=self.management_headers(), params=KeyInfoParams(key=key), response_type=KeyInfoResponse, ) @@ -477,7 +499,7 @@ class ProxyClient: return { url: transport.get( "/debug/memory/summary", - headers=transport.master, + headers=self.management_headers(transport=transport), params=NoBody(), response_type=MemorySummaryResponse, ) @@ -524,11 +546,12 @@ class ProxyClient: {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)} ) - @staticmethod def _body_poller[R: BaseModel]( - transport: Transport, path: str, params: BaseModel, response_type: type[R] + self, transport: Transport, path: str, params: BaseModel, response_type: type[R] ) -> Poller[Result[R]]: - return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type) + return lambda: transport.get( + path, headers=self.management_headers(transport=transport), params=params, response_type=response_type + ) def model_info(self) -> list[ModelInfoEntry]: """Every configured deployment with the price the proxy resolved for it @@ -536,7 +559,7 @@ class ProxyClient: return unwrap( self.transport.get( "/model/info", - headers=self.transport.master, + headers=self.management_headers(), params=NoBody(), response_type=ModelInfoResponse, ) @@ -546,7 +569,7 @@ class ProxyClient: return unwrap( self.transport.get( "/public/litellm_model_cost_map", - headers=self.transport.master, + headers=self.management_headers(), params=NoBody(), response_type=CostMap, ) @@ -607,7 +630,7 @@ class ProxyClient: model_id = unwrap( self.transport.post( "/model/new", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ModelNewResponse, ) @@ -623,7 +646,7 @@ class ProxyClient: def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None: """Block until every replica lists `model_name`, or fail at model_servable_timeout.""" - headers: Final = self.transport.master if listed_for is None else self.transport.bearer(listed_for) + headers: Final = self.management_headers(listed_for) outcome: Final = await_servable_everywhere( {url: self._models_poller(transport, headers) for url, transport in self.replicas.items()}, model_name=model_name, @@ -666,7 +689,7 @@ class ProxyClient: unwrap( self.transport.post( "/model/update", - headers=self.transport.master, + headers=self.management_headers(), json=ModelUpdateBody( litellm_params=litellm_params, model_info=ModelInfoBody(id=model_id), @@ -678,7 +701,7 @@ class ProxyClient: def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", - headers=self.transport.master, + headers=self.management_headers(), json=ModelDeleteBody(id=model_id), response_type=NoBody, ) @@ -747,11 +770,10 @@ class ProxyClient: f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}" ) - @staticmethod - def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: + def _reader[R: BaseModel](self, transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: return lambda request_timeout: transport.get( path, - headers=transport.master, + headers=self.management_headers(transport=transport), params=NoBody(), response_type=response_type, timeout=request_timeout, @@ -763,7 +785,7 @@ class ProxyClient: return unwrap( self.transport.post( "/v1/mcp/toolset", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ToolsetRow, ) @@ -775,7 +797,7 @@ class ProxyClient: return unwrap( self.transport.put( "/v1/mcp/toolset", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ToolsetRow, ) @@ -786,7 +808,7 @@ class ProxyClient: can unwrap it while a deferred teardown can ignore an already-deleted row.""" return self.transport.delete( f"/v1/mcp/toolset/{toolset_id}", - headers=self.transport.master, + headers=self.management_headers(), json=NoBody(), response_type=NoBody, ) @@ -795,7 +817,7 @@ class ProxyClient: unwrap( self.transport.post( "/credentials", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=CredentialCreateResponse, ) @@ -804,7 +826,7 @@ class ProxyClient: def delete_credential(self, credential_name: str) -> None: result = self.transport.delete( f"/credentials/{credential_name}", - headers=self.transport.master, + headers=self.management_headers(), json=NoBody(), response_type=NoBody, ) @@ -815,7 +837,7 @@ class ProxyClient: return unwrap( self.transport.post( "/team/new", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=TeamNewResponse, ) @@ -824,7 +846,7 @@ class ProxyClient: def delete_team(self, team_id: str) -> None: result = self.transport.post( "/team/delete", - headers=self.transport.master, + headers=self.management_headers(), json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) @@ -836,7 +858,7 @@ class ProxyClient: a user the proxy only upserts after a successful auth.""" result = self.transport.post( "/user/delete", - headers=self.transport.master, + headers=self.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=UserDeleteResponse, ) @@ -909,7 +931,7 @@ class ProxyClient: def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]: result = self.transport.get( "/spend/logs", - headers=self.transport.master, + headers=self.management_headers(), params=params, response_type=SpendLogs, ) @@ -924,7 +946,7 @@ class ProxyClient: return unwrap( self.transport.get( "/spend/logs/v2", - headers=self.transport.master, + headers=self.management_headers(), params=SpendLogsPageParams( start_date=start.strftime("%Y-%m-%d %H:%M:%S"), end_date=end.strftime("%Y-%m-%d %H:%M:%S"), @@ -977,7 +999,7 @@ class ProxyClient: # ---- route probe ---------------------------------------------------- def probe(self, path: str, *, params: NoBody) -> ProbeResult: - return self.transport.probe(path, params=params) + return self.transport.probe(path, params=params, headers=self.management_headers()) def build_proxy_client( diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 81cd6c8d3d1..7201da84924 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -29,6 +29,7 @@ from e2e_http import ( request_with_retry, streaming_outcome, wire_body, + without_retries, ) from pydantic import BaseModel, TypeAdapter @@ -56,6 +57,15 @@ def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse] class TestTransientRetryPolicy: + def test_qualification_disables_retries_and_restores_the_default(self) -> None: + responses: Final = (FakeResponse(529), FakeResponse(200)) + sleep: Final = SleepRecorder() + with without_retries(): + assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[0] + assert sleep.delays == () + assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[1] + assert sleep.delays == (0.5,) + def test_transient_set_is_only_statuses_the_proxy_cannot_emit(self) -> None: assert TRANSIENT_STATUSES == frozenset({529}) assert 429 not in TRANSIENT_STATUSES diff --git a/tests/e2e/test_idp.py b/tests/e2e/test_idp.py index 33a09a0f13a..8a92dfa527c 100644 --- a/tests/e2e/test_idp.py +++ b/tests/e2e/test_idp.py @@ -4,12 +4,19 @@ these carry no `e2e` marker and run everywhere.""" from __future__ import annotations +import os +import signal +import subprocess +import sys +import time +from builtins import ExceptionGroup from collections.abc import Callable, Generator from contextlib import ExitStack, contextmanager from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path from queue import SimpleQueue from threading import Thread -from typing import Final +from typing import Final, Literal import pytest from e2e_http import ExternalWrite @@ -18,6 +25,8 @@ from idp import ( KEYCLOAK_ADMIN_USER_ENV, KEYCLOAK_REALM_ENV, KEYCLOAK_URL_ENV, + BrowserClientBody, + Discovery, Keycloak, PasswordCredential, UserCreateBody, @@ -60,24 +69,48 @@ def _idp_server( ) -> Generator[tuple[Keycloak, SimpleQueue[str]]]: """Exercise provisioning failures through the same HTTP transport as live tests.""" deletions: SimpleQueue[str] = SimpleQueue() + clients: SimpleQueue[BrowserClientBody] = SimpleQueue() class Handler(BaseHTTPRequestHandler): def log_message(self, format: str, *args: object) -> None: pass def do_POST(self) -> None: - self.rfile.read(int(self.headers.get("Content-Length", "0"))) + body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) if self.path.endswith("/token"): self.send_response(admin_status) self.end_headers() self.wfile.write(b'{"access_token":"synthetic-harness-token"}') else: + if self.path.endswith("/clients"): + clients.put(BrowserClientBody.model_validate_json(body)) self.send_response(user_status if self.path.endswith("/users") else 201) self.send_header("Location", f"{self.path}/resource-1") self.end_headers() if user_status != 201 and self.path.endswith("/users"): self.wfile.write(b"injected create failure") + def do_GET(self) -> None: + self.send_response(200) + self.end_headers() + if "/clients/" in self.path: + client: Final = clients.get_nowait() + clients.put(client) + self.wfile.write(client.model_dump_json(by_alias=True).encode()) + else: + issuer: Final = f"http://127.0.0.1:{server.server_port}/realms/test" + self.wfile.write( + Discovery( + issuer=issuer, + authorization_endpoint=f"{issuer}/auth", + token_endpoint=f"{issuer}/token", + userinfo_endpoint=f"{issuer}/userinfo", + jwks_uri=f"{issuer}/certs", + ) + .model_dump_json() + .encode() + ) + def do_DELETE(self) -> None: deletions.put(self.path) self.send_response(delete_status) @@ -115,6 +148,54 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No assert deletions.empty() +@pytest.mark.parametrize("exit_mode", ("normal", "parent", "group")) +def test_oidc_launcher_removes_client_on_exit_and_termination( + tmp_path: Path, exit_mode: Literal["normal", "parent", "group"] +) -> None: + ready: Final = tmp_path / "ready" + child_command: Final = ( + "import os,time; from pathlib import Path; " + 'assert os.environ["GENERIC_CLIENT_SECRET"]; ' + 'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; ' + f"Path({str(ready)!r}).touch(); " + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") + ) + with _idp_server() as (idp, deletions): + with subprocess.Popen( + [ + sys.executable, + str(Path(__file__).with_name("idp.py")), + "http://127.0.0.1:9999", + sys.executable, + "-c", + child_command, + ], + env={ + **os.environ, + KEYCLOAK_URL_ENV: idp.base_url, + KEYCLOAK_REALM_ENV: idp.realm, + KEYCLOAK_ADMIN_USER_ENV: idp.admin_username, + KEYCLOAK_ADMIN_PASSWORD_ENV: idp.admin_password, + }, + start_new_session=True, + ) as process: + try: + deadline: Final = time.monotonic() + 15 + while not ready.exists() and time.monotonic() < deadline and process.poll() is None: + time.sleep(0.05) + assert ready.exists(), "OIDC child did not start" + if exit_mode == "parent": + process.terminate() + elif exit_mode == "group": + os.killpg(process.pid, signal.SIGTERM) + assert process.wait(timeout=10) == (7 if exit_mode == "normal" else 143) + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=5) + assert deletions.get(timeout=5) == "/admin/realms/test/clients/resource-1" + assert deletions.empty() + + def test_successful_provisioning_cleans_up_user_before_group() -> None: with _idp_server() as (idp, deletions): with ExitStack() as cleanup: @@ -134,6 +215,43 @@ def test_cleanup_failure_is_visible() -> None: idp.delete_group("group") +def test_strict_cleanup_reports_each_failure_and_continues() -> None: + from lifecycle import ResourceManager + from proxy_client import build_proxy_client + + with _idp_server(delete_status=500) as (idp, deletions): + resources: Final = ResourceManager(client=build_proxy_client(), strict_cleanup=True) + strict: Final = idp.with_strict_cleanup() + resources.defer(lambda: strict.delete_group("group")) + resources.defer(lambda: strict.delete_user("user")) + with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as error: + resources.teardown() + assert len(error.value.exceptions) == 2 + assert deletions.get_nowait() == "/admin/realms/test/users/user" + assert deletions.get_nowait() == "/admin/realms/test/groups/group" + + +@pytest.mark.parametrize("groups", ((), ("one",), ("one", "two"))) +def test_provisioning_records_zero_one_or_multiple_groups(groups: tuple[str, ...]) -> None: + with _idp_server() as (idp, deletions): + with ExitStack() as cleanup: + + def defer(callback: Callable[[], object]) -> None: + cleanup.callback(callback) + + identity: Final = idp.provision_groups( + marker="memberships", + groups=groups, + defer=defer, + ) + assert identity.groups == groups + assert len(identity.group_ids) == len(groups) + assert deletions.get_nowait() == "/admin/realms/test/users/resource-1" + for _ in groups: + assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1" + assert deletions.empty() + + def test_expired_admin_credentials_do_not_abort_remaining_cleanups() -> None: with _idp_server(admin_status=401) as (idp, _): cleanup: Final = ExitStack() diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 1b0133f12cb..879bd88980c 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -11,19 +11,53 @@ injected clock, so nothing here monkeypatches anything. from __future__ import annotations -from collections.abc import Iterable, Mapping +import json +from builtins import ExceptionGroup +from collections.abc import Callable, Generator, Iterable, Mapping +from contextlib import contextmanager from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import chain, repeat +from queue import SimpleQueue +from threading import Thread from types import MappingProxyType from typing import Final, cast import pytest from e2e_config import parse_replica_urls -from e2e_http import Result, Success -from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from e2e_http import NoBody, Result, Success, without_retries +from idp import Keycloak +from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory +from management.management_client import ManagementClient +from models import ( + ConnectionTestBody, + CredentialCreateBody, + KeyGenerateBody, + KeyInfo, + KeyInfoResponse, + KeyUpdateBody, + LiteLLMParamsBody, + McpServerCreateBody, + McpServerUpdateBody, + ModelListEntry, + ModelsListResponse, + OrgNewBody, + OrgUpdateBody, + SpendLogsParams, + TagNewBody, + TeamNewBody, + TeamUpdateBody, + ToolsetCreateBody, + ToolsetUpdateBody, + UserNewBody, + UserUpdateBody, +) from proxy_client import ( - ConvergeOutcome, + Caller, Converged, + ConvergeOutcome, + CredentialKind, EverywhereConverged, ModelsPoller, NeverConvergedOn, @@ -42,6 +76,108 @@ from proxy_client import ( ) from transport import Transport + +@contextmanager +def caller_boundary( + status: int = 200, bodies: SimpleQueue[bytes] | None = None, *, delete_status: int | None = None +) -> Generator[tuple[ManagementClient, SimpleQueue[str]]]: + received: Final[SimpleQueue[str]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: object) -> None: + pass + + def do_GET(self) -> None: + received.put(self.headers.get("Authorization", "")) + self.send_response(delete_status if self.path == "/key/delete" and delete_status is not None else status) + self.end_headers() + self.wfile.write( + b'{"key":"owned","info":{"key_alias":"owned"},"data":[{"id":"owned"}],"team_id":"owned","team_info":{},"model_id":"owned"}' + ) + + def do_POST(self) -> None: + body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + if bodies is not None: + bodies.put(body) + self.do_GET() + + do_PATCH = do_POST + do_PUT = do_POST + do_DELETE = do_POST + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + url: Final = f"http://127.0.0.1:{server.server_port}" + proxy: Final = build_proxy_client( + base_url=url, control_plane_base_url=url, replica_urls=(url,), master_key="bootstrap" + ) + try: + yield ManagementClient(proxy=proxy, master_key="bootstrap"), received + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +class TestBoundManagementCaller: + def test_actor_key_cleanup_reports_failure_and_continues(self) -> None: + with caller_boundary(delete_status=500) as (bootstrap, received), without_retries(): + resources: Final = ResourceManager(client=bootstrap.proxy, strict_cleanup=True) + remaining: SimpleQueue[str] = SimpleQueue() + resources.defer(lambda: remaining.put("cleaned")) + factory: Final = ActorFactory( + bootstrap=bootstrap, + idp=Keycloak(base_url="http://unused.test", realm="test", admin_username="test", admin_password="test"), + resources=resources, + ) + assert factory.key().key == "owned" + with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as failure: + resources.teardown() + assert len(failure.value.exceptions) == 1 + assert remaining.get_nowait() == "cleaned" + assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap") + + @pytest.mark.parametrize("kind", ("direct_jwt", "virtual_key", "dashboard_session")) + def test_direct_delegated_and_replica_reads_keep_the_bound_caller(self, kind: CredentialKind) -> None: + with caller_boundary() as (bootstrap, received): + caller: Final = Caller(credential="synthetic-caller", kind=kind, role="internal_user", tenant="tenant-a") + bound: Final = bootstrap.with_caller(caller) + bound.update_key(KeyUpdateBody(key="owned", key_alias="updated")) + bound.proxy.key_info("owned") + bound.proxy.read_back_everywhere( + "/key/info", + params=KeyUpdateBody(key="owned"), + response_type=KeyInfoResponse, + converged=lambda result: isinstance(result, Success), + ) + bound.proxy.read_body_back_everywhere( + "/key/info", KeyInfoResponse, settled=lambda result: result.info.key_alias == "owned" + ) + assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer synthetic-caller",) * 4 + assert received.empty() + bootstrap.proxy.key_info("owned") + assert received.get_nowait() == "Bearer bootstrap" + + def test_explicit_override_wins_without_rebinding_or_changing_master(self) -> None: + with caller_boundary() as (bootstrap, received): + bound: Final = bootstrap.with_caller(Caller(credential="bound", kind="direct_jwt", role="internal_user")) + bound.update_key(KeyUpdateBody(key="owned"), caller_key="override") + bound.proxy.key_info("owned") + assert received.get_nowait() == "Bearer override" + assert received.get_nowait() == "Bearer bound" + assert bound.master_key == "bootstrap" + + def test_credentials_are_absent_from_binding_and_header_diagnostics(self) -> None: + with caller_boundary() as (bootstrap, _): + caller: Final = Caller(credential="private-value", kind="direct_jwt", role="internal_user") + bound: Final = bootstrap.with_caller(caller) + assert "private-value" not in repr(caller) + assert "private-value" not in repr(bound) + assert "private-value" not in repr(bound.proxy.management_headers()) + assert "bootstrap" not in repr(bound) + + MODEL: Final = "gpt-under-test" _NO_TRANSPORTS: Final = cast(Transport, None) TIMEOUT: Final = 10.0 @@ -275,3 +411,166 @@ class TestReplicasFor: client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) with pytest.raises(AssertionError, match="no replica is configured"): _ = client.replicas_for("/v1/models") + + +MANAGEMENT_OPERATIONS: Final[tuple[tuple[str, Callable[[ManagementClient], object]], ...]] = ( + ("generate_key", lambda c: c.generate_key(KeyGenerateBody())), + ("llm_only_key", lambda c: c.llm_only_key()), + ("update_key", lambda c: c.update_key(KeyUpdateBody(key="owned"))), + ("update_key_models", lambda c: c.update_key_models("owned", [])), + ("key_info", lambda c: c.key_info_as("owned")), + ("delete_key_strict", lambda c: c.delete_key_strict("owned")), + ("delete_model_strict", lambda c: c.delete_model_strict("owned")), + ( + "connection_test", + lambda c: c.connection_test( + ConnectionTestBody(litellm_params=LiteLLMParamsBody(model="synthetic"), mode="chat") + ), + ), + ("block_key", lambda c: c.block_key("owned")), + ("regenerate_key", lambda c: c.regenerate_key("owned")), + ("reset_key_spend", lambda c: c.reset_key_spend("owned", 0)), + ("key_list", lambda c: c.key_list("owned")), + ("key_alias_count", lambda c: c.key_alias_count("owned")), + ("create_team", lambda c: c.create_team(TeamNewBody(team_alias="owned"))), + ("update_team", lambda c: c.update_team(TeamUpdateBody(team_id="owned", team_alias="updated"))), + ("delete_team", lambda c: c.delete_team("owned")), + ("team_info", lambda c: c.team_info("owned")), + ("team_list_ids", lambda c: c.team_list_ids()), + ("team_info_status", lambda c: c.team_info_status("owned")), + ("add_team_member", lambda c: c.add_team_member("owned", "user")), + ("delete_team_member", lambda c: c.delete_team_member("owned", "user")), + ("create_user", lambda c: c.create_user(UserNewBody(user_email="actor@example.com", user_role="internal_user"))), + ("create_customer", lambda c: c.create_customer("owned")), + ("customer_info", lambda c: c.customer_info("owned")), + ("delete_customer", lambda c: c.delete_customer("owned")), + ("update_user", lambda c: c.update_user(UserUpdateBody(user_id="owned", user_role="internal_user"))), + ("delete_user", lambda c: c.delete_user("owned")), + ("delete_user_strict", lambda c: c.delete_user_strict("owned")), + ("user_info", lambda c: c.user_info("owned")), + ("user_count", lambda c: c.user_count("owned")), + ("user_list_ids", lambda c: c.user_list_ids("owned")), + ("create_org", lambda c: c.create_org(OrgNewBody(organization_alias="owned"))), + ("update_org", lambda c: c.update_org(OrgUpdateBody(organization_id="owned", organization_alias="updated"))), + ("delete_org", lambda c: c.delete_org("owned")), + ("org_info", lambda c: c.org_info("owned")), + ("org_info_status", lambda c: c.org_info_status("owned")), + ("create_tag", lambda c: c.create_tag(TagNewBody(name="owned"))), + ("delete_tag", lambda c: c.delete_tag("owned")), + ("tag_list", lambda c: c.tag_list()), + ("create_mcp_server", lambda c: c.create_mcp_server(McpServerCreateBody(alias="owned", url="http://example.test"))), + ("update_mcp_server", lambda c: c.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None))), + ("delete_mcp_server", lambda c: c.delete_mcp_server("owned")), + ("proxy.generate_key", lambda c: c.proxy.generate_key(KeyGenerateBody())), + ("proxy.delete_key", lambda c: c.proxy.delete_key("owned")), + ("proxy.delete_customers", lambda c: c.proxy.delete_customers(["owned"])), + ("proxy.key_info", lambda c: c.proxy.key_info("owned")), + ("proxy.memory_summary", lambda c: c.proxy.memory_summary_everywhere()), + ("proxy.model_info", lambda c: c.proxy.model_info()), + ("proxy.model_cost_map", lambda c: c.proxy.model_cost_map()), + ("proxy.create_model", lambda c: c.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic"))), + ("proxy.update_model", lambda c: c.proxy.update_model("owned", LiteLLMParamsBody(model="synthetic"))), + ("proxy.delete_model", lambda c: c.proxy.delete_model("owned")), + ("proxy.create_toolset", lambda c: c.proxy.create_toolset(ToolsetCreateBody(toolset_name="owned", tools=[]))), + ("proxy.update_toolset", lambda c: c.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None))), + ("proxy.delete_toolset", lambda c: c.proxy.delete_toolset("owned")), + ( + "proxy.create_credential", + lambda c: c.proxy.create_credential(CredentialCreateBody(credential_name="owned", credential_values={})), + ), + ("proxy.delete_credential", lambda c: c.proxy.delete_credential("owned")), + ("proxy.create_team", lambda c: c.proxy.create_team(TeamNewBody(team_alias="owned"))), + ("proxy.delete_team", lambda c: c.proxy.delete_team("owned")), + ("proxy.delete_user", lambda c: c.proxy.delete_user("owned")), + ("proxy.spend_logs", lambda c: c.proxy.spend_logs(SpendLogsParams(api_key="owned"))), + ("proxy.probe", lambda c: c.proxy.probe("/user/info", params=NoBody())), +) + + +@pytest.mark.parametrize( + ("name", "operation"), MANAGEMENT_OPERATIONS, ids=tuple(name for name, _ in MANAGEMENT_OPERATIONS) +) +@pytest.mark.parametrize("kind", ("master", "direct_jwt", "virtual_key", "dashboard_session")) +def test_management_operations_send_the_selected_credential( + name: str, + operation: Callable[[ManagementClient], object], + kind: CredentialKind, +) -> None: + with caller_boundary(status=401) as (bootstrap, received), without_retries(): + client: Final = ( + bootstrap + if kind == "master" + else bootstrap.with_caller(Caller(credential=f"synthetic-{kind}", kind=kind, role="internal_user")) + ) + try: + operation(client) + except AssertionError: + pass + expected: Final = "Bearer bootstrap" if kind == "master" else f"Bearer synthetic-{kind}" + assert received.get_nowait() == expected, name + assert received.empty(), "an unauthorized request must not be retried" + + +class TestSplitCallerPropagation: + def test_control_and_data_replica_readers_keep_the_caller(self) -> None: + with caller_boundary() as (data, data_headers), caller_boundary() as (control, control_headers): + data_url: Final = next(iter(data.proxy.replicas)) + control_url: Final = next(iter(control.proxy.replicas)) + proxy: Final = build_proxy_client( + base_url=data_url, + control_plane_base_url=control_url, + replica_urls=(data_url,), + master_key="bootstrap", + ).with_caller(Caller(credential="tenant-token", kind="direct_jwt", role="team_member")) + proxy.key_info("owned") + proxy.read_body_back_everywhere( + "/key/info", KeyInfoResponse, settled=lambda info: info.info.key_alias == "owned" + ) + proxy.read_back_everywhere( + "/key/info", + params=NoBody(), + response_type=KeyInfoResponse, + converged=lambda result: isinstance(result, Success), + ) + assert control_headers.get_nowait() == "Bearer tenant-token" + assert control_headers.get_nowait() == "Bearer tenant-token" + assert data_headers.get_nowait() == "Bearer tenant-token" + assert control_headers.empty() and data_headers.empty() + + def test_successful_team_and_model_polling_uses_the_bound_caller(self) -> None: + with caller_boundary() as (bootstrap, received): + bound: Final = bootstrap.with_caller(Caller(credential="caller", kind="direct_jwt", role="proxy_admin")) + bound.create_team(TeamNewBody(team_alias="owned")) + bound.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic")) + assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer caller",) * 4 + assert received.empty() + + def test_expired_shaped_token_is_sent_once_without_renewal(self) -> None: + with caller_boundary(status=401) as (bootstrap, received): + bound: Final = bootstrap.with_caller( + Caller(credential="expired.payload.signature", kind="direct_jwt", role="internal_user") + ) + result: Final = bound.key_info_as("owned") + assert not isinstance(result, Success) + assert received.get_nowait() == "Bearer expired.payload.signature" + assert received.empty() + + +@pytest.mark.parametrize("operation", ("server", "toolset")) +def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation: str) -> None: + bodies: Final[SimpleQueue[bytes]] = SimpleQueue() + with caller_boundary(status=401, bodies=bodies) as (bootstrap, _): + try: + if operation == "server": + bootstrap.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None)) + else: + bootstrap.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None)) + except AssertionError: + pass + expected: Final = ( + {"server_id": "owned", "alias": None} + if operation == "server" + else {"toolset_id": "owned", "description": None} + ) + assert json.loads(bodies.get_nowait()) == expected + assert bodies.empty() diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index e8caa801467..037db0c340f 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -7,11 +7,9 @@ client touches requests.* or builds raw dicts; they pass pydantic models here. from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Protocol -from pydantic import BaseModel - import e2e_http from e2e_http import ( URL, @@ -21,6 +19,7 @@ from e2e_http import ( Result, StreamingResponse, ) +from pydantic import BaseModel class Transport(Protocol): @@ -85,7 +84,7 @@ class Transport(Protocol): self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] ) -> Result[R]: ... - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: ... + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: ... def upload[R: BaseModel]( self, @@ -113,7 +112,7 @@ class Transport(Protocol): @dataclass(frozen=True, slots=True) class HttpTransport: base_url: str - master_key: str + master_key: str = field(repr=False) request_timeout: float = 60.0 def _url(self, path: str) -> URL: @@ -245,10 +244,10 @@ class HttpTransport: timeout=self.request_timeout, ) - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return e2e_http.probe( self._url(path), - headers=self.master, + headers=self.master if headers is None else headers, params=params, timeout=self.request_timeout, ) @@ -434,8 +433,8 @@ class SplitTransport: path, headers=headers, json=json, params=params, stream=stream ) - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: - return self._route(path).probe(path, params=params) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: + return self._route(path).probe(path, params=params, headers=headers) def upload[R: BaseModel]( self, diff --git a/tests/e2e/ui/oidcSetup.ts b/tests/e2e/ui/oidcSetup.ts new file mode 100644 index 00000000000..943fa23cab3 --- /dev/null +++ b/tests/e2e/ui/oidcSetup.ts @@ -0,0 +1,30 @@ +import { chromium, expect } from "@playwright/test"; +import * as fs from "fs"; +import * as path from "path"; + +export default async function oidcSetup() { + const baseURL = process.env.E2E_OIDC_UI_URL; + const issuer = process.env.JWT_ISSUER; + const username = process.env.E2E_OIDC_USERNAME; + const password = process.env.E2E_OIDC_PASSWORD; + if (!baseURL || !issuer || !username || !password) { + throw new Error("The OIDC setup requires a running stack, issuer, and provisioned actor credentials"); + } + const artifactDir = process.env.E2E_UI_ARTIFACT_DIR || "."; + fs.mkdirSync(artifactDir, { recursive: true }); + const browser = await chromium.launch(); + try { + const page = await browser.newPage(); + await page.goto(`${baseURL.replace(/\/$/, "")}/sso/key/generate`); + await expect(page).toHaveURL(new RegExp(`^${issuer.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}/`)); + await page.getByLabel("Username or email").fill(username); + await page.getByLabel("Password", { exact: true }).fill(password); + await page.getByRole("button", { name: "Sign In", exact: true }).click(); + await page.waitForURL((url) => url.origin === new URL(baseURL).origin && url.pathname.startsWith("/ui")); + const statePath = path.join(artifactDir, "oidc.storageState.json"); + await page.context().storageState({ path: statePath }); + fs.chmodSync(statePath, 0o600); + } finally { + await browser.close(); + } +} diff --git a/tests/e2e/ui/playwright.oidc.config.ts b/tests/e2e/ui/playwright.oidc.config.ts new file mode 100644 index 00000000000..0fbe77e9bd2 --- /dev/null +++ b/tests/e2e/ui/playwright.oidc.config.ts @@ -0,0 +1,22 @@ +import { defineConfig, devices } from "@playwright/test"; +import * as path from "path"; + +const baseURL = process.env.E2E_OIDC_UI_URL; +if (!baseURL) throw new Error("E2E_OIDC_UI_URL must point to the running OIDC stack"); + +export default defineConfig({ + testDir: ".", + testMatch: "oidc/**/*.spec.ts", + retries: 0, + workers: 1, + outputDir: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc", "test-results"), + globalSetup: require.resolve("./oidcSetup"), + use: { + ...devices["Desktop Chrome"], + baseURL, + storageState: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc.storageState.json"), + trace: "off", + screenshot: "off", + video: "off", + }, +}); From c8bb54993e8b6db4eda84819b0015a1d9a85ca99 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:49:49 -0700 Subject: [PATCH 07/43] test: enforce isolated actors and stop OIDC process groups --- .../e2e/coverage_registry/management_cases.py | 2 +- tests/e2e/idp.py | 41 +++++++++++++++---- tests/e2e/management/jwt_actors.py | 2 +- tests/e2e/management/management_client.py | 17 ++++---- .../e2e/management/test_jwt_management_e2e.py | 20 ++++----- tests/e2e/test_idp.py | 25 ++++++++--- tests/e2e/test_proxy_client.py | 7 ++++ 7 files changed, 80 insertions(+), 34 deletions(-) diff --git a/tests/e2e/coverage_registry/management_cases.py b/tests/e2e/coverage_registry/management_cases.py index 15dc7d333c5..812dbfe5b8d 100644 --- a/tests/e2e/coverage_registry/management_cases.py +++ b/tests/e2e/coverage_registry/management_cases.py @@ -63,7 +63,7 @@ MANAGEMENT_CASES: Final = tuple( node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[virtual_key]", credential_kind="virtual_key", actor="proxy_admin", - profile="database_role", + profile="group_scoped", method="POST", path="/key/generate", operation_family="key_lifecycle", diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py index 12db91bbd88..2dc7c2ad71b 100644 --- a/tests/e2e/idp.py +++ b/tests/e2e/idp.py @@ -8,6 +8,7 @@ import secrets import signal import subprocess import sys +import time import warnings from collections.abc import Callable from contextlib import ExitStack @@ -426,6 +427,36 @@ def token_claims(token: str) -> TokenClaims: return TokenClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))) +def _signal_process_group(process_id: int, signum: int) -> bool: + try: + os.killpg(process_id, signum) + except ProcessLookupError: + return False + return True + + +def _stop_process_group(child: subprocess.Popen[bytes]) -> None: + _signal_process_group(child.pid, signal.SIGTERM) + deadline: Final = time.monotonic() + 5 + while _process_group_exists(child.pid): + child.poll() + if time.monotonic() >= deadline: + _signal_process_group(child.pid, signal.SIGKILL) + break + time.sleep(0.05) + child.wait() + + +def _process_group_exists(process_id: int) -> bool: + try: + os.killpg(process_id, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + def run_oidc_profile(proxy_url: str, command: list[str]) -> int: idp: Final = keycloak_from_env().with_strict_cleanup() with ExitStack() as cleanup: @@ -441,17 +472,11 @@ def run_oidc_profile(proxy_url: str, command: list[str]) -> int: client: Final = idp.browser_client(callback_url=f"{proxy_url.rstrip('/')}/sso/callback", defer=defer) environment: Final = {**os.environ, **client.environment(idp.discovery()), "PROXY_BASE_URL": proxy_url} - with subprocess.Popen(command, env=environment) as child: + with subprocess.Popen(command, env=environment, start_new_session=True) as child: try: return child.wait() finally: - if child.poll() is None: - child.terminate() - try: - child.wait(timeout=5) - except subprocess.TimeoutExpired: - child.kill() - child.wait() + _stop_process_group(child) if __name__ == "__main__": diff --git a/tests/e2e/management/jwt_actors.py b/tests/e2e/management/jwt_actors.py index 909d1652ada..2d23549fe71 100644 --- a/tests/e2e/management/jwt_actors.py +++ b/tests/e2e/management/jwt_actors.py @@ -84,7 +84,7 @@ class ActorFactory: ) ) ) - self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key)) + self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key, missing_ok=True)) return created def tenant(self) -> Tenant: diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index a0243e868e2..8470d318db8 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -167,17 +167,18 @@ class ManagementClient: response_type=KeyInfoResponse, ) - def delete_key_strict(self, key: str, *, caller_key: str | None = None) -> None: + def delete_key_strict(self, key: str, *, caller_key: str | None = None, missing_ok: bool = False) -> None: """Strict delete for the act phase of a test: a failed delete is a hard failure, unlike the warn-only ProxyClient.delete_key used at teardown.""" - _ = unwrap( - self.proxy.transport.post( - "/key/delete", - headers=self.proxy.management_headers(caller_key), - json=KeyDeleteBody(keys=[key]), - response_type=NoBody, - ) + result = self.proxy.transport.post( + "/key/delete", + headers=self.proxy.management_headers(caller_key), + json=KeyDeleteBody(keys=[key]), + response_type=NoBody, ) + if missing_ok and isinstance(result, UnknownApiError) and result.status_code == 404: + return + _ = unwrap(result) def delete_model_strict(self, model_id: str) -> None: """Strict delete for the act phase of a test: a failed delete is a hard diff --git a/tests/e2e/management/test_jwt_management_e2e.py b/tests/e2e/management/test_jwt_management_e2e.py index 0d23f954158..5898073a4e6 100644 --- a/tests/e2e/management/test_jwt_management_e2e.py +++ b/tests/e2e/management/test_jwt_management_e2e.py @@ -102,31 +102,29 @@ class TestJwtManagement: @pytest.mark.parametrize("credential_kind", ("direct_jwt", "virtual_key")) def test_admin_creates_reads_updates_clears_and_deletes_a_key( self, - client: ManagementClient, - idp: Keycloak, - jwt_identity: Identity, - resources: ResourceManager, actor_factory: ActorFactory, credential_kind: Literal["direct_jwt", "virtual_key"], ) -> None: - actor: Final = actor_factory.create("proxy_admin") + tenant: Final = actor_factory.tenant() + actor: Final = actor_factory.create("proxy_admin", tenants=(tenant,), profile="group_scoped") virtual_key: Final = ( actor_factory.key(user_id=actor.identity.user_id).key if credential_kind == "virtual_key" else None ) - admin: Final = ( - virtual_key if virtual_key is not None else idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + admin: Final = virtual_key if virtual_key is not None else actor.mint_caller(actor_factory.idp).credential + bound: Final = actor_factory.bootstrap.with_caller( + Caller(credential=admin, kind=credential_kind, role="proxy_admin") ) - bound: Final = client.with_caller(Caller(credential=admin, kind=credential_kind, role="proxy_admin")) + assert bound.user_info().user_id == actor.identity.user_id alias: Final = f"e2e-jwt-key-{unique_marker()}" created: Final = unwrap( bound.generate_key( - KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]), + KeyGenerateBody(key_alias=alias, team_id=tenant.team_id, models=[CHEAP_OPENAI_MODEL]), ) ) - resources.defer(lambda: client.proxy.delete_key(created.key)) + actor_factory.resources.defer(lambda: actor_factory.bootstrap.delete_key_strict(created.key, missing_ok=True)) original: Final = unwrap(bound.key_info_as(created.key)).info - assert original.key_alias == alias and original.team_id == jwt_identity.group + assert original.key_alias == alias and original.team_id == tenant.team_id assert original.models == [CHEAP_OPENAI_MODEL] updated_alias: Final = f"{alias}-updated" diff --git a/tests/e2e/test_idp.py b/tests/e2e/test_idp.py index 8a92dfa527c..cf2d4f3118a 100644 --- a/tests/e2e/test_idp.py +++ b/tests/e2e/test_idp.py @@ -6,6 +6,7 @@ from __future__ import annotations import os import signal +import socket import subprocess import sys import time @@ -148,16 +149,27 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No assert deletions.empty() -@pytest.mark.parametrize("exit_mode", ("normal", "parent", "group")) +@pytest.mark.parametrize( + ("exit_mode", "ignore_termination"), (("normal", False), ("parent", False), ("group", False), ("parent", True)) +) def test_oidc_launcher_removes_client_on_exit_and_termination( - tmp_path: Path, exit_mode: Literal["normal", "parent", "group"] + tmp_path: Path, exit_mode: Literal["normal", "parent", "group"], ignore_termination: bool ) -> None: ready: Final = tmp_path / "ready" + descendant_command: Final = ( + "import signal,socket,time; from pathlib import Path; " + + ("signal.signal(signal.SIGTERM, signal.SIG_IGN); " if ignore_termination else "") + + "listener=socket.socket(); listener.bind(('127.0.0.1',0)); listener.listen(); " + f"Path({str(ready)!r}).write_text(str(listener.getsockname()[1])); time.sleep(120)" + ) child_command: Final = ( - "import os,time; from pathlib import Path; " + "import os,subprocess,sys,time; from pathlib import Path; " 'assert os.environ["GENERIC_CLIENT_SECRET"]; ' 'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; ' - f"Path({str(ready)!r}).touch(); " + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") + f"subprocess.Popen([sys.executable, '-c', {descendant_command!r}]); " + f"ready=Path({str(ready)!r})\n" + "while not ready.exists(): time.sleep(0.05)\n" + + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") ) with _idp_server() as (idp, deletions): with subprocess.Popen( @@ -187,7 +199,10 @@ def test_oidc_launcher_removes_client_on_exit_and_termination( process.terminate() elif exit_mode == "group": os.killpg(process.pid, signal.SIGTERM) - assert process.wait(timeout=10) == (7 if exit_mode == "normal" else 143) + assert process.wait(timeout=15) == (7 if exit_mode == "normal" else 143) + with socket.socket() as connection: + connection.settimeout(1) + assert connection.connect_ex(("127.0.0.1", int(ready.read_text()))) != 0 finally: if process.poll() is None: os.killpg(process.pid, signal.SIGKILL) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 879bd88980c..0c4aed5bd65 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -121,6 +121,13 @@ def caller_boundary( class TestBoundManagementCaller: + def test_strict_key_cleanup_accepts_missing_only_when_requested(self) -> None: + with caller_boundary(delete_status=404) as (bootstrap, received), without_retries(): + with pytest.raises(AssertionError): + bootstrap.delete_key_strict("owned") + bootstrap.delete_key_strict("owned", missing_ok=True) + assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap") + def test_actor_key_cleanup_reports_failure_and_continues(self) -> None: with caller_boundary(delete_status=500) as (bootstrap, received), without_retries(): resources: Final = ResourceManager(client=bootstrap.proxy, strict_cleanup=True) From c1d0d29d01dbd4c28e5a80a52b5af99a73383dd0 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sun, 13 Sep 2026 04:50:29 +0000 Subject: [PATCH 08/43] fix(guardrails): import ModelResponse lazily to avoid cyclic import alert Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 34b79e295e1..1d00ad8c29a 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -31,7 +31,6 @@ from litellm.types.utils import ( GuardrailStatus, GuardrailTracingDetail, LLMResponseTypes, - ModelResponse, StandardLoggingGuardrailInformation, ) @@ -907,6 +906,8 @@ class CustomGuardrail(CustomLogger): response: Final = ( kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result ) + from litellm.types.utils import ModelResponse + output_translation: Final = ( get_guardrail_translation_mapping(CallTypes.acompletion)() if isinstance(response, ModelResponse) From 49b2d71057d4cff4e3a4baad843a7db7ac35c7c2 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sun, 13 Sep 2026 08:47:47 +0000 Subject: [PATCH 09/43] test(guardrails): type the native lifecycle logging_only test double Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_litellm/integrations/test_custom_guardrail.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 8095be3c0de..bb29bfed283 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2902,9 +2902,15 @@ class _NativeLifecycleLoggingGuardrail(CustomGuardrail): guardrail_name="native-logging-guardrail", event_hook=GuardrailEventHooks.logging_only, ) - self.calls: list = [] + self.calls: list[tuple[Literal["request", "response"], list[str]]] = [] - async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: self.calls.append((input_type, list(inputs.get("texts") or []))) return inputs From 073d4fe2b01500526829523f6596a60428fbead9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:11:06 -0700 Subject: [PATCH 10/43] fix(responses): route mid-stream error events through exception_type so content_policy_fallbacks fire Mid-stream error events on the streaming Responses API were all raised as APIError, so a content_policy_violation event never matched the router's content-policy fallback dispatch and the client got the raw error instead of the fallback model's answer. Map each error event's code and status through the existing exception_type mapping, matching the non-streaming path, and unwrap the typed ContentPolicyViolationError and ContextWindowExceededError so the router routes them to the configured content_policy_fallbacks and context_window_fallbacks. --- litellm/responses/streaming_iterator.py | 48 ++++-- litellm/router.py | 10 +- .../test_streaming_iterator_error_events.py | 157 ++++++++++++++++-- tests/test_litellm/test_router.py | 105 ++++++++++++ 4 files changed, 288 insertions(+), 32 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 40ff88fc557..b3426fbbfef 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( @@ -221,6 +222,13 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None ) +def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: + if isinstance(mapped_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError)): + return True + status_code: Final = getattr(mapped_exception, "status_code", None) + return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -521,15 +529,8 @@ class BaseResponsesAPIStreamingIterator: getattr(self.completed_response, "response", None) if self.completed_response else None ) error_info: Final = getattr(response_obj, "error", None) if response_obj else None - error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) - exception: Final = litellm.APIError( - status_code=_status_code_for_error_fields(error_type, error_code), - message=error_message, - llm_provider=self.custom_llm_provider or "", - model=self.model or "", - ) - self._handle_failure(exception) + self._handle_failure(self._map_error_event_exception(error_info)) def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: @@ -551,6 +552,26 @@ class BaseResponsesAPIStreamingIterator: self.logging_obj._response_cost_calculator(result=response_obj) or 0.0 ) + def _map_error_event_exception(self, error_obj: object) -> Exception: + error_message, error_type, error_code = _error_event_fields(error_obj) + status_code: Final = _status_code_for_error_fields(error_type, error_code) + error_body: Final = {"message": error_message, "type": error_type, "code": error_code} + provider_exception: Final = BaseLLMException( + status_code=status_code, + message=f"Error code: {status_code} - {{'error': {error_body}}}", + body=error_body, + ) + try: + return litellm.exception_type( + model=self.model or "", + custom_llm_provider=self.custom_llm_provider or "", + original_exception=provider_exception, + completion_kwargs={}, + extra_kwargs={}, + ) + except Exception as mapped_exception: + return mapped_exception + def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) if chunk_type not in ("error", "response.failed"): @@ -562,15 +583,8 @@ class BaseResponsesAPIStreamingIterator: else getattr(result, "error", None) ) - error_message, error_type, error_code = _error_event_fields(error_obj) - status_code: Final = _status_code_for_error_fields(error_type, error_code) - mapped_exception: Final = litellm.APIError( - status_code=status_code, - message=error_message, - llm_provider=self.custom_llm_provider or "", - model=self.model or "", - ) - if 400 <= status_code < 500 and status_code != 429: + mapped_exception: Final = self._map_error_event_exception(error_obj) + if not _mid_stream_fallback_eligible(mapped_exception): raise mapped_exception raise MidStreamFallbackError( message=str(mapped_exception), diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..79854a11150 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3271,8 +3271,16 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + fallback_trigger: Final[Exception] = ( + e.original_exception + if isinstance( + e.original_exception, + (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError), + ) + else e + ) fallback_response = await self.async_function_with_fallbacks_common_utils( - e=e, + e=fallback_trigger, disable_fallbacks=False, fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index ad74861c096..73afbb5e63a 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -1,9 +1,14 @@ """ Regression: in-stream error events (type="error", type="response.failed") must raise instead of being returned as benign chunks, mirroring chat streaming -semantics (_handle_stream_fallback_error): non-retriable 4xx (except 429) -raise litellm.APIError directly; 429 and 5xx are wrapped in -MidStreamFallbackError so the Router's mid-stream fallback machinery fires. +semantics (_handle_stream_fallback_error). The event's code, type and status go +through litellm.exception_type, so each event raises the same typed exception +the non-streaming path raises for that provider error: non-retriable 4xx +(except 429) raise that typed exception directly, while 429, 5xx, +ContentPolicyViolationError and ContextWindowExceededError are wrapped in +MidStreamFallbackError so the Router's mid-stream fallback machinery fires and +its content_policy_fallbacks / context_window_fallbacks dispatch sees the +trigger it matches on. Status mapping must consider both the OpenAI error `type` (e.g. "invalid_request_error") and `code` (e.g. "invalid_prompt", @@ -66,12 +71,12 @@ def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback( with pytest.raises(MidStreamFallbackError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 500 - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.InternalServerError) assert exc_info.value.original_exception.status_code == 500 def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fallback(): - """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped APIError.""" + """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped RateLimitError.""" iterator = _make_iterator() chunk = _make_error_chunk("tokens", "rate_limit_exceeded", "Too many requests") with pytest.raises(MidStreamFallbackError) as exc_info: @@ -79,15 +84,15 @@ def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fall assert exc_info.value.status_code == 429 assert exc_info.value.generated_content == "" assert exc_info.value.is_pre_first_chunk is True - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) assert exc_info.value.original_exception.status_code == 429 def test_maybe_raise_for_error_event_maps_invalid_request_type_to_400(): - """Client errors classified via the `type` field must raise APIError directly (no fallback).""" + """Client errors classified via the `type` field must raise BadRequestError directly (no fallback).""" iterator = _make_iterator() chunk = _make_error_chunk("invalid_request_error", "invalid_prompt", "bad request") - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) @@ -99,12 +104,84 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): chunk = Mock() chunk.type = "error" chunk.error = {"code": "context_length_exceeded", "message": "too long"} - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) +def test_maybe_raise_for_error_event_wraps_context_window_exceeded_for_context_window_fallbacks(): + """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming + path and is wrapped so the Router's context_window_fallbacks dispatch fires mid-stream.""" + iterator = _make_iterator() + chunk = _make_error_chunk( + "invalid_request_error", + "context_length_exceeded", + "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContextWindowExceededError) + assert exc_info.value.status_code == 400 + + +CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure"]) +def test_maybe_raise_for_error_event_wraps_content_policy_violation_for_content_policy_fallbacks( + custom_llm_provider: str, +): + """Regression: a content_policy_violation error event used to raise a bare APIError, so the Router's + content_policy_fallbacks never fired. It must map to ContentPolicyViolationError (the same exception the + non-streaming path raises) and be wrapped so the Router's mid-stream fallback catches it.""" + iterator = _make_iterator() + iterator.custom_llm_provider = custom_llm_provider + chunk = _make_error_chunk("invalid_request_error", "content_policy_violation", CONTENT_POLICY_MESSAGE) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + assert exc_info.value.original_exception.status_code == 400 + assert exc_info.value.status_code == 400 + assert exc_info.value.is_pre_first_chunk is True + assert CONTENT_POLICY_MESSAGE in str(exc_info.value.original_exception) + + +def test_maybe_raise_for_response_failed_event_wraps_content_policy_violation(): + iterator = _make_iterator() + chunk = _make_failed_chunk( + {"type": "invalid_request_error", "code": "content_policy_violation", "message": CONTENT_POLICY_MESSAGE} + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + + +@pytest.mark.parametrize( + "error_type,error_code,expected_exception", + [ + ("invalid_request_error", "content_policy_violation", litellm.ContentPolicyViolationError), + ("tokens", "rate_limit_exceeded", litellm.RateLimitError), + ("invalid_request_error", "insufficient_quota", litellm.RateLimitError), + ("server_error", "internal_error", litellm.InternalServerError), + ("invalid_request_error", "invalid_prompt", litellm.BadRequestError), + ("invalid_request_error", "model_not_found", litellm.NotFoundError), + ("server_error", "vector_store_timeout", litellm.Timeout), + ], +) +def test_error_event_raises_the_same_typed_exception_as_the_non_streaming_path( + error_type: str, error_code: str, expected_exception: type[Exception] +): + iterator = _make_iterator() + chunk = _make_error_chunk(error_type, error_code, "provider message") + with pytest.raises((MidStreamFallbackError, expected_exception)) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + raised = exc_info.value + typed_exception = raised.original_exception if isinstance(raised, MidStreamFallbackError) else raised + assert type(typed_exception) is expected_exception + assert "provider message" in str(typed_exception) + + def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429(): """OpenAI returns HTTP 429 for insufficient_quota; it must not map to 400 even though its type is invalid_request_error-adjacent, and it must be wrapped for fallback.""" @@ -113,6 +190,7 @@ def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429(): with pytest.raises(MidStreamFallbackError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 429 + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) def test_maybe_raise_for_error_event_passes_through_normal_chunk(): @@ -186,10 +264,43 @@ async def test_async_iterator_raises_mid_stream_fallback_on_rate_limit_error_eve assert exc_info.value.status_code == 429 assert exc_info.value.is_pre_first_chunk is True assert exc_info.value.generated_content == "" - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) assert exc_info.value.original_exception.status_code == 429 +@pytest.mark.asyncio +async def test_async_iterator_content_policy_violation_after_first_chunk_carries_generated_content(): + """The customer's case: text streams, then the provider halts the stream with a + content_policy_violation error event. The iterator must surface ContentPolicyViolationError + inside MidStreamFallbackError, together with the text already streamed.""" + iterator = _make_async_iterator_with_events( + [ + {"type": "response.output_text.delta", "delta": "partial "}, + { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "content_policy_violation", + "message": CONTENT_POLICY_MESSAGE, + }, + }, + ] + ) + + chunks = [] + + async def _drain(): + async for chunk in iterator: + chunks.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _drain() + assert len(chunks) == 1 + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + assert exc_info.value.is_pre_first_chunk is False + assert exc_info.value.generated_content == "partial " + + @pytest.mark.asyncio async def test_async_iterator_error_after_first_chunk_carries_generated_content(): """An error after streamed output must expose the accumulated text so the router's @@ -265,7 +376,7 @@ def test_handle_logging_failed_response_maps_rate_limit_to_429(): ): iterator._handle_logging_failed_response() logged_exception = mock_run_async.call_args.kwargs["exception"] - assert isinstance(logged_exception, litellm.APIError) + assert isinstance(logged_exception, litellm.RateLimitError) assert logged_exception.status_code == 429 assert "throttled" in str(logged_exception) @@ -282,10 +393,28 @@ def test_handle_logging_failed_response_maps_type_field_to_400(): ): iterator._handle_logging_failed_response() logged_exception = mock_run_async.call_args.kwargs["exception"] - assert isinstance(logged_exception, litellm.APIError) + assert isinstance(logged_exception, litellm.BadRequestError) assert logged_exception.status_code == 400 +def test_handle_logging_failed_response_logs_content_policy_violation(): + """Failure logging must record the same typed exception the stream raises, so logging + integrations see a content policy violation instead of a generic APIError.""" + iterator = _make_iterator() + iterator.completed_response = _make_failed_chunk( + {"type": "invalid_request_error", "code": "content_policy_violation", "message": CONTENT_POLICY_MESSAGE} + ) + with ( + patch.object(import_module("litellm.responses.streaming_iterator"), "run_async_function") as mock_run_async, + patch.object(import_module("litellm.responses.streaming_iterator"), "executor"), + ): + iterator._handle_logging_failed_response() + logged_exception = mock_run_async.call_args.kwargs["exception"] + assert isinstance(logged_exception, litellm.ContentPolicyViolationError) + assert logged_exception.status_code == 400 + assert CONTENT_POLICY_MESSAGE in str(logged_exception) + + def test_handle_logging_failed_response_records_usage_and_cost(): """Usage on a response.failed event must reach failure spend accounting via combined_usage_object.""" iterator = _make_iterator() @@ -357,7 +486,7 @@ def test_sync_iterator_raises_mid_stream_fallback_on_rate_limit_error_event(): for _ in iterator: pass assert exc_info.value.status_code == 429 - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) def test_every_openai_sdk_response_error_code_has_explicit_status_mapping(): @@ -413,7 +542,7 @@ def test_maybe_raise_for_response_failed_event_maps_image_code_to_400(): chunk = Mock() chunk.type = "response.failed" chunk.response = mock_response_obj - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b8a0d70f5bc..d74074dd149 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3654,6 +3654,111 @@ async def test_aresponses_streaming_iterator_fallback(): assert call_kwargs["disable_fallbacks"] is False +@pytest.mark.asyncio +async def test_aresponses_streaming_content_policy_error_event_routes_to_content_policy_fallback(): + """Regression: a mid-stream content_policy_violation error event never reached + content_policy_fallbacks. The iterator raised a bare APIError the wrapper does not + catch, and even once wrapped, the MidStreamFallbackError envelope was handed to the + fallback dispatch, whose isinstance branch on ContentPolicyViolationError never matched. + The stream below is the customer's shape: a raw OpenAI error event with code + content_policy_violation, transformed by the real OpenAI config, and the router must + call the content_policy_fallbacks target, not the general fallbacks one.""" + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + router = litellm.Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k1"}}, + { + "model_name": "content-fallback", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k2"}, + }, + {"model_name": "general-fallback", "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k3"}}, + ], + fallbacks=[{"primary": ["general-fallback"]}], + content_policy_fallbacks=[{"primary": ["content-fallback"]}], + ) + error_event = { + "type": "error", + "sequence_number": 2, + "error": { + "type": "invalid_request_error", + "code": "content_policy_violation", + "message": "This content was flagged for possible cybersecurity risk. The response was halted mid-stream.", + "param": None, + }, + } + + async def aiter_bytes(): + yield f"data: {json.dumps(error_event)}\n\n".encode() + + raw_response = MagicMock() + raw_response.headers = {} + raw_response.aiter_bytes = aiter_bytes + logging_obj = MagicMock(spec=LiteLLMLogging) + logging_obj.model_call_details = {"litellm_params": {}} + logging_obj.completion_start_time = None + source = ResponsesAPIStreamingIterator( + response=raw_response, + model="gpt-5.4", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + fallback_chunks = [MagicMock(type="response.output_text.delta"), MagicMock(type="response.completed")] + fallback_call = AsyncMock(return_value=_AsyncList(fallback_chunks)) + + wrapped = await router._aresponses_streaming_iterator( + response=source, + initial_kwargs={ + "model": "primary", + "stream": True, + "input": "Hi", + "original_generic_function": fallback_call, + }, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == fallback_chunks + fallback_call.assert_awaited_once() + assert fallback_call.await_args.kwargs["model"] == "gemini/gemini-2.5-flash" + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_fallback_dispatch(): + """The fallback dispatch matches on the trigger's own type, so the wrapper must hand it the + ContentPolicyViolationError carried inside MidStreamFallbackError, not the envelope.""" + router = _make_router_with_fallback("openai/gpt-5.4", "openai/gpt-5-mini") + content_policy_error = litellm.ContentPolicyViolationError( + message="flagged mid-stream", llm_provider="openai", model="openai/gpt-5.4" + ) + src = _make_responses_iterator( + chunks=[MagicMock(type="response.created")], + error=MidStreamFallbackError( + message=str(content_policy_error), + model="openai/gpt-5.4", + llm_provider="openai", + original_exception=content_policy_error, + is_pre_first_chunk=True, + ), + model="openai/gpt-5.4", + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_AsyncList([MagicMock(type="response.completed")])), + ) as mock_fallback_utils: + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={"model": "openai/gpt-5.4", "stream": True, "input": "Hi"}, + ) + [chunk async for chunk in wrapped] + + mock_fallback_utils.assert_awaited_once() + assert mock_fallback_utils.await_args.kwargs["e"] is content_policy_error + + @pytest.mark.asyncio @pytest.mark.parametrize( "fallback_headers", From c2463728593b935448bcb09f54ef4fd070e9f8bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:35:53 -0700 Subject: [PATCH 11/43] fix(responses): import BaseLLMException lazily and collect stream chunks via anext Move the BaseLLMException import into _map_error_event_exception so the module no longer imports it at load time, clearing the module-level cyclic import CodeQL flagged. The class is used only on the cold error path. Replace the mutable list-append test collector with aiter/anext so the regression tests read the stream immutably. --- litellm/responses/streaming_iterator.py | 3 ++- .../test_streaming_iterator_error_events.py | 22 ++++++++----------- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index b3426fbbfef..a15571acb7c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,7 +31,6 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) from litellm.litellm_core_utils.thread_pool_executor import executor -from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( @@ -553,6 +552,8 @@ class BaseResponsesAPIStreamingIterator: ) def _map_error_event_exception(self, error_obj: object) -> Exception: + from litellm.llms.base_llm.chat.transformation import BaseLLMException + error_message, error_type, error_code = _error_event_fields(error_obj) status_code: Final = _status_code_for_error_fields(error_type, error_code) error_body: Final = {"message": error_message, "type": error_type, "code": error_code} diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 73afbb5e63a..2f4fba45cee 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -287,15 +287,12 @@ async def test_async_iterator_content_policy_violation_after_first_chunk_carries ] ) - chunks = [] - - async def _drain(): - async for chunk in iterator: - chunks.append(chunk) + stream = aiter(iterator) + first_chunk = await anext(stream) + assert first_chunk is not None with pytest.raises(MidStreamFallbackError) as exc_info: - await _drain() - assert len(chunks) == 1 + await anext(stream) assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) assert exc_info.value.is_pre_first_chunk is False assert exc_info.value.generated_content == "partial " @@ -316,14 +313,13 @@ async def test_async_iterator_error_after_first_chunk_carries_generated_content( ] ) - chunks = [] - async def _drain(): - async for chunk in iterator: - chunks.append(chunk) + stream = aiter(iterator) + first_chunk = await anext(stream) + second_chunk = await anext(stream) + assert first_chunk is not None and second_chunk is not None with pytest.raises(MidStreamFallbackError) as exc_info: - await _drain() - assert len(chunks) == 2 + await anext(stream) assert exc_info.value.status_code == 500 assert exc_info.value.is_pre_first_chunk is False assert exc_info.value.generated_content == "hello world" From fff7a2cecfb113082681665d236525f2690aae45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:18:59 -0700 Subject: [PATCH 12/43] fix(responses): keep context-window events out of mid-stream fallback and fix stale exception assertions --- litellm/responses/streaming_iterator.py | 2 +- litellm/router.py | 7 +++---- .../test_openai_responses_api.py | 9 +++++---- ...est_router_aresponses_streaming_fallback.py | 2 +- .../test_streaming_iterator_error_events.py | 18 ++++++++++-------- 5 files changed, 20 insertions(+), 18 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index a15571acb7c..b39e130242d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -222,7 +222,7 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: - if isinstance(mapped_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError)): + if isinstance(mapped_exception, litellm.ContentPolicyViolationError): return True status_code: Final = getattr(mapped_exception, "status_code", None) return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 diff --git a/litellm/router.py b/litellm/router.py index 79854a11150..a9a8f3e2739 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3271,12 +3271,11 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + # The content-policy dispatch branch matches on the trigger's own type, so a refusal's + # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. fallback_trigger: Final[Exception] = ( e.original_exception - if isinstance( - e.original_exception, - (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError), - ) + if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e ) fallback_response = await self.async_function_with_fallbacks_common_utils( diff --git a/tests/llm_responses_api_testing/test_openai_responses_api.py b/tests/llm_responses_api_testing/test_openai_responses_api.py index 05bb9113835..c7712d96969 100644 --- a/tests/llm_responses_api_testing/test_openai_responses_api.py +++ b/tests/llm_responses_api_testing/test_openai_responses_api.py @@ -1627,9 +1627,10 @@ async def test_openai_responses_api_token_limit_error(): Parsing the in-stream ErrorEvent must not raise "pydantic_core._pydantic_core.ValidationError: 3 validation errors for ErrorEvent". - The iterator now surfaces the event as litellm.APIError with status 400 - (invalid_request_error is a non-retriable client error, so no - MidStreamFallbackError wrapping) carrying the provider's message. + The iterator routes the event through litellm.exception_type, so it surfaces as + the typed 400 client error the non-streaming path raises (litellm.BadRequestError) + carrying the provider's message. invalid_request_error is a non-retriable client + error, so there is no MidStreamFallbackError wrapping. """ litellm._turn_on_debug() @@ -1644,7 +1645,7 @@ async def test_openai_responses_api_token_limit_error(): async for event in response: print(event) - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: await _drain() assert exc_info.value.status_code == 400 diff --git a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py index ee4750e9db8..0d33435cf7a 100644 --- a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py +++ b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py @@ -372,7 +372,7 @@ async def test_aresponses_fallback_on_in_stream_error_event(): raised = mock_fallback.await_args.kwargs["e"] assert isinstance(raised, MidStreamFallbackError) assert raised.status_code == 429 - assert isinstance(raised.original_exception, litellm.APIError) + assert isinstance(raised.original_exception, litellm.RateLimitError) assert raised.original_exception.status_code == 429 assert mock_fallback.await_args.kwargs["kwargs"]["input"] == "original question" diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 2f4fba45cee..3d7c220804a 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -4,11 +4,11 @@ raise instead of being returned as benign chunks, mirroring chat streaming semantics (_handle_stream_fallback_error). The event's code, type and status go through litellm.exception_type, so each event raises the same typed exception the non-streaming path raises for that provider error: non-retriable 4xx -(except 429) raise that typed exception directly, while 429, 5xx, -ContentPolicyViolationError and ContextWindowExceededError are wrapped in +(except 429) raise that typed exception directly, so a context-length event +surfaces as ContextWindowExceededError(400) with no MidStreamFallbackError +wrapping, while 429, 5xx and ContentPolicyViolationError are wrapped in MidStreamFallbackError so the Router's mid-stream fallback machinery fires and -its content_policy_fallbacks / context_window_fallbacks dispatch sees the -trigger it matches on. +its content_policy_fallbacks dispatch sees the trigger it matches on. Status mapping must consider both the OpenAI error `type` (e.g. "invalid_request_error") and `code` (e.g. "invalid_prompt", @@ -110,19 +110,21 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): assert not isinstance(exc_info.value, MidStreamFallbackError) -def test_maybe_raise_for_error_event_wraps_context_window_exceeded_for_context_window_fallbacks(): +def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(): """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming - path and is wrapped so the Router's context_window_fallbacks dispatch fires mid-stream.""" + path and, being a non-retriable client error, is raised directly rather than wrapped for mid-stream + fallback, preserving the direct-SDK 400 contract from issue #15785.""" iterator = _make_iterator() chunk = _make_error_chunk( "invalid_request_error", "context_length_exceeded", "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", ) - with pytest.raises(MidStreamFallbackError) as exc_info: + with pytest.raises(litellm.ContextWindowExceededError) as exc_info: iterator._maybe_raise_for_error_event(chunk) - assert isinstance(exc_info.value.original_exception, litellm.ContextWindowExceededError) assert exc_info.value.status_code == 400 + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert "maximum context length" in str(exc_info.value) CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." From 03e6dd051cdc2d4fc307783b3a11ab618bc70216 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 06:59:10 +0000 Subject: [PATCH 13/43] fix(cli): drop enum.StrEnum so the CLI imports on Python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/pi.py | 4 ++-- tests/test_litellm/proxy/client/cli/test_pi.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 9810e81ae36..5c749959638 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -10,7 +10,7 @@ import os import tempfile from collections.abc import Callable, Mapping from dataclasses import dataclass -from enum import StrEnum +from enum import Enum from pathlib import Path from types import MappingProxyType from typing import Annotated, Final @@ -25,7 +25,7 @@ LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY" _REJECTED_STATUSES: Final = frozenset((401, 403)) -class ListingFailure(StrEnum): +class ListingFailure(str, Enum): """Why a proxy could not be listed, decided once where the HTTP outcome is classified. `unreachable` means no response at all; the other kinds prove the proxy answered, so callers diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 03e5d7dd197..3c343c9e268 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -1,3 +1,4 @@ +import inspect import json import os import stat @@ -7,6 +8,7 @@ from pathlib import Path import pytest import requests +from litellm.proxy.client.cli.commands import pi from litellm.proxy.client.cli.commands.pi import ( ListingFailure, ModelLimits, @@ -20,6 +22,18 @@ from litellm.proxy.client.cli.commands.pi import ( ) +def test_pi_module_has_no_strenum_import(): + """enum.StrEnum is Python 3.11+; pyproject allows 3.10, so pi.py must not import it.""" + assert "StrEnum" not in inspect.getsource(pi) + + +def test_listing_failure_is_str_enum(): + assert issubclass(ListingFailure, str) + assert ListingFailure.REJECTED.value == "rejected" + assert ListingFailure("rejected") is ListingFailure.REJECTED + assert str(ListingFailure.REJECTED.value) == "rejected" + + class _FakeResponse: def __init__(self, status_code, payload=None): self.status_code = status_code From 0679d799d483dabe74887a6a54b25b737676af8f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 07:07:06 +0000 Subject: [PATCH 14/43] ci: exercise the lite CLI on the Python 3.10 import smoke job Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-code-quality.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 9c7e0db7065..987f66773f2 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -178,7 +178,7 @@ jobs: version: "0.10.9" - name: Install dependencies - run: uv sync --frozen --extra proxy --python 3.10 + run: uv sync --frozen --extra proxy --extra cli --python 3.10 - run: uv run --no-sync python --version @@ -187,3 +187,6 @@ jobs: - name: Check litellm CLI run: uv run --no-sync litellm --version + + - name: Check lite CLI + run: uv run --no-sync lite version From 19c43eb875944e326d3c98743f4bcd0835298761 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 07:08:24 +0000 Subject: [PATCH 15/43] test(cli): drop structural StrEnum source check; smoke job covers the 3.10 import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/client/cli/test_pi.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 3c343c9e268..6bb49566580 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -1,4 +1,3 @@ -import inspect import json import os import stat @@ -8,7 +7,6 @@ from pathlib import Path import pytest import requests -from litellm.proxy.client.cli.commands import pi from litellm.proxy.client.cli.commands.pi import ( ListingFailure, ModelLimits, @@ -22,11 +20,6 @@ from litellm.proxy.client.cli.commands.pi import ( ) -def test_pi_module_has_no_strenum_import(): - """enum.StrEnum is Python 3.11+; pyproject allows 3.10, so pi.py must not import it.""" - assert "StrEnum" not in inspect.getsource(pi) - - def test_listing_failure_is_str_enum(): assert issubclass(ListingFailure, str) assert ListingFailure.REJECTED.value == "rejected" From eceb67b69a931d3006f3b9ad07a38967b19599bf Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Mon, 14 Sep 2026 18:03:22 +0000 Subject: [PATCH 16/43] test(model_armor): inject HTTP handler and token provider Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_armor/model_armor.py | 21 ++- .../guardrail_hooks/test_model_armor.py | 126 +++++++++++------- 2 files changed, 95 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index a0563a7a1c9..285d883bd28 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,5 +1,5 @@ import time -from collections.abc import AsyncGenerator, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from enum import Enum, auto from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_or_create_metadata_bucket, ) from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) @@ -146,6 +147,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): credentials: VERTEX_CREDENTIALS_TYPES | None = None, api_endpoint: str | None = None, sanitize_error_detail: "bool | None" = True, + async_handler: AsyncHTTPHandler | None = None, + access_token_provider: Callable[[], Awaitable[tuple[str, str]]] | None = None, **kwargs, ): # Set supported event hooks if not already provided @@ -162,7 +165,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): VertexBase.__init__(self) # Then set our attributes (this ensures project_id is not overwritten) - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = async_handler or get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.access_token_provider = access_token_provider self.template_id = template_id self.project_id = project_id self.location = location or "us-central1" @@ -286,11 +290,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): If file_bytes and file_type are provided, file prompt sanitization is performed. """ # Get access token using VertexBase auth - access_token, resolved_project_id = await self._ensure_access_token_async( - credentials=self.credentials, - project_id=self.project_id, - custom_llm_provider="vertex_ai", - ) + if self.access_token_provider is not None: + access_token, resolved_project_id = await self.access_token_provider() + else: + access_token, resolved_project_id = await self._ensure_access_token_async( + credentials=self.credentials, + project_id=self.project_id, + custom_llm_provider="vertex_ai", + ) # Use resolved project ID if not explicitly set if not self.project_id and resolved_project_id: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 126e162fec8..5aa3c7dfc17 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -2,6 +2,8 @@ import asyncio import base64 import io import json +from collections.abc import Iterator, Sequence +from typing import cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -14,7 +16,7 @@ import litellm import litellm.types.utils from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache -from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail @@ -4931,7 +4933,7 @@ def test_every_responses_delta_event_is_in_the_scanned_set(): assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES -def _clean_armor_response() -> dict: +def _clean_armor_response() -> dict[str, object]: return { "sanitizationResult": { "filterMatchState": "NO_MATCH_FOUND", @@ -4940,7 +4942,7 @@ def _clean_armor_response() -> dict: } -def _flagged_armor_response() -> dict: +def _flagged_armor_response() -> dict[str, object]: return { "sanitizationResult": { "filterMatchState": "MATCH_FOUND", @@ -4949,17 +4951,50 @@ def _flagged_armor_response() -> dict: } -def _logging_only_guardrail() -> ModelArmorGuardrail: - return ModelArmorGuardrail( +class _FakeArmorHandler(AsyncHTTPHandler): + def __init__(self, responses: Sequence[dict[str, object] | Exception]): + self.responses: Iterator[dict[str, object] | Exception] = iter(responses) + self.calls: list[dict[str, object]] = [] + self.raise_on_call: Exception | None = None + + async def post( + self, + url: str, + json: dict[str, object] | None = None, + headers: dict[str, str] | None = None, + **kwargs: object, + ) -> httpx.Response: + if self.raise_on_call is not None: + raise self.raise_on_call + if json is not None: + self.calls.append(json) + response: dict[str, object] | Exception = next(self.responses) + if isinstance(response, Exception): + raise response + return httpx.Response(200, json=response, request=httpx.Request("POST", url)) + + +async def _async_token_provider() -> tuple[str, str]: + return ("test-token", "test-project") + + +def _logging_only_guardrail( + responses: Sequence[dict[str, object] | Exception] = (_clean_armor_response(), _clean_armor_response()), +) -> ModelArmorGuardrail: + handler = _FakeArmorHandler(responses) + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-logging", event_hook=GuardrailEventHooks.logging_only, + async_handler=handler, + access_token_provider=_async_token_provider, ) + return guardrail -def _logged_kwargs() -> dict: +def _logged_kwargs() -> dict[str, object]: return { "model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], @@ -4990,8 +5025,10 @@ def _stream_chunk(text: str) -> litellm.ModelResponseStream: ) -def _metadata_entries(kwargs: dict) -> list: - return kwargs["standard_logging_object"].get("guardrail_information") or [] +def _metadata_entries(kwargs: dict[str, object]) -> list[dict[str, object]]: + standard_logging_object = cast(dict[str, object], kwargs["standard_logging_object"]) + entries = standard_logging_object.get("guardrail_information") or [] + return cast(list[dict[str, object]], entries) def test_logging_only_mode_is_accepted_and_keeps_native_hooks(): @@ -5014,10 +5051,11 @@ def test_logging_only_mode_is_accepted_and_keeps_native_hooks(): async def test_logging_only_stream_yields_chunks_without_waiting_for_scan(): """A logging_only guardrail must pass stream chunks straight through; the scan happens afterwards on the assembled response via async_logging_hook.""" - guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock( - side_effect=AssertionError("logging_only must not scan the stream") + guardrail = _logging_only_guardrail( + [_clean_armor_response(), _clean_armor_response()] ) + handler = cast(_FakeArmorHandler, guardrail.async_handler) + handler.raise_on_call = AssertionError("logging_only must not scan the stream") produced = 0 @@ -5038,9 +5076,9 @@ async def test_logging_only_stream_yields_chunks_without_waiting_for_scan(): async for chunk in hook_iter: chunks.append(chunk) assert len(chunks) == 3 - guardrail.make_model_armor_request.assert_not_awaited() + assert handler.calls == [] + handler.raise_on_call = None - guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) response = _chat_response("all clear") kwargs = _logged_kwargs() out_kwargs, out_result = await guardrail.async_logging_hook( @@ -5057,8 +5095,9 @@ async def test_logging_only_stream_yields_chunks_without_waiting_for_scan(): @pytest.mark.asyncio async def test_logging_only_records_flagged_verdict_without_altering_response(): - guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + guardrail = _logging_only_guardrail( + [_flagged_armor_response(), _flagged_armor_response()] + ) response = _chat_response("flagged output") kwargs = _logged_kwargs() @@ -5074,9 +5113,11 @@ async def test_logging_only_records_flagged_verdict_without_altering_response(): @pytest.mark.asyncio async def test_logging_only_records_model_armor_api_error(): - guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock( - side_effect=ModelArmorAPIError("Model Armor API error (upstream 500)") + guardrail = _logging_only_guardrail( + [ + ModelArmorAPIError("Model Armor API error (upstream 500)"), + ModelArmorAPIError("Model Armor API error (upstream 500)"), + ] ) response = _chat_response("some output") kwargs = _logged_kwargs() @@ -5124,7 +5165,6 @@ async def test_logging_only_scans_assembled_responses_api_stream(): ) guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) kwargs = _logged_kwargs() del kwargs["messages"] kwargs["input"] = "hello" @@ -5134,13 +5174,10 @@ async def test_logging_only_scans_assembled_responses_api_stream(): kwargs=kwargs, result=event, call_type="aresponses" ) - response_scans = [ - call - for call in guardrail.make_model_armor_request.await_args_list - if call.kwargs.get("source") == "model_response" - ] + handler = cast(_FakeArmorHandler, guardrail.async_handler) + response_scans = [call for call in handler.calls if "modelResponseData" in call] assert response_scans, "expected a model_response scan of the assembled response" - assert "assembled output text" in response_scans[0].kwargs["content"] + assert "assembled output text" in response_scans[0]["modelResponseData"]["text"] assert _metadata_entries(out_kwargs) @@ -5148,7 +5185,6 @@ async def test_logging_only_scans_assembled_responses_api_stream(): async def test_logging_only_scans_anthropic_messages_model_response(): """/v1/messages logs a ModelResponse; the output scan must extract the assistant text.""" guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) kwargs = _logged_kwargs() kwargs["messages"] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] response = _chat_response("anthropic assembled text") @@ -5158,38 +5194,36 @@ async def test_logging_only_scans_anthropic_messages_model_response(): ) assert out_result is response - response_scans = [ - call - for call in guardrail.make_model_armor_request.await_args_list - if call.kwargs.get("source") == "model_response" - ] + handler = cast(_FakeArmorHandler, guardrail.async_handler) + response_scans = [call for call in handler.calls if "modelResponseData" in call] assert response_scans - assert "anthropic assembled text" in response_scans[0].kwargs["content"] + assert "anthropic assembled text" in response_scans[0]["modelResponseData"]["text"] assert _metadata_entries(out_kwargs) @pytest.mark.asyncio async def test_logging_only_skips_output_scan_when_no_assembled_response(): guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) kwargs = _logged_kwargs() await guardrail.async_logging_hook(kwargs=kwargs, result=None, call_type="acompletion") - sources = [call.kwargs.get("source") for call in guardrail.make_model_armor_request.await_args_list] - assert "model_response" not in sources + handler = cast(_FakeArmorHandler, guardrail.async_handler) + assert all("modelResponseData" not in call for call in handler.calls) @pytest.mark.asyncio async def test_native_post_call_mode_ignores_logging_hook(): + handler = _FakeArmorHandler([_clean_armor_response()]) guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-post", event_hook=GuardrailEventHooks.post_call, + async_handler=handler, + access_token_provider=_async_token_provider, ) - guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response()) response = _chat_response("some output") kwargs = _logged_kwargs() @@ -5199,13 +5233,12 @@ async def test_native_post_call_mode_ignores_logging_hook(): assert out_kwargs is kwargs assert out_result is response - guardrail.make_model_armor_request.assert_not_awaited() + assert handler.calls == [] @pytest.mark.asyncio async def test_apply_guardrail_records_flagged_without_raising(): - guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + guardrail = _logging_only_guardrail([_flagged_armor_response()]) request_data = {"metadata": {}} inputs = {"texts": ["forbidden output"]} @@ -5222,8 +5255,7 @@ async def test_apply_guardrail_records_flagged_without_raising(): @pytest.mark.asyncio async def test_logging_only_records_transport_error(): - guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock(side_effect=httpx.ConnectError("boom")) + guardrail = _logging_only_guardrail([httpx.ConnectError("boom"), httpx.ConnectError("boom")]) response = _chat_response("some output") kwargs = _logged_kwargs() @@ -5241,8 +5273,9 @@ async def test_logging_only_records_transport_error(): @pytest.mark.asyncio async def test_logging_only_flagged_prompt_still_scans_response(): """A flagged input scan must not abort the output scan; both verdicts are recorded.""" - guardrail = _logging_only_guardrail() - guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) + guardrail = _logging_only_guardrail( + [_flagged_armor_response(), _flagged_armor_response()] + ) response = _chat_response("flagged output") kwargs = _logged_kwargs() @@ -5250,7 +5283,8 @@ async def test_logging_only_flagged_prompt_still_scans_response(): kwargs=kwargs, result=response, call_type="acompletion" ) - sources = [call.kwargs.get("source") for call in guardrail.make_model_armor_request.await_args_list] + handler = cast(_FakeArmorHandler, guardrail.async_handler) + sources = ["user_prompt" if "userPromptData" in call else "model_response" for call in handler.calls] assert sources == ["user_prompt", "model_response"] entries = _metadata_entries(out_kwargs) flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] @@ -5261,14 +5295,16 @@ async def test_logging_only_flagged_prompt_still_scans_response(): async def test_apply_guardrail_raises_on_flagged_when_not_logging_only(): """The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a non-logging_only instance must signal the block so flagged text is not returned as clean.""" + handler = _FakeArmorHandler([_flagged_armor_response()]) guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-pre", event_hook=GuardrailEventHooks.pre_call, + async_handler=handler, + access_token_provider=_async_token_provider, ) - guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response()) request_data = {"metadata": {}} with pytest.raises(HTTPException) as exc_info: From fb67d8335cc56540475f2768e1898595545d8f7a Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Mon, 14 Sep 2026 18:06:41 +0000 Subject: [PATCH 17/43] style(model_armor): format guardrail handler Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/guardrail_hooks/model_armor/model_armor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 285d883bd28..a7c93e63b32 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -165,7 +165,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): VertexBase.__init__(self) # Then set our attributes (this ensures project_id is not overwritten) - self.async_handler = async_handler or get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.access_token_provider = access_token_provider self.template_id = template_id self.project_id = project_id From 26e2e208ad347542f7ec4a5c34a2c4a02a5c9145 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:18:08 +0000 Subject: [PATCH 18/43] fix(bedrock guardrails): derive contextual grounding source and query from plain messages Bedrock only runs a contextualGroundingPolicy when the ApplyGuardrail payload carries grounding_source and query qualifiers. Callers sending ordinary system and user messages never got those, so a configured grounding threshold was silently skipped on /v1/chat/completions and /guardrails/apply_guardrail. When no explicit grounding_source or query tags are present, system and developer text is sent as grounding_source and the latest user message as query. The apply_guardrail response branch now forwards the request messages, which it previously dropped. Resolves LIT-4224 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/bedrock_guardrails.py | 32 +++- .../test_bedrock_guardrails.py | 154 ++++++++++++++++++ 2 files changed, 185 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 6a7ac4361b9..fe254803d04 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -493,6 +493,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): result carrying externally-influenced content can supply fake evidence for the contextual-grounding check to grade the response against. ``query`` is accepted from any role (it is the user's question). + + A request with no tagged blocks falls back to the plain messages: system / + developer text is the grounding source and the latest user message is the query. """ grounding: Final[list[QualifiedTextBlock]] = [] for message in messages or []: @@ -504,7 +507,33 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): and role in _GROUNDING_SOURCE_TRUSTED_ROLES ): grounding.append(block) - return grounding + return grounding or self._derive_grounding_blocks_from_plain_messages(messages) + + def _derive_grounding_blocks_from_plain_messages( + self, messages: list[AllMessageValues] | None + ) -> list[QualifiedTextBlock]: + """Bedrock scores grounding only when source, query and response are all present, + and rejects a source without a query, so return nothing unless both exist.""" + if not messages: + return [] + latest_user_index: Final = self._find_latest_message_index(messages, target_role="user") + if latest_user_index is None: + return [] + sources: Final = tuple( + QualifiedTextBlock(text=block.text, qualifier="grounding_source") + for message in messages + if message.get("role") in _GROUNDING_SOURCE_TRUSTED_ROLES + for block in self.get_content_items_for_message(message=message) or [] + if block.text + ) + queries: Final = tuple( + QualifiedTextBlock(text=block.text, qualifier="query") + for block in self.get_content_items_for_message(message=messages[latest_user_index]) or [] + if block.text + ) + if not sources or not queries: + return [] + return [*sources, *queries] def supports_scan_only_tool_results(self) -> bool: return self.experimental_use_latest_role_message_only is not True @@ -3210,6 +3239,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): bedrock_response = await self.make_bedrock_api_request( source="OUTPUT", response=synthetic_response, + messages=request_data.get("messages"), request_data=request_data, logging_event_type=_log_hook, ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index c0762edec92..a9a208e1807 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2474,6 +2474,122 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): assert actual_request == expected_request +def test_grounding_output_derives_source_and_query_from_plain_messages(): + """Untagged string system + user messages become grounding_source + query, so a + guardrail with a grounding threshold actually grades the response.""" + messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = { + "source": "OUTPUT", + "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], + } + + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + + assert actual_request == expected_request + + +def test_grounding_output_derived_query_is_latest_user_turn_only(): + """In a multi-turn chat only the latest user message is the query; earlier user + turns and assistant turns are not sent as query or source. Developer messages + count as source alongside system.""" + developer_text = "Answer in one sentence." + messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "developer", "content": [{"type": "text", "text": developer_text}]}, + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello, how can I help?"}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = { + "source": "OUTPUT", + "content": [ + _GROUNDING_SOURCE_BLOCK, + {"text": {"text": developer_text, "qualifiers": ["grounding_source"]}}, + _QUERY_BLOCK, + _GUARD_BLOCK, + ], + } + + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + + assert actual_request == expected_request + + +@pytest.mark.parametrize( + "messages", + [ + pytest.param([{"role": "system", "content": _GROUNDING_SOURCE_TEXT}], id="system-without-user"), + pytest.param( + [ + {"role": "tool", "content": _GROUNDING_SOURCE_TEXT, "tool_call_id": "c1"}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ], + id="tool-result-is-not-a-source", + ), + pytest.param( + [ + {"role": "system", "content": ""}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ], + id="empty-system-prompt", + ), + pytest.param( + [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://x.test/a.png"}}]}, + ], + id="image-only-user-turn", + ), + ], +) +def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(messages): + """Bedrock rejects a grounding_source without a query (and vice versa), so a + request that cannot supply both from trusted roles keeps the untagged payload.""" + expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]} + + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + + assert actual_request == expected_request + + +def test_grounding_output_explicit_tags_take_precedence_over_plain_messages(): + """A caller that tags blocks keeps full control: untagged system text is not + added as a second source and the untagged user text is not a second query.""" + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + *_grounding_messages(), + {"role": "user", "content": "Please be brief."}, + ] + expected_request = { + "source": "OUTPUT", + "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], + } + + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + + assert actual_request == expected_request + + +def test_grounding_input_ignores_plain_message_derivation(): + """Derivation is OUTPUT-only: an INPUT scan of plain system + user text stays an + untagged payload, so input policies keep scanning every block.""" + messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = { + "source": "INPUT", + "content": [{"text": {"text": _GROUNDING_SOURCE_TEXT}}, {"text": {"text": _GROUNDING_QUERY_TEXT}}], + } + + actual_request = _input_request(messages) + + assert actual_request == expected_request + + def test_grounding_output_combines_multiple_sources(): """Every grounding_source block is emitted; Bedrock combines them into one corpus.""" uk_source_text = "London is the capital of UK." @@ -2597,6 +2713,44 @@ async def test_grounding_output_blocked_raises_400(): assert exc_info.value.status_code == 400 +@pytest.mark.asyncio +async def test_apply_guardrail_response_forwards_request_messages_for_grounding(): + """/guardrails/apply_guardrail with input_type=response: the request messages + stored in request_data must reach the OUTPUT payload as grounding_source + query + around the guarded text. Before LIT-4224 the response branch dropped them, so + Bedrock never ran its contextual-grounding policy on this route.""" + guardrail = _grounding_guardrail() + request_messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = { + "source": "OUTPUT", + "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], + } + + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare, + ): + mock_post.return_value = _passing_bedrock_httpx_response(_GROUNDING_RESPONSE_TEXT) + + await guardrail.apply_guardrail( + inputs={"texts": [_GROUNDING_RESPONSE_TEXT]}, + request_data={"messages": request_messages}, + input_type="response", + ) + + assert mock_prepare.call_count == 1 + assert json.loads(json.dumps(mock_prepare.call_args.kwargs["data"])) == expected_request + + ############################################################################### # LIT-4186: disable_exception_on_block regression tests # From 442af3aab69f39d490829714025ffc2e49e7acc3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:04:42 +0000 Subject: [PATCH 19/43] fix(bedrock guardrails): gate plain-message grounding behind contextual_grounding_from_messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 6 +++ .../guardrail_hooks/bedrock_guardrails.py | 17 ++++--- .../guardrails/guardrail_initializers.py | 1 + litellm/types/guardrails.py | 10 ++++ .../test_bedrock_guardrails.py | 48 +++++++++++++------ .../proxy/guardrails/test_init_guardrails.py | 46 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 +++ 7 files changed, 112 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7a110eff080..fb2f014e3d8 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -11503,6 +11503,12 @@ "description": "Enable content moderation to check for harmful content (harassment, hate speech, etc.).", "title": "Content Moderation Check" }, + "contextual_grounding_from_messages": { + "default": false, + "description": "ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.", + "title": "Contextual Grounding From Messages", + "type": "boolean" + }, "credentials": { "anyOf": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index fe254803d04..2c407d91a48 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -244,6 +244,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + contextual_grounding_from_messages: bool = False, streaming_buffer_until_moderated: bool | None = None, streaming_sampling_rate: int | None = None, streaming_end_of_stream_only: bool | None = None, @@ -265,6 +266,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.guardrailVersion = guardrailVersion self.guardrail_provider = "bedrock" self.chunk_budget_chars = chunk_budget_chars + self.contextual_grounding_from_messages = contextual_grounding_from_messages self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only")) # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks` @@ -459,8 +461,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ Flatten a message into text blocks, preserving any contextual-grounding qualifier carried by the content-block ``type`` (grounding_source / query). - Untagged text keeps ``qualifier=None`` so the payload is unchanged for - callers that do not use grounding. + Untagged text keeps ``qualifier=None``; the OUTPUT scan decides whether to + derive grounding qualifiers from it. """ content: Final = message.get("content") if content is None: @@ -494,8 +496,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): contextual-grounding check to grade the response against. ``query`` is accepted from any role (it is the user's question). - A request with no tagged blocks falls back to the plain messages: system / - developer text is the grounding source and the latest user message is the query. + With ``contextual_grounding_from_messages`` on, a request with no tagged blocks + falls back to the plain messages: system / developer text is the grounding + source and the latest user message is the query. """ grounding: Final[list[QualifiedTextBlock]] = [] for message in messages or []: @@ -507,13 +510,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): and role in _GROUNDING_SOURCE_TRUSTED_ROLES ): grounding.append(block) - return grounding or self._derive_grounding_blocks_from_plain_messages(messages) + if grounding or not self.contextual_grounding_from_messages: + return grounding + return self._derive_grounding_blocks_from_plain_messages(messages) def _derive_grounding_blocks_from_plain_messages( self, messages: list[AllMessageValues] | None ) -> list[QualifiedTextBlock]: - """Bedrock scores grounding only when source, query and response are all present, - and rejects a source without a query, so return nothing unless both exist.""" if not messages: return [] latest_user_index: Final = self._find_latest_message_index(messages, target_role="user") diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 16369abbfb0..7858adeb55d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -23,6 +23,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): prompt_attack_threshold=litellm_params.prompt_attack_threshold, pii_confidence_threshold=litellm_params.pii_confidence_threshold, chunk_budget_chars=litellm_params.chunk_budget_chars, + contextual_grounding_from_messages=litellm_params.contextual_grounding_from_messages, default_on=litellm_params.default_on, disable_exception_on_block=litellm_params.disable_exception_on_block, mask_request_content=litellm_params.mask_request_content, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 69cb88bfa2f..ab2f9edea1d 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -552,6 +552,16 @@ class BedrockGuardrailConfigModel(BaseModel): "still rejects is bisected automatically, so this value only trades round trips against " "batch size and cannot fail a request on its own.", ) + contextual_grounding_from_messages: bool = Field( + default=False, + description="ApplyGuardrail: when True, post-call scans of a request with no grounding_source / " + "query content parts send the system and developer messages as the grounding source and " + "the latest user message as the query, so the guardrail's contextual grounding policy can " + "score the response. Bedrock bills contextual grounding units for these scans and rejects " + "queries, sources and responses over its contextual grounding length limits, so leave this " + "off for guardrails without a contextual grounding policy. Default False: plain messages " + "are never sent as grounding context.", + ) class BedrockGuardrailStreamingParams(BaseModel): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index a9a208e1807..d14ff34e0f4 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2375,8 +2375,12 @@ _GROUNDING_QUERY_TEXT = "What is the capital of Japan?" _GROUNDING_RESPONSE_TEXT = "The capital of Japan is Tokyo." -def _grounding_guardrail() -> BedrockGuardrail: - return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") +def _grounding_guardrail(from_messages: bool = False) -> BedrockGuardrail: + return BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + contextual_grounding_from_messages=from_messages, + ) def _grounding_messages() -> list: @@ -2418,9 +2422,11 @@ def _input_request(messages: list) -> dict: return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) -def _output_request(messages: list, response=None) -> dict: +def _output_request(messages: list, response=None, from_messages: bool = False) -> dict: """Arrange a guardrail and act: build the Bedrock OUTPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) + return _grounding_guardrail(from_messages).convert_to_bedrock_format( + source="OUTPUT", response=response, messages=messages + ) def test_grounding_input_strips_grounding_and_query_qualifiers(): @@ -2475,8 +2481,9 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): def test_grounding_output_derives_source_and_query_from_plain_messages(): - """Untagged string system + user messages become grounding_source + query, so a - guardrail with a grounding threshold actually grades the response.""" + """With contextual_grounding_from_messages on, untagged string system + user + messages become grounding_source + query, so a guardrail with a grounding + threshold grades the response.""" messages = [ {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, {"role": "user", "content": _GROUNDING_QUERY_TEXT}, @@ -2486,6 +2493,19 @@ def test_grounding_output_derives_source_and_query_from_plain_messages(): "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], } + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True) + + assert actual_request == expected_request + + +def test_grounding_output_plain_messages_stay_legacy_when_flag_is_off(): + """Default config: plain system + user text is never sent as grounding context.""" + messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]} + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2513,7 +2533,7 @@ def test_grounding_output_derived_query_is_latest_user_turn_only(): ], } - actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True) assert actual_request == expected_request @@ -2550,7 +2570,7 @@ def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(mes request that cannot supply both from trusted roles keeps the untagged payload.""" expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]} - actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True) assert actual_request == expected_request @@ -2568,7 +2588,7 @@ def test_grounding_output_explicit_tags_take_precedence_over_plain_messages(): "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], } - actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True) assert actual_request == expected_request @@ -2585,7 +2605,9 @@ def test_grounding_input_ignores_plain_message_derivation(): "content": [{"text": {"text": _GROUNDING_SOURCE_TEXT}}, {"text": {"text": _GROUNDING_QUERY_TEXT}}], } - actual_request = _input_request(messages) + actual_request = _grounding_guardrail(from_messages=True).convert_to_bedrock_format( + source="INPUT", messages=messages + ) assert actual_request == expected_request @@ -2715,11 +2737,7 @@ async def test_grounding_output_blocked_raises_400(): @pytest.mark.asyncio async def test_apply_guardrail_response_forwards_request_messages_for_grounding(): - """/guardrails/apply_guardrail with input_type=response: the request messages - stored in request_data must reach the OUTPUT payload as grounding_source + query - around the guarded text. Before LIT-4224 the response branch dropped them, so - Bedrock never ran its contextual-grounding policy on this route.""" - guardrail = _grounding_guardrail() + guardrail = _grounding_guardrail(from_messages=True) request_messages = [ {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, {"role": "user", "content": _GROUNDING_QUERY_TEXT}, diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index ceb084b4a4d..8377db57b6e 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -71,6 +71,52 @@ def test_initialize_bedrock_forwards_chunk_budget_chars(): assert initialized[-1].chunk_budget_chars == 60_000 +def test_initialize_bedrock_forwards_contextual_grounding_from_messages(): + """`contextual_grounding_from_messages: true` in config.yaml must make the post-call + payload carry the plain system prompt and user turn as grounding_source and query.""" + import litellm + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + from litellm.types.utils import Choices, Message, ModelResponse + + test_guardrail = { + "guardrail_name": "test_bedrock_grounding_from_messages", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.BEDROCK.value, + "mode": "post_call", + "guardrailIdentifier": "test-guardrail", + "guardrailVersion": "DRAFT", + "contextual_grounding_from_messages": True, + }, + } + messages = [ + {"role": "system", "content": "Returns are accepted for 30 days."}, + {"role": "user", "content": "How long is the return window?"}, + ] + response = ModelResponse( + choices=[Choices(index=0, message=Message(role="assistant", content="30 days."), finish_reason="stop")] + ) + expected_request = { + "source": "OUTPUT", + "content": [ + {"text": {"text": "Returns are accepted for 30 days.", "qualifiers": ["grounding_source"]}}, + {"text": {"text": "How long is the return window?", "qualifiers": ["query"]}}, + {"text": {"text": "30 days.", "qualifiers": ["guard_content"]}}, + ], + } + + guardrail_handler = InMemoryGuardrailHandler() + guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + initialized = [ + callback + for callback in litellm.callbacks + if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_grounding_from_messages" + ] + assert initialized, "bedrock guardrail was not registered as a callback" + actual_request = initialized[-1].convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) + assert json.loads(json.dumps(actual_request)) == expected_request + + def test_initialize_guardrail_preserves_guardrail_info(): """ Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0b0e3e18215..d76b256f752 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30875,6 +30875,12 @@ export interface components { * @description Enable content moderation to check for harmful content (harassment, hate speech, etc.). */ content_moderation_check?: boolean | null; + /** + * Contextual Grounding From Messages + * @description ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context. + * @default false + */ + contextual_grounding_from_messages: boolean; /** * Credentials * @description Path to Google Cloud credentials JSON file or JSON string From 2e11f7bc7b8eef524418dcf00305b2d80a556f71 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:09:59 +0000 Subject: [PATCH 20/43] test(bedrock guardrails): shorten grounding test docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/test_bedrock_guardrails.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index d14ff34e0f4..81bbaf86e46 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2481,9 +2481,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): def test_grounding_output_derives_source_and_query_from_plain_messages(): - """With contextual_grounding_from_messages on, untagged string system + user - messages become grounding_source + query, so a guardrail with a grounding - threshold grades the response.""" + """Flag on: untagged system + user text is sent as grounding_source + query.""" messages = [ {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, {"role": "user", "content": _GROUNDING_QUERY_TEXT}, @@ -2512,9 +2510,7 @@ def test_grounding_output_plain_messages_stay_legacy_when_flag_is_off(): def test_grounding_output_derived_query_is_latest_user_turn_only(): - """In a multi-turn chat only the latest user message is the query; earlier user - turns and assistant turns are not sent as query or source. Developer messages - count as source alongside system.""" + """Only the latest user turn is the query; system and developer turns are the source.""" developer_text = "Answer in one sentence." messages = [ {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, @@ -2566,8 +2562,7 @@ def test_grounding_output_derived_query_is_latest_user_turn_only(): ], ) def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(messages): - """Bedrock rejects a grounding_source without a query (and vice versa), so a - request that cannot supply both from trusted roles keeps the untagged payload.""" + """Bedrock rejects a source without a query and vice versa, so send neither.""" expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]} actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True) @@ -2576,8 +2571,7 @@ def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(mes def test_grounding_output_explicit_tags_take_precedence_over_plain_messages(): - """A caller that tags blocks keeps full control: untagged system text is not - added as a second source and the untagged user text is not a second query.""" + """Tagged blocks win: untagged text around them is not added as source or query.""" messages = [ {"role": "system", "content": "You are a helpful assistant."}, *_grounding_messages(), @@ -2594,8 +2588,7 @@ def test_grounding_output_explicit_tags_take_precedence_over_plain_messages(): def test_grounding_input_ignores_plain_message_derivation(): - """Derivation is OUTPUT-only: an INPUT scan of plain system + user text stays an - untagged payload, so input policies keep scanning every block.""" + """INPUT scans never derive grounding qualifiers from plain messages.""" messages = [ {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, {"role": "user", "content": _GROUNDING_QUERY_TEXT}, From d8d7210b5734c521a1d0fd0f651b620785492515 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:28:58 +0000 Subject: [PATCH 21/43] test(bedrock guardrails): cover tagged messages on the apply_guardrail response path with the flag off Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_bedrock_guardrails.py | 28 +++++++++++++++---- 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 81bbaf86e46..d173c5f5c70 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2729,12 +2729,28 @@ async def test_grounding_output_blocked_raises_400(): @pytest.mark.asyncio -async def test_apply_guardrail_response_forwards_request_messages_for_grounding(): - guardrail = _grounding_guardrail(from_messages=True) - request_messages = [ - {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, - {"role": "user", "content": _GROUNDING_QUERY_TEXT}, - ] +@pytest.mark.parametrize( + "from_messages, request_messages", + [ + ( + True, + [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ], + ), + ( + False, + [ + {"role": "system", "content": [{"type": "grounding_source", "text": _GROUNDING_SOURCE_TEXT}]}, + {"role": "user", "content": [{"type": "query", "text": _GROUNDING_QUERY_TEXT}]}, + ], + ), + ], + ids=["plain-messages-flag-on", "tagged-messages-flag-off"], +) +async def test_apply_guardrail_response_forwards_request_messages_for_grounding(from_messages, request_messages): + guardrail = _grounding_guardrail(from_messages=from_messages) expected_request = { "source": "OUTPUT", "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], From d4f2119b03faa175e790dd86cb3c8aa46f546293 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:41:46 +0000 Subject: [PATCH 22/43] fix(cost): bill gemini-embedding-2 per token and stop double charging audio Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/utils.py | 14 ++++- .../batch_embed_content_transformation.py | 60 ++----------------- ...odel_prices_and_context_window_backup.json | 15 ++--- model_prices_and_context_window.json | 15 ++--- .../llm_cost_calc/test_llm_cost_calc_utils.py | 33 ++++++++++ ...test_batch_embed_content_transformation.py | 42 +++++-------- 6 files changed, 75 insertions(+), 104 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 8fc428b38ae..a004f46b291 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -956,12 +956,17 @@ def _calculate_input_cost( ) ### AUDIO COST - if prompt_tokens_details["audio_tokens"]: + if prompt_tokens_details["audio_tokens"] and not ( + prompt_tokens_details["audio_length_seconds"] + and model_info.get("input_cost_per_audio_per_second") is not None + ): audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier) prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"]) ### IMAGE TOKEN COST - if prompt_tokens_details["image_tokens"]: + if prompt_tokens_details["image_tokens"] and not ( + prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None + ): # For image token costs: # First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token. image_token_cost_key = "input_cost_per_image_token" @@ -970,7 +975,10 @@ def _calculate_input_cost( prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"]) ### VIDEO TOKEN COST - if prompt_tokens_details["video_tokens"]: + if prompt_tokens_details["video_tokens"] and not ( + prompt_tokens_details["video_length_seconds"] + and model_info.get("input_cost_per_video_per_second") is not None + ): video_token_cost_key = "input_cost_per_video_token" if model_info.get(video_token_cost_key) is None: video_token_cost_key = "input_cost_per_token" diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index e7fd9a0d08b..8e120ab9fe6 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -297,9 +297,6 @@ def transform_openai_input_gemini_embed_content( return request_body -_IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"}) -_VIDEO_TOKENS_PER_SECOND: Final = 258.0 -_AUDIO_TOKENS_PER_SECOND: Final = 32.0 _usage_metadata_adapter: Final = TypeAdapter(UsageMetadata) @@ -312,40 +309,6 @@ def _parse_usage_metadata(raw_usage_metadata: object) -> UsageMetadata | None: return None -def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: - if isinstance(input, str): - return (input,) - return tuple(sub for element in input for sub in (element if isinstance(element, list) else [element])) - - -def _is_image_element( - element: str, - resolved_files: Mapping[str, Mapping[str, str]], -) -> bool: - if element.startswith("data:") and ";base64," in element: - try: - mime_type, _ = _parse_data_url(element) - except ValueError: - return False - return mime_type in _IMAGE_MIME_TYPES - if _is_gcs_url(element): - try: - return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES - except ValueError: - return False - if _is_file_reference(element): - file_info: Final = resolved_files.get(element) - return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES - return False - - -def _count_input_images( - input: GeminiEmbeddingInput, - resolved_files: Mapping[str, Mapping[str, str]], -) -> int: - return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files)) - - def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: return sum(detail["tokenCount"] for detail in details if detail["modality"] == modality) @@ -362,7 +325,6 @@ def _usage_from_embed_content_response( input: GeminiEmbeddingInput, model: str, raw_usage_metadata: object, - resolved_files: Mapping[str, Mapping[str, str]], ) -> Usage: usage_metadata: Final = _parse_usage_metadata(raw_usage_metadata) if usage_metadata is None: @@ -374,28 +336,17 @@ def _usage_from_embed_content_response( details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or () text_tokens: Final = _tokens_for_modality(details, "TEXT") audio_tokens: Final = _tokens_for_modality(details, "AUDIO") + image_tokens: Final = _tokens_for_modality(details, "IMAGE") video_tokens: Final = _tokens_for_modality(details, "VIDEO") - image_count: Final = _count_input_images(input, resolved_files) - - video_length_seconds: Final = video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0 - audio_length_seconds: Final = audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0 - - # generic_cost_per_token rewrites text_tokens to the full prompt minus - # other modalities when both text_tokens and image_count are zero. For - # video, that misallocates video tokens to text; a 1-token floor sidesteps - # the rewrite and keeps billing on input_cost_per_video_per_second. - needs_video_text_floor: Final = video_length_seconds > 0 and text_tokens == 0 and image_count == 0 - resolved_text_tokens: Final = 1 if needs_video_text_floor else text_tokens return Usage( prompt_tokens=prompt_tokens, total_tokens=total_tokens, prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=resolved_text_tokens, + text_tokens=text_tokens, audio_tokens=audio_tokens, - image_count=image_count, - video_length_seconds=video_length_seconds, - audio_length_seconds=audio_length_seconds, + image_tokens=image_tokens, + video_tokens=video_tokens, ), ) @@ -415,8 +366,6 @@ def process_embed_content_response( model_response: EmbeddingResponse to populate model: Model name response_json: Raw JSON response from embedContent endpoint - resolved_files: Mapping of file references (files/abc) to {mime_type, uri}, - used to bill resolved image references at the per-image rate Returns: EmbeddingResponse with single embedding @@ -438,7 +387,6 @@ def process_embed_content_response( input=input, model=model, raw_usage_metadata=response_json.get("usageMetadata"), - resolved_files=resolved_files or {}, ) return model_response diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0c60ff26635..74ca23fcc8a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25615,13 +25615,11 @@ "uses_embed_content": true }, "gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_token": 6.5e-06, - "input_cost_per_image": 0.00012, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25648,13 +25646,11 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_token": 6.5e-06, - "input_cost_per_image": 0.00012, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25709,10 +25705,11 @@ "tpm": 10000000 }, "gemini/gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0c60ff26635..74ca23fcc8a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25615,13 +25615,11 @@ "uses_embed_content": true }, "gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_token": 6.5e-06, - "input_cost_per_image": 0.00012, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25648,13 +25646,11 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_token": 6.5e-06, - "input_cost_per_image": 0.00012, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25709,10 +25705,11 @@ "tpm": 10000000 }, "gemini/gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 2289de9a951..4a02bb1a638 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -74,6 +74,39 @@ def test_missing_cache_read_policy_preserves_billing(prompt_tokens, read_rate, s assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4]) +def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: + model_info: ModelInfo = { + "key": "gemini-embedding-2", + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": 2e-7, + "input_cost_per_audio_token": 6.5e-6, + "input_cost_per_audio_per_second": 0.00016, + "output_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "embedding", + } + usage = Usage( + prompt_tokens=64, + completion_tokens=0, + total_tokens=64, + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=64, + audio_length_seconds=2, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-embedding-2", + usage=usage, + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(2 * 0.00016) + + def test_missing_cache_read_uses_off_peak_input_rate(): from datetime import datetime, timezone diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 86b3f0976ab..fbf86105e71 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -22,7 +22,6 @@ from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject from litellm.types.utils import EmbeddingResponse - IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" GCS_URL = "gs://my-bucket/image.png" @@ -324,7 +323,7 @@ class TestProcessEmbedContentResponseUsage: ) assert result.usage.prompt_tokens == 258 assert result.usage.total_tokens == 258 - assert result.usage.prompt_tokens_details.image_count == 1 + assert result.usage.prompt_tokens_details.image_tokens == 258 prompt_cost, _ = generic_cost_per_token( model=self.MODEL, @@ -358,7 +357,7 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost > 0 - def test_video_modality_derives_seconds_and_text_floor(self): + def test_video_modality_preserves_token_count(self): response_json = { "embedding": {"values": [0.1]}, "usageMetadata": { @@ -374,10 +373,8 @@ class TestProcessEmbedContentResponseUsage: response_json=response_json, ) assert result.usage.prompt_tokens == 516 - assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( - 2.0 - ) - assert result.usage.prompt_tokens_details.text_tokens == 1 + assert result.usage.prompt_tokens_details.video_tokens == 516 + assert result.usage.prompt_tokens_details.text_tokens == 0 def test_missing_usage_metadata_does_not_estimate_from_base64(self): response_json = {"embedding": {"values": [0.1, 0.2]}} @@ -400,8 +397,7 @@ class TestProcessEmbedContentResponseUsage: ) assert result.usage.prompt_tokens > 0 - def test_file_reference_image_billed_per_image_not_text(self): - """files/... image refs must bill per-image, not at the text token rate.""" + def test_file_reference_image_billed_per_image_token_rate(self): response_json = { "embedding": {"values": [0.1, 0.2, 0.3]}, "usageMetadata": { @@ -422,7 +418,7 @@ class TestProcessEmbedContentResponseUsage: } }, ) - assert result.usage.prompt_tokens_details.image_count == 1 + assert result.usage.prompt_tokens_details.image_tokens == 258 assert result.usage.prompt_tokens_details.text_tokens == 0 prompt_cost, _ = generic_cost_per_token( @@ -430,10 +426,10 @@ class TestProcessEmbedContentResponseUsage: usage=result.usage, custom_llm_provider="vertex_ai", ) - assert prompt_cost == pytest.approx(0.00012) + assert prompt_cost == pytest.approx(258 * 4.5e-7) def test_file_reference_non_image_not_counted_as_image(self): - """A files/... ref resolving to a non-image mime must not be image-counted.""" + """A files/... ref resolving to a non-image mime keeps audio token billing.""" response_json = { "embedding": {"values": [0.1, 0.2]}, "usageMetadata": { @@ -454,21 +450,18 @@ class TestProcessEmbedContentResponseUsage: } }, ) - assert result.usage.prompt_tokens_details.image_count == 0 assert result.usage.prompt_tokens_details.audio_tokens == 64 - assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( - 2.0 - ) + assert result.usage.prompt_tokens_details.image_tokens == 0 prompt_cost, _ = generic_cost_per_token( model=self.MODEL, usage=result.usage, custom_llm_provider="vertex_ai", ) - assert prompt_cost == pytest.approx(2.0 * 0.00016) + assert prompt_cost == pytest.approx(64 * 6.5e-6) def test_video_plus_audio_does_not_double_bill_text(self): - """Video+audio responses must not get video tokens reassigned to text.""" + """Video and audio responses are billed from their respective token counts.""" response_json = { "embedding": {"values": [0.1]}, "usageMetadata": { @@ -486,18 +479,13 @@ class TestProcessEmbedContentResponseUsage: model=self.MODEL, response_json=response_json, ) - assert result.usage.prompt_tokens_details.text_tokens == 1 - assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( - 2.0 - ) - assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( - 2.0 - ) + assert result.usage.prompt_tokens_details.text_tokens == 0 + assert result.usage.prompt_tokens_details.video_tokens == 516 + assert result.usage.prompt_tokens_details.audio_tokens == 64 prompt_cost, _ = generic_cost_per_token( model=self.MODEL, usage=result.usage, custom_llm_provider="vertex_ai", ) - # 1 floor text token at 2e-7 + 2s of video at 7.9e-4 + 2s of audio at 1.6e-4 - assert prompt_cost == pytest.approx(1 * 2e-7 + 2 * 0.00079 + 2 * 0.00016) + assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6) From e26a4970dd0ba5efe277a5f42b51854daaf5da6d Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:42:21 +0000 Subject: [PATCH 23/43] fix(test): complete synthetic model metadata Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 4a02bb1a638..97d04a03a78 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -86,6 +86,7 @@ def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: "output_cost_per_token": 0.0, "litellm_provider": "vertex_ai", "mode": "embedding", + "supported_openai_params": None, } usage = Usage( prompt_tokens=64, From 0c91d9157c43ba7728b58393b6088641aa824367 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:45:08 +0000 Subject: [PATCH 24/43] refactor(vertex): drop unused resolved_files from embed response parsing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../gemini_embeddings/batch_embed_content_handler.py | 2 -- .../batch_embed_content_transformation.py | 3 +-- .../test_batch_embed_content_transformation.py | 12 ------------ 3 files changed, 1 insertion(+), 16 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index f81d4ca777e..e09622ba236 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -268,7 +268,6 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, - resolved_files=resolved_files, ) else: _predictions: Final = VertexAIBatchEmbeddingsResponseObject(**_json_response) @@ -372,7 +371,6 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, - resolved_files=resolved_files, ) else: _predictions: Final = VertexAIBatchEmbeddingsResponseObject(**_json_response) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 8e120ab9fe6..f2cce775f3f 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,7 +4,7 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import Final from pydantic import TypeAdapter, ValidationError @@ -356,7 +356,6 @@ def process_embed_content_response( model_response: EmbeddingResponse, model: str, response_json: dict, - resolved_files: Mapping[str, Mapping[str, str]] | None = None, ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index fbf86105e71..0251a799b66 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -411,12 +411,6 @@ class TestProcessEmbedContentResponseUsage: model_response=EmbeddingResponse(), model=self.MODEL, response_json=response_json, - resolved_files={ - "files/img123": { - "mime_type": "image/png", - "uri": "https://example.com/img123", - } - }, ) assert result.usage.prompt_tokens_details.image_tokens == 258 assert result.usage.prompt_tokens_details.text_tokens == 0 @@ -443,12 +437,6 @@ class TestProcessEmbedContentResponseUsage: model_response=EmbeddingResponse(), model=self.MODEL, response_json=response_json, - resolved_files={ - "files/clip1": { - "mime_type": "audio/mpeg", - "uri": "https://example.com/clip1", - } - }, ) assert result.usage.prompt_tokens_details.audio_tokens == 64 assert result.usage.prompt_tokens_details.image_tokens == 0 From 6c9fe65608997dbfa85c35d01ad196f8b97c9a9d Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:47:30 +0000 Subject: [PATCH 25/43] style(cost): apply ruff formatting to modality guards Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/llm_cost_calc/utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index a004f46b291..88ea4b602cc 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -957,8 +957,7 @@ def _calculate_input_cost( ### AUDIO COST if prompt_tokens_details["audio_tokens"] and not ( - prompt_tokens_details["audio_length_seconds"] - and model_info.get("input_cost_per_audio_per_second") is not None + prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None ): audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier) prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"]) @@ -976,8 +975,7 @@ def _calculate_input_cost( ### VIDEO TOKEN COST if prompt_tokens_details["video_tokens"] and not ( - prompt_tokens_details["video_length_seconds"] - and model_info.get("input_cost_per_video_per_second") is not None + prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None ): video_token_cost_key = "input_cost_per_video_token" if model_info.get(video_token_cost_key) is None: From e5845c17ffde232ee1e460b648fb223ea4561348 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:51:49 +0000 Subject: [PATCH 26/43] fix(vertex): bill image inputs at the image rate when usage lacks modality details Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../batch_embed_content_handler.py | 2 + .../batch_embed_content_transformation.py | 52 +++++++++++++++- ...test_batch_embed_content_transformation.py | 60 +++++++++++++++++++ 3 files changed, 113 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index e09622ba236..f81d4ca777e 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -268,6 +268,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions: Final = VertexAIBatchEmbeddingsResponseObject(**_json_response) @@ -371,6 +372,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions: Final = VertexAIBatchEmbeddingsResponseObject(**_json_response) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index f2cce775f3f..b61fb47cf5c 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,7 +4,7 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final from pydantic import TypeAdapter, ValidationError @@ -297,6 +297,7 @@ def transform_openai_input_gemini_embed_content( return request_body +_IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"}) _usage_metadata_adapter: Final = TypeAdapter(UsageMetadata) @@ -309,6 +310,40 @@ def _parse_usage_metadata(raw_usage_metadata: object) -> UsageMetadata | None: return None +def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: + if isinstance(input, str): + return (input,) + return tuple(sub for element in input for sub in (element if isinstance(element, list) else [element])) + + +def _is_image_element( + element: str, + resolved_files: Mapping[str, Mapping[str, str]], +) -> bool: + if element.startswith("data:") and ";base64," in element: + try: + mime_type, _ = _parse_data_url(element) + except ValueError: + return False + return mime_type in _IMAGE_MIME_TYPES + if _is_gcs_url(element): + try: + return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES + except ValueError: + return False + if _is_file_reference(element): + file_info: Final = resolved_files.get(element) + return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES + return False + + +def _count_input_images( + input: GeminiEmbeddingInput, + resolved_files: Mapping[str, Mapping[str, str]], +) -> int: + return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files)) + + def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: return sum(detail["tokenCount"] for detail in details if detail["modality"] == modality) @@ -325,6 +360,7 @@ def _usage_from_embed_content_response( input: GeminiEmbeddingInput, model: str, raw_usage_metadata: object, + resolved_files: Mapping[str, Mapping[str, str]], ) -> Usage: usage_metadata: Final = _parse_usage_metadata(raw_usage_metadata) if usage_metadata is None: @@ -334,6 +370,17 @@ def _usage_from_embed_content_response( total_tokens: Final = usage_metadata.get("totalTokenCount") or prompt_tokens details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or () + if not details: + image_tokens: Final = prompt_tokens if _count_input_images(input, resolved_files) else 0 + return Usage( + prompt_tokens=prompt_tokens, + total_tokens=total_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=0, + image_tokens=image_tokens, + ), + ) + text_tokens: Final = _tokens_for_modality(details, "TEXT") audio_tokens: Final = _tokens_for_modality(details, "AUDIO") image_tokens: Final = _tokens_for_modality(details, "IMAGE") @@ -356,6 +403,7 @@ def process_embed_content_response( model_response: EmbeddingResponse, model: str, response_json: dict, + resolved_files: Mapping[str, Mapping[str, str]] | None = None, ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). @@ -365,6 +413,7 @@ def process_embed_content_response( model_response: EmbeddingResponse to populate model: Model name response_json: Raw JSON response from embedContent endpoint + resolved_files: Mapping of file references to resolved metadata Returns: EmbeddingResponse with single embedding @@ -386,6 +435,7 @@ def process_embed_content_response( input=input, model=model, raw_usage_metadata=response_json.get("usageMetadata"), + resolved_files=resolved_files or {}, ) return model_response diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 0251a799b66..df5903b9285 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -411,6 +411,12 @@ class TestProcessEmbedContentResponseUsage: model_response=EmbeddingResponse(), model=self.MODEL, response_json=response_json, + resolved_files={ + "files/img123": { + "mime_type": "image/png", + "uri": "https://example.com/img123", + } + }, ) assert result.usage.prompt_tokens_details.image_tokens == 258 assert result.usage.prompt_tokens_details.text_tokens == 0 @@ -437,6 +443,12 @@ class TestProcessEmbedContentResponseUsage: model_response=EmbeddingResponse(), model=self.MODEL, response_json=response_json, + resolved_files={ + "files/clip1": { + "mime_type": "audio/mpeg", + "uri": "https://example.com/clip1", + } + }, ) assert result.usage.prompt_tokens_details.audio_tokens == 64 assert result.usage.prompt_tokens_details.image_tokens == 0 @@ -477,3 +489,51 @@ class TestProcessEmbedContentResponseUsage: custom_llm_provider="vertex_ai", ) assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6) + + def test_image_without_modality_details_uses_image_rate(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + }, + } + result = process_embed_content_response( + input=IMAGE_DATA_URI, + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.image_tokens == 258 + assert result.usage.prompt_tokens_details.text_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(258 * 4.5e-7) + + def test_text_without_modality_details_uses_text_rate(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 12, + "totalTokenCount": 12, + }, + } + result = process_embed_content_response( + input="a short caption", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.text_tokens == 0 + assert result.usage.prompt_tokens_details.image_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(12 * 2e-7) From e31c64d2e038fc0895cf02e45e2c6a798bba3cf3 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:54:58 +0000 Subject: [PATCH 27/43] fix(schema): sync model price schema with cost map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- model_prices_and_context_window.schema.json | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index d1ac3e67b2b..eff3f192b3c 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -375,6 +375,10 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "input_cost_per_video_token": { + "type": "number", + "minimum": 0 + }, "input_dbu_cost_per_token": { "type": "number", "minimum": 0 From ac8e1a355cf69830ec989fd19afb42a2bef78efd Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:00:32 +0000 Subject: [PATCH 28/43] test(vertex): load local pricing in embedding billing tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_batch_embed_content_transformation.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index df5903b9285..49f5167fd6b 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -10,6 +10,7 @@ Covers: import pytest +import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( _build_part_for_input, @@ -26,6 +27,15 @@ IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+ GCS_URL = "gs://my-bucket/image.png" +@pytest.fixture(autouse=True) +def _local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + class TestIsMultimodalInput: def test_text_only_string(self): assert _is_multimodal_input("hello world") is False From 6cbed7b4c0ae643a429b0ebc7cb85a99e52b5b9e Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:10:27 +0000 Subject: [PATCH 29/43] fix(vertex): drop Final image_tokens redeclaration flagged by basedpyright Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../gemini_embeddings/batch_embed_content_transformation.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index b61fb47cf5c..b618f6e5165 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -371,13 +371,12 @@ def _usage_from_embed_content_response( details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or () if not details: - image_tokens: Final = prompt_tokens if _count_input_images(input, resolved_files) else 0 return Usage( prompt_tokens=prompt_tokens, total_tokens=total_tokens, prompt_tokens_details=PromptTokensDetailsWrapper( text_tokens=0, - image_tokens=image_tokens, + image_tokens=prompt_tokens if _count_input_images(input, resolved_files) else 0, ), ) From 6a18105275223a39170e24d8fec96d123af82b32 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:16:00 +0000 Subject: [PATCH 30/43] fix(vertex): only bill image rate without modality details when every input is an image Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../batch_embed_content_transformation.py | 9 ++++---- ...test_batch_embed_content_transformation.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index b618f6e5165..d669acecfd9 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -337,11 +337,12 @@ def _is_image_element( return False -def _count_input_images( +def _is_image_only_input( input: GeminiEmbeddingInput, resolved_files: Mapping[str, Mapping[str, str]], -) -> int: - return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files)) +) -> bool: + elements: Final = _flatten_input(input) + return bool(elements) and all(_is_image_element(element, resolved_files) for element in elements) def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: @@ -376,7 +377,7 @@ def _usage_from_embed_content_response( total_tokens=total_tokens, prompt_tokens_details=PromptTokensDetailsWrapper( text_tokens=0, - image_tokens=prompt_tokens if _count_input_images(input, resolved_files) else 0, + image_tokens=prompt_tokens if _is_image_only_input(input, resolved_files) else 0, ), ) diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 49f5167fd6b..5dfeac6f469 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -524,6 +524,29 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost == pytest.approx(258 * 4.5e-7) + def test_mixed_text_and_image_without_modality_details_not_billed_as_image(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 270, + "totalTokenCount": 270, + }, + } + result = process_embed_content_response( + input=["a short caption", IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.image_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(270 * 2e-7) + def test_text_without_modality_details_uses_text_rate(self): response_json = { "embedding": {"values": [0.1]}, From 4a8ec7b9d8c60896b28448b3bd87380617692d74 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:18:50 +0000 Subject: [PATCH 31/43] fix(cost): bill gemini-embedding-2-preview per token like the GA entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 18 +++++++-------- model_prices_and_context_window.json | 18 +++++++-------- ...test_batch_embed_content_transformation.py | 22 +++++++++++++++++++ 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 74ca23fcc8a..9299a1493f8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25601,10 +25601,10 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, "gemini-embedding-2-preview": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25631,10 +25631,10 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2-preview": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25689,10 +25689,10 @@ }, "gemini/gemini-embedding-2-preview": { "deprecation_date": "2026-08-10", - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 74ca23fcc8a..9299a1493f8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25601,10 +25601,10 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, "gemini-embedding-2-preview": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25631,10 +25631,10 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2-preview": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25689,10 +25689,10 @@ }, "gemini/gemini-embedding-2-preview": { "deprecation_date": "2026-08-10", - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 5dfeac6f469..926570d7929 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -500,6 +500,28 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6) + def test_preview_alias_bills_audio_per_token(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 64, + "totalTokenCount": 64, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}], + }, + } + result = process_embed_content_response( + input="audio", + model_response=EmbeddingResponse(), + model="gemini-embedding-2-preview", + response_json=response_json, + ) + prompt_cost, _ = generic_cost_per_token( + model="gemini-embedding-2-preview", + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(64 * 6.5e-6) + def test_image_without_modality_details_uses_image_rate(self): response_json = { "embedding": {"values": [0.1]}, From 27a486e4d320b4481c8aad6062f0c518d1c52ea4 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:34:17 +0000 Subject: [PATCH 32/43] test(cost): cover modality guards and image detection fallbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 68 +++++++++++++++++++ ...test_batch_embed_content_transformation.py | 39 +++++++++++ 2 files changed, 107 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 97d04a03a78..854bc9bbb81 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -108,6 +108,74 @@ def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: assert prompt_cost == pytest.approx(2 * 0.00016) +def test_generic_cost_per_token_prefers_image_per_image_rate() -> None: + model_info: ModelInfo = { + "key": "gemini-embedding-2", + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": 2e-7, + "input_cost_per_image_token": 4.5e-7, + "input_cost_per_image": 0.00012, + "output_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "embedding", + "supported_openai_params": None, + } + usage = Usage( + prompt_tokens=258, + completion_tokens=0, + total_tokens=258, + prompt_tokens_details=PromptTokensDetailsWrapper( + image_tokens=258, + image_count=1, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-embedding-2", + usage=usage, + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(0.00012) + + +def test_generic_cost_per_token_prefers_video_per_second_rate() -> None: + model_info: ModelInfo = { + "key": "gemini-embedding-2", + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": 2e-7, + "input_cost_per_video_token": 1.2e-5, + "input_cost_per_video_per_second": 0.00079, + "output_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "embedding", + "supported_openai_params": None, + } + usage = Usage( + prompt_tokens=516, + completion_tokens=0, + total_tokens=516, + prompt_tokens_details=PromptTokensDetailsWrapper( + video_tokens=516, + video_length_seconds=2, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-embedding-2", + usage=usage, + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(2 * 0.00079) + + def test_missing_cache_read_uses_off_peak_input_rate(): from datetime import datetime, timezone diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 926570d7929..fd8c2a9cf6a 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -546,6 +546,45 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost == pytest.approx(258 * 4.5e-7) + @pytest.mark.parametrize( + "input_value,resolved_files,expected_image_tokens", + [ + (GCS_URL, {}, 258), + ("gs://my-bucket/clip.mp4", {}, 0), + ("gs://my-bucket/unknown.bin", {}, 0), + ("files/image-123", {"files/image-123": {"mime_type": "image/jpeg"}}, 258), + ("files/missing", {}, 0), + ("data:application/octet-stream;base64,abc", {}, 0), + ([[IMAGE_DATA_URI]], {}, 258), + ([], {}, 0), + ], + ) + def test_missing_modality_details_classifies_image_inputs(self, input_value, resolved_files, expected_image_tokens): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + }, + } + result = process_embed_content_response( + input=input_value, + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files=resolved_files, + ) + assert result.usage.prompt_tokens_details.image_tokens == expected_image_tokens + assert result.usage.prompt_tokens_details.text_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + expected_rate = 4.5e-7 if expected_image_tokens else 2e-7 + assert prompt_cost == pytest.approx(258 * expected_rate) + def test_mixed_text_and_image_without_modality_details_not_billed_as_image(self): response_json = { "embedding": {"values": [0.1]}, From a28ea22ec131d1ce9f47af4dceb1faf7df7aa2f2 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:34:22 +0000 Subject: [PATCH 33/43] fix(cost): move gemini-embedding-2-preview to per-token rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 3 +++ model_prices_and_context_window.json | 3 +++ tests/test_litellm/test_utils.py | 11 +++++++---- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9299a1493f8..29243832c45 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25604,6 +25604,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, @@ -25634,6 +25635,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, @@ -25692,6 +25694,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9299a1493f8..29243832c45 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25604,6 +25604,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, @@ -25634,6 +25635,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, @@ -25692,6 +25694,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 04d2d35e05f..1da53fba923 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2946,7 +2946,7 @@ def test_model_info_for_openrouter_kimi_k2_5(): def test_gemini_embedding_2_ga_in_cost_map(): - """GA and Vertex preview gemini-embedding-2 entries align with multimodal unit pricing.""" + """GA and Vertex preview gemini-embedding-2 entries align with multimodal token pricing.""" import json from pathlib import Path @@ -2968,9 +2968,12 @@ def test_gemini_embedding_2_ga_in_cost_map(): assert info.get("mode") == "embedding" assert info.get("supports_multimodal") is True assert info.get("input_cost_per_token") == 2e-07 - assert info.get("input_cost_per_image") == 0.00012 - assert info.get("input_cost_per_audio_per_second") == 0.00016 - assert info.get("input_cost_per_video_per_second") == 0.00079 + assert info.get("input_cost_per_audio_token") == 6.5e-06 + assert info.get("input_cost_per_image_token") == 4.5e-07 + assert info.get("input_cost_per_video_token") == 1.2e-05 + assert "input_cost_per_image" not in info + assert "input_cost_per_audio_per_second" not in info + assert "input_cost_per_video_per_second" not in info if provider in ("vertex_ai-embedding-models", "vertex_ai"): assert ( info.get("uses_embed_content") is True From ce83fac3515c36c927ed133fe48abd7dc1a3ee74 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:50:21 +0000 Subject: [PATCH 34/43] fix(cost): bill batch embeddings per modality token rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/generate_model_prices_schema.py | 3 ++ litellm/cost_calculator.py | 30 ++++++++++- ...odel_prices_and_context_window_backup.json | 18 +++++++ litellm/types/utils.py | 6 +++ litellm/utils.py | 3 ++ model_prices_and_context_window.json | 18 +++++++ model_prices_and_context_window.schema.json | 15 ++++++ tests/test_litellm/test_cost_calculator.py | 51 +++++++++++++++++++ tests/test_litellm/test_utils.py | 3 ++ 9 files changed, 146 insertions(+), 1 deletion(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index ab29b70bdd4..ee0e7f22eb1 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -172,7 +172,10 @@ COST_DESCRIPTIONS: dict[str, str] = { ), "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", + "input_cost_per_audio_token_batches": "USD per audio prompt token via the provider's batch API.", + "input_cost_per_image_token_batches": "USD per image prompt token via the provider's batch API.", "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", + "input_cost_per_video_token_batches": "USD per video prompt token via the provider's batch API.", "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", } diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f5319776213..cbb9a45ada9 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2337,7 +2337,35 @@ def batch_cost_calculator( total_prompt_cost = 0.0 total_completion_cost = 0.0 if input_cost_per_token_batches is not None: - total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches + batch_details: Final = parse_prompt_tokens_details(usage) + audio_tokens, image_tokens, video_tokens = ( + batch_details["audio_tokens"], + batch_details["image_tokens"], + batch_details["video_tokens"], + ) + modality_rates: Final = ( + cast(float, model_info.get("input_cost_per_audio_token_batches")) + if model_info.get("input_cost_per_audio_token_batches") is not None + else input_cost_per_token_batches, + cast(float, model_info.get("input_cost_per_image_token_batches")) + if model_info.get("input_cost_per_image_token_batches") is not None + else input_cost_per_token_batches, + cast(float, model_info.get("input_cost_per_video_token_batches")) + if model_info.get("input_cost_per_video_token_batches") is not None + else input_cost_per_token_batches, + ) + total_prompt_cost = sum( + tokens * rate + for tokens, rate in zip( + ( + max(cast(int, usage.prompt_tokens) - audio_tokens - image_tokens - video_tokens, 0), + audio_tokens, + image_tokens, + video_tokens, + ), + (input_cost_per_token_batches, *modality_rates), + ) + ) elif input_cost_per_token: details: Final = parse_prompt_tokens_details(usage) cache_read_tokens: Final = details["cache_hit_tokens"] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 29243832c45..9f91cf82f41 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25602,10 +25602,13 @@ }, "gemini-embedding-2-preview": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25617,10 +25620,13 @@ }, "gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25633,10 +25639,13 @@ }, "vertex_ai/gemini-embedding-2-preview": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25649,10 +25658,13 @@ }, "vertex_ai/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25692,10 +25704,13 @@ "gemini/gemini-embedding-2-preview": { "deprecation_date": "2026-08-10", "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25709,10 +25724,13 @@ }, "gemini/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index fbfcc678de9..2e8b20edf7a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -283,8 +283,11 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_video_token: float | None # for gemini omni models with video input input_cost_per_audio_per_second: float | None # only for vertex ai models input_cost_per_video_per_second: float | None # only for vertex ai models + input_cost_per_audio_token_batches: ReadOnly[float | None] + input_cost_per_image_token_batches: ReadOnly[float | None] input_cost_per_second: float | None # for OpenAI Speech models input_cost_per_token_batches: float | None + input_cost_per_video_token_batches: ReadOnly[float | None] output_cost_per_token_batches: float | None output_cost_per_token: Required[float | None] output_cost_per_token_flex: float | None # OpenAI flex service tier pricing @@ -3583,7 +3586,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_video_per_second_above_128k_tokens: float | None = None input_cost_per_video_per_second_above_15s_interval: float | None = None input_cost_per_video_per_second_above_8s_interval: float | None = None + input_cost_per_audio_token_batches: float | None = None + input_cost_per_image_token_batches: float | None = None input_cost_per_token_batches: float | None = None + input_cost_per_video_token_batches: float | None = None output_cost_per_token_batches: float | None = None output_cost_per_token_flex: float | None = None output_cost_per_token_priority: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index d4e3d58ba9f..b2715b41739 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5923,10 +5923,13 @@ def _get_model_info_helper( input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None), input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None), input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None), + input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None), + input_cost_per_image_token_batches=_model_info.get("input_cost_per_image_token_batches", None), input_cost_per_image=_model_info.get("input_cost_per_image", None), input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None), input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None), input_cost_per_token_batches=_model_info.get("input_cost_per_token_batches"), + input_cost_per_video_token_batches=_model_info.get("input_cost_per_video_token_batches", None), output_cost_per_token_batches=_model_info.get("output_cost_per_token_batches"), output_cost_per_token=_output_cost_per_token, output_cost_per_token_flex=_model_info.get("output_cost_per_token_flex", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 29243832c45..9f91cf82f41 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25602,10 +25602,13 @@ }, "gemini-embedding-2-preview": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25617,10 +25620,13 @@ }, "gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25633,10 +25639,13 @@ }, "vertex_ai/gemini-embedding-2-preview": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25649,10 +25658,13 @@ }, "vertex_ai/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25692,10 +25704,13 @@ "gemini/gemini-embedding-2-preview": { "deprecation_date": "2026-08-10", "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25709,10 +25724,13 @@ }, "gemini/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index eff3f192b3c..b7e0a9fd414 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -249,6 +249,11 @@ "type": "number", "minimum": 0 }, + "input_cost_per_audio_token_batches": { + "type": "number", + "minimum": 0, + "description": "USD per audio prompt token via the provider's batch API." + }, "input_cost_per_audio_token_priority": { "type": "number", "minimum": 0, @@ -276,6 +281,11 @@ "type": "number", "minimum": 0 }, + "input_cost_per_image_token_batches": { + "type": "number", + "minimum": 0, + "description": "USD per image prompt token via the provider's batch API." + }, "input_cost_per_pixel": { "type": "number", "minimum": 0 @@ -379,6 +389,11 @@ "type": "number", "minimum": 0 }, + "input_cost_per_video_token_batches": { + "type": "number", + "minimum": 0, + "description": "USD per video prompt token via the provider's batch API." + }, "input_dbu_cost_per_token": { "type": "number", "minimum": 0 diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 68e9b6143a0..8c3436d3108 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3909,6 +3909,57 @@ def _batch_cache_usage() -> Usage: ) +def test_batch_cost_calculator_prices_multimodal_tokens_at_modality_rates(): + from litellm.cost_calculator import batch_cost_calculator + + model_info: ModelInfo = { + "input_cost_per_token_batches": 1e-7, + "input_cost_per_audio_token_batches": 3.25e-6, + "input_cost_per_image_token_batches": 2.25e-7, + "input_cost_per_video_token_batches": 6e-6, + } + usage = Usage( + prompt_tokens=100, + completion_tokens=0, + total_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=64, + image_tokens=10, + video_tokens=6, + ), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-embedding-2", + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(20 * 1e-7 + 64 * 3.25e-6 + 10 * 2.25e-7 + 6 * 6e-6) + + +def test_batch_cost_calculator_falls_back_to_text_batch_rate_for_modalities(): + from litellm.cost_calculator import batch_cost_calculator + + model_info: ModelInfo = {"input_cost_per_token_batches": 1e-7} + usage = Usage( + prompt_tokens=100, + completion_tokens=0, + total_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=64), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-embedding-2", + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(100 * 1e-7) + + def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate(): """ LIT-4008 regression: anthropic batch usage is dominated by cache tokens. diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1da53fba923..02ffaee0543 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2971,6 +2971,9 @@ def test_gemini_embedding_2_ga_in_cost_map(): assert info.get("input_cost_per_audio_token") == 6.5e-06 assert info.get("input_cost_per_image_token") == 4.5e-07 assert info.get("input_cost_per_video_token") == 1.2e-05 + assert info.get("input_cost_per_audio_token_batches") == 3.25e-06 + assert info.get("input_cost_per_image_token_batches") == 2.25e-07 + assert info.get("input_cost_per_video_token_batches") == 6e-06 assert "input_cost_per_image" not in info assert "input_cost_per_audio_per_second" not in info assert "input_cost_per_video_per_second" not in info From 2b32f586c087ca94c3427c31eb83513c2a03c599 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:51:41 +0000 Subject: [PATCH 35/43] refactor(cost): extract batch modality rate lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index cbb9a45ada9..dce6b7299a2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2278,6 +2278,19 @@ def default_video_cost_calculator( return 0.0 +def _batch_rate( + model_info: ModelInfo, + key: Literal[ + "input_cost_per_audio_token_batches", + "input_cost_per_image_token_batches", + "input_cost_per_video_token_batches", + ], + fallback: float, +) -> float: + rate: Final = model_info.get(key) + return fallback if rate is None else cast(float, rate) + + def batch_cost_calculator( usage: Usage, model: str, @@ -2344,15 +2357,9 @@ def batch_cost_calculator( batch_details["video_tokens"], ) modality_rates: Final = ( - cast(float, model_info.get("input_cost_per_audio_token_batches")) - if model_info.get("input_cost_per_audio_token_batches") is not None - else input_cost_per_token_batches, - cast(float, model_info.get("input_cost_per_image_token_batches")) - if model_info.get("input_cost_per_image_token_batches") is not None - else input_cost_per_token_batches, - cast(float, model_info.get("input_cost_per_video_token_batches")) - if model_info.get("input_cost_per_video_token_batches") is not None - else input_cost_per_token_batches, + _batch_rate(model_info, "input_cost_per_audio_token_batches", input_cost_per_token_batches), + _batch_rate(model_info, "input_cost_per_image_token_batches", input_cost_per_token_batches), + _batch_rate(model_info, "input_cost_per_video_token_batches", input_cost_per_token_batches), ) total_prompt_cost = sum( tokens * rate From c0c5044c45bac0a696cd270c442b00d29cb2756e Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:57:50 +0000 Subject: [PATCH 36/43] fix(batches): keep modality token details in raw vertex batch usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 34 +++++++++++++++++-- .../test_litellm/batches/test_batch_utils.py | 32 +++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +++++++ 3 files changed, 76 insertions(+), 2 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 87f8fd3946e..397bc0a35a2 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -3,14 +3,14 @@ from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass from dataclasses import replace as dataclasses_replace from enum import Enum -from typing import Any, Final, Literal +from typing import Any, Final, Literal, cast import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch -from litellm.types.utils import ModelInfo, Usage +from litellm.types.utils import ModelInfo, PromptTokensDetailsWrapper, Usage from litellm.utils import token_counter @@ -310,6 +310,35 @@ def _aggregate_batch_cost_usage_models( ) +def _vertex_prompt_tokens_details( + usage_metadata: Mapping[str, object], +) -> PromptTokensDetailsWrapper | None: + raw_details: Final = usage_metadata.get("promptTokensDetails") + if not isinstance(raw_details, list): + return None + + raw_list: Final = cast(list[object], raw_details) + if not all(isinstance(detail, Mapping) for detail in raw_list): + return None + + details: Final = tuple(cast(Mapping[str, object], detail) for detail in raw_list) + normalized: Final = tuple( + (modality.upper(), token_count) + for detail in details + if isinstance(modality := detail.get("modality"), str) + and isinstance(token_count := detail.get("tokenCount"), int) + ) + if len(normalized) != len(details): + return None + + return PromptTokensDetailsWrapper( + text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), + audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), + image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), + video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), + ) + + def calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses: list[dict], model_name: str | None = None, @@ -356,6 +385,7 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=_prompt, completion_tokens=_completion, total_tokens=_total, + prompt_tokens_details=_vertex_prompt_tokens_details(cast(Mapping[str, object], usage_metadata)), ) try: diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 768ea332677..8b04d7af70a 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -695,6 +695,38 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch): assert result.failed_requests == 0 +def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "vertex_ai/gemini-embedding-2", + { + "input_cost_per_token_batches": 1e-7, + "input_cost_per_audio_token_batches": 3.25e-6, + "input_cost_per_image_token_batches": 2.25e-7, + "input_cost_per_video_token_batches": 6e-6, + }, + ) + responses = [ + { + "response": { + "usageMetadata": { + "promptTokenCount": 84, + "candidatesTokenCount": 0, + "totalTokenCount": 84, + "promptTokensDetails": [ + {"modality": "AUDIO", "tokenCount": 64}, + {"modality": "TEXT", "tokenCount": 20}, + ], + } + } + } + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2") + + assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7) + + def test_vertex_cost_skips_none_response_body(monkeypatch): import litellm.cost_calculator as cc diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8aa05cf8c7c..73245806bb2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29857,6 +29857,8 @@ export interface components { input_cost_per_audio_per_second_above_128k_tokens?: number | null; /** Input Cost Per Audio Token */ input_cost_per_audio_token?: number | null; + /** Input Cost Per Audio Token Batches */ + input_cost_per_audio_token_batches?: number | null; /** Input Cost Per Character */ input_cost_per_character?: number | null; /** Input Cost Per Character Above 128K Tokens */ @@ -29867,6 +29869,8 @@ export interface components { input_cost_per_image_above_128k_tokens?: number | null; /** Input Cost Per Image Token */ input_cost_per_image_token?: number | null; + /** Input Cost Per Image Token Batches */ + input_cost_per_image_token_batches?: number | null; /** Input Cost Per Pixel */ input_cost_per_pixel?: number | null; /** Input Cost Per Query */ @@ -29909,6 +29913,8 @@ export interface components { input_cost_per_video_per_second_above_8s_interval?: number | null; /** Input Cost Per Video Token */ input_cost_per_video_token?: number | null; + /** Input Cost Per Video Token Batches */ + input_cost_per_video_token_batches?: number | null; /** Itpm */ itpm?: number | null; /** Keepalive Seconds */ @@ -40071,6 +40077,8 @@ export interface components { input_cost_per_audio_per_second_above_128k_tokens?: number | null; /** Input Cost Per Audio Token */ input_cost_per_audio_token?: number | null; + /** Input Cost Per Audio Token Batches */ + input_cost_per_audio_token_batches?: number | null; /** Input Cost Per Character */ input_cost_per_character?: number | null; /** Input Cost Per Character Above 128K Tokens */ @@ -40081,6 +40089,8 @@ export interface components { input_cost_per_image_above_128k_tokens?: number | null; /** Input Cost Per Image Token */ input_cost_per_image_token?: number | null; + /** Input Cost Per Image Token Batches */ + input_cost_per_image_token_batches?: number | null; /** Input Cost Per Pixel */ input_cost_per_pixel?: number | null; /** Input Cost Per Query */ @@ -40123,6 +40133,8 @@ export interface components { input_cost_per_video_per_second_above_8s_interval?: number | null; /** Input Cost Per Video Token */ input_cost_per_video_token?: number | null; + /** Input Cost Per Video Token Batches */ + input_cost_per_video_token_batches?: number | null; /** Itpm */ itpm?: number | null; /** Keepalive Seconds */ From ca7364fb0568a1d7f6b085529d45da2487fcb62c Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:00:48 +0000 Subject: [PATCH 37/43] fix(batches): avoid strict lint violation in usage parser Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 397bc0a35a2..acc0036f27d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -3,7 +3,7 @@ from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass from dataclasses import replace as dataclasses_replace from enum import Enum -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger @@ -317,18 +317,18 @@ def _vertex_prompt_tokens_details( if not isinstance(raw_details, list): return None - raw_list: Final = cast(list[object], raw_details) - if not all(isinstance(detail, Mapping) for detail in raw_list): - return None + def _normalize(detail: object) -> tuple[str, int] | None: + if not isinstance(detail, Mapping): + return None + modality: Final = detail.get("modality") + token_count: Final = detail.get("tokenCount") + if not isinstance(modality, str) or not isinstance(token_count, int): + return None + return modality.upper(), token_count - details: Final = tuple(cast(Mapping[str, object], detail) for detail in raw_list) - normalized: Final = tuple( - (modality.upper(), token_count) - for detail in details - if isinstance(modality := detail.get("modality"), str) - and isinstance(token_count := detail.get("tokenCount"), int) - ) - if len(normalized) != len(details): + parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) + normalized: Final = tuple(detail for detail in parsed_details if detail is not None) + if len(normalized) != len(parsed_details): return None return PromptTokensDetailsWrapper( @@ -385,7 +385,7 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=_prompt, completion_tokens=_completion, total_tokens=_total, - prompt_tokens_details=_vertex_prompt_tokens_details(cast(Mapping[str, object], usage_metadata)), + prompt_tokens_details=_vertex_prompt_tokens_details(usage_metadata), ) try: From a5fc880c907356a79843987eecc6aa170263271c Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:03:32 +0000 Subject: [PATCH 38/43] refactor(vertex): move batch usage modality parsing under llms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 34 ++----------------- .../llms/vertex_ai/batches/transformation.py | 32 ++++++++++++++++- 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index acc0036f27d..26b4318da2d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -9,8 +9,9 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details +from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details from litellm.types.llms.openai import Batch -from litellm.types.utils import ModelInfo, PromptTokensDetailsWrapper, Usage +from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -310,35 +311,6 @@ def _aggregate_batch_cost_usage_models( ) -def _vertex_prompt_tokens_details( - usage_metadata: Mapping[str, object], -) -> PromptTokensDetailsWrapper | None: - raw_details: Final = usage_metadata.get("promptTokensDetails") - if not isinstance(raw_details, list): - return None - - def _normalize(detail: object) -> tuple[str, int] | None: - if not isinstance(detail, Mapping): - return None - modality: Final = detail.get("modality") - token_count: Final = detail.get("tokenCount") - if not isinstance(modality, str) or not isinstance(token_count, int): - return None - return modality.upper(), token_count - - parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) - normalized: Final = tuple(detail for detail in parsed_details if detail is not None) - if len(normalized) != len(parsed_details): - return None - - return PromptTokensDetailsWrapper( - text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), - audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), - image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), - video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), - ) - - def calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses: list[dict], model_name: str | None = None, @@ -385,7 +357,7 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=_prompt, completion_tokens=_completion, total_tokens=_total, - prompt_tokens_details=_vertex_prompt_tokens_details(usage_metadata), + prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata), ) try: diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index e63c80dd3cf..f5f1ab2068a 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final from urllib.parse import unquote @@ -8,7 +9,36 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest from litellm.types.llms.vertex_ai import * -from litellm.types.utils import LiteLLMBatch +from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper + + +def vertex_prompt_tokens_details( + usage_metadata: Mapping[str, object], +) -> PromptTokensDetailsWrapper | None: + raw_details: Final = usage_metadata.get("promptTokensDetails") + if not isinstance(raw_details, list): + return None + + def _normalize(detail: object) -> tuple[str, int] | None: + if not isinstance(detail, Mapping): + return None + modality: Final = detail.get("modality") + token_count: Final = detail.get("tokenCount") + if not isinstance(modality, str) or not isinstance(token_count, int): + return None + return modality.upper(), token_count + + parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) + normalized: Final = tuple(detail for detail in parsed_details if detail is not None) + if len(normalized) != len(parsed_details): + return None + + return PromptTokensDetailsWrapper( + text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), + audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), + image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), + video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), + ) class VertexAIBatchTransformation: From bf1bdb3045670322350e1789f210da8e41c5a8e9 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:07:44 +0000 Subject: [PATCH 39/43] fix(ci): keep cost map schema generated by the base branch generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/generate_model_prices_schema.py | 3 --- model_prices_and_context_window.schema.json | 9 +++------ 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index ee0e7f22eb1..ab29b70bdd4 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -172,10 +172,7 @@ COST_DESCRIPTIONS: dict[str, str] = { ), "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", - "input_cost_per_audio_token_batches": "USD per audio prompt token via the provider's batch API.", - "input_cost_per_image_token_batches": "USD per image prompt token via the provider's batch API.", "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", - "input_cost_per_video_token_batches": "USD per video prompt token via the provider's batch API.", "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index b7e0a9fd414..c2490041cf7 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -251,8 +251,7 @@ }, "input_cost_per_audio_token_batches": { "type": "number", - "minimum": 0, - "description": "USD per audio prompt token via the provider's batch API." + "minimum": 0 }, "input_cost_per_audio_token_priority": { "type": "number", @@ -283,8 +282,7 @@ }, "input_cost_per_image_token_batches": { "type": "number", - "minimum": 0, - "description": "USD per image prompt token via the provider's batch API." + "minimum": 0 }, "input_cost_per_pixel": { "type": "number", @@ -391,8 +389,7 @@ }, "input_cost_per_video_token_batches": { "type": "number", - "minimum": 0, - "description": "USD per video prompt token via the provider's batch API." + "minimum": 0 }, "input_dbu_cost_per_token": { "type": "number", From a5f00b9189fa1dad1515800b0e0797a2fa411099 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:15:22 +0000 Subject: [PATCH 40/43] fix(cost): drop unnecessary cast in batch rate lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index dce6b7299a2..39af725a231 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2288,7 +2288,7 @@ def _batch_rate( fallback: float, ) -> float: rate: Final = model_info.get(key) - return fallback if rate is None else cast(float, rate) + return fallback if rate is None else rate def batch_cost_calculator( From 168d0d3d997e11b39240479cf47ddd1675ba22c2 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:26:19 +0000 Subject: [PATCH 41/43] fix(cost): drop remaining unnecessary cast in batch cost calculator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 39af725a231..3dc6d81256b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2365,7 +2365,7 @@ def batch_cost_calculator( tokens * rate for tokens, rate in zip( ( - max(cast(int, usage.prompt_tokens) - audio_tokens - image_tokens - video_tokens, 0), + max((usage.prompt_tokens or 0) - audio_tokens - image_tokens - video_tokens, 0), audio_tokens, image_tokens, video_tokens, From 7761d044509e0ab0dae37ed6ded7835d4d06d7e9 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:37:51 +0000 Subject: [PATCH 42/43] test(vertex): cover malformed batch usage details Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_ai/batches/test_transformation.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 232c6413e78..e6126b02790 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -17,9 +17,9 @@ from unittest.mock import patch import pytest - from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, + vertex_prompt_tokens_details, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 VertexAIError, @@ -41,6 +41,22 @@ ENDPOINT_INPUT_FILE = ( ) +def test_vertex_prompt_tokens_details_rejects_malformed_details(): + assert vertex_prompt_tokens_details({"promptTokensDetails": [1]}) is None + assert vertex_prompt_tokens_details({"promptTokensDetails": [{"modality": "AUDIO"}]}) is None + assert ( + vertex_prompt_tokens_details( + { + "promptTokensDetails": [ + {"modality": "AUDIO", "tokenCount": 1}, + "malformed", + ] + } + ) + is None + ) + + # =========================================================================== # # transform_openai_batch_request_to_vertex_ai_batch_request # =========================================================================== # From 954dfa6ba74882eecb3109438c993d769c636c92 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:38:53 +0000 Subject: [PATCH 43/43] test(utils): allow modality batch cost fields in cost map schema test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 02ffaee0543..d5feda6f892 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -892,7 +892,10 @@ def validate_model_cost_values(model_data, exceptions=None): "input_cost_per_video_per_second_above_8s_interval", "input_cost_per_video_per_second_above_15s_interval", "input_cost_per_video_per_second_above_128k_tokens", + "input_cost_per_audio_token_batches", + "input_cost_per_image_token_batches", "input_cost_per_token_batches", + "input_cost_per_video_token_batches", "output_cost_per_token_batches", "input_cost_per_token_cache_hit", "cache_creation_input_token_cost", @@ -1041,7 +1044,10 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_second": {"type": "number"}, "input_cost_per_token": {"type": "number"}, "input_cost_per_token_above_128k_tokens": {"type": "number"}, + "input_cost_per_audio_token_batches": {"type": "number"}, + "input_cost_per_image_token_batches": {"type": "number"}, "input_cost_per_token_batches": {"type": "number"}, + "input_cost_per_video_token_batches": {"type": "number"}, "input_cost_per_token_cache_hit": {"type": "number"}, "input_cost_per_video_per_second": {"type": "number"}, "input_cost_per_video_per_second_above_8s_interval": {"type": "number"},